diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e4861d8fa..1d31c5288 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -642,6 +642,8 @@ jobs: - uses: actions/checkout@v4 - name: Providers doc generator drift (docs/api/providers.md must match generator output) run: python3 scripts/gen_providers_doc.py --check + - name: Provider boundaries (no call-level reads or provider-name branches in providers, removed names stay removed) + run: bash scripts/check_provider_boundaries.sh - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 with: diff --git a/.gitignore b/.gitignore index 6e6fabbd4..6567f9180 100644 --- a/.gitignore +++ b/.gitignore @@ -53,3 +53,6 @@ quirks-corpus/ # Generated audit artifacts — keep conclusions in Markdown, not raw output. /docs/quality-audit/**/lcov.info /docs/quality-audit/**/*.log + +# AI SDK protocol fixture recorder (scripts/aisdk-fixtures) — npm install output +scripts/aisdk-fixtures/node_modules/ diff --git a/CHANGELOG.md b/CHANGELOG.md index ea313a12b..e4dc50ace 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,14 +9,85 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Breaking +- Middleware wraps generation and streaming with callbacks for both operations, + parameter transforms and provider/model/supported-URL overrides. Explicit + model and provider ids take precedence over middleware overrides. + `NoSuchProvider` carries `provider_id`, `model_id`, `model_type` and + `available_providers`. Node and Python OpenAI constructors use Responses + by default. + +- 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, + build_aggregator_prompt, extract_text}`, `recording::new_call_id`, + `util::rfc3339_now`, two session helpers, `fetch_error_to_ai_mux_error`, + `extract_response_header_pairs`, the SigV4 signing internals, the logging + internals and `DEFAULT_MAX_JSON_RESPONSE_SIZE`. + +- Proxy settings belong to a transport instance, not to the process + (`aimux-provider-utils`). `init_proxy` and its set-once global are + removed; `ProxyConfig` moves to `fetch`. `ReqwestFetch::new(proxy)` is a + transport with its own proxy and connection pool, given to a provider as + its `fetch` setting, as the AI SDK leaves routing to `fetch`; the + WebSocket counterpart is `ws::TungsteniteConnector::new(proxy)`, given as + `web_socket`. `set_default_transport` replaces the default that providers + without `fetch`, built-in WebSocket connects and download hops read per + request. `PinnedFetch::new` takes the hop's `ProxyConfig`. + `aimux_init_proxy` (C ABI unchanged) and the Node and Python `initProxy` / + `init_proxy` now replace the default transport for later requests, every + time they are called, instead of taking effect only before the first + request. + +- Removed items the AI SDK has no counterpart for: `aimux_core::ModelId` + (and `model_id`, the Node `ModelId` type), the hidden + `aimux_provider_utils::download_guard` re-export (use + `aimux_core::download_guard`; `same_origin` stays re-exported), the + deprecated panicking converters `convert_prompt_to_anthropic`, + `convert_prompt_to_anthropic_full`, `convert_prompt_to_openai_messages` + and `convert_prompt_to_openai_messages_with_mode` (use the `_fallible` + forms), and the `ResponsesSystemMessageMode` alias (use + `SystemMessageMode`). + +- `aimux_provider_utils::create_standard_json_error_response_handler` is + removed. Failed-response handlers live in the provider modules, as the AI + SDK keeps `groqFailedResponseHandler` and DeepSeek's error schema in their + packages; build one with `create_json_error_response_handler`. The Groq + handler now reads only `error.type` as the provider code, and DeepSeek's + prefers `error.code` (a string or a number) over `error.type`. + +- `LanguageModel::supported_urls` and + `LanguageModelMiddleware::override_supported_urls` are async, as the AI + SDK's `supportedUrls` and `overrideSupportedUrls` may be a `PromiseLike`. + A wrapped model reads the override, else the inner model, when asked + instead of once at wrap time. The OpenAI-compatible `SupportedUrlsFn` + setting returns a `BoxFuture<'static, SupportedUrls>`. + +- The user operations download as the AI SDK's do, through a transport the + caller can replace (`aimux-core`, new public `download` module). + `generate_text` and `stream_text` read the model's `supported_urls()` + before the call and download each prompt file URL it does not cover (user + `FileUrl` parts and `file` items of tool outputs; not `data:` URLs), which + then reaches the provider as inline data; they used to pass every URL + through. `GenerateTextOptions::download` takes a `DownloadFunction` (the AI + SDK's `experimental_download`) instead. `AudioInput::Url` is new: + `transcribe` downloads it with `TranscriptionCallOptions::download`, and a + transcription model given one returns `AudioInput::url_error()`. + `generate_video` downloads every video returned as a URL with + `VideoCallOptions::download` and returns it as `VideoData::Binary` (it + returned the URL before). Without a download each operation uses + `create_download(None)`, the guarded download (public HTTP(S) URLs only, + pinned hops, no provider credentials, 2 GiB limit), so a Vertex `gs://` + video result needs an injected download. + - Provider results carry request and response information as the AI SDK does (`aimux-core`). `GenerateResult` has `request: Option` and `response: Option`; `ResponseInfo` gains `id` and holds the timestamp, model id, headers and body. `StreamResult` has `request` and `response: Option`. The flat `request_body` and `response_headers` fields are removed and `response` is optional; the call - layer fills a missing id, timestamp and model id. The user-facing results and - their JSON are unchanged. + layer fills a missing id, timestamp and model id. `GenerateTextResult` + exposes `request: RequestInfo` and `response: ResponseInfo`, including + headers and body. - The prompt a provider receives is modelled by role (`aimux-core`), as the AI SDK's `LanguageModelV4Message`. `LanguageModelPrompt` is @@ -139,11 +210,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 (`aimux-core`). A provider emits `tool::RawToolCall` (`input` is the raw argument text) in `GenerateContent::ToolCall(..)` and `StreamPart::ToolCall(..)`, and `tool::ToolResult` in the `ToolResult` - variants; these variants are newtype variants now. `StreamPart` is generic - over the tool-call type: a provider's `do_stream` yields `StreamPart` - (= `StreamPart`), `stream_text` and everything after it yield - `TextStreamPart` (= `StreamPart`, parsed input with `invalid` / - `error`). `RawToolCall` moved from `parse_tool_call` to `tool`. Wire format: + variants; these variants are newtype variants now. A provider's `do_stream` + yields `StreamPart` with raw tool-call input. `stream_text` and everything + after it yield `TextStreamPart` with parsed input and `invalid` / `error`. + Call-layer streams consume provider response-metadata events internally; + `TextStreamPart` has no `ResponseMetadata` variant. `RawToolCall` moved from `parse_tool_call` to `tool`. Wire format: enum variants keep their JSON shape; a provider-layer `StreamPart::ToolCall` has a string `input` and no `invalid` / `error`; a `GenerateContent::ToolCall` with a non-string `input` no longer @@ -230,8 +301,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Removed the generated `ProviderName` type in every binding (Rust enum, TS const object, Go/Java/Kotlin consts, Swift enum, Dart consts, Python `Literal`) and `scripts/gen_provider_names.py`. Provider names are plain - strings: `provider("groq", ...)`. Built-in and overlay-registered names now - share one string path. The provider list lives in + strings: Rust uses `create_provider("groq", PresetSettings)`; bindings retain + `provider("groq", ...)`, including overlay-registered names. The provider list lives in [docs/api/providers.md](docs/api/providers.md); Rust also gains `provider_names()`. @@ -244,9 +315,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `generateVideo` `poll` option (`aimux-core`, `aimux-providers`). `VideoModel::poll_config()` and `VideoPollConfig` are removed; unset `VideoCallOptions::poll` fields use the AI SDK defaults (5000 ms interval, - 600000 ms timeout). `RunwaymlConfig` loses `poll_interval`, `timeout`, - `with_poll_interval` and `with_timeout` (RunwayML used 2 s / 300 s); pass - `poll` per call instead. Each status check, its retries included, now races + 600000 ms timeout). RunwayML no longer paces its own polling (it used + 2 s / 300 s, first as settings, then as package constants); pass `poll` per + call instead. Each status check, its retries included, now races the time left in the poll budget and a hanging check fails with `Timeout` "Video generation timed out after {timeoutMs}ms." as upstream. @@ -406,6 +477,265 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `InvalidResponseData` (previously it started a call with an empty name). `DeltaToolCall.index` is now `Option`. +**Rust (provider factory, RFC-0036)** + +- The native Decision/Jev extension from master is retained. Its typed + `JevConfig` / `JevProvider::new` entry point and `Provider::decision_model` + remain available. Decision replay now takes the caller's registry; + recordings contain provider identity only and keep decision capabilities + separately under input. Decision retries use the core call-level policy. + +- `OpenAIProvider::language_model(id)` (and so `create_provider("openai", ..)` + and the registry id `openai:`) returns the **Responses** model, as + `@ai-sdk/openai` does. It returned the Chat Completions model. Call + `.chat(id)` for Chat Completions, for example against an endpoint that has + no `/responses` route. + The bindings' `openai(...)` constructors and the C ABI's `aimux_openai_new*` + are unchanged: they still create the Chat Completions model. + +- Removed crate-root exports `provider()`, `provider_handle`, + `provider_from_env`, `provider_discovery`, `provider_registry_entry`, + `ProviderOptions`, `ProviderProfile`, `register_provider` and + `load_providers_from_json`. Binding compatibility lives in + `aimux_providers::provider`; binding APIs are unchanged. Rust by-name + creation is `create_provider(name, PresetSettings)`; `provider_names()` + lists names, and `default_providers()` supplies the map for + `aimux_core::provider_registry::create_provider_registry(providers, options)` + with `provider:model` ids by default. + +- Redirects are handled in the request helper layer, not in the transport. + An ordinary API call follows every redirect, as the AI SDK's `postToApi` + does under fetch's default `redirect: "follow"`. A hop that leaves the + request's origin keeps only `User-Agent`, the rule the AI SDK's + `fetchWithValidatedRedirects` applies on the server, so no credential + header (`authorization`, `x-api-key`, `x-goog-api-key`, `api-key`, ...) + reaches another origin, and dropped headers stay dropped if a later hop + returns. `FetchRequest.credentials` (`RequestCredentials::Include` / + `Omit`, after the fetch `credentials` member) is `Omit` on such a hop, and + `SigV4Fetch` then sends it unsigned, without `x-amz-security-token`. Every + followed hop goes through the request's transport again, so SigV4 signs + the URL it sends. A `Fetch` never follows a redirect itself: + `RedirectPolicy`, `FetchRequest.redirect` and `PinnedFetch::unpinned` are + removed. Validated downloads keep their hop-by-hop guard. +- `Provider` follows the AI SDK's `ProviderV4`: `language_model`, + `embedding_model` and `image_model` are required and return + `NoSuchModel { model_id, model_type }` when a vendor has no such modality; + `transcription_model`, `speech_model`, `reranking_model` and `files` are + optional (`None` = not offered); video and search stay aimux extensions. + `Provider` gains `specification_version()`, returning `"v4"` like + `ProviderV4` and the model traits. + AI SDK model constructors return an `Arc`. Removed: `Provider::name()`, + `LanguageModel::config_snapshot()` and `list_models` on `Provider`. + Discovery is the separate `ProviderDiscovery` trait, accessed through + `Provider::discovery()`; `supported_urls()` is added to `LanguageModel`. +- AI SDK-aligned packages remove `XxxConfig`, config builders, `from_env()`, + `with_*` methods and `XxxProvider::new(...)`. Each package + has package-specific `XxxProviderSettings`, `create_xxx(settings)` + (validating package-specific settings, including base URLs, conflicting + credentials and template parameters) and a default instance `xxx()` where + the package provides one. + Models read a crate-private per-model config; there are no getters and no + snapshot. `Fetch` (HTTP) and package-specific WebSocket settings are + transport injection points. OpenAI, Anthropic and Google use fixed string + credentials and header maps, without a body-transform setting. The + OpenAI-compatible package exposes its upstream `transform_request_body`, + `supported_urls` and `convert_usage` hooks. Cohere and Mistral expose + `generate_id` settings. Package settings omit `name` and body transforms + when the upstream package does not offer them. Every request carries + `ai-sdk-/` in its user-agent. +- For packages using the API-key loader, environment fallbacks are evaluated + on every request: `api_key: None` reads the package's environment variable, `Some("")` is + sent verbatim and never falls back to the environment, and a missing key + fails the call (not the factory) with `AiMuxError::LoadApiKey { env_var, + description }`. A missing required setting (AWS region, Azure resource name, + Vertex project, …) fails the call with `AiMuxError::LoadSetting { env_var, + name }`. Credentials and custom headers otherwise follow each package's + settings; the generic compatible factory sends no auth for a missing or empty + key. Headers merge case-insensitively in package-specific order; custom and + call headers can override credential headers. +- Retry is a call-level concern only. `RetryConfig`, `retry_config()`, + `with_retry_config`, the `max_retries` fields on every provider config and + `ProviderRecord.max_retries` are deleted; the nine core operations call + `prepare_retries(max_retries, abort)` with the constant defaults 2 retries, + 2000 ms initial delay, factor 2 (jitter, `Retry-After` and the error history + are unchanged). Providers no longer retry anything themselves: `list_models` + and files uploads are single exchanges, and job-creating requests are sent + exactly once. +- `body_overrides` is removed everywhere (provider configs, builders, + `CallOptions`, `ProviderOptions`, `config_json`, binding `ProviderConfig`). + OpenAI-compatible provider-level body rewrites use its upstream + `transform_request_body` closure on the finished JSON body. `ProviderOptions`, + the C ABI `config_json` and the Node `ProviderConfig` that carry `max_retries` or + `body_overrides` now fail with `InvalidArgument` instead of ignoring them + (`reject_removed_provider_options` for other shapes). +- Recording: `RECORDING_SCHEMA` stays 2, and recordings written by aimux 0.5 + and earlier (schemas 1 and 2) still load; schema 3 and above are refused. + `ProviderRecord` holds identity only (`provider_id`, `provider`, + `model_id`); base URL, key source, profile, options and retry settings are + no longer recorded, and when read from an older recording they are + ignored. What this version cannot replay from an older recording fails at + replay time with an error that names it: its empty `provider_id` + (`ProviderRecord::registry_id`) for a rebuild from a registry, a prompt in + the older format (`InputRecord::unreadable_prompt`, loaded as empty), + older call options, or decision capabilities that only the older + `provider.profile` holds. `rebuild_provider(record, registry)` resolves + `provider_id` through the caller-supplied registry, retaining its settings. + If the default model method differs from the recorded one, pass the model + to `replay_with_model`. Recorded provider strings change with the next item. +- `model.provider()` follows the package's fixed identity or, where offered, + its `settings.name` override. Defaults: `openai.chat` / + `openai.responses` / `openai.embedding` / `openai.image` / `openai.speech` / + `openai.transcription` / `openai.files`; `anthropic.messages` (a custom name + is used verbatim; `{name}.files` derives from it), `anthropic-aws`, + `googleVertex.anthropic.messages`; `google.generative-ai` (+ `.video`, + `.files`), `google.vertex` (+ `.video`, `.transcription`), + `amazon-bedrock`; `azure.chat` / `.responses` / `.embeddings` / `.image` / + `.transcription` / `.speech`; `xai.responses`, `huggingface.responses`, + `mistral.chat` / `.embedding`, `cohere.chat` / `.textEmbedding` / + `.reranking`, `codex.responses`, `{name}.responses` (open_responses), + `voyage.embedding` / `.reranking`, `elevenlabs.speech` / `.transcription`, + `groq.chat`, `deepseek.chat`; compat and preset providers + `{name}.chat` / `.embedding` / `.image`; single-modality packages + `luma.image`, `deepgram.transcription`, `tavily.search`, `amazon-polly.speech`, + `blackForestLabs.image`, `cartesia.transcription`, … . Consumers that + compared `provider` to `"openai"` must match the prefix (the CLI and web + probes filter by `model.provider()`). +- providerOptions namespaces: `googleVertex` and `amazonBedrock` are the + canonical keys; the legacy `vertex` / `bedrock` keys are still read, and + provider metadata is written under both, as the upstream packages do + (Bedrock also reads `anthropic` for its Anthropic models). The + Anthropic package reads `anthropic` merged with the first segment of a + custom `name` (custom wins) and writes metadata under the custom key; + compat providers no longer read the `openai` key; `providerOptions.deepseek` + is honoured (`reasoningEffort`, `thinking`). +- Presets: `provider_registry.json` has 282 rows; Groq and DeepSeek are + vendor packages instead of preset rows. Former thin wrapper types are + deleted. The JSON is embedded and parsed once into a runtime descriptor + table in `preset.rs`, held in a `OnceLock` without leaking memory: + `PresetDescriptor`, `ParamSpec` and `DeriveSpec` own their data (`String`, + `Vec`) instead of `&'static` slices, and lose `Copy`; there is no `family` + field, generated presets source or per-name Rust function. Create presets + or vendor packages by name via `create_provider(name, PresetSettings)`. + Rows carry `auth: api_key | none` (`none` sends no + `Authorization` unless the row's optional key variable, such as + `VLLM_API_KEY`, is set; `PLACEHOLDER_API_KEY` is gone), `base_url_env` and + template `params` (env, default, derived host maps). Only declared + parameters are accepted, host parameters reject `/ @ : ?`, and an + unexpanded placeholder is an error. Unknown names are `NoSuchProvider`; + nothing falls back to OpenAI. `litellm_proxy` reads `LITELLM_PROXY_BASE_URL` + (the old wrapper read its API-key variable as a URL). Registry validation + runs in the table test, and + `scripts/check_provider_boundaries.sh` runs in CI. +- Decisions (#226) come from the providers: `OpenAIProvider::decision` and + `Provider::decision_model` on OpenAI, and `decision_model` on the official + runtime presets (ollama, llamacpp, localai, laya, vllm, sglang, + cloudflare). `laya` is a registry row. Decision recordings store the + provider identity only, and `rebuild_decision_provider` takes the + caller's `ProviderRegistry` instead of an API key. +- Groq and DeepSeek are standalone packages (`create_groq`, `create_deepseek`) + with their own chat models; the native OpenAI package no longer knows any + other vendor. Compat providers expose chat, embedding and image only (the + upstream set); the native OpenAI package warns on and drops `top_k`; + `OpenAICompatProfile` and `with_profile` are deleted in favour of dialect + hooks. +- Anthropic: the base URL includes `/v1` (endpoint `{base}/messages`; only the + bare `https://api.anthropic.com` is rewritten); `ANTHROPIC_BASE_URL` and + `OPENAI_BASE_URL` supply the base URL when the corresponding setting is + absent; `stream` is omitted on non-streaming calls; `providerOptions.anthropic.metadata.userId` maps to `metadata.user_id`; `anthropic-beta` is + sent on every host. Result-level `providerMetadata` is now populated + (`usage`, `stopSequence`, `iterations`, `container`, `contextManagement`, + under `anthropic` and the custom name), `usage.raw` is the provider's own + usage object, and `usage.iterations` now feeds the token totals. Vertex- + Anthropic and `anthropic_aws` run on the Anthropic core. +- Google: request bodies follow the AI SDK (`generationConfig` is always + sent, function tools use `parametersJsonSchema`, `providerOptions.google` + maps `thinkingConfig`, `responseModalities`, `safetySettings`, + `cachedContent`, …; response `modelId` comes from `modelVersion`). + Vertex: Gemini uses `v1beta1`, Anthropic-on-Vertex `v1`; a tuned model in + express mode now fails at request time instead of at construction. + Bedrock signs the final request bytes through `SigV4Fetch`; + `list_models` now calls the control-plane host + `bedrock.{region}.amazonaws.com` (it used `….api.amazonaws.com`). +- Azure defaults to `api_version = "v1"` with `/v1{path}`; a dated + `api_version` selects the deployment URL form. xAI, Hugging Face and Azure + default to the Responses API. **xAI and Hugging Face no longer have a Chat + Completions model**: the AI SDK packages serve Responses only, so `XaiModel` + and the Hugging Face chat model are removed (no `chat` / `chat_completions` + method). To keep calling those endpoints, use + `create_openai_compatible` with `https://api.x.ai/v1` or + `https://router.huggingface.co/v1`; the xAI-specific chat response + handling (citations, search parameters, 200-status errors) is not available + on that path. The xAI Responses model now sends `top_k` and warns for + `frequencyPenalty` and `presencePenalty`, as upstream does. The OpenAI default `language_model` and + `call` use Responses, matching upstream; `chat` remains an explicit accessor. Codex is `codex.responses`, its ChatGPT-account base URL + defaults to `https://chatgpt.com/backend-api/codex` (it was missing + `/codex`), and `store: false` is a package rule. `open_responses`: `url` → + `base_url`. The Azure deepseek / completion / MAI models and the Foundry item + type are not ported. +- Single-modality vendors (28 packages) run async jobs on package constants: + a fixed interval and attempt cap (previously unbounded loops now stop after + 6000 × 100 ms), a transient poll error spends one attempt, and the interval + is overridden through the package namespace (`pollIntervalMs`, luma + `pollIntervalMillis`) which replaces the old per-provider poll settings. + `cartesia` `version` and `runwayml` `poll_interval` / `timeout` settings are + removed (headers / constants); `google_pse` `cx` resolves settings, then + providerOptions, then `GOOGLE_CSE_ID` per request; DataForSEO without + credentials fails with `LoadApiKey`; searxng accepts an optional bearer + token; the user-agent suffix is gone. +- Removed from `aimux-providers`: `body_merge`, `openai_legacy`, `AzureAuth`, + `TokenProvider`, `VertexAuth`, `BedrockAuth`, `StaticBearerConfig` and the + `hmac` dependency (`sha2` / `hex` are dev-dependencies). +- Bedrock regions must be a single DNS label; invalid values fail the call + with `InvalidArgument` before a request is sent. +- `ExternalProviderEntry` debug output reports only whether `api_key` is + present, keeping literal keys out of debug logs. + +**C ABI** + +- New error codes `AIMUX_E_LOAD_API_KEY = 19` and `AIMUX_E_LOAD_SETTING = 20` + (appended; retired 4 is not reused). `aimux_error_provider_code` returns the + environment variable that was consulted and `aimux_error_provider_message` + the key's description (`"OpenAI"`) or the setting name (`"region"`); no new + symbol. `aimux_error_model_type` carries `NoSuchModel.model_type`. These two + failures used to surface as `AIMUX_E_INVALID_ARGUMENT`. +- `aimux_provider_new`, `aimux_provider_handle_new` and `config_json` resolve + names through the binding compatibility module (overlays, vendor packages + and presets); `config_json` accepts + `base_url`, `headers`, `organization`, `project` and `params`, and rejects + `max_retries` and `body_overrides` with `AIMUX_E_INVALID_ARGUMENT`. A + provider handle is provider + discovery. `aimux_register_providers` + entries reject the same two keys. + +**Node / Python** + +- Node `ProviderConfig`: `maxRetries` and `bodyOverrides` are kept in the type + only so that passing them throws `InvalidArgumentError`, for every native + constructor (`openai`, `anthropic`, + `google`, `cohere`, `mistral`, `xai`, `deepseek`, `bedrock`, `vertex`, + `anthropicAws`, `azure`, `provider`, `createProvider`), not only OpenAI. + `headers` now applies to all of them. New `params: Record` + fills a preset's template parameters. Python `config` / + `config_json` accept `params` and reject `max_retries` / `body_overrides` + the same way. +- New error classes `LoadAPIKeyError` (`envVar`, `description`) and + `LoadSettingError` (`envVar`, `settingName`) in Node; `LoadAPIKeyError` + (`env_var`, `description`) and `LoadSettingError` (`env_var`, + `setting_name`) in Python. Both extend `AimuxError`, no longer + `InvalidArgumentError`. +- Recorded and traced provider strings follow `"{name}.{method}"` + (`openai.chat`, not `openai`). + +**Go / Java / Kotlin / Swift / Dart** + +- New error codes 19 / 20 (`CodeLoadAPIKey` / `CodeLoadSetting`, + `LoadAPIKeyError` / `LoadSettingError`, `.loadApiKey` / `.loadSetting`), + with the consulted environment variable as `EnvVar` / `envVar`. +- The call-level `body_overrides` / `bodyOverrides` field is deleted from + `GenerateTextOptions` (the core no longer accepts the key; leaving it would + have been silently ignored). The provider `config_json` helpers + (`ProviderConfig` in Go and Dart) drop `max_retries` / `body_overrides` and + gain `params`; call-level `max_retries` is unchanged. + **All bindings** - Typed provider tools in Java, Kotlin, Swift, Python and Flutter carry @@ -428,8 +758,72 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `ToolApprovalResponse` are decoded. Tool-call repair reads the top-level `tool_calls` only. +### Added + +- `aimux_core::wrap_provider`, the AI SDK's `wrapProvider`: wraps every + language model of a provider with language-model middleware and, when + given, every image model with image-model middleware; all other members + are the provider's own. +- C ABI: `void aimux_drop_handle_ptr(void *handle)`, `aimux_drop_handle` with + the handle passed as a pointer-sized value (NULL is a no-op), for a + finalizer that can only register a `void (*)(void *)` callback. +- `aimux_providers::provider_api_key_env(name)`: the environment variable a + built-in provider reads its key from when `create_provider` gets none, and + `provider_default_base_url(name)`, the endpoint it calls when given no base + URL (`None` when that is built from settings). + ### Fixed +- `default_providers()`, `openai()` and `anthropic()` no longer panic on an + invalid `OPENAI_BASE_URL` / `ANTHROPIC_BASE_URL`. The variable is still + read when the provider is created, as the AI SDK does; an invalid value now + fails each call of that provider with `InvalidArgument` naming the + variable. An invalid explicit `base_url` still fails the factory. +- Anthropic Messages streams honor `include_raw_chunks`: each SSE chunk is + sent as `StreamPart::Raw` before the parts made from it, a chunk that is + not JSON as `Raw` with `null`, as the AI SDK does. +- OpenAI Responses: reasoning parts of one assistant message that share an + `itemId` become one input item, as the AI SDK does: one `item_reference` + with `store`, otherwise one `reasoning` item whose summary collects the + parts and whose `encrypted_content` is the last one given. Each part used + to become its own item, which the API rejects as a duplicate id. +- Recordings and logs redact Prodia's `X-Prodia-Key` and any header ending in + `subscription-key` (Azure's `Ocp-Apim-Subscription-Key`), with or without + separators; both were stored in clear text. +- Flutter: the `NativeFinalizer` that releases a garbage-collected `Model` or + `ProviderHandle` calls `aimux_drop_handle_ptr`. It called + `aimux_drop_handle(uint64_t)` through a `void (*)(void *)` pointer, which on + 32-bit ARM (armeabi-v7a) passes a garbage high word and leaks the handle. +- Python: every blocking native call (`generate_text`, `generate_object`, + `consume_stream_text`, `generate_text_as_openai`, + `ProviderHandle.list_models`, `get_model_specs`, the multimodal calls and + the recording flushes) releases the GIL while it waits, so other Python + threads keep running. +- aimux-web: an `env:VAR` key reference is accepted only for the variable the + provider's key is registered under (for example `env:OPENAI_API_KEY` for + `openai`); any other name is refused before it is read, on the calls, + cache-probe and replay paths. It read any variable of the console process. +- aimux-web: a call or cache probe whose `base_url` is not the provider's + default endpoint gets only a key carried in the request (plaintext, loopback + binds only): the Settings key and the environment are not read, `env:` is + refused, and with no key the call goes out without one. Providers that read + other credentials from the environment (Bedrock, Vertex, Codex, Polly, + DataForSEO, Anthropic on AWS) refuse a custom `base_url`. A client could + point `base_url` at its own host and receive the console's key. +- Recordings: a call that lost a recording event (the JSONL recorder's queue + was full or its writer gone) or was marked inconsistent (a duplicate + attempt, an exchange update that matches no exchange) is written with + `complete: false`, by the JSONL and the ring recorder, whether its last + event or a flush writes it. Such calls used to be marked complete. +- `aimux-replay` replays a decision recording live for any provider with a + decision model (OpenAI, the runtime presets, Jev), rebuilt from the same + registry as language recordings: the built-in providers, the recorded one + with `--api-key` when given. It accepted only Jev. +- aimux-web: a call at a custom `base_url` gets its own cache-probe trace + scope, named after the provider and that endpoint, so two endpoints of one + provider and model no longer match each other's prefix history. The trace + scope isolates by provider name and model id; `TraceLayer::with_scope_name` + names a scope for a host that calls one provider at several endpoints. - Gated aimux-core's ts-rs exports behind `ts-export`, so plain `cargo test` no longer rewrites tracked TypeScript types. `scripts/gen_ts_types.py` enables exports and rejects ungated export attributes with `--check`. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 959bcfa67..32e07985f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -9,8 +9,8 @@ to set up a development environment, run the tests, and submit changes. aimux/ ├── aimux-core/ # Core abstractions: LanguageModel / Provider / Message / StreamPart ├── aimux-providers/ # provider implementations + cassettes (counts: docs/api/providers.md) -├── aimux-stream/ # SSE parsing -├── aimux-provider-utils/ # HTTP utilities: retry, backoff, error parsing, API-key loading, streamed tool-call tracking +├── aimux-stream/ # SSE decoding +├── aimux-provider-utils/ # HTTP utilities: Fetch transport, Resolvable settings, error parsing, API-key loading, streamed tool-call tracking ├── aimux-ffi/ # C ABI (opaque handle + JSON + push callback) for non-native bindings ├── bindings/ # Node, Python, Swift, Kotlin, Flutter, Go, C — share one Rust core ├── rfc/ # Design docs (RFCs) @@ -38,6 +38,20 @@ cargo fmt --all -- --check cargo clippy --workspace --all-targets -- -D warnings ``` +Provider changes also run the generator and boundary checks CI runs in the +`contract-tests` job: + +```bash +python3 scripts/gen_providers_doc.py --check # docs/api/providers.md +python3 scripts/gen_ts_types.py --check # ts-rs types for the Node binding +bash scripts/check_provider_boundaries.sh # factory-shape rules for provider packages +``` + +The Node and Python bindings are separate workspaces: run `cargo fmt` and +`cargo clippy -- -D warnings` inside each, then `npm ci && npm run build && +npm test` (Node) or `maturin build --release`, `pip install` the wheel and +`python -m pytest tests/` (Python). + ## Testing Tests run entirely on **cassette playback** — no network access and no API keys @@ -62,10 +76,13 @@ aimux distinguishes three kinds of providers: 1. **Native protocol** — providers with their own request/response model (OpenAI, Anthropic, Google, Bedrock, …). Implement the full `convert`/model path and handle provider-specific differences. -2. **OpenAI-compatible thin wrapper** — the majority. Describe differences via - `OpenAICompatProfile` (top_k, tools, response_format, streaming usage, - request-body post-processing) so the thin wrapper does not erase - provider-specific behavior. +2. **OpenAI-compatible preset** — the majority. One row of + `aimux-providers/src/provider_registry.json` (base URL, key variable or + `auth: none`, optional URL variable and template parameters) that + is embedded and parsed once into a runtime table. Create presets by name + through `provider(name, ...)`. Vendors with behavior of their own get a package on + the shared compatible chat model (`groq/`, `deepseek/`) instead of a + branch on their name. 3. **Modality-specific** — speech, image, video, transcription, etc. Step-by-step checklists (with the generator and CI rules) live in diff --git a/README.md b/README.md index ba5627084..d12a9f0c3 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ aimux banner

-> **A unified LLM access layer written in Rust. One API for [327 AI providers](docs/api/providers.md).** +> **A unified LLM access layer written in Rust. One API for [327 provider names](docs/api/providers.md).** [![CI](https://github.com/arcships/aimux/actions/workflows/ci.yml/badge.svg)](https://github.com/arcships/aimux/actions/workflows/ci.yml) [![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE) @@ -32,13 +32,10 @@ For planned milestones and their current status, see the [roadmap](ROADMAP.md). ## Why aimux -- **327 providers** (as of 2026-09-24) — 251 registry-backed OpenAI-compatible - (unified `provider(name, ...)` entry) + 76 typed providers (native protocol - implementations such as OpenAI/Anthropic/Google/Bedrock/Vertex, local - engines like Ollama/vLLM, and speech/image/video/search modality providers). - Counts by category and the full list live in - [docs/api/providers.md](docs/api/providers.md) — generated, and the single - source of truth for provider counts (only the badge above repeats the total). +- **327 provider names** — 282 registry-backed OpenAI-compatible presets + + 45 vendor packages, covering text and other modalities. Local engines + such as Ollama/vLLM are presets. `provider_names()` lists the Rust names; + the generated catalogue lives in [docs/api/providers.md](docs/api/providers.md). - **Unified, object-safe interface** — the `LanguageModel` trait supports `Box` so providers are interchangeable without changing call sites. - **Full multimodal** — text, streaming, tool calling, embeddings, image, @@ -51,10 +48,12 @@ For planned milestones and their current status, see the [roadmap](ROADMAP.md). - **Composite models** (new in 0.3.0) — `RouterModel` routes each call with fallback (RFC-0021); `MoaModel` aggregates parallel reference models mixture-of-agents style (RFC-0022). Both are plain `LanguageModel`s. -- **Config-driven provider registry** — `provider-registry.json` describes - each registry-backed OpenAI-compatible provider (base URL, env var, profile - quirks: top_k, tools, response_format, streaming usage, max_tokens key); - one unified `provider(name, ...)` entry in every binding. +- **Config-driven provider registry** — `provider_registry.json` describes + each registry-backed OpenAI-compatible provider (base URL, env var, auth + mode, template parameters); the embedded JSON is parsed once into a runtime + descriptor table. Rust uses `create_provider(name, PresetSettings)` for + presets and vendor packages; bindings retain `provider(name, ...)`. + Packages follow the AI SDK shape: `XxxProviderSettings` + `create_xxx()` + a default `xxx()`. - **Fast and small** — Rust core, release profile tuned for binary size (`lto`, `codegen-units=1`, `panic="abort"`, `strip`, `opt-level="z"`). - **8 language bindings** from one core: Node, Python, Swift, Kotlin, Flutter, @@ -150,14 +149,12 @@ console.log(result.text) ```rust use aimux_core::prelude::*; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::openai::openai; #[tokio::main] async fn main() -> Result<(), AiMuxError> { - let provider = OpenAIProvider::new( - OpenAIConfig::new(std::env::var("OPENAI_API_KEY")?) - ); - let model = provider.model("gpt-4o"); + // The default provider loads OPENAI_API_KEY when a request is made. + let model = openai().chat("gpt-4o"); let result = generate_text( &model, @@ -222,9 +219,10 @@ use aimux_core::moa::{MoaConfig, MoaModel}; use aimux_core::router::{FallbackPolicy, RouterConfig, RouterModel, RuleRouter}; // Router: pick one model per call, fall back on failure (RFC-0021). +let provider = aimux_providers::openai::openai(); let children: Vec = vec![ - Arc::new(provider.model("gpt-4o")) as ChildModel, - Arc::new(provider.model("gpt-4o-mini")) as ChildModel, + Arc::new(provider.chat("primary-model-id")) as ChildModel, + Arc::new(provider.chat("secondary-model-id")) as ChildModel, ]; let routed = RouterModel::new( children.clone(), @@ -236,7 +234,7 @@ let routed = RouterModel::new( // MoA: fan out to references, aggregate their answers (RFC-0022). let moa = MoaModel::new( children, - Arc::new(provider.model("gpt-4o")) as ChildModel, + Arc::new(provider.chat("primary-model-id")) as ChildModel, MoaConfig::default(), ); @@ -275,35 +273,36 @@ session.close() ## Switch providers ```rust -// OpenAI → DeepSeek: only the provider name changes (registry-backed; -// key read from the provider's env var) -use aimux_providers::{provider, provider_from_env}; +use aimux_core::provider_registry::create_provider_registry; +use aimux_providers::{create_provider, default_providers, provider_names, PresetSettings}; -let model = provider("deepseek", None, "deepseek-chat", None)?; -let model = provider_from_env("deepseek", "deepseek-chat", None)?; +// Keys are loaded when a request is made. +let provider = create_provider("deepseek", PresetSettings::default())?; +let model = provider.language_model("deepseek-chat")?; + +// Or assemble the built-in providers and resolve a provider:model id. +let registry = create_provider_registry(default_providers(), Default::default()); +let model = registry.language_model("deepseek:deepseek-chat")?; +let names = provider_names(); // model usage is identical — it's all dyn LanguageModel ``` -All registry-backed OpenAI-compatible providers share one entry: -`provider(name, ...)` in every binding. Provider names are runtime strings; -the generated provider list is the discovery surface. -The retired per-provider shell types (`XxxConfig`/`XxxProvider`) are gone — -see [docs/API.md](docs/API.md#providers). - -> **Scope:** OpenAI-compatible → `provider(name, ...)`; others (Anthropic, -> multimodal, local…) → their constructors. List: -> [docs/api/providers.md](docs/api/providers.md). +Rust by-name creation covers vendor packages and all 282 presets. Each vendor +also has `create_xxx(XxxProviderSettings {..})` and a default instance `xxx()`. +Bindings keep their own `provider(name, ...)` API. Runtime model listing uses +`Provider::discovery()` and then `ProviderDiscovery::list_models()` — see +[docs/api/rust.md](docs/api/rust.md#providers). ## Provider coverage -Counts by category live in -[docs/api/providers.md](docs/api/providers.md#totals) — generated from the -registry and `lib.rs`, and the single source of truth. The rough shape of +The generated catalogue is [docs/api/providers.md](docs/api/providers.md). +The accepted Rust names come from `provider_names()`. The rough shape of the roster: -- **Native protocol** — OpenAI, Anthropic, Google, Bedrock, Vertex, Azure, Cohere, Mistral, xAI, Anthropic-AWS, Voyage, Codex, OpenRouter -- **OpenAI-compatible (registry-backed)** — Groq, Fireworks, Together, Perplexity, Ollama Cloud, DeepSeek, Alibaba Tongyi, Zhipu, Baidu, Tencent, Moonshot, SiliconFlow… -- **OpenAI-compatible (standalone + Vertex-hosted)** — Hugging Face, Ollama, vLLM, SGLang, Llama.cpp, LiteLLM Proxy, Vertex-hosted DeepSeek/Qwen/Llama… +- **Native protocol** — OpenAI, Anthropic, Google, Bedrock, Vertex, Azure, Cohere, Mistral, Groq, DeepSeek, xAI, Anthropic-AWS, Voyage +- **OpenAI-compatible (registry-backed)** — Fireworks, Together, Perplexity, OpenRouter, Ollama Cloud, Alibaba Tongyi, Zhipu, Baidu, Tencent, Moonshot, SiliconFlow… +- **Local and Vertex-hosted presets** — Ollama, vLLM, SGLang, Llama.cpp, LiteLLM Proxy, Vertex-hosted DeepSeek/Qwen/Llama… +- **Responses API for text** — xAI and Hugging Face (no chat-completions entry) - **Speech / transcription** — ElevenLabs, Deepgram, AssemblyAI, AWS Polly, Cartesia, Hume, Gladia, RevAI, LMNT, Fal - **Image / video** — Black Forest Labs, Replicate, Luma, Prodia, KlingAI, Recraft, Stability, RunwayML - **Embeddings / rerank / search** — Voyage, Jina, Tavily, Exa, Firecrawl, Serper, SearXNG, You.com… diff --git a/ROADMAP.md b/ROADMAP.md index 1added1d2..1ff8fa7fd 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -35,7 +35,7 @@ aimux 是 **provider 接入与治理的运行时**： │ ├─→ #174 Auth L1 ──→ #175 Auth L2 ← 依赖 B1 │ └─→ #167 transport-level replay ──→ #179 replay 子命令 ← 依赖 B1 ├─ C 轨 FFI（现 123 个 extern "C"，#166 基线 109）→ ops 协议（C1 升级为传输无关协议：op 表 / 二进制帧 / 版本协商） -│ └─ C2 构造器转发 shim 是 B4/B5 前置 +│ └─ C2 构造器转发 shim 取消：B4/B5 的直接构造调用点随 provider 工厂切换一次迁移（调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)） ├─ D 轨 8 种绑定重写为 ops 薄封装（依赖 C1）+ D8 类型镜像生成 ├─ E 轨 内部清理（#164 已合入，改为 master 上独立 PR）+ 文档三层重组 │ @@ -75,12 +75,12 @@ aimux 是 **provider 接入与治理的运行时**： | 项 | 内容 | |---|---| | B2–B8 | 17 个 LanguageModel 实现 → ~7 个协议；33 个 wrapper 退役；responses 家族合并（遵循 #166 修正 3/6） | -| #174 / #175 | registry `auth` schema + `apply_auth()`；Credential 解析序 + CredentialStore + TokenRefresher | +| #174 / #175 | registry `auth` 列作为 preset descriptor 数据，鉴权在各厂商包内以 headers 闭包 / fetch 装饰器实现；CredentialStore + TokenRefresher 随宿主回调后置（调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)） | | L2 补齐 | `StreamPart::Raw` / `provider_metadata` 在各协议接线——投影不下的不得静默丢弃 | | **S3** | `cargo-bloat` 审计 0.5.0 增重（Android `.so` 13.4 → 21 MB）；feature gating 仅作可选小体积路径，默认全量 | | ~~调研 RFC~~ | ~~L2 数据模型选型（AI SDK 形态 / Open Responses items / 自有），指标见 RFC-0036 §3~~ 撤销：选型已定为 AI SDK V4 形态（调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)） | -版本语义：Rust API 大 breaking（33 wrapper 退役）；C ABI 经 C2 转发不受影响。 +版本语义：Rust API、binding wire 与 C ABI 在全链路对齐切换时一次性 breaking（33 wrapper 退役、构造器替换），不经 C2 转发（调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)）。 ### 0.8.0 ——「ops 协议落地」（~4-5 周） @@ -92,13 +92,13 @@ aimux 是 **provider 接入与治理的运行时**： | D8 | 类型镜像生成：serde → JSON Schema → 各语言，`--check` 门禁扩到全部输出 | | P | stdio 入口往返开销纳入性能门禁 | -版本语义：新 ABI 加入，旧导出并存（#166 要求共存至少一个 minor 版本）；绑定按各自节奏迁移。“共存至少一个 minor”不再作为承诺：全链路对齐切换本身没有共存期，之后 ops ABI 引入时的旧符号也按切换门一次替换；C2 转发 shim 与 C4 旧导出清理随之取消（调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)）。 +版本语义：新 ABI 按切换门一次替换旧符号，不提供共存期（#166 原要求的“共存至少一个 minor”撤销）；全链路对齐切换本身同样没有共存期，C2 转发 shim 与 C4 旧导出清理随之取消（调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)）。 ### 0.9.0 ——「治理」（~3-4 周） | 项 | 内容 | |---|---| -| #167 | transport-level replay：mock 挂 HTTP 层跑真协议代码，覆盖全协议 / 全模态；`ProviderRecord` = registry 行 + protocol | +| #167 | transport-level replay：mock 挂 HTTP 层跑真协议代码，覆盖全协议 / 全模态；录制只存身份，不保存可重建的 `ProviderRecord`（调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)） | | 漂移检测 | #170 扩展：定期重录，与 cassette 字节级 diff | | 能力矩阵 | provider × 能力，每格由探测或录制支撑；后置到全链路对齐迁移之后（见 §6.1） | | CLI | `aimux probe / replay / diff`（#179 replay 子命令并入）；#181 debug CLI 并入 | @@ -107,7 +107,7 @@ aimux 是 **provider 接入与治理的运行时**： | 项 | 内容 | |---|---| -| C4 | 8 种绑定全部迁移后删除旧导出、旧头文件、旧 FFI 测试 | +| ~~C4~~ | ~~8 种绑定全部迁移后删除旧导出、旧头文件、旧 FFI 测试~~ 取消：旧导出在切换时随绑定一起删除，1.0 不再补删（调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)） | | 冻结范围 | ops 协议与 L0 / L1 契约。**L2 独立版本化，不随 1.0 冻结** | | 发布判据 | #166 ledger 全勾（B9 / #171 后半除外）；错误模型 #95 与请求管线 #164 稳定一个周期；各门禁（S2 / P）连续绿 | diff --git a/aimux-core/Cargo.toml b/aimux-core/Cargo.toml index 9530e33c6..9e0043176 100644 --- a/aimux-core/Cargo.toml +++ b/aimux-core/Cargo.toml @@ -32,6 +32,7 @@ httpdate = { workspace = true } rand = { workspace = true } base64 = "0.22" jsonschema = { workspace = true } +regex = "1" [dev-dependencies] tokio = { workspace = true, features = ["macros", "rt", "test-util"] } diff --git a/aimux-core/src/composite.rs b/aimux-core/src/composite.rs index a11e5712f..9d3098404 100644 --- a/aimux-core/src/composite.rs +++ b/aimux-core/src/composite.rs @@ -10,9 +10,9 @@ //! - [`ChildModel`] — the child-handle shape (`Arc`). //! - `step_context` — the recording context of one child step. //! - `child_max_retries` — the operation's retry count for a child call. -//! - [`add_usage`] — per-field `Usage` accumulation. -//! - [`extract_text`] — pull text out of a `GenerateContent` list. -//! - [`build_aggregator_prompt`] — assemble the MoA aggregator prompt. +//! - `add_usage` — per-field `Usage` accumulation. +//! - `extract_text` — pull text out of a `GenerateContent` list. +//! - `build_aggregator_prompt` — assemble the MoA aggregator prompt. //! //! See [`crate::trace::layer::TraceLayer`] for the established decorator pattern //! (`inner: Arc` implementing `LanguageModel`); composite @@ -43,7 +43,8 @@ pub(crate) fn step_context(step: impl Into) -> Option /// The retry count of a child call: the `max_retries` of the operation the /// composite runs in ([`OperationSettings::current`]). `None` outside an -/// operation, or when the operation sets none, leaves the child's default. +/// operation, or when the operation sets none, retries twice, the AI SDK +/// default. pub(crate) fn child_max_retries() -> Option { OperationSettings::current().and_then(|settings| settings.max_retries) } @@ -54,7 +55,7 @@ pub(crate) fn child_max_retries() -> Option { /// Each usage detail field is `Option`; `None` is treated as zero so a /// child that doesn't report a breakdown doesn't erase the other child's data. #[must_use] -pub fn add_usage(a: Usage, b: &Usage) -> Usage { +pub(crate) fn add_usage(a: Usage, b: &Usage) -> Usage { Usage { input_tokens: InputTokenUsage { total: opt_add(a.input_tokens.total, b.input_tokens.total), @@ -84,7 +85,7 @@ fn opt_add(a: Option, b: Option) -> Option { /// are kept; `Reasoning` / `ToolCall` / `Source` / `File` are dropped (MoA /// thin version — references contribute analysis text, not tool calls). #[must_use] -pub fn extract_text(content: &[GenerateContent]) -> String { +pub(crate) fn extract_text(content: &[GenerateContent]) -> String { let mut out = String::new(); for c in content { if let GenerateContent::Text { text, .. } = c { @@ -103,7 +104,7 @@ pub fn extract_text(content: &[GenerateContent]) -> String { /// The reference list may be empty — in that case the aggregator just runs the /// original prompt (degenerates to a single-model call). #[must_use] -pub fn build_aggregator_prompt( +pub(crate) fn build_aggregator_prompt( prompt: &LanguageModelPrompt, instructions: Option<&str>, references: &[(String, String)], diff --git a/aimux-core/src/decision_model.rs b/aimux-core/src/decision_model.rs index 36a867904..f7c3289ec 100644 --- a/aimux-core/src/decision_model.rs +++ b/aimux-core/src/decision_model.rs @@ -503,21 +503,13 @@ impl DecisionResult { #[async_trait] pub trait DecisionModel: Send + Sync { + /// Specification version (always `"v4"`). fn specification_version(&self) -> &'static str { "v4" } fn provider(&self) -> &str; fn model_id(&self) -> &str; fn capabilities(&self) -> DecisionCapabilities; - fn config_snapshot(&self) -> crate::recording::ProviderRecord { - let mut snapshot = - crate::recording::ProviderRecord::minimal(self.provider(), self.model_id()); - snapshot.profile = Some(serde_json::json!({"decision_capabilities": self.capabilities()})); - snapshot - } - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig::default() - } /// Perform one provider attempt. /// # Errors /// Returns provider, transport or invalid response errors. @@ -602,11 +594,7 @@ pub async fn decide( } let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; let abort_signal = options.abort_signal.clone(); - let retries = retry::prepare_retries( - options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, abort_signal.clone()); let context = crate::recording::recorder().map(|recorder| { let ctx = crate::recording::RecordingContext::new(crate::recording::new_call_id(), recorder); @@ -615,9 +603,12 @@ pub async fn decide( &options, model.provider(), model.model_id(), + &capabilities, + ); + ctx.recorder.record_provider( + &ctx.call_id, + &crate::recording::ProviderRecord::from_model(model.provider(), model.model_id()), ); - ctx.recorder - .record_provider(&ctx.call_id, &model.config_snapshot()); ctx }); let result = timeout::run( diff --git a/aimux-core/src/decision_replay.rs b/aimux-core/src/decision_replay.rs index 1b3c0ddb8..90bfdcd8e 100644 --- a/aimux-core/src/decision_replay.rs +++ b/aimux-core/src/decision_replay.rs @@ -51,16 +51,13 @@ impl MockDecisionReplayModel { "decision replay: no matching provider/model recording".into(), ) })?; - let capabilities = serde_json::from_value( - first - .provider - .profile - .as_ref() - .and_then(|v| v.get("decision_capabilities")) - .cloned() - .unwrap_or(Value::Null), - ) - .map_err(|e| AiMuxError::InvalidArgument(format!("decision replay capabilities: {e}")))?; + let capabilities = first.input.decision_capabilities.clone().ok_or_else(|| { + AiMuxError::InvalidArgument( + "decision replay: missing decision capabilities (recordings written by aimux \ + 0.5 or earlier keep them in provider.profile, which this version does not read)" + .into(), + ) + })?; Ok(Self { provider, model_id, diff --git a/aimux-core/src/download.rs b/aimux-core/src/download.rs index cfa7d0f45..17699b879 100644 --- a/aimux-core/src/download.rs +++ b/aimux-core/src/download.rs @@ -1,8 +1,31 @@ -//! Resolve generated file URLs before exposing their file data. +//! Downloads made by the user operations, as the AI SDK's `download` +//! options (ai@7.0.127 `util/download/`). +//! +//! - `transcribe` and `generate_video` take a [`Download`], which fetches +//! one URL (`transcribe/transcribe.ts:56`, `:109-117`; +//! `generate-video/generate-video.ts:136`, `:253-261`). +//! - `generate_text` and `stream_text` take a [`DownloadFunction`] for the +//! file URLs of the prompt (`util/download/download-function.ts`; +//! `prompt/convert-to-language-model-prompt.ts:36-55`). +//! +//! An operation given neither uses the guarded download ([`create_download`]): +//! every hop is validated and its connection pinned to the validated +//! addresses ([`crate::download_guard`]), and the provider's fetch and +//! credentials are never used for these URLs. A caller that needs another +//! transport (offline replay, a proxy, an authenticated store) passes its +//! own. +//! +//! The URLs of generated files in a text result are always fetched with the +//! guarded download: upstream's `resolveGeneratedFileData` takes no download +//! option either (`generate-text/resolve-generated-file-data.ts:14-41`). +use std::fmt; +use std::future::Future; use std::net::SocketAddr; +use std::sync::Arc; use futures::StreamExt; +use futures::future::BoxFuture; use crate::abort_signal::AbortSignal; use crate::download_guard::{validate_download_target, validate_download_url}; @@ -12,6 +35,177 @@ use crate::shared::{FileBytes, GeneratedFileData}; const DEFAULT_MAX_DOWNLOAD_SIZE: usize = 2 * 1024 * 1024 * 1024; const MAX_DOWNLOAD_REDIRECTS: usize = 10; +/// A downloaded file: its bytes and the media type the response reported +/// (`{ data, mediaType }`). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DownloadedFile { + /// The file's bytes. + pub data: Vec, + /// The response's `Content-Type`, when it had one. + pub media_type: Option, +} + +/// What a [`Download`] is asked to fetch (`{ url, abortSignal }`). +#[derive(Debug, Clone)] +pub struct DownloadRequest { + /// The URL to fetch. + pub url: String, + /// The operation's abort signal. + pub abort_signal: Option, +} + +type DownloadCallback = + dyn Fn(DownloadRequest) -> BoxFuture<'static, Result> + Send + Sync; + +/// Fetches one URL: the `download` option of `transcribe` and +/// `generate_video` (ai@7.0.127 `transcribe/transcribe.ts:109-117`, +/// `generate-video/generate-video.ts:253-261`). [`create_download`] builds +/// the default one. +#[derive(Clone)] +pub struct Download(Arc); + +impl Download { + /// Wrap a download callback. + #[must_use] + pub fn new(download: F) -> Self + where + F: Fn(DownloadRequest) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + Self(Arc::new(move |request| Box::pin(download(request)))) + } + + /// Fetch `url`. + /// + /// # Errors + /// + /// The callback's error. + pub async fn download( + &self, + url: &str, + abort_signal: Option<&AbortSignal>, + ) -> Result { + (self.0)(DownloadRequest { + url: url.to_owned(), + abort_signal: abort_signal.cloned(), + }) + .await + } +} + +impl fmt::Debug for Download { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("Download()") + } +} + +/// The guarded download, the AI SDK's `createDownload({ maxBytes })` +/// (`util/download/create-download.ts`): it fails for a file larger than +/// `max_bytes` (default 2 GiB) and for a URL the download guard rejects. +#[must_use] +pub fn create_download(max_bytes: Option) -> Download { + Download::new(move |request: DownloadRequest| async move { + let (data, media_type) = + download(&request.url, max_bytes, request.abort_signal.as_ref()).await?; + Ok(DownloadedFile { data, media_type }) + }) +} + +/// A file URL of a prompt, as a [`DownloadFunction`] receives it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PlannedDownload { + /// The file's URL. + pub url: String, + /// Whether the model fetches this URL itself: the part's media type and + /// URL match the model's + /// [`supported_urls`](crate::language_model::LanguageModel::supported_urls). + pub is_url_supported_by_model: bool, +} + +type DownloadFunctionCallback = dyn Fn(Vec) -> BoxFuture<'static, Result>, AiMuxError>> + + Send + + Sync; + +/// The prompt download of `generate_text` and `stream_text`: the AI SDK's +/// experimental `DownloadFunction` (`util/download/download-function.ts`), +/// given as `experimental_download`. +/// +/// It is called once per model call with every file URL of the prompt, in +/// prompt order, and returns one entry per URL in the same order: the +/// downloaded file, or `None` to send the URL to the model as it is. +/// [`create_default_download_function`] builds the default one. +#[derive(Clone)] +pub struct DownloadFunction(Arc); + +impl DownloadFunction { + /// Wrap a prompt download callback. + #[must_use] + pub fn new(download: F) -> Self + where + F: Fn(Vec) -> Fut + Send + Sync + 'static, + Fut: Future>, AiMuxError>> + Send + 'static, + { + Self(Arc::new(move |downloads| Box::pin(download(downloads)))) + } + + /// Download the planned URLs. + /// + /// # Errors + /// + /// The callback's error, or `InvalidArgument` when it does not return + /// one entry per URL. + pub async fn download( + &self, + downloads: Vec, + ) -> Result>, AiMuxError> { + let requested = downloads.len(); + let files = (self.0)(downloads).await?; + if files.len() != requested { + return Err(AiMuxError::InvalidArgument(format!( + "the download function returned {} results for {requested} URLs", + files.len() + ))); + } + Ok(files) + } +} + +impl fmt::Debug for DownloadFunction { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("DownloadFunction()") + } +} + +/// The default prompt download, the AI SDK's +/// `createDefaultDownloadFunction(download, abortSignal)` +/// (`util/download/download-function.ts`): it leaves the URLs the model +/// fetches itself to the model and downloads the others concurrently with +/// `download` (default [`create_download`]`(None)`). +#[must_use] +pub fn create_default_download_function( + download: Option, + abort_signal: Option, +) -> DownloadFunction { + let download = download.unwrap_or_else(|| create_download(None)); + DownloadFunction::new(move |downloads: Vec| { + let files = downloads.into_iter().map(|planned| { + let download = download.clone(); + let abort_signal = abort_signal.clone(); + async move { + if planned.is_url_supported_by_model { + Ok(None) + } else { + download + .download(&planned.url, abort_signal.as_ref()) + .await + .map(Some) + } + } + }); + futures::future::try_join_all(files) + }) +} + pub(crate) async fn resolve_generated_file_data( data: &GeneratedFileData, abort_signal: Option<&AbortSignal>, @@ -128,3 +322,94 @@ fn size_error(max_bytes: usize) -> AiMuxError { "download exceeded maximum size of {max_bytes} bytes" )) } + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Mutex; + + fn recording_download(requests: Arc>>) -> Download { + Download::new(move |request: DownloadRequest| { + requests.lock().unwrap().push(request.url.clone()); + async move { + Ok(DownloadedFile { + data: request.url.into_bytes(), + media_type: Some("image/png".into()), + }) + } + }) + } + + fn planned(url: &str, is_url_supported_by_model: bool) -> PlannedDownload { + PlannedDownload { + url: url.into(), + is_url_supported_by_model, + } + } + + #[tokio::test] + async fn the_default_function_downloads_only_the_urls_the_model_does_not_fetch() { + let requests = Arc::new(Mutex::new(Vec::new())); + let function = + create_default_download_function(Some(recording_download(requests.clone())), None); + let files = function + .download(vec![ + planned("https://a.test/1.png", false), + planned("https://a.test/2.png", true), + planned("https://a.test/3.png", false), + ]) + .await + .unwrap(); + assert_eq!( + files, + [ + Some(DownloadedFile { + data: b"https://a.test/1.png".to_vec(), + media_type: Some("image/png".into()), + }), + None, + Some(DownloadedFile { + data: b"https://a.test/3.png".to_vec(), + media_type: Some("image/png".into()), + }), + ] + ); + assert_eq!( + *requests.lock().unwrap(), + ["https://a.test/1.png", "https://a.test/3.png"] + ); + } + + #[tokio::test] + async fn a_function_must_return_one_entry_per_url() { + let function = DownloadFunction::new(|_| async { Ok(vec![None]) }); + let error = function + .download(vec![ + planned("https://a.test/1", false), + planned("https://a.test/2", false), + ]) + .await + .unwrap_err(); + assert!( + matches!(&error, AiMuxError::InvalidArgument(message) + if message == "the download function returned 1 results for 2 URLs"), + "{error:?}" + ); + } + + /// The default transport is the guarded download: a loopback URL is + /// refused before any connection is made. + #[tokio::test] + async fn the_default_download_is_guarded() { + let error = create_download(None) + .download("http://127.0.0.1:9/file", None) + .await + .unwrap_err(); + assert!( + matches!(&error, AiMuxError::InvalidArgument(message) + if message == "download URL targets a non-public address: 127.0.0.1"), + "{error:?}" + ); + } +} diff --git a/aimux-core/src/download_guard.rs b/aimux-core/src/download_guard.rs index 005f4c1b3..4c3f02915 100644 --- a/aimux-core/src/download_guard.rs +++ b/aimux-core/src/download_guard.rs @@ -37,7 +37,7 @@ pub fn hop_trusted_origin<'a>( } /// The URL as an error message may show it: no query, no userinfo. -fn without_query(parsed: &url::Url) -> String { +fn redacted(parsed: &url::Url) -> String { let mut redacted = parsed.clone(); redacted.set_query(None); // Both fail only for URLs that cannot carry userinfo at all. @@ -66,7 +66,7 @@ pub fn validate_download_url(url: &str) -> Result { if !matches!(parsed.scheme(), "http" | "https") { return Err(AiMuxError::InvalidArgument(format!( "download URL must use HTTP or HTTPS: {}", - without_query(&parsed) + redacted(&parsed) ))); } let host = parsed @@ -257,6 +257,16 @@ pub fn retain_user_agent(headers: &mut Vec<(String, String)>) { mod tests { use super::*; + #[test] + fn rejected_url_error_omits_userinfo_and_query() { + let error = validate_download_url("ftp://user:secret@example.com/file?token=t0k") + .expect_err("ftp must be blocked"); + let message = error.to_string(); + assert!(message.contains("ftp://example.com/file"), "{message}"); + assert!(!message.contains("secret"), "{message}"); + assert!(!message.contains("t0k"), "{message}"); + } + fn public(address: &str) -> bool { is_public_download_address(address.parse().unwrap()) } @@ -354,16 +364,6 @@ mod tests { } } - #[test] - fn rejected_url_error_omits_userinfo_and_query() { - let error = validate_download_url("ftp://user:secret@example.com/file?token=t0k") - .expect_err("ftp must be blocked"); - let message = error.to_string(); - assert!(message.contains("ftp://example.com/file"), "{message}"); - assert!(!message.contains("secret"), "{message}"); - assert!(!message.contains("t0k"), "{message}"); - } - #[test] fn resolved_address_validation_fails_closed_and_preserves_all_answers() { let empty = validate_resolved_download_addresses("empty.example", Vec::new()) diff --git a/aimux-core/src/embedding_model.rs b/aimux-core/src/embedding_model.rs index 63471e1bb..1bb94a110 100644 --- a/aimux-core/src/embedding_model.rs +++ b/aimux-core/src/embedding_model.rs @@ -36,9 +36,7 @@ pub struct EmbeddingCallOptions { pub abort_signal: Option, /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - /// retrying. `None` uses 2, the AI SDK default, unless the model's - /// `retry_config()` sets another count (model-level retry settings are - /// removed in the target design, docs/aisdk-architecture-alignment.md D14). + /// retrying. `None` uses 2, the AI SDK default. pub max_retries: Option, /// Per-call operation timeout. @@ -131,10 +129,6 @@ pub trait EmbeddingModel: Send + Sync { /// Provider-specific model ID, e.g. `"text-embedding-3-small"`. fn model_id(&self) -> &str; - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig::default() - } - /// Limit of how many embeddings can be generated in a single API call. /// /// `None` means the model has no fixed limit. The TS spec allows this to @@ -162,11 +156,7 @@ pub async fn embed( ) -> Result { let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; let abort_signal = options.abort_signal.clone(); - let retries = retry::prepare_retries( - options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, abort_signal.clone()); timeout::run( retries.retry(|| model.do_embed(&options)), abort_signal.as_ref(), diff --git a/aimux-core/src/error.rs b/aimux-core/src/error.rs index 25458f9d8..daea2d073 100644 --- a/aimux-core/src/error.rs +++ b/aimux-core/src/error.rs @@ -243,6 +243,34 @@ pub enum AiMuxError { #[error("invalid prompt: {0}")] InvalidPrompt(String), + /// No API key was passed and the fallback environment variable is unset + /// (the AI SDK's `LoadAPIKeyError`). Pre-HTTP: no request was made. + #[error( + "No API key found for {description}. \ + Please provide it via the `api_key` parameter \ + or set the `{env_var}` environment variable." + )] + LoadApiKey { + /// The environment variable consulted as the fallback. + env_var: String, + /// What the key is for, e.g. `"OpenAI"`. + description: String, + }, + + /// A required provider setting was not passed and its fallback + /// environment variable is unset (the AI SDK's `LoadSettingError`). + #[error( + "No value found for setting `{name}`. \ + Please provide it via the `{name}` parameter \ + or set the `{env_var}` environment variable." + )] + LoadSetting { + /// The environment variable consulted as the fallback. + env_var: String, + /// The setting's parameter name, e.g. `"region"`. + name: String, + }, + /// The prompt has tool calls without tool results (the AI SDK's /// `MissingToolResultsError`, `ai` `error/missing-tool-result-error.ts`). /// Message verbatim from upstream (`:15-20`). @@ -282,8 +310,13 @@ pub enum AiMuxError { /// Registry-level "provider name does not resolve" (the AI SDK's /// `NoSuchProviderError`). - #[error("No such provider: {provider_id}")] - NoSuchProvider { provider_id: String }, + #[error("No such provider: {provider_id} (available providers: {})", .available_providers.join(","))] + NoSuchProvider { + provider_id: String, + model_id: String, + model_type: String, + available_providers: Vec, + }, #[error("request timed out: {0}")] Timeout(String), @@ -367,6 +400,18 @@ pub fn is_retryable_status(status: u16) -> bool { } impl AiMuxError { + /// Build a [`NoSuchModel`](Self::NoSuchModel) for `model_id`. + /// + /// `model_type` is the AI SDK's `modelType` (`"languageModel"`, + /// `"embeddingModel"`, `"imageModel"`, …). + #[must_use] + pub fn no_such_model(model_id: &str, model_type: &str) -> Self { + Self::NoSuchModel { + model_id: model_id.to_string(), + model_type: model_type.to_string(), + } + } + /// Returns `true` for a malformed individual stream frame that does not /// prove the underlying transport has failed. Stream reducers may report /// these errors and continue polling for later, independently framed data. diff --git a/aimux-core/src/generate.rs b/aimux-core/src/generate.rs index d85a6e344..043fd8341 100644 --- a/aimux-core/src/generate.rs +++ b/aimux-core/src/generate.rs @@ -1,7 +1,9 @@ //! User-facing API: `generate_text` and `stream_text`. //! //! These are the primary entry points for users. They: -//! 1. Convert user prompt (`ModelPrompt`) to `LanguageModelPrompt` (provider-facing). +//! 1. Convert user prompt (`ModelPrompt`) to `LanguageModelPrompt` +//! (provider-facing), first downloading the file URLs the model does not +//! fetch itself. //! 2. Build `CallOptions` from user-facing options. //! 3. Call `LanguageModel::do_generate` / `do_stream`. //! 4. Record the model call as a [`StepResult`] and aggregate the steps into @@ -20,9 +22,12 @@ use serde_json::Value; use tracing::Instrument; use ts_rs::TS; +use crate::download::{DownloadFunction, create_default_download_function}; use crate::error::AiMuxError; use crate::language_model::LanguageModel; -use crate::language_model_message::convert_to_language_model_prompt; +use crate::language_model_message::{ + LanguageModelPrompt, convert_with_downloaded_assets, download_assets, +}; use crate::message::{ModelMessage, ModelPrompt}; use crate::operation::OperationSettings; use crate::options::{CallOptions, ToolChoice, prepare_tool_choice}; @@ -126,12 +131,8 @@ pub struct GenerateTextOptions { pub reasoning: Option, /// System instructions prepended to the prompt. pub instructions: Option, - /// Per-call request body overrides (deep-merged). See RFC-0017. - pub body_overrides: Option, /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - /// retrying. `None` uses 2, the AI SDK default, unless the model's - /// `retry_config()` sets another count (model-level retry settings are - /// removed in the target design, docs/aisdk-architecture-alignment.md D14). + /// retrying. `None` uses 2, the AI SDK default. pub max_retries: Option, /// Per-call timeout configuration (total / first-chunk / chunk idle). pub timeout: Option, @@ -155,6 +156,19 @@ pub struct GenerateTextOptions { #[serde(skip)] #[ts(skip)] pub repair_tool_call: Option, + /// The download of the prompt's file URLs, the AI SDK's + /// `experimental_download` (ai@7.0.127 `generate-text.ts:382-386`, + /// `stream-text.ts:439`). Before the model call it receives every file + /// URL of the prompt, each marked with whether the model's + /// `supported_urls` cover it; the files it returns are sent inline. + /// `None` downloads the URLs the model does not fetch itself with the + /// guarded download + /// ([`create_default_download_function`]). + /// + /// Runtime handle; never crosses the JSON boundary. + #[serde(skip)] + #[ts(skip)] + pub download: Option, } impl GenerateTextOptions { @@ -198,7 +212,6 @@ impl GenerateTextOptions { headers: self.headers, provider_options: self.provider_options, reasoning: self.reasoning, - body_overrides: self.body_overrides, abort_signal: self.abort_signal, include_raw_chunks: self.include_raw_chunks, } @@ -893,17 +906,23 @@ pub async fn generate_text( prompt: impl Into, options: GenerateTextOptions, ) -> Result { - let (messages, instructions) = - standardize_prompt(prompt.into(), options.instructions.as_deref())?; - let lm_prompt = convert_to_language_model_prompt(&messages, instructions, provider_id(model))?; - generate_text_from_language_model_prompt(model, lm_prompt, messages, options).await + let settings = options.operation_settings(); + settings + .clone() + .scope(generate_in_operation( + model, + OperationPrompt::User(prompt.into()), + options, + settings, + )) + .await } // Replay already has a provider prompt. Keep it intact instead of converting // through user content, which cannot represent every provider part. pub(crate) async fn generate_text_from_language_model_prompt( model: &dyn LanguageModel, - lm_prompt: crate::language_model_message::LanguageModelPrompt, + lm_prompt: LanguageModelPrompt, messages: Vec, options: GenerateTextOptions, ) -> Result { @@ -911,17 +930,66 @@ pub(crate) async fn generate_text_from_language_model_prompt( settings .clone() .scope(generate_in_operation( - model, lm_prompt, messages, options, settings, + model, + OperationPrompt::Model { + prompt: lm_prompt, + messages, + }, + options, + settings, )) .await } +/// The prompt of a text operation. +enum OperationPrompt { + /// A user prompt, prepared inside the operation ([`prepare_prompt`]). + User(ModelPrompt), + /// A model prompt sent as it is (replay), with the user messages it was + /// made from. + Model { + prompt: LanguageModelPrompt, + messages: Vec, + }, +} + +/// Standardize a user prompt, download the file URLs the model does not +/// fetch itself, and convert it for the model: the AI SDK's +/// `convertToLanguageModelPrompt` with the model's `supportedUrls` and the +/// operation's `experimental_download` (ai@7.0.127 +/// `prompt/convert-to-language-model-prompt.ts:36-55`, +/// `generate-text/generate-text.ts:963-971`). +/// +/// It runs inside the operation, before the model call: the download +/// observes the caller's abort signal and the operation timeout. +async fn prepare_prompt( + model: &dyn LanguageModel, + prompt: ModelPrompt, + instructions: Option<&str>, + download: Option, + abort_signal: Option<&AbortSignal>, + operation_timeout: timeout::OperationTimeout, +) -> Result<(LanguageModelPrompt, Vec), AiMuxError> { + let (messages, instructions) = standardize_prompt(prompt, instructions)?; + let supported_urls = model.supported_urls().await; + let download = + download.unwrap_or_else(|| create_default_download_function(None, abort_signal.cloned())); + let assets = timeout::run( + download_assets(&messages, &download, &supported_urls), + abort_signal, + operation_timeout, + ) + .await?; + let lm_prompt = + convert_with_downloaded_assets(&messages, instructions, provider_id(model), &assets)?; + Ok((lm_prompt, messages)) +} + // The operation body runs inside its settings' scope, where the model, the // composite children and the recorder read them. async fn generate_in_operation( model: &dyn LanguageModel, - lm_prompt: crate::language_model_message::LanguageModelPrompt, - messages: Vec, + prompt: OperationPrompt, options: GenerateTextOptions, settings: OperationSettings, ) -> Result { @@ -932,16 +1000,28 @@ async fn generate_in_operation( if let Some(output) = &output { output.validate()?; } + let operation_timeout = timeout::OperationTimeout::new(settings.timeout.unwrap_or_default())?; + let abort_signal = options.abort_signal.clone(); + + // 1. Convert the user prompt to the provider-facing prompt. + let (lm_prompt, messages) = match prompt { + OperationPrompt::User(prompt) => { + prepare_prompt( + model, + prompt, + operation_instructions.as_deref(), + options.download.clone(), + abort_signal.as_ref(), + operation_timeout, + ) + .await? + } + OperationPrompt::Model { prompt, messages } => (prompt, messages), + }; // 2. Build CallOptions. let call_options = options.into_call_options(lm_prompt); - let operation_timeout = timeout::OperationTimeout::new(settings.timeout.unwrap_or_default())?; - let abort_signal = call_options.abort_signal.clone(); - let retries = retry::prepare_retries( - settings.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(settings.max_retries, abort_signal.clone()); // 2a. RFC-0023: 关闭时零成本(M2 评审)——仅在录制开启时生成 call_id // 并绑定 recorder 快照;传输封闭由层 A 收尾声明(P1 无层 B,barrier @@ -955,8 +1035,10 @@ async fn generate_in_operation( model.provider(), model.model_id(), ); - ctx.recorder - .record_provider(&ctx.call_id, &model.config_snapshot()); + ctx.recorder.record_provider( + &ctx.call_id, + &crate::recording::ProviderRecord::from_model(model.provider(), model.model_id()), + ); // 层 A 早发 closed(defense-in-depth):无 HTTP 的调用(mock/local)也 // 能完成 barrier;真实 HTTP 的骨架 finalized=false 仍会挡写,由层 B // 结束再发 closed 完成。双发幂等。 @@ -1433,27 +1515,31 @@ async fn stream_in_operation( if let Some(output) = &output { output.validate()?; } - let (messages, instructions) = standardize_prompt(prompt, operation_instructions.as_deref())?; - let lm_prompt = convert_to_language_model_prompt(&messages, instructions, provider_id(model))?; + let stream_timeout = settings.timeout.unwrap_or_default(); + let operation_timeout = timeout::OperationTimeout::new(stream_timeout)?; + let (lm_prompt, messages) = prepare_prompt( + model, + prompt, + operation_instructions.as_deref(), + options.download.clone(), + options.abort_signal.as_ref(), + operation_timeout, + ) + .await?; // 2. Build CallOptions. let call_options = options.into_call_options(lm_prompt); - let stream_timeout = settings.timeout.unwrap_or_default(); - let operation_timeout = timeout::OperationTimeout::new(stream_timeout)?; - // Armed at operation start: providers await their first SSE event inside - // do_stream, so the first-chunk budget must already be counting there — - // a 200-then-silence server is otherwise unbounded when only - // first_chunk_ms is configured. + // Armed before the model call (after the prompt's downloads, which only + // the total and step deadlines bound): providers await their first SSE + // event inside do_stream, so the first-chunk budget must already be + // counting there — a 200-then-silence server is otherwise unbounded when + // only first_chunk_ms is configured. let first_chunk_deadline = stream_timeout .first_chunk_ms .map(|duration_ms| timeout::TimeoutDeadline::from_now("First chunk", duration_ms)) .transpose()?; let abort_signal = call_options.abort_signal.clone(); - let retries = retry::prepare_retries( - settings.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(settings.max_retries, abort_signal.clone()); // 2a. RFC-0023: 关闭时零成本(M2 评审)——仅在录制开启时生成 call_id // 并绑定 recorder 快照;传输封闭由层 A 收尾声明(P1 无层 B,barrier @@ -1467,8 +1553,10 @@ async fn stream_in_operation( model.provider(), model.model_id(), ); - ctx.recorder - .record_provider(&ctx.call_id, &model.config_snapshot()); + ctx.recorder.record_provider( + &ctx.call_id, + &crate::recording::ProviderRecord::from_model(model.provider(), model.model_id()), + ); // 层 A 早发 closed(defense-in-depth):无 HTTP 的调用(mock/local)也 // 能完成 barrier;真实 HTTP 的骨架 finalized=false 仍会挡写,由层 B // 结束再发 closed 完成。双发幂等。 diff --git a/aimux-core/src/generate/tests.rs b/aimux-core/src/generate/tests.rs index 104c53879..8fb143292 100644 --- a/aimux-core/src/generate/tests.rs +++ b/aimux-core/src/generate/tests.rs @@ -1318,6 +1318,181 @@ async fn generated_and_streamed_output_replay_the_same_messages() { ); } +/// A scripted model that fetches `https://model.test/` images itself and +/// records the prompt of each call. +struct UrlModel { + inner: ScriptedModel, + prompts: Mutex>, +} + +#[async_trait] +impl LanguageModel for UrlModel { + fn provider(&self) -> &str { + self.inner.provider() + } + + fn model_id(&self) -> &str { + self.inner.model_id() + } + + async fn do_generate(&self, options: &CallOptions) -> Result { + self.prompts.lock().unwrap().push(options.prompt.clone()); + self.inner.do_generate(options).await + } + + async fn do_stream(&self, options: &CallOptions) -> Result { + self.prompts.lock().unwrap().push(options.prompt.clone()); + self.inner.do_stream(options).await + } + + async fn supported_urls(&self) -> crate::language_model::SupportedUrls { + crate::language_model::SupportedUrls(HashMap::from([( + "image/*".to_owned(), + vec![regex::Regex::new(r"^https://model\.test/").unwrap()], + )])) + } +} + +fn url_prompt() -> Vec { + serde_json::from_value(json!([ + { "role": "user", "content": [ + { "type": "file_url", "url": "https://model.test/a.png", "media_type": "image/png" }, + { "type": "file_url", "url": "https://cdn.test/b.png", "media_type": "image/png" } + ] } + ])) + .unwrap() +} + +/// What the model received for the two URLs of [`url_prompt`]. +fn sent_files(model: &UrlModel) -> Vec { + let prompts = model.prompts.lock().unwrap(); + let [prompt] = &prompts[..] else { + panic!("one model call: {prompts:?}") + }; + let [crate::language_model_message::LanguageModelMessage::User { content, .. }] = &prompt[..] + else { + panic!("one user message: {prompt:?}") + }; + content + .iter() + .map(|part| match part { + crate::language_model_message::UserPart::File(file) => file.data.clone(), + crate::language_model_message::UserPart::Text(_) => panic!("file parts only"), + }) + .collect() +} + +/// A prompt download that records its requests and downloads every URL the +/// model does not fetch itself. +fn recording_prompt_download( + requests: std::sync::Arc>>, +) -> crate::download::DownloadFunction { + crate::download::DownloadFunction::new(move |planned: Vec| { + requests.lock().unwrap().extend(planned.clone()); + async move { + Ok(planned + .into_iter() + .map(|planned| { + (!planned.is_url_supported_by_model).then(|| crate::download::DownloadedFile { + data: vec![1, 2, 3], + media_type: None, + }) + }) + .collect()) + } + }) +} + +fn expected_files() -> Vec { + vec![ + crate::shared::FileData::Url { + url: "https://model.test/a.png".into(), + original_url: None, + }, + crate::shared::FileData::Data { + data: crate::shared::FileBytes::Binary(vec![1, 2, 3]), + }, + ] +} + +/// generate-text.ts:963-971 and convert-to-language-model-prompt.ts:36-55: +/// the prompt's URLs go to `experimental_download`, marked with the model's +/// `supportedUrls`; downloaded files are sent inline, the others as URLs. +#[tokio::test] +async fn generate_text_downloads_the_urls_the_model_does_not_fetch() { + let model = UrlModel { + inner: ScriptedModel::generating(generated( + vec![text_content("ok")], + FinishReasonUnified::Stop, + )), + prompts: Mutex::new(Vec::new()), + }; + let requests = std::sync::Arc::new(Mutex::new(Vec::new())); + let options = GenerateTextOptions { + download: Some(recording_prompt_download(requests.clone())), + ..Default::default() + }; + generate_text(&model, url_prompt(), options).await.unwrap(); + assert_eq!( + *requests.lock().unwrap(), + [ + crate::download::PlannedDownload { + url: "https://model.test/a.png".into(), + is_url_supported_by_model: true, + }, + crate::download::PlannedDownload { + url: "https://cdn.test/b.png".into(), + is_url_supported_by_model: false, + }, + ] + ); + assert_eq!(sent_files(&model), expected_files()); +} + +/// stream-text.ts:1343: `streamText` passes the same download. +#[tokio::test] +async fn stream_text_downloads_the_urls_the_model_does_not_fetch() { + let model = UrlModel { + inner: ScriptedModel::streaming(text_stream(&["ok"], FinishReasonUnified::Stop)), + prompts: Mutex::new(Vec::new()), + }; + let options = GenerateTextOptions { + download: Some(recording_prompt_download(std::sync::Arc::default())), + ..Default::default() + }; + let result = stream_text(&model, url_prompt(), options).await.unwrap(); + collect(result).await; + assert_eq!(sent_files(&model), expected_files()); +} + +/// Without `download`, the default downloads with the guarded download, +/// which refuses a loopback URL before the model is called. +#[tokio::test] +async fn the_default_prompt_download_is_guarded() { + let model = UrlModel { + inner: ScriptedModel::generating(generated( + vec![text_content("ok")], + FinishReasonUnified::Stop, + )), + prompts: Mutex::new(Vec::new()), + }; + let prompt: Vec = serde_json::from_value(json!([ + { "role": "user", "content": [ + { "type": "file_url", "url": "http://127.0.0.1:9/a.png", "media_type": "image/png" } + ] } + ])) + .unwrap(); + let error = generate_text(&model, prompt, GenerateTextOptions::default()) + .await + .unwrap_err(); + assert!( + matches!(&error, AiMuxError::InvalidArgument(message) + if message == "download URL targets a non-public address: 127.0.0.1"), + "{error:?}" + ); + assert!(model.prompts.lock().unwrap().is_empty()); +} + /// stream-text.ts:1725-1757: an abort after a step finished keeps the step; /// the result resolves with the recorded step and, without a `finish` part, /// the null total usage. Only a stream that recorded no step fails with the diff --git a/aimux-core/src/image_model.rs b/aimux-core/src/image_model.rs index e41ab3fe8..b2af51a1b 100644 --- a/aimux-core/src/image_model.rs +++ b/aimux-core/src/image_model.rs @@ -113,9 +113,7 @@ pub struct ImageCallOptions { pub abort_signal: Option, /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - /// retrying. `None` uses 2, the AI SDK default, unless the model's - /// `retry_config()` sets another count (model-level retry settings are - /// removed in the target design, docs/aisdk-architecture-alignment.md D14). + /// retrying. `None` uses 2, the AI SDK default. pub max_retries: Option, /// Per-call operation timeout. @@ -195,10 +193,6 @@ pub trait ImageModel: Send + Sync { /// Provider-specific model ID, e.g. `"dall-e-3"`. fn model_id(&self) -> &str; - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig::default() - } - /// Limit of how many images can be generated in a single API call. /// /// `None` means there is no fixed limit. The TS spec allows this to be a @@ -206,6 +200,16 @@ pub trait ImageModel: Send + Sync { /// `Option`. fn max_images_per_call(&self) -> Option; + /// Whether image file inputs are supported; `None` means unknown. + fn supports_file_inputs(&self) -> Option { + None + } + + /// Whether mask inputs are supported; `None` means unknown. + fn supports_mask_inputs(&self) -> Option { + None + } + /// Generate an array of images. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. @@ -223,11 +227,7 @@ pub async fn generate_image( ) -> Result { let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; let abort_signal = options.abort_signal.clone(); - let retries = retry::prepare_retries( - options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, abort_signal.clone()); timeout::run( retries.retry(|| model.do_generate(&options)), abort_signal.as_ref(), diff --git a/aimux-core/src/image_model_middleware.rs b/aimux-core/src/image_model_middleware.rs new file mode 100644 index 000000000..c1789aa8c --- /dev/null +++ b/aimux-core/src/image_model_middleware.rs @@ -0,0 +1,137 @@ +//! Image-model middleware and wrapping, aligned with the upstream V4 hooks. + +use std::sync::Arc; + +use async_trait::async_trait; +use futures::future::BoxFuture; + +use crate::error::AiMuxError; +use crate::image_model::{ImageCallOptions, ImageModel, ImageResult}; + +/// Optional hooks for overriding model properties and wrapping generation. +#[async_trait] +pub trait ImageModelMiddleware: Send + Sync { + fn override_provider(&self, _model: &dyn ImageModel) -> Option { + None + } + + fn override_model_id(&self, _model: &dyn ImageModel) -> Option { + None + } + + fn override_max_images_per_call(&self, _model: &dyn ImageModel) -> Option { + None + } + + /// `Some(None)` overrides support to unknown; `None` leaves it unchanged. + fn override_supports_file_inputs(&self, _model: &dyn ImageModel) -> Option> { + None + } + + /// `Some(None)` overrides support to unknown; `None` leaves it unchanged. + fn override_supports_mask_inputs(&self, _model: &dyn ImageModel) -> Option> { + None + } + + /// Transform the parameters passed to the wrapped operation. + /// + /// # Errors + /// Returns an error if the parameters cannot be transformed. + async fn transform_params( + &self, + params: ImageCallOptions, + _model: &dyn ImageModel, + ) -> Result { + Ok(params) + } + + /// Wrap generation, with a callback using the transformed parameters. + /// + /// # Errors + /// Returns a middleware error or the underlying generation error. + async fn wrap_generate( + &self, + _params: &ImageCallOptions, + _model: &dyn ImageModel, + do_generate: &(dyn Fn() -> BoxFuture<'_, Result> + Send + Sync), + ) -> Result { + do_generate().await + } +} + +struct WrappedImageModel { + model: Arc, + middleware: Arc, + provider: String, + model_id: String, + max_images_per_call: Option, + supports_file_inputs: Option, + supports_mask_inputs: Option, +} + +/// Apply middleware in input order: the first entry is the outermost wrapper. +#[must_use] +pub fn wrap_image_model( + model: Arc, + middleware: &[Arc], + model_id: Option<&str>, + provider_id: Option<&str>, +) -> Arc { + middleware.iter().rev().fold(model, |model, middleware| { + Arc::new(WrappedImageModel { + provider: provider_id + .map(str::to_string) + .or_else(|| middleware.override_provider(model.as_ref())) + .unwrap_or_else(|| model.provider().to_string()), + model_id: model_id + .map(str::to_string) + .or_else(|| middleware.override_model_id(model.as_ref())) + .unwrap_or_else(|| model.model_id().to_string()), + max_images_per_call: middleware + .override_max_images_per_call(model.as_ref()) + .or_else(|| model.max_images_per_call()), + supports_file_inputs: middleware + .override_supports_file_inputs(model.as_ref()) + .unwrap_or_else(|| model.supports_file_inputs()), + supports_mask_inputs: middleware + .override_supports_mask_inputs(model.as_ref()) + .unwrap_or_else(|| model.supports_mask_inputs()), + model, + middleware: Arc::clone(middleware), + }) + }) +} + +#[async_trait] +impl ImageModel for WrappedImageModel { + fn provider(&self) -> &str { + &self.provider + } + + fn model_id(&self) -> &str { + &self.model_id + } + + fn max_images_per_call(&self) -> Option { + self.max_images_per_call + } + + fn supports_file_inputs(&self) -> Option { + self.supports_file_inputs + } + + fn supports_mask_inputs(&self) -> Option { + self.supports_mask_inputs + } + + async fn do_generate(&self, options: &ImageCallOptions) -> Result { + let params = self + .middleware + .transform_params(options.clone(), self.model.as_ref()) + .await?; + let do_generate = || self.model.do_generate(¶ms); + self.middleware + .wrap_generate(¶ms, self.model.as_ref(), &do_generate) + .await + } +} diff --git a/aimux-core/src/language_model.rs b/aimux-core/src/language_model.rs index c6a37f50f..731f93ffa 100644 --- a/aimux-core/src/language_model.rs +++ b/aimux-core/src/language_model.rs @@ -3,12 +3,68 @@ //! Every LLM provider implements this trait. Users never call `do_generate` / //! `do_stream` directly; they call `generate_text` / `stream_text` instead. +use std::collections::HashMap; + use async_trait::async_trait; +use regex::Regex; use crate::error::AiMuxError; use crate::options::CallOptions; use crate::result::{GenerateResult, StreamResult}; +/// URL patterns a language model can fetch itself, keyed by media-type +/// pattern (`"image/*"`, `"application/pdf"`, …). Each value is a list of +/// regular expressions matched against the full URL. +/// +/// The empty default means the model fetches nothing itself. +#[derive(Debug, Clone, Default)] +pub struct SupportedUrls(pub HashMap>); + +impl SupportedUrls { + /// Whether no media type has a supported URL pattern. + #[must_use] + pub fn is_empty(&self) -> bool { + self.0.values().all(Vec::is_empty) + } + + /// Whether the model fetches `url` itself for a file of `media_type`, + /// the AI SDK's `isUrlSupported` (provider-utils@5.0.55 + /// `is-url-supported.ts`). + /// + /// Media types and the URL are compared in lower case. A `*` or `*/*` + /// key covers every media type and a `type/*` key every subtype of + /// `type`. A top-level `media_type` such as `image` matches only an + /// `image/*` key, since its subtype is unknown. The URL is supported when + /// one pattern of a matching key finds a match in it. + #[must_use] + pub fn is_supported(&self, url: &str, media_type: &str) -> bool { + let url = url.to_lowercase(); + let media_type = media_type.to_lowercase(); + let top_level_only = !media_type.contains('/'); + self.0 + .iter() + .filter(|(key, _)| { + let key = key.to_lowercase(); + let prefix = if key == "*" || key == "*/*" { + String::new() + } else { + key.replacen('*', "", 1) + }; + if prefix.is_empty() { + true + } else if top_level_only { + format!("{media_type}/") == prefix + } else if prefix.ends_with('/') { + media_type.starts_with(&prefix) + } else { + media_type == prefix + } + }) + .flat_map(|(_, patterns)| patterns) + .any(|pattern| pattern.is_match(&url)) + } +} + /// The unified language model trait (provider-facing). /// /// Aligned with Vercel AI SDK `LanguageModelV4`. @@ -34,22 +90,70 @@ pub trait LanguageModel: Send + Sync { /// Model identifier, e.g. `"gpt-4o"`. fn model_id(&self) -> &str; - /// Provider/model retry settings used when the caller does not override - /// the retry count. - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig::default() - } - /// Generate a complete (non-streaming) response. async fn do_generate(&self, options: &CallOptions) -> Result; /// Generate a streaming response. async fn do_stream(&self, options: &CallOptions) -> Result; - /// RFC-0023: 配置侧快照供录制。默认返回最小信息(provider/model_id); - /// 各 provider 覆盖以提供 base_url/profile/api_key 来源等完整配置。 - /// 透明 decorator(TraceLayer 等)必须覆盖并转发 inner 快照。 - fn config_snapshot(&self) -> crate::recording::ProviderRecord { - crate::recording::ProviderRecord::minimal(self.provider(), self.model_id()) + /// URLs the model can fetch on its own, keyed by media-type pattern + /// (e.g. `"image/*"`), AI SDK `supportedUrls`. Async, as the AI SDK's + /// value may be a `PromiseLike` (provider@4.0.21 + /// `language-model-v4.ts:36-38`). + /// + /// Default: empty, meaning every file part is sent as inline data. + /// Transparent decorators must forward the inner model's value. + /// `generate_text` and `stream_text` read it before each model call and + /// download the file URLs it does not cover + /// ([`crate::download::DownloadFunction`]). + async fn supported_urls(&self) -> SupportedUrls { + SupportedUrls::default() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn urls(entries: &[(&str, &str)]) -> SupportedUrls { + let mut map: HashMap> = HashMap::new(); + for (media_type, pattern) in entries { + map.entry((*media_type).to_owned()) + .or_default() + .push(Regex::new(pattern).unwrap()); + } + SupportedUrls(map) + } + + #[test] + fn a_subtype_wildcard_covers_its_subtypes_and_a_top_level_type() { + let supported = urls(&[("image/*", r"^https://")]); + assert!(supported.is_supported("https://a.test/x.png", "image/png")); + assert!(supported.is_supported("https://a.test/x", "image")); + assert!(!supported.is_supported("http://a.test/x.png", "image/png")); + assert!(!supported.is_supported("https://a.test/x.pdf", "application/pdf")); + } + + #[test] + fn an_exact_key_needs_the_same_media_type() { + let supported = urls(&[("application/pdf", r"^https://")]); + assert!(supported.is_supported("https://a.test/x.pdf", "application/pdf")); + assert!(!supported.is_supported("https://a.test/x.pdf", "application/pdfx")); + // A top-level type cannot say which subtype it is. + assert!(!supported.is_supported("https://a.test/x.pdf", "application")); + } + + #[test] + fn a_star_key_covers_every_media_type_and_case_is_ignored() { + let supported = urls(&[("*", r"^https://cdn\.test/")]); + assert!(supported.is_supported("HTTPS://CDN.TEST/X", "AUDIO/WAV")); + assert!(!supported.is_supported("https://other.test/x", "audio/wav")); + let supported = urls(&[("IMAGE/*", r"^https://")]); + assert!(supported.is_supported("https://a.test/x", "Image/PNG")); + } + + #[test] + fn no_entry_supports_nothing() { + assert!(!SupportedUrls::default().is_supported("https://a.test/x", "image/png")); } } diff --git a/aimux-core/src/language_model_message.rs b/aimux-core/src/language_model_message.rs index 61374048e..c732a8a99 100644 --- a/aimux-core/src/language_model_message.rs +++ b/aimux-core/src/language_model_message.rs @@ -7,7 +7,9 @@ //! the user-facing `ModelMessage`s. use crate::content::{self as user, AssistantContentPart, ToolContentPart, UserContentPart}; +use crate::download::{DownloadFunction, DownloadedFile, PlannedDownload}; use crate::error::AiMuxError; +use crate::language_model::SupportedUrls; use crate::message::{AssistantContent, ModelMessage, UserContent}; use crate::shared::{ FileBytes, FileData, GeneratedFileData, SharedProviderOptions, SharedProviderReference, @@ -279,12 +281,139 @@ fn text_part(text: &str, provider_options: Option<&SharedProviderOptions>) -> Te } } +/// The files downloaded for a prompt's URLs, by URL. +pub(crate) type DownloadedAssets = HashMap; + +/// A file part holding a file downloaded for its URL (upstream +/// `convertPartToLanguageModelPart`, `:567-578`): the downloaded media type +/// replaces the part's when the part's is not a full `type/subtype`. +fn downloaded_file_part( + file: &DownloadedFile, + media_type: &str, + filename: Option, + provider_options: Option, +) -> FilePart { + let media_type = match &file.media_type { + Some(downloaded) if !is_full_media_type(media_type) => downloaded, + _ => media_type, + }; + FilePart { + data: bytes(&file.data), + media_type: media_type.to_owned(), + filename, + provider_options, + } +} + +/// Upstream `isFullMediaType` (provider-utils@5.0.55 +/// `detect-media-type.ts:381-388`): `type/subtype` with a subtype that is not +/// empty and not `*`. +fn is_full_media_type(media_type: &str) -> bool { + media_type + .split_once('/') + .is_some_and(|(_, subtype)| !subtype.is_empty() && subtype != "*") +} + +/// Download the prompt's file URLs, the AI SDK's `downloadAssets` +/// (`ai` `prompt/convert-to-language-model-prompt.ts:444-540`). +/// +/// It plans the URL of every user `file` part and of every `file` item of +/// a `content` tool output (in tool and assistant messages), in prompt +/// order, marks the URLs the model fetches itself (`supported_urls`), and +/// calls `download` once with all of them. As upstream, assistant file parts +/// and the deprecated `file-url` / `image-url` tool output items are not +/// downloaded. Neither is a `data:` URL, which upstream turns into inline +/// data before planning and aimux sends as it is. +/// +/// # Errors +/// +/// The download function's error. +pub(crate) async fn download_assets( + messages: &[ModelMessage], + download: &DownloadFunction, + supported_urls: &SupportedUrls, +) -> Result { + fn plan_output(output: &user::ToolResultOutput, plan: &mut impl FnMut(&str, &str)) { + if let user::ToolResultOutput::Content { value } = output { + for item in value { + if let user::ToolResultContentPart::File { + data: FileData::Url { url, .. }, + media_type, + .. + } = item + { + plan(url, media_type); + } + } + } + } + + let mut planned = Vec::new(); + let mut plan = |url: &str, media_type: &str| { + if url + .get(..5) + .is_some_and(|scheme| scheme.eq_ignore_ascii_case("data:")) + { + return; + } + planned.push(PlannedDownload { + url: url.to_owned(), + is_url_supported_by_model: supported_urls.is_supported(url, media_type), + }); + }; + for message in messages { + match message { + ModelMessage::User { + content: UserContent::Parts(parts), + .. + } => { + for part in parts { + if let UserContentPart::FileUrl { + url, media_type, .. + } = part + { + plan(url, media_type); + } + } + } + ModelMessage::Tool { content, .. } => { + for part in content { + if let ToolContentPart::ToolResult(result) = part { + plan_output(&result.output, &mut plan); + } + } + } + ModelMessage::Assistant { + content: AssistantContent::Parts(parts), + .. + } => { + for part in parts { + if let AssistantContentPart::ToolResult(result) = part { + plan_output(&result.output, &mut plan); + } + } + } + _ => {} + } + } + let files = download.download(planned.clone()).await?; + Ok(planned + .into_iter() + .zip(files) + .filter_map(|(planned, file)| file.map(|file| (planned.url, file))) + .collect()) +} + /// A user part as the provider sees it: text, or one `file` part. /// /// An `image` part is deprecated upstream (`provider-utils` /// `types/content-part.ts:29`); it becomes a `file` part with a deprecation /// warning (`ai` `prompt/convert-to-language-model-prompt.ts:229-238`). -fn user_part(part: &UserContentPart, warnings: &mut Vec) -> Result { +fn user_part( + part: &UserContentPart, + assets: &DownloadedAssets, + warnings: &mut Vec, +) -> Result { Ok(match part { UserContentPart::Text { text, @@ -332,12 +461,10 @@ fn user_part(part: &UserContentPart, warnings: &mut Vec) -> Result UserPart::File(file_part( - url(file_url), - media_type, - None, - provider_options.as_ref(), - )), + } => UserPart::File(match assets.get(file_url) { + Some(file) => downloaded_file_part(file, media_type, None, provider_options.clone()), + None => file_part(url(file_url), media_type, None, provider_options.as_ref()), + }), UserContentPart::FileReference { media_type, reference: file_reference, @@ -460,6 +587,8 @@ struct Context<'a> { provider: Option<&'a str>, /// Deprecation warnings of the message being converted. warnings: Vec, + /// The files downloaded for the prompt's URLs. + assets: &'a DownloadedAssets, } fn tool_result_part( @@ -479,7 +608,12 @@ fn tool_result_part( "tool result {tool_call_id} has no tool name or preceding tool call" )) })?, - output: tool_result_output(&part.output, context.provider, &mut context.warnings)?, + output: tool_result_output( + &part.output, + context.provider, + context.assets, + &mut context.warnings, + )?, provider_options: part.provider_options.clone(), }) } @@ -491,6 +625,7 @@ fn tool_result_part( fn tool_result_output( output: &user::ToolResultOutput, provider: Option<&str>, + assets: &DownloadedAssets, warnings: &mut Vec, ) -> Result { Ok(match output.clone() { @@ -532,7 +667,7 @@ fn tool_result_output( user::ToolResultOutput::Content { value } => ToolResultOutput::Content { value: value .into_iter() - .map(|item| tool_result_content(item, provider, warnings)) + .map(|item| tool_result_content(item, provider, assets, warnings)) .collect::>()?, }, }) @@ -577,6 +712,7 @@ fn file_reference(reference: SharedProviderReference) -> FileData { fn tool_result_content( item: user::ToolResultContentPart, provider: Option<&str>, + assets: &DownloadedAssets, warnings: &mut Vec, ) -> Result { use user::ToolResultContentPart as Item; @@ -601,7 +737,12 @@ fn tool_result_content( media_type, filename, provider_options, - } => file(data, &media_type, filename, provider_options), + } => match &data { + FileData::Url { url, .. } if assets.contains_key(url) => ToolResultContent::File( + downloaded_file_part(&assets[url], &media_type, filename, provider_options), + ), + _ => file(data, &media_type, filename, provider_options), + }, Item::FileData { data, media_type, @@ -793,7 +934,7 @@ fn convert_message( UserContent::Text(text) => vec![UserPart::Text(text_part(text, None))], UserContent::Parts(parts) => parts .iter() - .map(|part| user_part(part, &mut context.warnings)) + .map(|part| user_part(part, context.assets, &mut context.warnings)) .filter(|part| { // remove empty text parts (`:242-243`) !matches!(part, Ok(UserPart::Text(TextPart { text, .. })) if text.is_empty()) @@ -875,6 +1016,9 @@ fn convert_message( /// - Consecutive tool messages become one; the earlier message's /// `provider_options` are merged into its last part (`:104-130`). /// - A tool message left empty is dropped. +/// - File URLs are sent as they are. `generate_text` and `stream_text` first +/// download the ones the model does not fetch itself +/// ([`crate::download::DownloadFunction`]). /// /// # Errors /// @@ -889,7 +1033,19 @@ pub fn convert_to_language_model_prompt( instructions: Option<&str>, provider: Option<&str>, ) -> Result { - convert_with_warnings(messages, instructions, provider).map(|(prompt, _)| prompt) + convert_with_warnings(messages, instructions, provider, &DownloadedAssets::new()) + .map(|(prompt, _)| prompt) +} + +/// [`convert_to_language_model_prompt`] with the files [`download_assets`] +/// downloaded for the prompt's URLs: a downloaded URL becomes inline data. +pub(crate) fn convert_with_downloaded_assets( + messages: &[ModelMessage], + instructions: Option<&str>, + provider: Option<&str>, + assets: &DownloadedAssets, +) -> Result { + convert_with_warnings(messages, instructions, provider, assets).map(|(prompt, _)| prompt) } /// [`convert_to_language_model_prompt`], also returning the deprecation @@ -898,6 +1054,7 @@ fn convert_with_warnings( messages: &[ModelMessage], instructions: Option<&str>, provider: Option<&str>, + assets: &DownloadedAssets, ) -> Result<(LanguageModelPrompt, Vec), AiMuxError> { let mut result = Vec::new(); if let Some(instructions) = instructions { @@ -911,6 +1068,7 @@ fn convert_with_warnings( tool_names: HashMap::new(), provider, warnings: Vec::new(), + assets, }; let mut warnings = Vec::new(); for message in messages { @@ -1096,7 +1254,12 @@ mod tests { items: Value, provider: Option<&str>, ) -> Result<(Vec, Vec), AiMuxError> { - let (prompt, warnings) = convert_with_warnings(&tool_output(items), None, provider)?; + let (prompt, warnings) = convert_with_warnings( + &tool_output(items), + None, + provider, + &DownloadedAssets::new(), + )?; let [LanguageModelMessage::Tool { content, .. }] = &prompt[..] else { panic!("one tool message: {prompt:?}") }; @@ -1310,7 +1473,8 @@ mod tests { ] } ])) .unwrap(); - let (prompt, warnings) = convert_with_warnings(&messages, None, None).unwrap(); + let (prompt, warnings) = + convert_with_warnings(&messages, None, None, &DownloadedAssets::new()).unwrap(); let [LanguageModelMessage::User { content, .. }] = &prompt[..] else { panic!() }; @@ -1326,3 +1490,156 @@ mod tests { ); } } + +#[cfg(test)] +mod download_tests { + //! Prompt downloads, against `ai@7.0.127` + //! `prompt/convert-to-language-model-prompt.ts` (`downloadAssets`, + //! `convertPartToLanguageModelPart`). + + use super::*; + use std::sync::{Arc, Mutex}; + + use regex::Regex; + use serde_json::json; + + fn messages() -> Vec { + serde_json::from_value(json!([ + { "role": "user", "content": [ + { "type": "file_url", "url": "https://model.test/a.png", "media_type": "image/png" }, + { "type": "file_url", "url": "https://cdn.test/b", "media_type": "image" }, + { "type": "file_url", "url": "data:image/png;base64,AAAA", "media_type": "image/png" } + ] }, + { "role": "assistant", "content": [ + { "type": "file_url", "url": "https://cdn.test/assistant.png", "media_type": "image/png" }, + { "type": "tool_call", "tool_call_id": "c1", "tool_name": "t", "input": {} } + ] }, + { "role": "tool", "content": [ + { "type": "tool_result", "tool_call_id": "c1", "tool_name": "t", + "output": { "type": "content", "value": [ + { "type": "file", "data": { "Url": { "url": "https://cdn.test/c.pdf" } }, + "media_type": "application/pdf" }, + { "type": "file-url", "url": "https://cdn.test/c.pdf", "media_type": "application/pdf" } + ] } } + ] } + ])) + .expect("valid messages") + } + + fn supported_urls() -> SupportedUrls { + SupportedUrls(HashMap::from([( + "image/*".to_owned(), + vec![Regex::new(r"^https://model\.test/").unwrap()], + )])) + } + + /// A download function that records what it is asked for and downloads + /// every URL the model does not fetch itself. + fn recording(planned: Arc>>) -> DownloadFunction { + DownloadFunction::new(move |downloads: Vec| { + planned.lock().unwrap().extend(downloads.clone()); + async move { + Ok(downloads + .into_iter() + .map(|download| { + (!download.is_url_supported_by_model).then(|| DownloadedFile { + data: download.url.into_bytes(), + media_type: Some("image/jpeg".into()), + }) + }) + .collect()) + } + }) + } + + #[tokio::test] + async fn user_file_urls_and_tool_output_files_are_planned_in_prompt_order() { + let planned = Arc::new(Mutex::new(Vec::new())); + download_assets(&messages(), &recording(planned.clone()), &supported_urls()) + .await + .unwrap(); + let planned = planned.lock().unwrap().clone(); + assert_eq!( + planned, + [ + PlannedDownload { + url: "https://model.test/a.png".into(), + is_url_supported_by_model: true, + }, + PlannedDownload { + url: "https://cdn.test/b".into(), + is_url_supported_by_model: false, + }, + PlannedDownload { + url: "https://cdn.test/c.pdf".into(), + is_url_supported_by_model: false, + }, + ], + "no data: URL, no assistant file part, no deprecated file-url item" + ); + } + + #[tokio::test] + async fn downloaded_urls_become_inline_data_and_the_others_stay_urls() { + let assets = download_assets( + &messages(), + &recording(Arc::new(Mutex::new(Vec::new()))), + &supported_urls(), + ) + .await + .unwrap(); + let prompt = convert_with_downloaded_assets(&messages(), None, None, &assets).unwrap(); + + let LanguageModelMessage::User { content, .. } = &prompt[0] else { + panic!("user message: {:?}", prompt[0]) + }; + let files: Vec<_> = content + .iter() + .map(|part| match part { + UserPart::File(file) => (file.data.clone(), file.media_type.as_str()), + UserPart::Text(_) => panic!("file parts only"), + }) + .collect(); + assert_eq!( + files, + [ + // The model fetches it: the URL is sent. + (url("https://model.test/a.png"), "image/png"), + // Downloaded: a top-level media type takes the downloaded one. + (bytes(b"https://cdn.test/b"), "image/jpeg"), + (url("data:image/png;base64,AAAA"), "image/png"), + ] + ); + + let LanguageModelMessage::Assistant { content, .. } = &prompt[1] else { + panic!("assistant message: {:?}", prompt[1]) + }; + assert!( + matches!(&content[0], AssistantPart::File(file) + if file.data == url("https://cdn.test/assistant.png")), + "{content:?}" + ); + + let LanguageModelMessage::Tool { content, .. } = &prompt[2] else { + panic!("tool message: {:?}", prompt[2]) + }; + let [ToolPart::ToolResult(result)] = &content[..] else { + panic!("one tool result: {content:?}") + }; + let ToolResultOutput::Content { value } = &result.output else { + panic!("content output: {:?}", result.output) + }; + let [ + ToolResultContent::File(file), + ToolResultContent::File(deprecated), + ] = &value[..] + else { + panic!("two files: {value:?}") + }; + // A full media type is kept. + assert_eq!(file.data, bytes(b"https://cdn.test/c.pdf")); + assert_eq!(file.media_type, "application/pdf"); + // The deprecated file-url item is not downloaded. + assert_eq!(deprecated.data, url("https://cdn.test/c.pdf")); + } +} diff --git a/aimux-core/src/language_model_middleware.rs b/aimux-core/src/language_model_middleware.rs new file mode 100644 index 000000000..bebd88973 --- /dev/null +++ b/aimux-core/src/language_model_middleware.rs @@ -0,0 +1,256 @@ +//! Language-model middleware and wrapping, aligned with the upstream V4 hooks. + +use std::sync::Arc; + +use async_trait::async_trait; +use futures::future::BoxFuture; + +use crate::error::AiMuxError; +use crate::language_model::{LanguageModel, SupportedUrls}; +use crate::options::CallOptions; +use crate::result::{GenerateResult, StreamResult}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LanguageModelOperation { + Generate, + Stream, +} + +/// Optional hooks for transforming parameters and wrapping either operation. +/// The supplied callbacks invoke either operation using the transformed parameters. +#[async_trait] +pub trait LanguageModelMiddleware: Send + Sync { + fn override_provider(&self, _model: &dyn LanguageModel) -> Option { + None + } + + fn override_model_id(&self, _model: &dyn LanguageModel) -> Option { + None + } + + /// Replaces the wrapped model's supported URLs. Async, as the AI SDK's + /// `overrideSupportedUrls` may return a `PromiseLike` (provider@4.0.21 + /// `language-model-v4-middleware.ts:33-35`). + async fn override_supported_urls(&self, _model: &dyn LanguageModel) -> Option { + None + } + + async fn transform_params( + &self, + _operation: LanguageModelOperation, + params: CallOptions, + _model: &dyn LanguageModel, + ) -> Result { + Ok(params) + } + + async fn wrap_generate( + &self, + _params: &CallOptions, + _model: &dyn LanguageModel, + do_generate: &(dyn Fn() -> BoxFuture<'_, Result> + Send + Sync), + _do_stream: &(dyn Fn() -> BoxFuture<'_, Result> + Send + Sync), + ) -> Result { + do_generate().await + } + + async fn wrap_stream( + &self, + _params: &CallOptions, + _model: &dyn LanguageModel, + _do_generate: &( + dyn Fn() -> BoxFuture<'_, Result> + Send + Sync + ), + do_stream: &(dyn Fn() -> BoxFuture<'_, Result> + Send + Sync), + ) -> Result { + do_stream().await + } +} + +struct WrappedLanguageModel { + model: Arc, + middleware: Arc, + provider: String, + model_id: String, +} + +/// Apply middleware in input order: the first entry is the outermost wrapper. +#[must_use] +pub fn wrap_language_model( + model: Arc, + middleware: &[Arc], + model_id: Option<&str>, + provider_id: Option<&str>, +) -> Arc { + middleware.iter().rev().fold(model, |model, middleware| { + let provider = provider_id + .map(str::to_string) + .or_else(|| middleware.override_provider(model.as_ref())) + .unwrap_or_else(|| model.provider().to_string()); + let model_id = model_id + .map(str::to_string) + .or_else(|| middleware.override_model_id(model.as_ref())) + .unwrap_or_else(|| model.model_id().to_string()); + Arc::new(WrappedLanguageModel { + model, + middleware: Arc::clone(middleware), + provider, + model_id, + }) + }) +} + +#[async_trait] +impl LanguageModel for WrappedLanguageModel { + fn provider(&self) -> &str { + &self.provider + } + + fn model_id(&self) -> &str { + &self.model_id + } + + /// The middleware's override, else the wrapped model's value, as + /// `overrideSupportedUrls?.({ model }) ?? model.supportedUrls` + /// (ai@7.0.127 `wrap-language-model.ts:79`), resolved when read. + async fn supported_urls(&self) -> SupportedUrls { + match self + .middleware + .override_supported_urls(self.model.as_ref()) + .await + { + Some(urls) => urls, + None => self.model.supported_urls().await, + } + } + + async fn do_generate(&self, options: &CallOptions) -> Result { + let params = self + .middleware + .transform_params( + LanguageModelOperation::Generate, + options.clone(), + self.model.as_ref(), + ) + .await?; + let do_generate = || self.model.do_generate(¶ms); + let do_stream = || self.model.do_stream(¶ms); + self.middleware + .wrap_generate(¶ms, self.model.as_ref(), &do_generate, &do_stream) + .await + } + + async fn do_stream(&self, options: &CallOptions) -> Result { + let params = self + .middleware + .transform_params( + LanguageModelOperation::Stream, + options.clone(), + self.model.as_ref(), + ) + .await?; + let do_generate = || self.model.do_generate(¶ms); + let do_stream = || self.model.do_stream(¶ms); + self.middleware + .wrap_stream(¶ms, self.model.as_ref(), &do_generate, &do_stream) + .await + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use regex::Regex; + + use super::*; + + fn urls(media_type: &str, pattern: &str) -> SupportedUrls { + SupportedUrls(HashMap::from([( + media_type.to_string(), + vec![Regex::new(pattern).unwrap()], + )])) + } + + fn keys(urls: &SupportedUrls) -> Vec<&str> { + let mut keys: Vec<&str> = urls.0.keys().map(String::as_str).collect(); + keys.sort_unstable(); + keys + } + + struct Model; + + #[async_trait] + impl LanguageModel for Model { + fn provider(&self) -> &str { + "mock" + } + + fn model_id(&self) -> &str { + "mock-model" + } + + async fn supported_urls(&self) -> SupportedUrls { + tokio::task::yield_now().await; + urls("image/*", "^https://") + } + + async fn do_generate(&self, _: &CallOptions) -> Result { + unreachable!("not called") + } + + async fn do_stream(&self, _: &CallOptions) -> Result { + unreachable!("not called") + } + } + + struct Override(Option, AtomicUsize); + + #[async_trait] + impl LanguageModelMiddleware for Override { + async fn override_supported_urls( + &self, + model: &dyn LanguageModel, + ) -> Option { + self.1.fetch_add(1, Ordering::SeqCst); + assert_eq!(model.model_id(), "mock-model"); + tokio::task::yield_now().await; + self.0.clone() + } + } + + /// `overrideSupportedUrls?.({ model }) ?? model.supportedUrls` + /// (ai@7.0.127 `wrap-language-model.ts:79`): an override wins, its + /// absence forwards the wrapped model's value, and both may be async. + #[tokio::test] + async fn supported_urls_come_from_the_override_else_the_model() { + let overriding = Arc::new(Override( + Some(urls("application/pdf", "^https://docs/")), + AtomicUsize::new(0), + )); + let middleware: Arc = overriding.clone(); + let wrapped = wrap_language_model(Arc::new(Model), &[middleware], None, None); + assert_eq!(keys(&wrapped.supported_urls().await), ["application/pdf"]); + assert_eq!(overriding.1.load(Ordering::SeqCst), 1); + + let passing: Arc = + Arc::new(Override(None, AtomicUsize::new(0))); + let wrapped = wrap_language_model(Arc::new(Model), &[passing], None, None); + assert_eq!(keys(&wrapped.supported_urls().await), ["image/*"]); + + // The first middleware is the outermost: its override wins. + let inner: Arc = + Arc::new(Override(Some(urls("audio/*", ".*")), AtomicUsize::new(0))); + let wrapped = wrap_language_model( + Arc::new(Model), + &[ + overriding.clone() as Arc, + inner, + ], + None, + None, + ); + assert_eq!(keys(&wrapped.supported_urls().await), ["application/pdf"]); + } +} diff --git a/aimux-core/src/lib.rs b/aimux-core/src/lib.rs index 60dc4d30d..50ba2edef 100644 --- a/aimux-core/src/lib.rs +++ b/aimux-core/src/lib.rs @@ -22,7 +22,7 @@ pub mod composite; pub mod content; pub mod decision_model; mod decision_replay; -pub(crate) mod download; +pub mod download; #[doc(hidden)] pub mod download_guard; pub mod embedding_model; @@ -30,20 +30,22 @@ pub mod error; pub mod files_model; pub mod generate; pub mod image_model; +pub mod image_model_middleware; pub mod json_repair; pub mod language_model; pub mod language_model_message; +pub mod language_model_middleware; pub mod math; pub mod message; pub mod moa; pub mod model_catalogue; -pub mod model_id; pub mod openai_output; pub mod operation; pub mod options; pub mod output; pub mod parse_tool_call; pub mod provider; +pub mod provider_registry; pub mod recording; pub mod replay; pub mod reranking_model; @@ -54,6 +56,7 @@ pub mod router; pub mod search_model; pub mod session; pub mod shared; +pub mod skills_model; pub mod speech_model; pub mod step_result; pub mod stream_part; @@ -64,6 +67,7 @@ pub mod transcription_model; pub mod types; pub mod util; pub mod video_model; +pub mod wrap_provider; /// Convenience re-exports. pub mod prelude { @@ -92,8 +96,12 @@ pub mod prelude { generate_text_result_to_chat_completion, stream_text, stream_text_as_openai, }; pub use crate::image_model::{ImageCallOptions, ImageModel, ImageResult, generate_image}; - pub use crate::language_model::LanguageModel; + pub use crate::image_model_middleware::{ImageModelMiddleware, wrap_image_model}; + pub use crate::language_model::{LanguageModel, SupportedUrls}; pub use crate::language_model_message::LanguageModelPrompt; + pub use crate::language_model_middleware::{ + LanguageModelMiddleware, LanguageModelOperation, wrap_language_model, + }; pub use crate::message::{ AssistantContent, ModelMessage, ModelPrompt, ToolContent, UserContent, }; @@ -102,7 +110,6 @@ pub mod prelude { CatalogueSource, Modality, ModelCapabilities, ModelCost, ModelLimits, ModelModalities, ModelSpec, ModelType, ReasoningMode, ReasoningSpec, ReasoningVisibility, RuntimeModel, }; - pub use crate::model_id::ModelId; pub use crate::openai_output::{ ChatCompletion, ChatCompletionChunk, ChatCompletionStream, DONE_FRAME, OpenAiStreamOptions, encode_chunk_sse, to_chat_completion, to_chat_completion_stream, @@ -139,6 +146,7 @@ pub mod prelude { SharedProviderMetadata, SharedProviderOptions, SharedProviderReference, Size, provider_namespace, }; + pub use crate::skills_model::{SkillFile, Skills, UploadSkillCallOptions, UploadSkillResult}; pub use crate::speech_model::{SpeechCallOptions, SpeechModel, SpeechResult, generate_speech}; pub use crate::step_result::{ OutputChunkTimingStats, StepModel, StepResponse, StepResponseMetadata, StepResult, @@ -162,6 +170,7 @@ pub mod prelude { VideoCallOptions, VideoModel, VideoOperationStart, VideoOperationStatus, VideoPollOptions, VideoResult, generate_video, }; + pub use crate::wrap_provider::wrap_provider; } // Root-level re-exports for convenience. @@ -173,8 +182,10 @@ pub use files_model::Files; pub use image_model::ImageModel; pub use language_model::LanguageModel; pub use provider::Provider; +pub use provider_registry::{ProviderRegistry, ProviderRegistryOptions, create_provider_registry}; pub use reranking_model::RerankingModel; pub use search_model::SearchModel; +pub use skills_model::Skills; pub use speech_model::SpeechModel; pub use transcription_model::TranscriptionModel; pub use video_model::VideoModel; diff --git a/aimux-core/src/moa.rs b/aimux-core/src/moa.rs index fc1c8929c..71953a024 100644 --- a/aimux-core/src/moa.rs +++ b/aimux-core/src/moa.rs @@ -136,11 +136,8 @@ impl MoaModel { let results = join_all(self.references.iter().enumerate().map(|(i, m)| { let context = step_context(format!("moa.ref[{i}]:{}/{}", m.provider(), m.model_id())); let ref_opts = &ref_opts; - let retries = retry::prepare_retries( - child_max_retries(), - m.retry_config(), - ref_opts.abort_signal.clone(), - ); + let retries = + retry::prepare_retries(child_max_retries(), ref_opts.abort_signal.clone()); async move { RecordingContext::scope( context.clone(), @@ -197,15 +194,6 @@ impl LanguageModel for MoaModel { &self.config.model_id } - fn retry_config(&self) -> retry::RetryConfig { - // Retrying the composite would rerun the entire reference fanout; - // references and aggregator are retried independently instead. - retry::RetryConfig { - max_retries: 0, - ..retry::RetryConfig::default() - } - } - async fn do_generate(&self, options: &CallOptions) -> Result { // 1. Fan out references (non-streaming). let (texts, ref_usage, warnings) = self.run_references_nonstream(options).await?; @@ -220,11 +208,7 @@ impl LanguageModel for MoaModel { agg_opts.prompt = agg_prompt; // 3. Run the aggregator. - let retries = retry::prepare_retries( - child_max_retries(), - self.aggregator.retry_config(), - agg_opts.abort_signal.clone(), - ); + let retries = retry::prepare_retries(child_max_retries(), agg_opts.abort_signal.clone()); let context = self.aggregator_context(); let mut agg = RecordingContext::scope( context.clone(), @@ -260,11 +244,7 @@ impl LanguageModel for MoaModel { // 3. Aggregator streams; we emit our own StreamStart and add reference // usage onto its Finish. We swallow the aggregator's StreamStart // (we've already emitted ours). - let retries = retry::prepare_retries( - child_max_retries(), - self.aggregator.retry_config(), - agg_opts.abort_signal.clone(), - ); + let retries = retry::prepare_retries(child_max_retries(), agg_opts.abort_signal.clone()); let context = self.aggregator_context(); let agg = RecordingContext::scope( context.clone(), @@ -573,7 +553,6 @@ mod tests { let (ref_b, ref_b_calls) = retry_child("ref-b", "B", 0); let (aggregator, aggregator_calls) = retry_child("aggregator", "final", 1); let moa = MoaModel::new(vec![ref_a, ref_b], aggregator, MoaConfig::default()); - assert_eq!(moa.retry_config().max_retries, 0); let result = crate::generate::generate_text( &moa, diff --git a/aimux-core/src/model_catalogue.rs b/aimux-core/src/model_catalogue.rs index 4412080a9..efd9f7613 100644 --- a/aimux-core/src/model_catalogue.rs +++ b/aimux-core/src/model_catalogue.rs @@ -2,7 +2,7 @@ //! //! Two independent data surfaces: //! - [`RuntimeModel`] — what a provider's `/models` endpoint returns (sparse: -//! id, owned_by, created). Produced by `Provider::list_models`. +//! id, owned_by, created). Produced by `ProviderDiscovery::list_models`. //! - [`ModelSpec`] — community knowledge portrait (context length, capabilities, //! reasoning). Produced by `get_model_specs` (anya2a). //! @@ -219,8 +219,7 @@ pub struct ModelCost { /// The full static portrait for a model, supplemented from community knowledge. /// /// All fields are advisory — aimux never auto-applies them in the request path. -/// The caller reads them to decide how to configure `GenerateTextOptions` / -/// `bodyOverrides`. +/// The caller reads them to decide how to configure `GenerateTextOptions`. #[derive(Debug, Clone, Default, Serialize, Deserialize, TS, PartialEq)] #[cfg_attr(feature = "ts-export", ts(export))] pub struct ModelSpec { diff --git a/aimux-core/src/model_id.rs b/aimux-core/src/model_id.rs deleted file mode 100644 index 1a8f6ae73..000000000 --- a/aimux-core/src/model_id.rs +++ /dev/null @@ -1,48 +0,0 @@ -//! Model identifier helpers. - -use serde::{Deserialize, Serialize}; -use std::fmt; -use ts_rs::TS; - -/// A model identifier in `"provider/model-name"` format, e.g. `"openai/gpt-4o"`. -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, TS)] -#[cfg_attr(feature = "ts-export", ts(export))] -pub struct ModelId { - provider: String, - model: String, -} - -impl ModelId { - pub fn new(provider: impl Into, model: impl Into) -> Self { - Self { - provider: provider.into(), - model: model.into(), - } - } - - #[must_use] - pub fn provider(&self) -> &str { - &self.provider - } - - #[must_use] - pub fn model(&self) -> &str { - &self.model - } -} - -impl fmt::Display for ModelId { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}/{}", self.provider, self.model) - } -} - -impl std::str::FromStr for ModelId { - type Err = String; - fn from_str(s: &str) -> Result { - let (provider, model) = s - .split_once('/') - .ok_or_else(|| format!("expected 'provider/model', got '{s}'"))?; - Ok(Self::new(provider, model)) - } -} diff --git a/aimux-core/src/options.rs b/aimux-core/src/options.rs index 3553c7858..902576d8f 100644 --- a/aimux-core/src/options.rs +++ b/aimux-core/src/options.rs @@ -128,11 +128,6 @@ pub struct CallOptions { /// Anthropic `thinking` config. pub reasoning: Option, - /// Per-call request body overrides. Deep-merged into the provider-built - /// request body (after any built-in vendor override) before sending. - /// `null` values delete the corresponding key. See RFC-0017. - pub body_overrides: Option, - /// Abort signal for cancelling the call. /// /// Runtime handle — never crosses the JSON boundary (bindings that only @@ -141,14 +136,16 @@ pub struct CallOptions { #[ts(skip)] pub abort_signal: Option, - /// Emit raw provider stream chunks as `StreamPart::Raw` (debugging aid). - /// When `Some(true)`, streaming providers yield one `Raw` part per JSON - /// SSE event, carrying the **parsed JSON payload** of the event, emitted - /// before the parsed parts. Excludes the `[DONE]` sentinel; unparsable - /// chunks emit only `Error`. `null`/`Some(false)` = off. - /// Currently honored by the OpenAI-compatible family (openai / azure / - /// openai-compatible registry providers); other provider families ignore - /// it for now (RFC-0016 M2). + /// Emit raw provider stream chunks as `StreamPart::Raw` (the AI SDK's + /// `includeRawChunks`). When `Some(true)`, a streaming model yields one + /// `Raw` part per SSE event, carrying the event's parsed JSON, before the + /// parts made from it; the `[DONE]` sentinel is excluded. A chunk that is + /// not JSON yields `Raw` with `null` in Groq and Anthropic, as the AI SDK + /// does, and only `Error` elsewhere. `None` / `Some(false)` = off. + /// Honored by OpenAI chat, OpenAI-compatible chat and completion (registry + /// presets and Azure included), Groq, DeepSeek, Anthropic Messages, and + /// the OpenAI and ElevenLabs transcription streams; the other model + /// families do not emit `Raw` yet. pub include_raw_chunks: Option, } @@ -178,7 +175,6 @@ impl CallOptions { headers: None, provider_options: None, reasoning: None, - body_overrides: None, abort_signal: None, include_raw_chunks: None, } diff --git a/aimux-core/src/provider.rs b/aimux-core/src/provider.rs index 2883e0169..8a06e8c21 100644 --- a/aimux-core/src/provider.rs +++ b/aimux-core/src/provider.rs @@ -1,70 +1,208 @@ -//! The `Provider` trait — a factory for model task interfaces. +//! The `Provider` trait — a factory that creates model instances for every +//! modality a vendor offers. +//! +//! Mirrors the AI SDK's `ProviderV4`: three required model constructors +//! (language / embedding / image), four optional ones that default to "not +//! offered", and a `files()` accessor. Runtime model discovery lives on the +//! separate [`ProviderDiscovery`] trait. -use std::future::Future; -use std::pin::Pin; +use std::sync::Arc; +use futures::future::BoxFuture; + +use crate::embedding_model::EmbeddingModel; use crate::error::AiMuxError; +use crate::files_model::Files; +use crate::image_model::ImageModel; use crate::language_model::LanguageModel; use crate::model_catalogue::RuntimeModel; +use crate::reranking_model::RerankingModel; +use crate::search_model::SearchModel; +use crate::skills_model::Skills; +use crate::speech_model::SpeechModel; +use crate::transcription_model::TranscriptionModel; +use crate::video_model::VideoModel; /// A provider factory. /// -/// Holds connection configuration and creates task interfaces by model name. -/// A deployment may expose both language generation and decisions. +/// Holds API keys / config and creates model instances by model id. The +/// shape follows AI SDK `ProviderV4`: +/// +/// - [`language_model`](Self::language_model), +/// [`embedding_model`](Self::embedding_model) and +/// [`image_model`](Self::image_model) are required. A vendor that does not +/// offer a modality returns [`AiMuxError::NoSuchModel`] carrying the +/// matching `model_type` (`"languageModel"`, `"embeddingModel"`, +/// `"imageModel"`). Credential or settings errors are never reported as +/// `NoSuchModel`. +/// - [`transcription_model`](Self::transcription_model), +/// [`speech_model`](Self::speech_model), +/// [`reranking_model`](Self::reranking_model) and [`files`](Self::files) +/// are optional: `None` means the vendor does not offer that modality. +/// - [`video_model`](Self::video_model), +/// [`search_model`](Self::search_model) and +/// [`decision_model`](Self::decision_model) are aimux extensions; they are +/// not members of `ProviderV4`. +/// +/// A provider has no `name()`: the name it was registered under belongs to +/// whoever holds the registry, and each model reports its own +/// [`LanguageModel::provider`] string. pub trait Provider: Send + Sync { - /// Unique provider name (e.g. `"openai"`). - fn name(&self) -> &str; + /// Specification version (always `"v4"`). + fn specification_version(&self) -> &'static str { + "v4" + } - /// Create a typed decision model. + /// Create a typed decision model (aimux extension). /// # Errors - /// Returns `UnsupportedFunctionality` unless the provider implements decisions. - fn decision_model(&self, _model_id: &str) -> Result, AiMuxError> { - Err(AiMuxError::UnsupportedFunctionality(format!( - "provider '{}' does not provide decision models", - self.name() - ))) + /// Returns `UnsupportedFunctionality` when decisions are not offered. + fn decision_model(&self, _id: &str) -> Result, AiMuxError> { + Err(AiMuxError::UnsupportedFunctionality( + "provider does not provide decision models".into(), + )) } - /// Create a model instance by its name string (e.g. `"gpt-4o"`). + /// Create a language model by its id (e.g. `"gpt-4o"`). /// - /// Non-language-model providers (image/video/speech/search/… — e.g. Tavily, - /// Stability, Recraft) do not implement this and get the default - /// `Unsupported` error. Only providers that actually expose a language - /// model override it (issue M9). + /// # Errors + /// + /// Returns [`AiMuxError::NoSuchModel`] with `model_type: "languageModel"` + /// when the vendor offers no language models; implementations may also + /// return lookup/auth errors for ids they cannot serve. + fn language_model(&self, id: &str) -> Result, AiMuxError>; + + /// Create an embedding model by its id (e.g. `"text-embedding-3-large"`). /// /// # Errors /// - /// The default implementation returns `AiMuxError::UnsupportedFunctionality`; - /// implementations return lookup/auth errors for model ids they cannot serve. - fn language_model(&self, _model_id: &str) -> Result, AiMuxError> { - Err(AiMuxError::UnsupportedFunctionality(format!( - "provider '{}' does not provide language models", - self.name() - ))) + /// Returns [`AiMuxError::NoSuchModel`] with `model_type: "embeddingModel"` + /// when the vendor offers no embedding models. + fn embedding_model(&self, id: &str) -> Result, AiMuxError>; + + /// Create an image model by its id (e.g. `"dall-e-3"`). + /// + /// # Errors + /// + /// Returns [`AiMuxError::NoSuchModel`] with `model_type: "imageModel"` + /// when the vendor offers no image models. + fn image_model(&self, id: &str) -> Result, AiMuxError>; + + /// Create a transcription (STT) model; `None` when not offered. + fn transcription_model( + &self, + id: &str, + ) -> Option, AiMuxError>> { + let _ = id; + None } + /// Create a speech (TTS) model; `None` when not offered. + fn speech_model(&self, id: &str) -> Option, AiMuxError>> { + let _ = id; + None + } + + /// Create a reranking model; `None` when not offered. + fn reranking_model(&self, id: &str) -> Option, AiMuxError>> { + let _ = id; + None + } + + /// The vendor's file-management interface; `None` when not offered. + fn files(&self) -> Option> { + None + } + + /// The vendor's skill-management interface; `None` when not offered. + fn skills(&self) -> Option> { + None + } + + /// Create a video-generation model; `None` when not offered. + /// + /// aimux extension — not a member of AI SDK `ProviderV4`. + fn video_model(&self, id: &str) -> Option, AiMuxError>> { + let _ = id; + None + } + + /// Create a web-search model; `None` when not offered. + /// + /// aimux extension — not a member of AI SDK `ProviderV4`. + fn search_model(&self, id: &str) -> Option, AiMuxError>> { + let _ = id; + None + } + + /// aimux extension: this provider's model-discovery side, when it has one + /// (see [`ProviderDiscovery`]). `None` for vendors without a `/models` + /// endpoint. + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + None + } +} + +/// Runtime model discovery, implemented by providers that expose a +/// model-list endpoint (RFC-0027). +/// +/// Separate from [`Provider`] because `ProviderV4` has no equivalent and most +/// vendors cannot list models. +pub trait ProviderDiscovery: Send + Sync { /// List the models this account can call on this provider, via the - /// provider's `/models` endpoint (runtime discovery). + /// provider's `/models` endpoint. /// /// Returns **only the provider's official data** (`RuntimeModel`: id, /// owned_by, created) — no community catalogue enrichment. To supplement /// with model specs (context length, capabilities, reasoning portrait), /// call `get_model_specs` separately and merge in the host (RFC-0027). - /// - /// Default returns [`AiMuxError::UnsupportedFunctionality`] — providers that expose a - /// model-list endpoint override this. - /// - /// Implemented as a `Pin>` (rather than `#[async_trait]`) so that - /// the dozens of existing `Provider` impls need no changes — only providers - /// that actually support `/models` override this. - fn list_models( + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>>; +} + +/// A reference to a provider is a provider, so a `&'static` default instance +/// (`openai()`, `anthropic()`, ...) can be registered as it is. +impl Provider for &P { + fn specification_version(&self) -> &'static str { + (**self).specification_version() + } + + fn decision_model(&self, id: &str) -> Result, AiMuxError> { + (**self).decision_model(id) + } + + fn language_model(&self, id: &str) -> Result, AiMuxError> { + (**self).language_model(id) + } + fn embedding_model(&self, id: &str) -> Result, AiMuxError> { + (**self).embedding_model(id) + } + fn image_model(&self, id: &str) -> Result, AiMuxError> { + (**self).image_model(id) + } + fn transcription_model( &self, - ) -> Pin, AiMuxError>> + Send + '_>> { - let name = self.name().to_string(); - Box::pin(async move { - Err(AiMuxError::UnsupportedFunctionality(format!( - "list_models not implemented for provider '{name}'" - ))) - }) + id: &str, + ) -> Option, AiMuxError>> { + (**self).transcription_model(id) + } + fn speech_model(&self, id: &str) -> Option, AiMuxError>> { + (**self).speech_model(id) + } + fn reranking_model(&self, id: &str) -> Option, AiMuxError>> { + (**self).reranking_model(id) + } + fn files(&self) -> Option> { + (**self).files() + } + fn skills(&self) -> Option> { + (**self).skills() + } + fn video_model(&self, id: &str) -> Option, AiMuxError>> { + (**self).video_model(id) + } + fn search_model(&self, id: &str) -> Option, AiMuxError>> { + (**self).search_model(id) + } + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + (**self).discovery() } } diff --git a/aimux-core/src/provider_registry.rs b/aimux-core/src/provider_registry.rs new file mode 100644 index 000000000..d2040efc7 --- /dev/null +++ b/aimux-core/src/provider_registry.rs @@ -0,0 +1,241 @@ +//! The provider registry: models addressed as `"{provider}:{model}"`. +//! +//! Mirrors the AI SDK's `createProviderRegistry`: the caller assembles a map +//! of provider ids to providers and the registry resolves a combined id by +//! splitting it at the first separator. The registry knows no provider on its +//! own and never falls back to another one. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use crate::embedding_model::EmbeddingModel; +use crate::error::AiMuxError; +use crate::files_model::Files; +use crate::image_model::ImageModel; +use crate::image_model_middleware::{ImageModelMiddleware, wrap_image_model}; +use crate::language_model::LanguageModel; +use crate::language_model_middleware::{LanguageModelMiddleware, wrap_language_model}; +use crate::provider::Provider; +use crate::reranking_model::RerankingModel; +use crate::search_model::SearchModel; +use crate::skills_model::Skills; +use crate::speech_model::SpeechModel; +use crate::transcription_model::TranscriptionModel; +use crate::video_model::VideoModel; + +/// Options of [`create_provider_registry`]. +#[derive(Clone)] +pub struct ProviderRegistryOptions { + /// Separates the provider id from the model id. Defaults to `":"`. + pub separator: String, + /// Middleware applied to each language model in input order. + pub language_model_middleware: Vec>, + /// Middleware applied to each image model in input order. + pub image_model_middleware: Vec>, +} + +impl Default for ProviderRegistryOptions { + fn default() -> Self { + Self { + separator: ":".to_string(), + language_model_middleware: Vec::new(), + image_model_middleware: Vec::new(), + } + } +} + +/// A set of providers addressed by id; see [`create_provider_registry`]. +pub struct ProviderRegistry { + providers: BTreeMap>, + separator: String, + language_model_middleware: Vec>, + image_model_middleware: Vec>, +} + +/// Create a registry over `providers` (the AI SDK's `createProviderRegistry`). +#[must_use] +pub fn create_provider_registry( + providers: BTreeMap>, + options: ProviderRegistryOptions, +) -> ProviderRegistry { + ProviderRegistry { + providers, + separator: options.separator, + language_model_middleware: options.language_model_middleware, + image_model_middleware: options.image_model_middleware, + } +} + +impl ProviderRegistry { + /// The provider registered as `id`. + /// + /// # Errors + /// + /// [`AiMuxError::NoSuchProvider`] for an id that is not registered. + pub fn provider(&self, id: &str) -> Result<&Arc, AiMuxError> { + self.get_provider(id, "") + } + + fn get_provider(&self, id: &str, model_type: &str) -> Result<&Arc, AiMuxError> { + self.providers + .get(id) + .ok_or_else(|| AiMuxError::NoSuchProvider { + provider_id: id.to_string(), + model_id: id.to_string(), + model_type: model_type.to_string(), + available_providers: self.providers.keys().cloned().collect(), + }) + } + + /// The registered provider ids, in order. + pub fn provider_ids(&self) -> impl Iterator { + self.providers.keys().map(String::as_str) + } + + /// Split `id` at the first separator and look the provider up. + fn resolve<'a>( + &self, + id: &'a str, + model_type: &str, + ) -> Result<(&Arc, &'a str), AiMuxError> { + let (provider_id, model_id) = id + .split_once(self.separator.as_str()) + .ok_or_else(|| AiMuxError::no_such_model(id, model_type))?; + Ok((self.get_provider(provider_id, model_type)?, model_id)) + } + + /// The language model `"{provider}{separator}{model}"`. + /// + /// # Errors + /// + /// [`AiMuxError::NoSuchModel`] for an id without the separator, + /// [`AiMuxError::NoSuchProvider`] for an unregistered provider, and + /// whatever the provider returns for the model id. + pub fn language_model(&self, id: &str) -> Result, AiMuxError> { + let (provider, model_id) = self.resolve(id, "languageModel")?; + Ok(wrap_language_model( + provider.language_model(model_id)?, + &self.language_model_middleware, + None, + None, + )) + } + + /// The embedding model for a combined id; errors as [`Self::language_model`]. + /// + /// # Errors + /// + /// See [`Self::language_model`]. + pub fn embedding_model(&self, id: &str) -> Result, AiMuxError> { + let (provider, model_id) = self.resolve(id, "embeddingModel")?; + provider.embedding_model(model_id) + } + + /// The image model for a combined id; errors as [`Self::language_model`]. + /// + /// # Errors + /// + /// See [`Self::language_model`]. + pub fn image_model(&self, id: &str) -> Result, AiMuxError> { + let (provider, model_id) = self.resolve(id, "imageModel")?; + Ok(wrap_image_model( + provider.image_model(model_id)?, + &self.image_model_middleware, + None, + None, + )) + } + + /// The transcription model for a combined id. A provider that offers none + /// is [`AiMuxError::NoSuchModel`]. + /// + /// # Errors + /// + /// See [`Self::language_model`]. + pub fn transcription_model(&self, id: &str) -> Result, AiMuxError> { + let (provider, model_id) = self.resolve(id, "transcriptionModel")?; + offered( + provider.transcription_model(model_id), + id, + "transcriptionModel", + ) + } + + /// The speech model for a combined id. A provider that offers none is + /// [`AiMuxError::NoSuchModel`]. + /// + /// # Errors + /// + /// See [`Self::language_model`]. + pub fn speech_model(&self, id: &str) -> Result, AiMuxError> { + let (provider, model_id) = self.resolve(id, "speechModel")?; + offered(provider.speech_model(model_id), id, "speechModel") + } + + /// The reranking model for a combined id. A provider that offers none is + /// [`AiMuxError::NoSuchModel`]. + /// + /// # Errors + /// + /// See [`Self::language_model`]. + pub fn reranking_model(&self, id: &str) -> Result, AiMuxError> { + let (provider, model_id) = self.resolve(id, "rerankingModel")?; + offered(provider.reranking_model(model_id), id, "rerankingModel") + } + + /// The video model for a combined id. A provider that offers none is + /// [`AiMuxError::NoSuchModel`]. + /// + /// # Errors + /// + /// See [`Self::language_model`]. + pub fn video_model(&self, id: &str) -> Result, AiMuxError> { + let (provider, model_id) = self.resolve(id, "videoModel")?; + offered(provider.video_model(model_id), id, "videoModel") + } + + /// The search model for a combined id (an aimux extension). A provider + /// that offers none is [`AiMuxError::NoSuchModel`]. + /// + /// # Errors + /// + /// See [`Self::language_model`]. + pub fn search_model(&self, id: &str) -> Result, AiMuxError> { + let (provider, model_id) = self.resolve(id, "searchModel")?; + offered(provider.search_model(model_id), id, "searchModel") + } + + /// The skills interface of a registered provider. + /// + /// # Errors + /// Returns an unknown provider or unsupported functionality error. + pub fn skills(&self, provider_id: &str) -> Result, AiMuxError> { + self.provider(provider_id)?.skills().ok_or_else(|| { + AiMuxError::UnsupportedFunctionality(format!( + "provider '{provider_id}' does not expose skills" + )) + }) + } + + /// The files interface of the provider `provider_id`. + /// + /// # Errors + /// + /// [`AiMuxError::NoSuchProvider`] for an unregistered provider, + /// [`AiMuxError::UnsupportedFunctionality`] when it exposes no files. + pub fn files(&self, provider_id: &str) -> Result, AiMuxError> { + self.provider(provider_id)?.files().ok_or_else(|| { + AiMuxError::UnsupportedFunctionality(format!( + "provider '{provider_id}' does not expose files" + )) + }) + } +} + +fn offered( + model: Option, AiMuxError>>, + id: &str, + model_type: &str, +) -> Result, AiMuxError> { + model.unwrap_or_else(|| Err(AiMuxError::no_such_model(id, model_type))) +} diff --git a/aimux-core/src/recording.rs b/aimux-core/src/recording.rs index 331607214..bb2b88b8b 100644 --- a/aimux-core/src/recording.rs +++ b/aimux-core/src/recording.rs @@ -24,7 +24,7 @@ use std::io::{BufWriter, Write}; use std::path::PathBuf; use std::pin::Pin; use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; -use std::sync::mpsc::{Receiver, SyncSender, sync_channel}; +use std::sync::mpsc::{Receiver, SyncSender, TrySendError, sync_channel}; use std::sync::{Arc, Mutex, RwLock}; use std::task::{Context, Poll}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -33,8 +33,41 @@ use ts_rs::TS; // ── 数据模型(三层 + call_id 关联;schema 版本)──────────────────────────── /// 录制格式版本(用于未来字段迁移与绑定层兼容)。 +/// +/// Schema 2, as released aimux writes it: the V4 layers keep it until the +/// full schema 3 of the alignment doc (§4.3) lands with the last layers. pub const RECORDING_SCHEMA: u32 = 2; +/// Reads `schema`: 1 and 2, the schemas released aimux versions write, load; +/// any other value is refused with a clear error. No schema 3 recording was +/// ever released. +fn deserialize_known_schema<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + let schema = u32::deserialize(deserializer)?; + if (1..=RECORDING_SCHEMA).contains(&schema) { + Ok(schema) + } else { + Err(serde::de::Error::custom(format!( + "unsupported recording schema {schema}: only schemas 1 to {RECORDING_SCHEMA} \ + are readable" + ))) + } +} + +/// Reads `InputRecord.prompt`. A prompt this version cannot read, such as +/// one written by aimux 0.5 or earlier with a system message whose content is +/// a part array, loads as empty so the recording still loads; +/// [`InputRecord::unreadable_prompt`] then says why, and replay refuses it. +fn deserialize_prompt<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + let value = serde_json::Value::deserialize(deserializer)?; + Ok(serde_json::from_value(value).unwrap_or_default()) +} + /// Operation stored in a unified recording. Older recordings are language calls. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, TS)] #[serde(rename_all = "snake_case")] @@ -55,7 +88,11 @@ impl RecordingOperation { #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[cfg_attr(feature = "ts-export", ts(export))] pub struct Recording { - /// 格式版本。 + /// 格式版本。Written as [`RECORDING_SCHEMA`]; schemas 1 and 2 are read. + /// Recordings written by aimux 0.5 and earlier also carry a configuration + /// snapshot under `provider` (base URL, key source, profile, provider + /// options), which is ignored. + #[serde(deserialize_with = "deserialize_known_schema")] pub schema: u32, /// 暴露逻辑调用 ID,关联三层。 pub call_id: String, @@ -74,7 +111,9 @@ pub struct Recording { /// 最终结果摘要(状态 + finish_reason + usage)。 pub outcome: OutcomeRecord, - /// wire 是否完整(流式 body 未补全时为 false;由 writer 标记)。 + /// wire 是否完整(流式 body 未补全时为 false;由 writer 标记)。A call + /// that lost an event (full queue, writer gone) or was marked + /// inconsistent (duplicate attempt, unmatched update) is never complete. #[serde(default)] pub complete: bool, /// 传输层封闭信号:层 B 声明"不再有 exchange"(P1:层 A 收尾自动发; @@ -138,13 +177,19 @@ impl Recording { #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[cfg_attr(feature = "ts-export", ts(export))] pub struct InputRecord { + /// Decision contract used for validation and deterministic offline replay. + /// Connection settings remain outside the recording. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub decision_capabilities: Option, #[serde(default, skip_serializing_if = "RecordingOperation::is_language_model")] pub operation: RecordingOperation, - /// 完整 prompt(消息数组,含 ContentPart::Image 等多模态)。 + /// 完整 prompt(消息数组,含 ContentPart::Image 等多模态)。Loads as + /// empty when this version cannot read it ([`Self::unreadable_prompt`]). + #[serde(deserialize_with = "deserialize_prompt")] pub prompt: LanguageModelPrompt, /// 序列化的 CallOptions(abort_signal 已 serde skip), /// 加上所在操作的 `max_retries` / `timeout` / `session_id` - /// (`OperationSettings`);headers/provider_options/body_overrides 已递归脱敏。 + /// (`OperationSettings`);headers/provider_options 已递归脱敏。 pub options: serde_json::Value, } @@ -168,6 +213,7 @@ impl InputRecord { map.extend(settings); } Self { + decision_capabilities: None, operation: RecordingOperation::LanguageModel, prompt: options.prompt.clone(), options: redact_json(value), @@ -175,46 +221,93 @@ impl InputRecord { } #[must_use] - pub fn from_decision_options(options: &crate::decision_model::DecisionCallOptions) -> Self { + pub fn from_decision_options( + options: &crate::decision_model::DecisionCallOptions, + capabilities: &crate::decision_model::DecisionCapabilities, + ) -> Self { Self { + decision_capabilities: Some(capabilities.clone()), operation: RecordingOperation::Decision, prompt: Vec::new(), options: redact_json(serde_json::to_value(options).unwrap_or_default()), } } + + /// Why the recorded prompt cannot be replayed, when it cannot: the + /// recorded options hold a prompt this version does not read (one written + /// by aimux 0.5 or earlier, for example with a system message whose + /// content is a part array). `prompt` then loaded as empty. + #[must_use] + pub fn unreadable_prompt(&self) -> Option { + if !self.prompt.is_empty() { + return None; + } + let recorded = self.options.get("prompt")?; + if recorded.as_array().is_none_or(Vec::is_empty) { + return None; + } + serde_json::from_value::(recorded.clone()) + .err() + .map(|e| e.to_string()) + } } /// ② 配置侧:provider 身份(请求回放重建用)。 +/// +/// 只记身份,不记配置:回放时按 `provider_id` + `model_id` 重新构建 model, +/// 凭证与 base URL 只来自环境/registry(RFC-0036 §3.3)。 #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[cfg_attr(feature = "ts-export", ts(export))] pub struct ProviderRecord { - /// `model.provider()`,如 "openai"。 + /// provider 名(registry 名),回放重建的查找键,如 "openai"、"deepseek"。 + /// Empty in a recording written by aimux 0.5 or earlier, which has none + /// ([`Self::registry_id`]). + #[serde(default)] + pub provider_id: String, + /// `model.provider()`,如 "openai.chat"。 pub provider: String, /// `model.model_id()`,如 "gpt-4o"。 pub model_id: String, - /// provider 的 base_url。 - pub base_url: Option, - /// api_key 来源(不存明文):"env:OPENAI_API_KEY" / "explicit" / "none" / "unknown"。 - pub api_key_source: String, - /// OpenAICompatProfile(能力差异)。 - pub profile: Option, - /// ProviderOptions(headers/org/project/...;已脱敏)。 - pub provider_options: Option, } impl ProviderRecord { - /// 最小快照(provider/model_id)——`config_snapshot` 默认实现;cover 的部分放结构方法。 + /// 显式给出三项身份。 #[must_use] - pub fn minimal(provider: &str, model_id: &str) -> Self { + pub fn new(provider_id: &str, provider: &str, model_id: &str) -> Self { Self { + provider_id: provider_id.to_string(), provider: provider.to_string(), model_id: model_id.to_string(), - base_url: None, - api_key_source: "unknown".to_string(), - profile: None, - provider_options: None, } } + + /// `provider_id`, the registry name live replay rebuilds the model from. + /// + /// # Errors + /// + /// `InvalidArgument` when the recording has none: aimux 0.5 and earlier + /// record the model's `provider()` and its configuration instead, so such + /// a recording is replayed by passing the model to `replay_with_model`. + pub fn registry_id(&self) -> Result<&str, crate::error::AiMuxError> { + if self.provider_id.is_empty() { + return Err(crate::error::AiMuxError::InvalidArgument(format!( + "replay: the recording of `{}` has no provider_id (recordings written by \ + aimux 0.5 or earlier carry none), so the model cannot be rebuilt from a \ + registry; pass the model to replay_with_model", + self.provider + ))); + } + Ok(&self.provider_id) + } + + /// 由 model 的 `provider()` / `model_id()` 构造:`provider_id` 暂取 + /// `provider()` 字符串的首段(`"openai.chat"` → `"openai"`);registry + /// 名接入后由其填写。 + #[must_use] + pub fn from_model(provider: &str, model_id: &str) -> Self { + let provider_id = provider.split('.').next().unwrap_or_default(); + Self::new(provider_id, provider, model_id) + } } /// ③ HTTP 侧:单次 attempt 的 wire 交换。 @@ -408,6 +501,7 @@ pub trait Recorder: Send + Sync { _options: &crate::decision_model::DecisionCallOptions, _provider: &str, _model_id: &str, + _capabilities: &crate::decision_model::DecisionCapabilities, ) { } /// 录制会话归组信息(RFC-0024 P3):session_id + 会话内步号。 @@ -415,7 +509,7 @@ pub trait Recorder: Send + Sync { /// 默认空实现(不参与录制的 Recorder 无需关心);开启录制且调用被归组 /// 时,`generate_text`/`stream_text` 入口调用它。 fn record_session(&self, _call_id: &str, _session_id: &str, _step: u32) {} - /// 录制配置侧完整快照。 + /// 录制配置侧身份(`ProviderRecord`:provider_id / provider / model_id)。 fn record_provider(&self, call_id: &str, snapshot: &ProviderRecord); /// 录制单次 HTTP 交换(层 B http.rs 调用,per-attempt 一条)。 fn record_exchange(&self, call_id: &str, exchange: &HttpExchange); @@ -587,18 +681,6 @@ pub fn init_recording(recorder: Option>) { } } -/// 从环境变量初始化:`AIMUX_RECORD=1` 开启,`AIMUX_RECORD_DIR` 指定目录(默认 `./recordings`)。 -/// 未开启返回 false(静默)。 -#[must_use] -pub fn init_recording_from_env() -> bool { - if std::env::var("AIMUX_RECORD").as_deref() != Ok("1") { - return false; - } - let dir = std::env::var("AIMUX_RECORD_DIR").unwrap_or_else(|_| "./recordings".to_string()); - init_recording(Some(Arc::new(JsonlRecorder::new(dir)))); - true -} - /// 从当前全局 recorder 生成一次调用的 context(关闭时 None)。 /// 层 A 入口:先读一次,再用 `context && ctx.start()` 记录 ①+②。 pub fn context(call_id: impl Into) -> Option { @@ -618,7 +700,7 @@ pub fn recorder() -> Option> { static CALL_SEQ: AtomicU64 = AtomicU64::new(0); /// 生成进程级唯一 call_id(`call-{ns}-{seq}`;与 session.rs 同构)。 -pub fn new_call_id() -> String { +pub(crate) fn new_call_id() -> String { let ns = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_nanos()) @@ -664,7 +746,7 @@ pub fn is_sensitive_key(name: &str) -> bool { } /// 递归脱敏(JSON 中含敏感键的项值替换为 `[REDACTED]`)。 -/// 覆盖 `CallOptions.headers`/`provider_options`/`body_overrides` 任意层级。 +/// 覆盖 `CallOptions.headers`/`provider_options` 任意层级。 pub(crate) fn redact_json(value: serde_json::Value) -> serde_json::Value { match value { serde_json::Value::Object(map) => { @@ -729,6 +811,23 @@ enum RecordEvent { }, } +impl RecordEvent { + /// The call this event belongs to; `None` for `Flush`, which is never + /// dropped (it is sent with a blocking send). + fn call_id(&self) -> Option<&str> { + match self { + Self::Input { call_id, .. } + | Self::Session { call_id, .. } + | Self::Provider { call_id, .. } + | Self::Exchange { call_id, .. } + | Self::ExchangeUpdate { call_id, .. } + | Self::Outcome { call_id, .. } + | Self::TransportClosed { call_id } => Some(call_id), + Self::Flush { .. } => None, + } + } +} + /// JsonlRecorder 事件通道容量(有界,A4):防止 writer 跟不上时录制热路径 /// 无界堆积内存。满时 **drop-newest** 并递增 `dropped` 计数(经 /// [`JsonlRecorder::dropped_count`] 可查),保持 `send_ev` 非阻塞;writer @@ -743,6 +842,9 @@ pub struct JsonlRecorder { thread: Option>, /// 有界通道满 / writer 断开时丢弃的事件计数(drop-newest,可查)。 dropped: AtomicU64, + /// Calls that lost an event to `dropped`. Shared with the writer, which + /// never marks such a call `complete` (RFC-0023 completion barrier). + lossy: Arc>>, /// 被标记为不一致的 call_id(C4-7:重复 attempt / update 未命中或歧义)。 /// 存于 recorder 侧而非 `Recording` 字段——后者被 `replay.rs` 以全字段 /// 字面量构造,新增 serde 字段会破坏其它模块编译,故用旁路集合标记。 @@ -779,18 +881,23 @@ impl JsonlRecorder { let w = BufWriter::new(file); let (tx, rx) = sync_channel::(JSONL_CHANNEL_CAPACITY); let inconsistent = Arc::new(Mutex::new(HashSet::new())); - let writer_inconsistent = inconsistent.clone(); + let lossy = Arc::new(Mutex::new(HashSet::new())); + let damaged = DamagedCalls { + inconsistent: inconsistent.clone(), + lossy: lossy.clone(), + }; let write_error = Arc::new(Mutex::new(None)); let writer_write_error = write_error.clone(); let handle = std::thread::Builder::new() .name("aimux-recording".into()) - .spawn(move || writer_loop(rx, w, writer_inconsistent, writer_write_error)) + .spawn(move || writer_loop(rx, w, damaged, writer_write_error)) .map_err(|source| RecordingError::Spawn { source })?; Ok(Self { tx: Some(tx), dir, thread: Some(handle), dropped: AtomicU64::new(0), + lossy, inconsistent, write_error, }) @@ -812,6 +919,7 @@ impl JsonlRecorder { dir, thread: None, dropped: AtomicU64::new(0), + lossy: Arc::new(Mutex::new(HashSet::new())), inconsistent: Arc::new(Mutex::new(HashSet::new())), write_error: Arc::new(Mutex::new(None)), } @@ -820,7 +928,7 @@ impl JsonlRecorder { /// 事件入队(有界 `try_send`;满或 writer 断开 → drop-newest 并计数)。 fn send_ev(&self, ev: RecordEvent) { if let Some(tx) = &self.tx { - send_or_drop(tx, ev, &self.dropped); + send_or_drop(tx, ev, &self.dropped, &self.lossy); } } @@ -898,18 +1006,19 @@ impl Recorder for JsonlRecorder { options: &crate::decision_model::DecisionCallOptions, provider: &str, model_id: &str, + capabilities: &crate::decision_model::DecisionCapabilities, ) { self.send_ev(RecordEvent::Input { call_id: call_id.to_owned(), - input: InputRecord::from_decision_options(options), - provider: ProviderRecord::minimal(provider, model_id), + input: InputRecord::from_decision_options(options, capabilities), + provider: ProviderRecord::from_model(provider, model_id), }); } fn record_input(&self, call_id: &str, options: &CallOptions, provider: &str, model_id: &str) { self.send_ev(RecordEvent::Input { call_id: call_id.to_string(), input: InputRecord::from_call_options(options), - provider: ProviderRecord::minimal(provider, model_id), + provider: ProviderRecord::from_model(provider, model_id), }); } @@ -922,14 +1031,9 @@ impl Recorder for JsonlRecorder { } fn record_provider(&self, call_id: &str, snapshot: &ProviderRecord) { - // 核心边界统一强制脱敏 provider_options/profile(B6 评审), - // 不依赖 provider 端自觉。 - let mut snap = snapshot.clone(); - snap.provider_options = snap.provider_options.take().map(redact_json); - snap.profile = snap.profile.take().map(redact_json); self.send_ev(RecordEvent::Provider { call_id: call_id.to_string(), - provider: snap, + provider: snapshot.clone(), }); } @@ -990,11 +1094,12 @@ fn entry_or_init<'a>( Recording::new( call_id, InputRecord { + decision_capabilities: None, operation: Default::default(), prompt: Vec::new(), options: serde_json::Value::Null, }, - ProviderRecord::minimal("", ""), + ProviderRecord::from_model("", ""), ) }) } @@ -1002,12 +1107,13 @@ fn entry_or_init<'a>( /// writer 线程:按 call_id 合并分片;completion barrier 后写一行。 /// /// 文件已由 [`JsonlRecorder::try_new`] 打开并作为 `BufWriter` 传入(A9/N4: -/// 打开失败在 `try_new` 走 `Result`,而非此处静默早退)。`inconsistent` 为与 -/// recorder 共享的旁路集合,用于标记 C4-7 检测到的不一致 call_id。 +/// 打开失败在 `try_new` 走 `Result`,而非此处静默早退)。`damaged` 为与 +/// recorder 共享的旁路集合:C4-7 检测到的不一致 call_id 由这里标记,丢过事件的 +/// call_id 由发送端标记;两者写出时都不标 `complete`。 fn writer_loop( rx: Receiver, mut w: impl Write, - inconsistent: Arc>>, + damaged: DamagedCalls, write_error: Arc>>, ) { let mut pending: HashMap = HashMap::new(); @@ -1037,7 +1143,7 @@ fn writer_loop( let rec = entry_or_init(&mut pending, &call_id); rec.session_id = Some(session_id); rec.step = Some(step); - try_finalize(&mut w, &mut pending, &call_id, &write_error); + try_finalize(&mut w, &mut pending, &call_id, &damaged, &write_error); } // C4-6:用 entry_or_init(而非 get_mut)兜底建条目,保证乱序 // (先于 Input 到达)的 Provider 快照不丢。 @@ -1048,9 +1154,9 @@ fn writer_loop( // C4-7:同 (call_id, attempt) 重复 → 合并骨架并标记 inconsistent。 let rec = entry_or_init(&mut pending, &call_id); if insert_exchange(rec, exchange) { - mark_inconsistent(&inconsistent, &call_id); + mark_inconsistent(&damaged.inconsistent, &call_id); } - try_finalize(&mut w, &mut pending, &call_id, &write_error); + try_finalize(&mut w, &mut pending, &call_id, &damaged, &write_error); } RecordEvent::ExchangeUpdate { call_id, @@ -1066,20 +1172,20 @@ fn writer_loop( apply_exchange_update(rec, attempt, exchange_index, response, error), UpdateMatch::Patched ) { - mark_inconsistent(&inconsistent, &call_id); + mark_inconsistent(&damaged.inconsistent, &call_id); } - try_finalize(&mut w, &mut pending, &call_id, &write_error); + try_finalize(&mut w, &mut pending, &call_id, &damaged, &write_error); } RecordEvent::TransportClosed { call_id } => { entry_or_init(&mut pending, &call_id).transport_closed = true; - try_finalize(&mut w, &mut pending, &call_id, &write_error); + try_finalize(&mut w, &mut pending, &call_id, &damaged, &write_error); } RecordEvent::Outcome { call_id, outcome } => { entry_or_init(&mut pending, &call_id).outcome = outcome; - try_finalize(&mut w, &mut pending, &call_id, &write_error); + try_finalize(&mut w, &mut pending, &call_id, &damaged, &write_error); } RecordEvent::Flush { ack } => { - write_ready_all(&mut w, &mut pending, &write_error); + write_ready_all(&mut w, &mut pending, &damaged, &write_error); if let Err(e) = w.flush() { note_write_error(&write_error, &e); } @@ -1108,26 +1214,56 @@ fn writer_loop( } } -/// 若该 call 的所有分片到齐(barrier 满足)则写出,并标记 complete=true。 +/// Calls whose recording cannot be complete, shared between a +/// [`JsonlRecorder`] and its writer: `inconsistent` (C4-7, marked by the +/// writer) and `lossy` (an event was dropped, marked by the sender). +struct DamagedCalls { + inconsistent: Arc>>, + lossy: Arc>>, +} + +impl DamagedCalls { + /// Whether `call_id` lost an event or was marked inconsistent. Its + /// `lossy` mark is cleared: the call is being written now. + fn take(&self, call_id: &str) -> bool { + let lossy = self + .lossy + .lock() + .map(|mut set| set.remove(call_id)) + .unwrap_or(true); + let inconsistent = self + .inconsistent + .lock() + .map(|set| set.contains(call_id)) + .unwrap_or(true); + lossy || inconsistent + } +} + +/// 若该 call 的所有分片到齐(barrier 满足)则写出;丢过事件或被标记不一致的 +/// call 写为 `complete=false`,其余标记 complete=true。 fn try_finalize( w: &mut impl Write, pending: &mut HashMap, call_id: &str, + damaged: &DamagedCalls, write_error: &Mutex>, ) { if pending.get(call_id).map(Recording::ready).unwrap_or(false) && let Some(mut rec) = pending.remove(call_id) { if !rec.exchanges.is_empty() || rec.transport_closed { - rec.complete = true; + rec.complete = !damaged.take(call_id); } write_line(w, rec, write_error); } } +/// flush:写出所有 barrier 已满足的 call,`complete` 规则同 [`try_finalize`]。 fn write_ready_all( w: &mut impl Write, pending: &mut HashMap, + damaged: &DamagedCalls, write_error: &Mutex>, ) { let ids: Vec = pending @@ -1137,7 +1273,7 @@ fn write_ready_all( .collect(); for id in ids { if let Some(mut rec) = pending.remove(&id) { - rec.complete = true; + rec.complete = !damaged.take(&id); write_line(w, rec, write_error); } } @@ -1250,11 +1386,22 @@ fn mark_inconsistent(set: &Mutex>, call_id: &str) { // ── 有界 channel send(A4:drop-newest + 计数)───────────────────────────── /// 有界 `try_send`:满(drop-newest)或 writer 断开时丢弃事件并递增 `dropped`, -/// 保持热路径非阻塞且失败可查(而非静默 `let _ = tx.send(...)`)。 -fn send_or_drop(tx: &SyncSender, ev: RecordEvent, dropped: &AtomicU64) { +/// 保持热路径非阻塞且失败可查(而非静默 `let _ = tx.send(...)`)。丢弃事件所属 +/// 的 call 记入 `lossy`,writer 不会把它写成 `complete=true`。 +fn send_or_drop( + tx: &SyncSender, + ev: RecordEvent, + dropped: &AtomicU64, + lossy: &Mutex>, +) { if let Err(err) = tx.try_send(ev) { - // err 持有未入队事件,随 err 析构丢弃;仅计数,保持非阻塞。 - drop(err); + // err 持有未入队事件:记下它的 call_id 后丢弃,保持非阻塞。 + let (TrySendError::Full(ev) | TrySendError::Disconnected(ev)) = err; + if let Some(call_id) = ev.call_id() + && let Ok(mut set) = lossy.lock() + { + set.insert(call_id.to_string()); + } dropped.fetch_add(1, Ordering::Relaxed); } } @@ -1462,11 +1609,12 @@ impl RingInner { Recording::new( call_id, InputRecord { + decision_capabilities: None, operation: Default::default(), prompt: Vec::new(), options: serde_json::Value::Null, }, - ProviderRecord::minimal("", ""), + ProviderRecord::from_model("", ""), ), ); } @@ -1501,12 +1649,13 @@ impl RingInner { } /// 合并一个已 ready 的分片(barrier 满足则入 ring)。 + /// A call marked inconsistent is stored with `complete=false`. fn finalize(&mut self, call_id: &str) { if self.pending.get(call_id).is_some_and(Recording::ready) && let Some(mut rec) = self.pending.remove(call_id) { if !rec.exchanges.is_empty() || rec.transport_closed { - rec.complete = true; + rec.complete = !self.inconsistent.contains(call_id); } self.push(rec); } @@ -1539,12 +1688,13 @@ impl Recorder for RingRecorder { options: &crate::decision_model::DecisionCallOptions, provider: &str, model_id: &str, + capabilities: &crate::decision_model::DecisionCapabilities, ) { let mut inner = self.inner.lock().unwrap(); let rec = inner.entry_or_init_bounded(call_id); - rec.input = InputRecord::from_decision_options(options); + rec.input = InputRecord::from_decision_options(options, capabilities); if rec.provider.provider.is_empty() { - rec.provider = ProviderRecord::minimal(provider, model_id); + rec.provider = ProviderRecord::from_model(provider, model_id); } } fn record_input(&self, call_id: &str, options: &CallOptions, provider: &str, model_id: &str) { @@ -1554,7 +1704,7 @@ impl Recorder for RingRecorder { let rec = inner.entry_or_init_bounded(call_id); rec.input = InputRecord::from_call_options(options); if rec.provider.provider.is_empty() { - rec.provider = ProviderRecord::minimal(provider, model_id); + rec.provider = ProviderRecord::from_model(provider, model_id); } } @@ -1568,13 +1718,9 @@ impl Recorder for RingRecorder { } fn record_provider(&self, call_id: &str, snapshot: &ProviderRecord) { - // 与 JsonlRecorder 一致的边界强制脱敏,不依赖 provider 端自觉。 - let mut snap = snapshot.clone(); - snap.provider_options = snap.provider_options.take().map(redact_json); - snap.profile = snap.profile.take().map(redact_json); let mut inner = self.inner.lock().unwrap(); // C4-6:entry_or_init_bounded 兜底建条目,保证乱序 Provider 快照不丢。 - inner.entry_or_init_bounded(call_id).provider = snap; + inner.entry_or_init_bounded(call_id).provider = snapshot.clone(); } fn record_exchange(&self, call_id: &str, exchange: &HttpExchange) { @@ -1631,7 +1777,7 @@ impl Recorder for RingRecorder { .collect(); for id in ids { if let Some(mut rec) = inner.pending.remove(&id) { - rec.complete = true; + rec.complete = !inner.inconsistent.contains(&id); inner.push(rec); } } @@ -2083,11 +2229,13 @@ mod tests { "x-goog-api-key": "gkey", "x-amz-security-token": "sts-tok" }, - "provider_options": { "headers": { "Cookie": "s=1" } }, - "body_overrides": { - "api_key": "sk-secret", - "token": "bearer-secret", - "access_token": "oauth-secret" + "provider_options": { + "headers": { "Cookie": "s=1" }, + "openai": { + "api_key": "sk-secret", + "token": "bearer-secret", + "access_token": "oauth-secret" + } }, // 用量字段名含 "token" 子串,但非凭据——contains("token") 曾误伤。 "usage": { @@ -2103,9 +2251,12 @@ mod tests { assert_eq!(r["headers"]["x-amz-security-token"], "[REDACTED]"); assert_eq!(r["headers"]["X-Key"], "ok"); assert_eq!(r["provider_options"]["headers"]["Cookie"], "[REDACTED]"); - assert_eq!(r["body_overrides"]["api_key"], "[REDACTED]"); - assert_eq!(r["body_overrides"]["token"], "[REDACTED]"); - assert_eq!(r["body_overrides"]["access_token"], "[REDACTED]"); + assert_eq!(r["provider_options"]["openai"]["api_key"], "[REDACTED]"); + assert_eq!(r["provider_options"]["openai"]["token"], "[REDACTED]"); + assert_eq!( + r["provider_options"]["openai"]["access_token"], + "[REDACTED]" + ); // 含 "token" 子串的用量字段不再被误脱敏(回归 contains("token"))。 assert_eq!(r["usage"]["max_output_tokens"], 4096); assert_eq!(r["usage"]["prompt_tokens"], 10); @@ -2118,7 +2269,7 @@ mod tests { let rec = Recording::new( "call-1", InputRecord::from_call_options(&sample_options()), - ProviderRecord::minimal("openai", "gpt-4o"), + ProviderRecord::from_model("openai", "gpt-4o"), ); let json = serde_json::to_string(&rec).unwrap(); let back: Recording = serde_json::from_str(&json).unwrap(); @@ -2127,6 +2278,70 @@ mod tests { assert_eq!(back.provider.provider, "openai"); } + /// Recordings are written as schema 2, as master writes them. Schemas 1 + /// and 2 (what released versions write) are read; any other schema, + /// including the never-released 3, is refused with a clear error. + #[test] + fn schema_two_is_written_and_released_schemas_are_read() { + let rec = Recording::new( + "call-1", + InputRecord::from_call_options(&sample_options()), + ProviderRecord::from_model("openai.chat", "gpt-4o"), + ); + let json = serde_json::to_value(&rec).unwrap(); + assert_eq!(json["schema"], 2); + for schema in [1, 2] { + let mut value = json.clone(); + value["schema"] = schema.into(); + let back: Recording = serde_json::from_value(value).unwrap(); + assert_eq!(back.schema, schema); + } + for schema in [0, 3, 4] { + let mut value = json.clone(); + value["schema"] = schema.into(); + let err = serde_json::from_value::(value) + .unwrap_err() + .to_string(); + assert!( + err.contains(&format!("unsupported recording schema {schema}")), + "{err}" + ); + } + } + + /// Fields only this version writes are optional when reading, fields + /// only master writes are ignored, and a prompt this version cannot read + /// loads as empty with the reason kept for replay. + #[test] + fn a_master_shaped_record_loads() { + let system = serde_json::json!([{ + "role": "system", + "content": [{"type": "text", "text": "Be brief."}], + "provider_options": null + }]); + let line = serde_json::json!({ + "schema": 2, + "call_id": "c", + "recorded_at": "2026-10-13T00:00:00.000Z", + "input": {"prompt": system, "options": {"prompt": system}}, + "provider": { + "provider": "openai", + "model_id": "gpt-4o", + "base_url": "https://api.openai.com/v1", + "profile": null, + "provider_options": null + }, + "exchanges": [], + "outcome": {"status": "success", "finish_reason": "stop", "error": null, "usage": null} + }); + let rec: Recording = serde_json::from_value(line).unwrap(); + assert_eq!(rec.provider.provider_id, ""); + assert!(rec.provider.registry_id().is_err()); + assert!(rec.input.prompt.is_empty()); + assert!(rec.input.unreadable_prompt().is_some()); + assert!(!rec.complete); + } + #[test] fn jsonl_flush_writes_complete_line_blocking() { let dir = std::env::temp_dir().join(format!("aimux-rec-test-{}", std::process::id())); @@ -2167,11 +2382,12 @@ mod tests { let mut rec = Recording::new( "call-x", InputRecord { + decision_capabilities: None, operation: Default::default(), prompt: Vec::new(), options: serde_json::Value::Null, }, - ProviderRecord::minimal("openai", "gpt-4o"), + ProviderRecord::from_model("openai", "gpt-4o"), ); rec.transport_closed = true; rec.outcome = OutcomeRecord { @@ -2238,15 +2454,14 @@ mod tests { let (tx, rx) = sync_channel(16); let write_error = Arc::new(Mutex::new(None)); let handle = { - let inconsistent = Arc::new(Mutex::new(HashSet::new())); let w_err = write_error.clone(); - std::thread::spawn(move || writer_loop(rx, FailingWriter, inconsistent, w_err)) + std::thread::spawn(move || writer_loop(rx, FailingWriter, damaged_calls(), w_err)) }; tx.send(RecordEvent::Input { call_id: "c1".into(), input: InputRecord::from_call_options(&sample_options()), - provider: ProviderRecord::minimal("openai", "gpt-4o"), + provider: ProviderRecord::from_model("openai", "gpt-4o"), }) .unwrap(); tx.send(RecordEvent::TransportClosed { @@ -2322,7 +2537,7 @@ mod tests { let rec = Recording::new( "call-1", InputRecord::from_call_options(&sample_options()), - ProviderRecord::minimal("openai", "gpt-4o"), + ProviderRecord::from_model("openai", "gpt-4o"), ); let json = serde_json::to_string(&rec).unwrap(); assert!(!json.contains("session_id"), "ungrouped must omit: {json}"); @@ -2776,17 +2991,212 @@ mod tests { // 容量 1 + 不消费:第 1 条入队,后两条 Full → drop-newest 并计数。 let (tx, rx) = sync_channel::(1); let dropped = AtomicU64::new(0); - let ev = || RecordEvent::TransportClosed { - call_id: "c".into(), + let lossy = Mutex::new(HashSet::new()); + let ev = |call_id: &str| RecordEvent::TransportClosed { + call_id: call_id.into(), }; - send_or_drop(&tx, ev(), &dropped); - send_or_drop(&tx, ev(), &dropped); - send_or_drop(&tx, ev(), &dropped); + send_or_drop(&tx, ev("a"), &dropped, &lossy); + send_or_drop(&tx, ev("b"), &dropped, &lossy); + send_or_drop(&tx, ev("c"), &dropped, &lossy); assert_eq!(dropped.load(Ordering::Relaxed), 2); + // The calls that lost an event are remembered; the queued one is not. + let mut lossy: Vec = lossy.into_inner().unwrap().into_iter().collect(); + lossy.sort(); + assert_eq!(lossy, ["b", "c"]); assert!(rx.try_recv().is_ok()); // 第 1 条仍在 assert!(rx.try_recv().is_err()); // 无更多 } + fn damaged_calls() -> DamagedCalls { + DamagedCalls { + inconsistent: Arc::new(Mutex::new(HashSet::new())), + lossy: Arc::new(Mutex::new(HashSet::new())), + } + } + + /// A writer target the test can read back while the writer thread owns it. + #[derive(Clone, Default)] + struct SharedBuf(Arc>>); + + impl std::io::Write for SharedBuf { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + impl SharedBuf { + /// `complete` of each written recording, by call id. + fn complete_by_call(&self) -> HashMap { + String::from_utf8(self.0.lock().unwrap().clone()) + .unwrap() + .lines() + .map(|line| { + let rec: Recording = serde_json::from_str(line).unwrap(); + (rec.call_id, rec.complete) + }) + .collect() + } + } + + fn input_event(call_id: &str) -> RecordEvent { + RecordEvent::Input { + call_id: call_id.into(), + input: InputRecord::from_call_options(&sample_options()), + provider: ProviderRecord::from_model("openai", "gpt-4o"), + } + } + + /// The rest of a call after its input: one finished exchange, the + /// outcome and the transport close. The writer finalizes the call on the + /// last of them when the input arrived first. + fn rest_of_call(call_id: &str) -> Vec { + vec![ + RecordEvent::Exchange { + call_id: call_id.into(), + exchange: sample_exchange(), + }, + RecordEvent::Outcome { + call_id: call_id.into(), + outcome: success_outcome(), + }, + RecordEvent::TransportClosed { + call_id: call_id.into(), + }, + ] + } + + /// S2-5 (doc §4.3 completion barrier): a call that lost an event to the + /// full queue is written with `complete=false`, both when its last event + /// finalizes it (`try_finalize`) and when a flush writes it + /// (`write_ready_all`, the input arriving last). An intact call in the + /// same writer stays `complete=true`. + #[test] + fn a_call_that_lost_an_event_is_never_complete() { + let (tx, rx) = sync_channel::(1); + let damaged = damaged_calls(); + let dropped = AtomicU64::new(0); + // The queue holds one event and the writer is not running yet, so the + // session events of "finalized" and "flushed" are dropped. + tx.send(input_event("finalized")).unwrap(); + for call_id in ["finalized", "flushed"] { + let session = RecordEvent::Session { + call_id: call_id.into(), + session_id: "s".into(), + step: 0, + }; + send_or_drop(&tx, session, &dropped, &damaged.lossy); + } + assert_eq!(dropped.load(Ordering::Relaxed), 2); + + let out = SharedBuf::default(); + let writer = { + let out = out.clone(); + let write_error = Arc::new(Mutex::new(None)); + std::thread::spawn(move || writer_loop(rx, out, damaged, write_error)) + }; + for ev in rest_of_call("finalized") { + tx.send(ev).unwrap(); + } + tx.send(input_event("intact")).unwrap(); + for ev in rest_of_call("intact") { + tx.send(ev).unwrap(); + } + for ev in rest_of_call("flushed") { + tx.send(ev).unwrap(); + } + tx.send(input_event("flushed")).unwrap(); + let (ack_tx, ack_rx) = sync_channel(0); + tx.send(RecordEvent::Flush { ack: ack_tx }).unwrap(); + ack_rx.recv().unwrap().unwrap(); + + let complete = out.complete_by_call(); + assert_eq!(complete.get("finalized"), Some(&false), "{complete:?}"); + assert_eq!(complete.get("flushed"), Some(&false), "{complete:?}"); + assert_eq!(complete.get("intact"), Some(&true), "{complete:?}"); + drop(tx); + writer.join().unwrap(); + } + + /// S2-5: a call marked inconsistent (C4-7: a duplicate attempt, or an + /// update that matches no exchange) is written with `complete=false` on + /// both write paths. + #[test] + fn an_inconsistent_call_is_never_complete() { + let (tx, rx) = sync_channel::(64); + let damaged = damaged_calls(); + let out = SharedBuf::default(); + let writer = { + let out = out.clone(); + let write_error = Arc::new(Mutex::new(None)); + std::thread::spawn(move || writer_loop(rx, out, damaged, write_error)) + }; + let duplicate = |call_id: &str| RecordEvent::Exchange { + call_id: call_id.into(), + exchange: sample_exchange(), + }; + // Finalized by its last event: a duplicate exchange first. + tx.send(input_event("finalized")).unwrap(); + tx.send(duplicate("finalized")).unwrap(); + for ev in rest_of_call("finalized") { + tx.send(ev).unwrap(); + } + // Written by the flush: an update that matches no exchange, input last. + tx.send(RecordEvent::ExchangeUpdate { + call_id: "flushed".into(), + attempt: 7, + exchange_index: 0, + response: resp_status(200), + error: None, + }) + .unwrap(); + for ev in rest_of_call("flushed") { + tx.send(ev).unwrap(); + } + tx.send(input_event("flushed")).unwrap(); + let (ack_tx, ack_rx) = sync_channel(0); + tx.send(RecordEvent::Flush { ack: ack_tx }).unwrap(); + ack_rx.recv().unwrap().unwrap(); + + let complete = out.complete_by_call(); + assert_eq!(complete.get("finalized"), Some(&false), "{complete:?}"); + assert_eq!(complete.get("flushed"), Some(&false), "{complete:?}"); + drop(tx); + writer.join().unwrap(); + } + + /// S2-5 for the ring: an inconsistent call is stored with + /// `complete=false` when it finalizes and when a flush collects it. + #[test] + fn ring_inconsistent_call_is_never_complete() { + let ring = RingRecorder::new(); + // Finalized by its last event, after a duplicate exchange. + ring.record_input("finalized", &sample_options(), "openai", "gpt-4o"); + ring.record_exchange("finalized", &sample_exchange()); + ring.record_exchange("finalized", &sample_exchange()); + ring.record_outcome("finalized", &success_outcome()); + ring.record_transport_closed("finalized"); + // Collected by flush: an update that matches nothing, input last. + ring.record_exchange_update("flushed", 7, 0, &resp_status(200), None); + ring.record_outcome("flushed", &success_outcome()); + ring.record_transport_closed("flushed"); + ring.record_input("flushed", &sample_options(), "openai", "gpt-4o"); + ring.flush(); + drive_full_call(&ring, "intact"); + + let complete: HashMap = ring + .completed() + .into_iter() + .map(|rec| (rec.call_id, rec.complete)) + .collect(); + assert_eq!(complete.get("finalized"), Some(&false), "{complete:?}"); + assert_eq!(complete.get("flushed"), Some(&false), "{complete:?}"); + assert_eq!(complete.get("intact"), Some(&true), "{complete:?}"); + } + #[test] fn jsonl_dropped_count_starts_zero() { let dir = std::env::temp_dir().join(format!("aimux-rec-drop-{}", std::process::id())); @@ -2864,14 +3274,7 @@ mod tests { // ── C4-6:Provider 先于 Input 不丢 ───────────────────────────────────── fn full_provider() -> ProviderRecord { - ProviderRecord { - provider: "openai".into(), - model_id: "gpt-4o".into(), - base_url: Some("https://api.openai.com".into()), - api_key_source: "env:OPENAI_API_KEY".into(), - profile: None, - provider_options: None, - } + ProviderRecord::new("openai", "openai.chat", "gpt-4o") } #[test] @@ -2886,12 +3289,11 @@ mod tests { ring.export_jsonl(&mut buf).unwrap(); let parsed: Recording = serde_json::from_str(std::str::from_utf8(&buf).unwrap().trim()).unwrap(); - assert_eq!(parsed.provider.provider, "openai"); assert_eq!( - parsed.provider.base_url.as_deref(), - Some("https://api.openai.com"), - "full snapshot preserved, not overwritten by minimal" + parsed.provider.provider, "openai.chat", + "explicit record preserved, not overwritten by the input-derived one" ); + assert_eq!(parsed.provider.provider_id, "openai"); assert_eq!(parsed.provider.model_id, "gpt-4o"); } @@ -2907,11 +3309,8 @@ mod tests { rec.flush(); let content = std::fs::read_to_string(rec.path()).unwrap(); let parsed: Recording = serde_json::from_str(content.trim()).unwrap(); - assert_eq!( - parsed.provider.base_url.as_deref(), - Some("https://api.openai.com") - ); - assert_eq!(parsed.provider.provider, "openai"); + assert_eq!(parsed.provider.provider_id, "openai"); + assert_eq!(parsed.provider.provider, "openai.chat"); let _ = std::fs::remove_dir_all(&dir); } @@ -2952,11 +3351,12 @@ mod tests { Recording::new( call_id, InputRecord { + decision_capabilities: None, operation: Default::default(), prompt: Vec::new(), options: serde_json::Value::Null, }, - ProviderRecord::minimal("", ""), + ProviderRecord::from_model("", ""), ) } diff --git a/aimux-core/src/replay.rs b/aimux-core/src/replay.rs index 110a4032d..cbebcf287 100644 --- a/aimux-core/src/replay.rs +++ b/aimux-core/src/replay.rs @@ -52,10 +52,10 @@ pub trait ReplayMatcher: Send + Sync { /// 精确匹配:provider/model_id 相同,且 prompt + 影响响应的可重放选项 /// (temperature/max_output_tokens/seed/response_format/tools/tool_choice) -/// 以及 headers/provider_options/body_overrides 规范相同才命中。运行时字段 +/// 以及 headers/provider_options 规范相同才命中。运行时字段 /// (call_id/abort_signal)不参与比较。 /// -/// headers/provider_options/body_overrides 用**脱敏感知比较**:录制侧这些 +/// headers/provider_options 用**脱敏感知比较**:录制侧这些 /// 字段经 `recording::redact_json` 脱敏 /// (敏感键值→`"[REDACTED]"`),脱敏值视为通配,匹配任意显式请求值;非脱敏 /// 部分仍精确比较。规则见 `redaction_aware_eq`。 @@ -97,9 +97,9 @@ impl ReplayMatcher for ExactMatcher { /// `CallOptions` 中影响响应、可重放选项的规范键(ExactMatcher 用)。 /// /// 含 prompt + temperature/max_output_tokens/seed/response_format/tools/ -/// tool_choice + headers/provider_options/body_overrides;排除运行时字段 +/// tool_choice + headers/provider_options;排除运行时字段 /// (call_id/abort_signal)。序列化为 `Value` 后比较, -/// 等价于"稳定规范哈希"但无碰撞风险。headers/provider_options/body_overrides +/// 等价于"稳定规范哈希"但无碰撞风险。headers/provider_options /// 在比较时走脱敏感知语义(见 [`canonical_keys_match`])。 fn canonical_call_key(opts: &CallOptions) -> serde_json::Value { serde_json::json!({ @@ -112,7 +112,6 @@ fn canonical_call_key(opts: &CallOptions) -> serde_json::Value { "tool_choice": serde_json::to_value(&opts.tool_choice).unwrap_or_default(), "headers": serde_json::to_value(&opts.headers).unwrap_or_default(), "provider_options": serde_json::to_value(&opts.provider_options).unwrap_or_default(), - "body_overrides": serde_json::to_value(&opts.body_overrides).unwrap_or_default(), }) } @@ -129,11 +128,10 @@ fn canonical_recording_key(rec: &Recording) -> serde_json::Value { "tool_choice": o.get("tool_choice").cloned().unwrap_or_default(), "headers": o.get("headers").cloned().unwrap_or_default(), "provider_options": o.get("provider_options").cloned().unwrap_or_default(), - "body_overrides": o.get("body_overrides").cloned().unwrap_or_default(), }) } -/// 脱敏感知比较(用于 headers/provider_options/body_overrides)。 +/// 脱敏感知比较(用于 headers/provider_options)。 /// /// 录制侧这些字段经 `recording::redact_json` /// 脱敏:敏感键(authorization/api-key/apikey/key/token/cookie/...)的**值** @@ -175,7 +173,7 @@ pub(crate) fn redaction_aware_eq(rec: &serde_json::Value, req: &serde_json::Valu } } -/// 比较规范键:headers/provider_options/body_overrides 用脱敏感知比较 +/// 比较规范键:headers/provider_options 用脱敏感知比较 /// (录制侧可能已脱敏);其余字段(prompt/temperature/...)精确比较。 /// /// 注意:`max_output_tokens` 等含 "token" 子串的字段在录制侧也会被 @@ -183,7 +181,7 @@ pub(crate) fn redaction_aware_eq(rec: &serde_json::Value, req: &serde_json::Valu /// 被脱敏为 `"[REDACTED]"` 而请求值为数值,精确比较会 miss(保守:宁可 miss /// 不可误 hit,因为 max_output_tokens 影响响应,通配会引入伪命中)。 fn canonical_keys_match(rec_key: &serde_json::Value, call_key: &serde_json::Value) -> bool { - const REDACTED_FIELDS: [&str; 3] = ["headers", "provider_options", "body_overrides"]; + const REDACTED_FIELDS: [&str; 2] = ["headers", "provider_options"]; let (ro, co) = match (rec_key.as_object(), call_key.as_object()) { (Some(a), Some(b)) => (a, b), _ => return rec_key == call_key, @@ -442,8 +440,9 @@ impl MockReplayModel { /// /// # Errors /// - /// Returns `AiMuxError::InvalidArgument` when the file cannot be read or - /// contains no recordings, and `JsonParse` when a line is not valid JSON. + /// Returns `AiMuxError::InvalidArgument` when the file cannot be read, + /// contains no recordings, or holds a recording of an unsupported schema, + /// and `JsonParse` when a line is not valid JSON. pub fn from_jsonl(path: &str) -> Result { let content = std::fs::read_to_string(path) .map_err(|e| AiMuxError::InvalidArgument(format!("mock replay: {e}")))?; @@ -452,8 +451,16 @@ impl MockReplayModel { if line.trim().is_empty() { continue; } - let rec: Recording = serde_json::from_str(line) - .map_err(|e| AiMuxError::JsonParse(format!("mock replay line {}: {e}", idx + 1)))?; + let rec: Recording = serde_json::from_str(line).map_err(|e| { + let msg = format!("mock replay line {}: {e}", idx + 1); + // Well-formed JSON that is not a current-schema `Recording` + // (including an old `schema`) is a bad argument, not a parse error. + if e.classify() == serde_json::error::Category::Data { + AiMuxError::InvalidArgument(msg) + } else { + AiMuxError::JsonParse(msg) + } + })?; recordings.push(rec); } if recordings.is_empty() { @@ -471,6 +478,32 @@ impl MockReplayModel { pub fn recordings(&self) -> &[Recording] { &self.recordings } + + /// The matcher's recording, or its error naming the recordings of this + /// provider and model whose prompt this version cannot read (written by + /// aimux 0.5 or earlier), which never match. + fn matching(&self, options: &CallOptions) -> Result<&Recording, AiMuxError> { + self.matcher + .r#match(options, &self.recordings) + .map_err(|error| { + let unreadable = self + .recordings + .iter() + .filter(|r| { + r.provider.provider == self.provider + && r.provider.model_id == self.model_id + && r.input.unreadable_prompt().is_some() + }) + .count(); + if unreadable == 0 { + return error; + } + AiMuxError::InvalidArgument(format!( + "{error}; {unreadable} recording(s) of this model hold a prompt this \ + version cannot read (written by aimux 0.5 or earlier) and never match" + )) + }) + } } #[async_trait] @@ -484,13 +517,11 @@ impl LanguageModel for MockReplayModel { } async fn do_generate(&self, options: &CallOptions) -> Result { - let rec = self.matcher.r#match(options, &self.recordings)?; - rebuild_generate_result(rec) + rebuild_generate_result(self.matching(options)?) } async fn do_stream(&self, options: &CallOptions) -> Result { - let rec = self.matcher.r#match(options, &self.recordings)?; - rebuild_stream_result(rec) + rebuild_stream_result(self.matching(options)?) } } @@ -946,13 +977,13 @@ pub struct ReplayOverrides { /// - `recording.input.options` 中的 headers/provider_options 已脱敏, /// 重发会用脱敏后的 `[REDACTED]` 值——需要真实头时用 `overrides` 或 /// 重建 provider 时补充。 -/// - 逐消息 `provider_options`(如 anthropic cacheControl)不参与重建。 /// /// # Errors /// -/// Returns `AiMuxError::JsonParse` when the recorded call options cannot be -/// deserialized, `AiMuxError::InvalidPrompt` for inline text file data without -/// a user-facing representation, and propagates errors from `generate_text`. +/// Returns `AiMuxError::InvalidArgument` for a recording whose prompt this +/// version cannot read ([`InputRecord::unreadable_prompt`](crate::recording::InputRecord::unreadable_prompt)), +/// `AiMuxError::JsonParse` when the recorded call options cannot be +/// deserialized, and propagates errors from the generation pipeline. pub async fn replay_with_model( recording: &Recording, model: &dyn LanguageModel, @@ -963,12 +994,29 @@ pub async fn replay_with_model( "language replay requires a language-model recording".into(), )); } + if let Some(error) = recording.input.unreadable_prompt() { + return Err(AiMuxError::InvalidArgument(format!( + "language replay: this version cannot read the recorded prompt (recordings \ + written by aimux 0.5 or earlier use another prompt format): {error}" + ))); + } + // A recording without a provider_id was written by aimux 0.5 or earlier; + // its options may use fields this version reads differently. + let older = if recording.provider.provider_id.is_empty() { + " (the recording was written by aimux 0.5 or earlier)" + } else { + "" + }; // 1. 从录制输入重建 CallOptions(round-trip:录制时由 CallOptions 序列化) // and the operation settings recorded next to them. let mut call_options: CallOptions = serde_json::from_value(recording.input.options.clone()) - .map_err(|e| AiMuxError::JsonParse(format!("mock replay: input options invalid: {e}")))?; + .map_err(|e| { + AiMuxError::JsonParse(format!("mock replay: input options invalid{older}: {e}")) + })?; let settings: OperationSettings = serde_json::from_value(recording.input.options.clone()) - .map_err(|e| AiMuxError::JsonParse(format!("mock replay: input options invalid: {e}")))?; + .map_err(|e| { + AiMuxError::JsonParse(format!("mock replay: input options invalid{older}: {e}")) + })?; // 2. 应用 overrides。 if let Some(o) = overrides { @@ -1005,7 +1053,8 @@ fn user_tool_choice(choice: LanguageModelV4ToolChoice) -> ToolChoice { /// 侧到用户侧的反向映射)。 /// /// `instructions` 已烘焙进 prompt(system 消息),重建时置 None; -/// `abort_signal` 运行时句柄不跨 JSON,重建时置 None。 +/// `abort_signal` 运行时句柄不跨 JSON,重建时置 None;回放原样发送录制的 +/// prompt,不再下载其中的 URL,`download` 置 None。 fn generate_options_from_call_options( o: CallOptions, settings: OperationSettings, @@ -1026,13 +1075,13 @@ fn generate_options_from_call_options( provider_options: o.provider_options, reasoning: o.reasoning, instructions: None, - body_overrides: o.body_overrides, max_retries: settings.max_retries, timeout: settings.timeout, session_id: settings.session_id, abort_signal: None, include_raw_chunks: o.include_raw_chunks, repair_tool_call: None, + download: None, } } @@ -1100,6 +1149,7 @@ mod tests { call_id: trace_id.to_string(), recorded_at: "2026-08-06T00:00:00Z".to_string(), input: InputRecord { + decision_capabilities: None, operation: Default::default(), prompt: input_prompt, // A recorded call always carries the prepared tool choice. @@ -1108,14 +1158,7 @@ mod tests { "tool_choice": { "type": "auto" }, }), }, - provider: ProviderRecord { - provider: "openai".into(), - model_id: "gpt-4o".into(), - base_url: None, - api_key_source: "none".into(), - profile: None, - provider_options: None, - }, + provider: ProviderRecord::new("openai", "openai", "gpt-4o"), exchanges: vec![HttpExchange { step: None, attempt: 0, diff --git a/aimux-core/src/reranking_model.rs b/aimux-core/src/reranking_model.rs index 022141bed..d4677b756 100644 --- a/aimux-core/src/reranking_model.rs +++ b/aimux-core/src/reranking_model.rs @@ -44,9 +44,7 @@ pub struct RerankingCallOptions { pub abort_signal: Option, /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - /// retrying. `None` uses 2, the AI SDK default, unless the model's - /// `retry_config()` sets another count (model-level retry settings are - /// removed in the target design, docs/aisdk-architecture-alignment.md D14). + /// retrying. `None` uses 2, the AI SDK default. pub max_retries: Option, /// Per-call operation timeout. @@ -138,10 +136,6 @@ pub trait RerankingModel: Send + Sync { /// Provider-specific model ID, e.g. `"rerank-english-v3.0"`. fn model_id(&self) -> &str; - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig::default() - } - /// Rerank a list of documents using the query. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. @@ -162,11 +156,7 @@ pub async fn rerank( ) -> Result { let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; let abort_signal = options.abort_signal.clone(); - let retries = retry::prepare_retries( - options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, abort_signal.clone()); timeout::run( retries.retry(|| model.do_rerank(&options)), abort_signal.as_ref(), diff --git a/aimux-core/src/retry.rs b/aimux-core/src/retry.rs index 7e3fc53f7..a16e11926 100644 --- a/aimux-core/src/retry.rs +++ b/aimux-core/src/retry.rs @@ -8,40 +8,19 @@ use rand::Rng; use crate::AbortSignal; use crate::error::{AiMuxError, RetryError, RetryErrorReason}; +/// Default maximum number of retries after the initial attempt. const DEFAULT_MAX_RETRIES: u32 = 2; -const INITIAL_DELAY_MS: u64 = 2_000; -const BACKOFF_FACTOR: u32 = 2; - -/// Default retry settings for model operations. -/// -/// Per-call `max_retries` overrides only [`Self::max_retries`]; the configured -/// delay and backoff factor remain in effect. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct RetryConfig { - /// Maximum number of retries after the initial attempt. - pub max_retries: u32, - /// Initial delay between attempts. - pub initial_delay: Duration, - /// Multiplier applied to the delay after each retry. - pub backoff_factor: u32, -} - -impl Default for RetryConfig { - fn default() -> Self { - Self { - max_retries: DEFAULT_MAX_RETRIES, - initial_delay: Duration::from_millis(INITIAL_DELAY_MS), - backoff_factor: BACKOFF_FACTOR, - } - } -} +/// Default delay before the first retry, in milliseconds. +const DEFAULT_INITIAL_DELAY_MS: u64 = 2_000; +/// Default multiplier applied to the delay after each retry. +const DEFAULT_BACKOFF_FACTOR: u32 = 2; /// A prepared operation retry function and its resolved retry count. /// /// This is the Rust representation of the AI SDK's /// `prepareRetries()` result: `{ maxRetries, retry }`. #[derive(Debug, Clone)] -pub struct PreparedRetries { +pub(crate) struct PreparedRetries { /// The effective maximum after applying the default. pub max_retries: u32, initial_delay_ms: u64, @@ -49,21 +28,20 @@ pub struct PreparedRetries { abort_signal: Option, } -/// Bind retry settings, applying a per-call retry-count override. +/// Bind the call-level retry count, applying [`DEFAULT_MAX_RETRIES`] when the +/// caller gave none. The delay and backoff factor are the fixed +/// [`DEFAULT_INITIAL_DELAY_MS`] / [`DEFAULT_BACKOFF_FACTOR`] constants: retry +/// is a call-level concern with no provider-level settings. #[must_use] -pub fn prepare_retries( +pub(crate) fn prepare_retries( max_retries: Option, - mut config: RetryConfig, - abort_signal: Option, + abort: Option, ) -> PreparedRetries { - if let Some(max_retries) = max_retries { - config.max_retries = max_retries; - } PreparedRetries { - max_retries: config.max_retries, - initial_delay_ms: u64::try_from(config.initial_delay.as_millis()).unwrap_or(u64::MAX), - backoff_factor: u64::from(config.backoff_factor), - abort_signal, + max_retries: max_retries.unwrap_or(DEFAULT_MAX_RETRIES), + initial_delay_ms: DEFAULT_INITIAL_DELAY_MS, + backoff_factor: u64::from(DEFAULT_BACKOFF_FACTOR), + abort_signal: abort, } } @@ -299,14 +277,11 @@ mod tests { #[tokio::test] async fn prepared_retries_resolves_the_default_and_binds_abort() { - assert_eq!( - prepare_retries(None, RetryConfig::default(), None).max_retries, - 2 - ); + assert_eq!(prepare_retries(None, None).max_retries, 2); let signal = AbortSignal::new(); signal.abort(); - let retries = prepare_retries(Some(4), RetryConfig::default(), Some(signal)); + let retries = prepare_retries(Some(4), Some(signal)); let attempts = AtomicUsize::new(0); let error = retries @@ -323,20 +298,14 @@ mod tests { } #[test] - fn per_call_max_preserves_the_model_delay_and_backoff() { - let retries = prepare_retries( - Some(4), - RetryConfig { - max_retries: 7, - initial_delay: Duration::from_millis(123), - backoff_factor: 3, - }, - None, - ); + 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, 123); - assert_eq!(retries.backoff_factor, 3); + 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] diff --git a/aimux-core/src/router.rs b/aimux-core/src/router.rs index 6d97cc585..910cecf3d 100644 --- a/aimux-core/src/router.rs +++ b/aimux-core/src/router.rs @@ -178,11 +178,7 @@ impl RouterModel { if i == exclude { continue; } - let retries = retry::prepare_retries( - child_max_retries(), - m.retry_config(), - options.abort_signal.clone(), - ); + let retries = retry::prepare_retries(child_max_retries(), options.abort_signal.clone()); let context = step_context(self.step_label(i)); match RecordingContext::scope( context.clone(), @@ -232,24 +228,11 @@ impl LanguageModel for RouterModel { &self.config.model_id } - fn retry_config(&self) -> retry::RetryConfig { - // Retrying the composite would rerun routing and previously attempted - // children; each child is retried at its own execution boundary. - retry::RetryConfig { - max_retries: 0, - ..retry::RetryConfig::default() - } - } - async fn do_generate(&self, options: &CallOptions) -> Result { let raw = self.router.route(&options.prompt, &self.models)?; let idx = self.check_index(raw, "router")?; let model = &self.models[idx]; - let retries = retry::prepare_retries( - child_max_retries(), - model.retry_config(), - options.abort_signal.clone(), - ); + let retries = retry::prepare_retries(child_max_retries(), options.abort_signal.clone()); let context = step_context(self.step_label(idx)); match RecordingContext::scope( context.clone(), @@ -277,11 +260,7 @@ impl LanguageModel for RouterModel { let raw = self.router.route(&options.prompt, &self.models)?; let idx = self.check_index(raw, "router")?; let model = &self.models[idx]; - let retries = retry::prepare_retries( - child_max_retries(), - model.retry_config(), - options.abort_signal.clone(), - ); + let retries = retry::prepare_retries(child_max_retries(), options.abort_signal.clone()); let context = step_context(self.step_label(idx)); // Only setup can be retried. The returned stream is passed through, so // failures after setup never invoke another child or duplicate output. @@ -674,7 +653,6 @@ mod tests { FallbackPolicy::OnError, RouterConfig::default(), ); - assert_eq!(router.retry_config().max_retries, 0); let result = crate::generate::generate_text( &router, diff --git a/aimux-core/src/search_model.rs b/aimux-core/src/search_model.rs index fc5516763..544da00f2 100644 --- a/aimux-core/src/search_model.rs +++ b/aimux-core/src/search_model.rs @@ -46,9 +46,7 @@ pub struct SearchCallOptions { pub abort_signal: Option, /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - /// retrying. `None` uses 2, the AI SDK default, unless the model's - /// `retry_config()` sets another count (model-level retry settings are - /// removed in the target design, docs/aisdk-architecture-alignment.md D14). + /// retrying. `None` uses 2, the AI SDK default. pub max_retries: Option, /// Per-call operation timeout. @@ -153,10 +151,6 @@ pub trait SearchModel: Send + Sync { /// `"tavily-search"`; others accept endpoint-specific names). fn model_id(&self) -> &str; - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig::default() - } - /// Execute a search query and return results. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. @@ -174,11 +168,7 @@ pub async fn search( ) -> Result { let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; let abort_signal = options.abort_signal.clone(); - let retries = retry::prepare_retries( - options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, abort_signal.clone()); timeout::run( retries.retry(|| model.do_search(&options)), abort_signal.as_ref(), diff --git a/aimux-core/src/session.rs b/aimux-core/src/session.rs index e728fd286..dc5b4180a 100644 --- a/aimux-core/src/session.rs +++ b/aimux-core/src/session.rs @@ -348,7 +348,7 @@ pub fn session_inferer() -> Option>> { /// Lazily honor `AIMUX_SESSION_INFER=1` (checked once, like RFC-0014's env /// auto-init). No-op when the inferer is already registered programmatically /// (explicit configuration wins) or the env var is not set. -pub fn ensure_inferer_from_env() { +fn ensure_inferer_from_env() { if INFER_ENV_CHECKED.is_completed() { return; } @@ -365,7 +365,7 @@ pub fn ensure_inferer_from_env() { /// Returns `(session_id, source)` or `None` when no explicit id is present /// and inference is off (or not registered). #[must_use] -pub fn resolve_session_id( +pub(crate) fn resolve_session_id( explicit: Option<&str>, prompt: &LanguageModelPrompt, ) -> Option<(String, SessionSource)> { diff --git a/aimux-core/src/shared.rs b/aimux-core/src/shared.rs index 7322a4142..bf0785db3 100644 --- a/aimux-core/src/shared.rs +++ b/aimux-core/src/shared.rs @@ -376,6 +376,16 @@ impl From for ResponseInfo { } } +impl From for crate::types::ResponseMetadata { + fn from(response: ResponseInfo) -> Self { + Self { + id: response.id, + timestamp: response.timestamp, + model_id: response.model_id, + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/aimux-core/src/skills_model.rs b/aimux-core/src/skills_model.rs new file mode 100644 index 000000000..a619a3c36 --- /dev/null +++ b/aimux-core/src/skills_model.rs @@ -0,0 +1,51 @@ +//! Provider-facing skill upload interface, aligned with SkillsV4. + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use ts_rs::TS; + +use crate::error::AiMuxError; +use crate::files_model::UploadFileData; +use crate::shared::{ + SharedProviderMetadata, SharedProviderOptions, SharedProviderReference, Warning, +}; + +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[cfg_attr(feature = "ts-export", ts(export))] +pub struct SkillFile { + pub path: String, + pub data: UploadFileData, +} + +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[cfg_attr(feature = "ts-export", ts(export))] +pub struct UploadSkillCallOptions { + pub files: Vec, + pub display_title: Option, + pub provider_options: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[cfg_attr(feature = "ts-export", ts(export))] +pub struct UploadSkillResult { + pub provider_reference: SharedProviderReference, + pub display_title: Option, + pub name: Option, + pub description: Option, + pub latest_version: Option, + pub provider_metadata: Option, + pub warnings: Vec, +} + +#[async_trait] +pub trait Skills: Send + Sync { + /// Specification version (always `"v4"`). + fn specification_version(&self) -> &'static str { + "v4" + } + fn provider(&self) -> &str; + async fn upload_skill( + &self, + options: &UploadSkillCallOptions, + ) -> Result; +} diff --git a/aimux-core/src/speech_model.rs b/aimux-core/src/speech_model.rs index e07c4a2bd..8070bf8aa 100644 --- a/aimux-core/src/speech_model.rs +++ b/aimux-core/src/speech_model.rs @@ -60,9 +60,7 @@ pub struct SpeechCallOptions { pub abort_signal: Option, /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - /// retrying. `None` uses 2, the AI SDK default, unless the model's - /// `retry_config()` sets another count (model-level retry settings are - /// removed in the target design, docs/aisdk-architecture-alignment.md D14). + /// retrying. `None` uses 2, the AI SDK default. pub max_retries: Option, /// Per-call operation timeout. @@ -151,10 +149,6 @@ pub trait SpeechModel: Send + Sync { /// Provider-specific model ID, e.g. `"tts-1"`. fn model_id(&self) -> &str; - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig::default() - } - /// Generate speech audio from text. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. @@ -172,11 +166,7 @@ pub async fn generate_speech( ) -> Result { let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; let abort_signal = options.abort_signal.clone(); - let retries = retry::prepare_retries( - options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, abort_signal.clone()); timeout::run( retries.retry(|| model.do_generate(&options)), abort_signal.as_ref(), diff --git a/aimux-core/src/trace/layer.rs b/aimux-core/src/trace/layer.rs index 18847158b..03b5ac1ce 100644 --- a/aimux-core/src/trace/layer.rs +++ b/aimux-core/src/trace/layer.rs @@ -76,6 +76,8 @@ pub struct TraceLayer { default_session: Option, /// Scope salt (HMAC-style key for fingerprint scoping). scope_salt: u64, + /// Name the scope key isolates by; `None` uses the model's `provider()`. + scope_name: Option, session_tracker: Arc>>, } @@ -116,10 +118,24 @@ impl TraceLayer { strict: false, default_session: None, scope_salt: hash::mix(0x5eed_5eed ^ std::process::id() as u64), + scope_name: None, session_tracker: Arc::new(Mutex::new(HashMap::new())), } } + /// Name this layer's trace scope instead of the model's `provider()`. + /// + /// The scope key isolates LCP history by (scope name, model id), and the + /// model carries no endpoint, so two endpoints reached under one provider + /// name share a scope. A host that calls one provider at several base + /// URLs gives each endpoint its own name here. Records still report the + /// model's `provider()`. + #[must_use] + pub fn with_scope_name(mut self, name: impl Into) -> Self { + self.scope_name = Some(name.into()); + self + } + /// Default session id (used when the operation sets no session id) + /// session key (32-byte HMAC-style salt for fingerprint scoping). #[must_use] @@ -163,24 +179,21 @@ impl TraceLayer { let session_id = crate::operation::OperationSettings::current() .and_then(|settings| settings.session_id) .or_else(|| self.default_session.clone()); - // B1: scope_key isolates by (provider, base_url, model_id) — not just - // model_id — so two TraceLayers over the same model_id but different - // providers (or base_urls) sharing one RingTraceStore never cross-match - // their LCP history. `base_url` comes from the inner model's config - // snapshot; providers that only report a minimal snapshot omit it, in - // which case provider + model_id still isolate the scope. The default - // `scope_salt` stays process-level (set in `new` / `with_default_session`). + // B1: scope_key isolates by (scope name, model_id), so two TraceLayers + // over the same model_id but different providers sharing one + // RingTraceStore never cross-match their LCP history. The scope name + // is the model's `provider()` string, which carries the settings-level + // name, unless the host set one (`with_scope_name`). Two endpoints of + // one vendor under the same name are NOT isolated: the model reports + // no base URL, so the host names each endpoint. The default + // `scope_salt` stays process-level (set in `new` / + // `with_default_session`). let provider = self.inner.provider(); let model_id = self.inner.model_id(); - let base_url = self.inner.config_snapshot().base_url; - let mut scope_bytes = Vec::with_capacity( - provider.len() + model_id.len() + base_url.as_deref().map_or(0, str::len), - ); - scope_bytes.extend_from_slice(provider.as_bytes()); + let scope_name = self.scope_name.as_deref().unwrap_or(provider); + let mut scope_bytes = Vec::with_capacity(scope_name.len() + model_id.len()); + scope_bytes.extend_from_slice(scope_name.as_bytes()); scope_bytes.extend_from_slice(model_id.as_bytes()); - if let Some(b) = &base_url { - scope_bytes.extend_from_slice(b.as_bytes()); - } let scope_key = hash::hash64(self.scope_salt, &scope_bytes); let sent_at_unix_ms = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -505,10 +518,6 @@ impl Drop for TraceRecordingStream { #[async_trait] impl LanguageModel for TraceLayer { - fn specification_version(&self) -> &'static str { - "v4" - } - fn provider(&self) -> &str { self.inner.provider() } @@ -517,14 +526,8 @@ impl LanguageModel for TraceLayer { self.inner.model_id() } - fn retry_config(&self) -> crate::retry::RetryConfig { - self.inner.retry_config() - } - - /// RFC-0023 §3.3: transparent decorators must forward the inner snapshot - /// (otherwise recording sees the decorator's minimal record). - fn config_snapshot(&self) -> crate::recording::ProviderRecord { - self.inner.config_snapshot() + async fn supported_urls(&self) -> crate::language_model::SupportedUrls { + self.inner.supported_urls().await } async fn do_generate(&self, options: &CallOptions) -> Result { @@ -616,3 +619,59 @@ impl LanguageModel for TraceLayer { }) } } + +#[cfg(test)] +mod tests { + use super::*; + + struct Model(&'static str); + + #[async_trait] + impl LanguageModel for Model { + fn provider(&self) -> &str { + self.0 + } + + fn model_id(&self) -> &str { + "m" + } + + async fn do_generate(&self, _: &CallOptions) -> Result { + unreachable!("not called") + } + + async fn do_stream(&self, _: &CallOptions) -> Result { + unreachable!("not called") + } + } + + fn scope_key(layer: &TraceLayer) -> u64 { + let ctx = layer.make_record_ctx(None, None, Arc::new(AtomicU64::new(0))); + ctx.ctx.scope_key + } + + /// B18 (#232): the scope key isolates by provider name and model, so two + /// endpoints under one provider name share a scope unless the host names + /// each one. Records keep the model's provider. + #[test] + fn a_scope_name_separates_endpoints_of_one_provider() { + let sink: Arc = Arc::new(RingTraceStore::new()); + let layer = |provider| TraceLayer::new(Arc::new(Model(provider)), sink.clone()); + + let first = scope_key(&layer("openai.chat")); + assert_eq!(first, scope_key(&layer("openai.chat"))); + assert_ne!(first, scope_key(&layer("deepseek.chat"))); + + let named = layer("openai.chat").with_scope_name("openai.chat @ http://a/v1"); + assert_ne!(first, scope_key(&named)); + let other = layer("openai.chat").with_scope_name("openai.chat @ http://b/v1"); + assert_ne!(scope_key(&named), scope_key(&other)); + assert_eq!( + scope_key(&named), + scope_key(&layer("openai.chat").with_scope_name("openai.chat @ http://a/v1")) + ); + + let ctx = named.make_record_ctx(None, None, Arc::new(AtomicU64::new(0))); + assert_eq!(ctx.provider, "openai.chat"); + } +} diff --git a/aimux-core/src/trace/mod.rs b/aimux-core/src/trace/mod.rs index 606c23b2b..f8d60fde2 100644 --- a/aimux-core/src/trace/mod.rs +++ b/aimux-core/src/trace/mod.rs @@ -25,5 +25,5 @@ pub use store::{ }; pub use verdict::{ JudgmentInput, LcpInput, ProviderAuditSpec, ProviderFamily, SessionStats, Verdict, - VerdictConfidence, VerdictKind, judge, matrix, quantize_down, + VerdictConfidence, VerdictKind, judge, matrix, }; diff --git a/aimux-core/src/trace/verdict.rs b/aimux-core/src/trace/verdict.rs index 7e2eea1ee..1ba4ea8bb 100644 --- a/aimux-core/src/trace/verdict.rs +++ b/aimux-core/src/trace/verdict.rs @@ -337,7 +337,7 @@ pub struct JudgmentInput { } #[must_use] -pub fn quantize_down(x: u64, gran: Option) -> u64 { +fn quantize_down(x: u64, gran: Option) -> u64 { match gran { Some(g) if g > 0 => x / g * g, _ => x, diff --git a/aimux-core/src/transcription_model.rs b/aimux-core/src/transcription_model.rs index 29622bd36..80ec7b03b 100644 --- a/aimux-core/src/transcription_model.rs +++ b/aimux-core/src/transcription_model.rs @@ -16,11 +16,12 @@ use futures::Stream; use serde::{Deserialize, Serialize}; use ts_rs::TS; +use crate::download::{Download, create_download}; use crate::error::AiMuxError; use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning}; use crate::{AbortSignal, retry, timeout}; -/// Audio input: raw bytes or a base64-encoded string. +/// Audio input: raw bytes, a base64-encoded string, or a URL. #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[cfg_attr(feature = "ts-export", ts(export))] pub enum AudioInput { @@ -28,6 +29,21 @@ pub enum AudioInput { Binary(Vec), /// A base64-encoded string. Base64(String), + /// A URL that [`transcribe`] downloads before calling the model, the AI + /// SDK's `audio: URL` (ai@7.0.127 `transcribe/transcribe.ts:192-196`). + /// A model takes audio data and never receives a URL. + Url(String), +} + +impl AudioInput { + /// The error a transcription model returns for an [`AudioInput::Url`]: + /// only [`transcribe`] downloads a URL. + #[must_use] + pub fn url_error() -> AiMuxError { + AiMuxError::InvalidArgument( + "a transcription model takes audio data; `transcribe` downloads an audio URL".into(), + ) + } } /// A chunk of audio in a streaming transcription request. @@ -58,10 +74,13 @@ pub struct TranscriptionSegment { #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[cfg_attr(feature = "ts-export", ts(export))] pub struct TranscriptionCallOptions { - /// Audio data to transcribe (raw bytes or a base64-encoded string). + /// Audio to transcribe: raw bytes, a base64-encoded string, or a URL + /// that [`transcribe`] downloads. pub audio: AudioInput, - /// The IANA media type of the audio data, e.g. `"audio/mp3"`. + /// The IANA media type of the audio data, e.g. `"audio/mp3"`. For a URL, + /// the media type the download reports replaces it, as the AI SDK uses + /// the downloaded media type (`transcribe.ts:227-233`). pub media_type: String, /// Additional provider-specific options, keyed by provider name. @@ -73,9 +92,7 @@ pub struct TranscriptionCallOptions { pub abort_signal: Option, /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - /// retrying. `None` uses 2, the AI SDK default, unless the model's - /// `retry_config()` sets another count (model-level retry settings are - /// removed in the target design, docs/aisdk-architecture-alignment.md D14). + /// retrying. `None` uses 2, the AI SDK default. pub max_retries: Option, /// Per-call operation timeout. @@ -83,6 +100,15 @@ pub struct TranscriptionCallOptions { /// Additional HTTP headers to send with the request. pub headers: Option, + + /// The download of an [`AudioInput::Url`], the AI SDK's `transcribe` + /// `download` option (`transcribe.ts:56`, `:109-117`). `None` uses the + /// guarded download ([`create_download`]`(None)`, 2 GiB limit). + /// + /// Runtime handle; never crosses the JSON boundary. + #[serde(skip)] + #[ts(skip)] + pub download: Option, } impl TranscriptionCallOptions { @@ -96,6 +122,7 @@ impl TranscriptionCallOptions { max_retries: None, timeout: None, headers: None, + download: None, } } } @@ -298,10 +325,6 @@ pub trait TranscriptionModel: Send + Sync { /// Provider-specific model ID, e.g. `"whisper-1"`. fn model_id(&self) -> &str; - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig::default() - } - /// Generate a transcript. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. @@ -328,26 +351,37 @@ pub trait TranscriptionModel: Send + Sync { /// User-facing non-streaming transcription with Core-owned retry and timeout. /// +/// An [`AudioInput::Url`] is downloaded first, once, with +/// [`TranscriptionCallOptions::download`] (AI SDK `transcribe`, +/// `transcribe.ts:192-199`); the download counts against the operation +/// timeout and observes the abort signal. +/// /// # Errors /// -/// Returns the provider failure, retry exhaustion, timeout, or caller abort. +/// Returns the download failure, the provider failure, retry exhaustion, +/// timeout, or caller abort. pub async fn transcribe( model: &dyn TranscriptionModel, - options: TranscriptionCallOptions, + mut options: TranscriptionCallOptions, ) -> Result { let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; let abort_signal = options.abort_signal.clone(); - let retries = retry::prepare_retries( - options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); - timeout::run( - retries.retry(|| model.do_generate(&options)), - abort_signal.as_ref(), - timeout, - ) - .await + let retries = retry::prepare_retries(options.max_retries, abort_signal.clone()); + let operation = async { + if let AudioInput::Url(url) = &options.audio { + let download = options + .download + .clone() + .unwrap_or_else(|| create_download(None)); + let file = download.download(url, abort_signal.as_ref()).await?; + options.audio = AudioInput::Binary(file.data); + if let Some(media_type) = file.media_type { + options.media_type = media_type; + } + } + retries.retry(|| model.do_generate(&options)).await + }; + timeout::run(operation, abort_signal.as_ref(), timeout).await } /// Start user-facing live transcription. @@ -365,3 +399,115 @@ pub async fn stream_transcribe( ) -> Result { model.do_stream(options).await } + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::{Arc, Mutex}; + + use crate::download::{DownloadRequest, DownloadedFile}; + + /// Records the audio and media type each call receives. + #[derive(Default)] + struct RecordingModel { + calls: Mutex, String)>>, + } + + #[async_trait] + impl TranscriptionModel for RecordingModel { + fn provider(&self) -> &str { + "test" + } + + fn model_id(&self) -> &str { + "recording" + } + + async fn do_generate( + &self, + options: &TranscriptionCallOptions, + ) -> Result { + let AudioInput::Binary(audio) = &options.audio else { + return Err(AudioInput::url_error()); + }; + self.calls + .lock() + .unwrap() + .push((audio.clone(), options.media_type.clone())); + Ok(TranscriptionResult { + text: "hello".into(), + segments: Vec::new(), + language: None, + duration_in_seconds: None, + warnings: Vec::new(), + request: None, + response: TranscriptionResponse::default(), + provider_metadata: None, + }) + } + } + + fn url_options( + media_type: Option<&'static str>, + ) -> (TranscriptionCallOptions, Arc>>) { + let requests = Arc::new(Mutex::new(Vec::new())); + let mut options = TranscriptionCallOptions::new( + AudioInput::Url("https://cdn.test/a.mp3".into()), + "audio/wav", + ); + let recorded = requests.clone(); + options.download = Some(Download::new(move |request: DownloadRequest| { + recorded.lock().unwrap().push(request.url); + async move { + Ok(DownloadedFile { + data: vec![1, 2, 3], + media_type: media_type.map(str::to_owned), + }) + } + })); + (options, requests) + } + + #[tokio::test] + async fn a_url_is_downloaded_with_the_injected_download_and_sent_as_data() { + let model = RecordingModel::default(); + let (options, requests) = url_options(Some("audio/mpeg")); + let result = transcribe(&model, options).await.unwrap(); + assert_eq!(result.text, "hello"); + assert_eq!(*requests.lock().unwrap(), ["https://cdn.test/a.mp3"]); + assert_eq!( + *model.calls.lock().unwrap(), + [(vec![1, 2, 3], "audio/mpeg".to_owned())] + ); + } + + #[tokio::test] + async fn without_a_downloaded_media_type_the_given_one_stays() { + let model = RecordingModel::default(); + let (options, _) = url_options(None); + transcribe(&model, options).await.unwrap(); + assert_eq!( + *model.calls.lock().unwrap(), + [(vec![1, 2, 3], "audio/wav".to_owned())] + ); + } + + /// Without an injected download, the guarded download refuses a + /// loopback URL before the model is called. + #[tokio::test] + async fn the_default_download_is_guarded() { + let model = RecordingModel::default(); + let options = TranscriptionCallOptions::new( + AudioInput::Url("http://127.0.0.1:9/a.wav".into()), + "audio/wav", + ); + let error = transcribe(&model, options).await.unwrap_err(); + assert!( + matches!(&error, AiMuxError::InvalidArgument(message) + if message == "download URL targets a non-public address: 127.0.0.1"), + "{error:?}" + ); + assert!(model.calls.lock().unwrap().is_empty()); + } +} diff --git a/aimux-core/src/util.rs b/aimux-core/src/util.rs index 7ac22001c..39e4212f7 100644 --- a/aimux-core/src/util.rs +++ b/aimux-core/src/util.rs @@ -19,7 +19,7 @@ pub use crate::math::{UtilError, cosine_similarity}; /// Current time as RFC 3339 UTC with millisecond precision /// (`2026-08-05T04:52:30.123Z`). #[must_use] -pub fn rfc3339_now() -> String { +pub(crate) fn rfc3339_now() -> String { let d = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default(); diff --git a/aimux-core/src/video_model.rs b/aimux-core/src/video_model.rs index 45684b633..0970bdd03 100644 --- a/aimux-core/src/video_model.rs +++ b/aimux-core/src/video_model.rs @@ -9,6 +9,7 @@ use ts_rs::TS; use std::time::{Duration, Instant}; +use crate::download::{Download, create_download}; use crate::error::AiMuxError; use crate::shared::{ AspectRatio, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Size, Warning, @@ -140,9 +141,7 @@ pub struct VideoCallOptions { pub abort_signal: Option, /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - /// retrying. `None` uses 2, the AI SDK default, unless the model's - /// `retry_config()` sets another count (model-level retry settings are - /// removed in the target design, docs/aisdk-architecture-alignment.md D14). + /// retrying. `None` uses 2, the AI SDK default. pub max_retries: Option, /// Poll pacing for the start/status flow, the AI SDK's `generateVideo` @@ -155,6 +154,16 @@ pub struct VideoCallOptions { /// Additional HTTP headers to send with the request. pub headers: Option, + + /// The download of the videos a model returns as URLs, the AI SDK's + /// `generateVideo` `download` option (`generate-video.ts:136`, + /// `:253-261`). `None` uses the guarded download + /// ([`create_download`]`(None)`, 2 GiB limit). + /// + /// Runtime handle; never crosses the JSON boundary. + #[serde(skip)] + #[ts(skip)] + pub download: Option, } fn default_video_n() -> u32 { @@ -182,6 +191,7 @@ impl VideoCallOptions { poll: None, timeout: None, headers: None, + download: None, } } } @@ -325,10 +335,6 @@ pub trait VideoModel: Send + Sync { /// Provider-specific model ID, e.g. `"kling-video"`. fn model_id(&self) -> &str; - fn retry_config(&self) -> retry::RetryConfig { - retry::RetryConfig::default() - } - /// Limit of how many videos can be generated in a single API call. /// /// `None` means no fixed limit. Most video models only support `1`. @@ -368,11 +374,20 @@ pub trait VideoModel: Send + Sync { /// deep-merged the same way start/completion metadata is within a batch (see /// `merge_provider_metadata`). /// +/// Every video returned as a URL is then downloaded, in order, with +/// [`VideoCallOptions::download`] and returned as [`VideoData::Binary`], as +/// the AI SDK returns generated files (`generate-video.ts:383-410`). Its +/// media type is the video's own unless that is empty or +/// `application/octet-stream`, else the downloaded one under the same rule, +/// else `video/mp4`; aimux has no byte sniffing in Core, so upstream's +/// `detectMediaType` step is skipped. +/// /// # Errors /// /// Returns `InvalidArgument` when `options.n == 0` (checked before any /// network call), or the first batch's provider failure, retry exhaustion, -/// poll or operation timeout, or caller abort — Rust drops the other +/// poll or operation timeout, a video download failure, or caller abort — +/// Rust drops the other /// in-flight batches' futures on the first error, unlike `Promise.all`, /// which lets sibling settles run to completion; their partial network /// effects (e.g. an already-started but abandoned batch) are the same @@ -402,11 +417,7 @@ async fn start_and_poll( } let abort_signal = options.abort_signal.clone(); - let retries = retry::prepare_retries( - options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, abort_signal.clone()); let poll = PollPacing::resolve(options.poll); @@ -449,9 +460,37 @@ async fn start_and_poll( merged.provider_metadata = merge_provider_metadata(merged.provider_metadata, result.provider_metadata); } + + let download = options + .download + .clone() + .unwrap_or_else(|| create_download(None)); + for video in &mut merged.videos { + if let VideoData::Url { url, media_type } = video { + let file = download.download(url, abort_signal.as_ref()).await?; + let media_type = downloaded_video_media_type(media_type, file.media_type.as_deref()); + *video = VideoData::Binary { + data: file.data, + media_type, + }; + } + } Ok(merged) } +/// The media type of a downloaded video (generate-video.ts:395-406): the +/// video's own, else the downloaded one, else `video/mp4`, skipping empty +/// and `application/octet-stream` values (`isUsableMediaType`). +fn downloaded_video_media_type(own: &str, downloaded: Option<&str>) -> String { + let usable = + |media_type: &&str| !media_type.is_empty() && *media_type != "application/octet-stream"; + Some(own) + .filter(usable) + .or_else(|| downloaded.filter(usable)) + .unwrap_or("video/mp4") + .to_owned() +} + /// Run one `do_start` + poll cycle for a single provider-sized batch. /// /// Each batch mints its own idempotency key (AI SDK batching semantics): a @@ -543,8 +582,8 @@ fn merge_provider_metadata( #[cfg(test)] mod tests { - use std::sync::Mutex; use std::sync::atomic::{AtomicU32, Ordering}; + use std::sync::{Arc, Mutex}; use super::*; use crate::error::ApiCallError; @@ -608,13 +647,6 @@ mod tests { "scripted" } - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig { - initial_delay: Duration::from_millis(1), - ..crate::retry::RetryConfig::default() - } - } - fn max_videos_per_call(&self) -> Option { Some(1) } @@ -681,8 +713,10 @@ mod tests { }))), StatusStep::Pending => Ok(VideoOperationStatus::Pending), StatusStep::Complete => Ok(VideoOperationStatus::Completed(VideoResult { + // A loopback URL: an injected download fetches it, the + // guarded default download refuses it. videos: vec![VideoData::Url { - url: "https://cdn/video.mp4".to_string(), + url: "http://127.0.0.1:9/video.mp4".to_string(), media_type: "video/mp4".to_string(), }], warnings: vec![Warning::Other { @@ -701,6 +735,12 @@ mod tests { interval_ms: Some(1), timeout_ms: Some(2_000), }); + options.download = Some(Download::new(|_| async { + Ok(crate::download::DownloadedFile { + data: b"video".to_vec(), + media_type: None, + }) + })); options } @@ -729,7 +769,8 @@ mod tests { assert_eq!(result.videos.len(), 1); } - #[tokio::test] + // Paused time: the fixed 2s retry delay must not make the test wait. + #[tokio::test(start_paused = true)] async fn start_retry_reuses_one_idempotency_key() { let model = ScriptedVideoModel::new(vec![StatusStep::Complete]).with_start_failures(1); let mut options = fast_poll_options(); @@ -918,7 +959,9 @@ mod tests { /// The race covers the status retries too (upstream passes the status /// timeout signal to `prepareRetries`, generate-video.ts:602-605): a /// retry delay longer than the remaining budget ends with the polling - /// timeout, not after the delay. + /// timeout, not after the delay. The failed check asks for a 30 s delay + /// (`retry-after-ms`, honoured exactly below 60 s), so the delay does not + /// depend on the jittered exponential backoff. #[tokio::test] async fn status_retries_race_the_remaining_poll_timeout() { struct SlowRetries(ScriptedVideoModel); @@ -930,12 +973,6 @@ mod tests { fn model_id(&self) -> &str { "slow-retries" } - fn retry_config(&self) -> crate::retry::RetryConfig { - crate::retry::RetryConfig { - initial_delay: Duration::from_secs(60), - ..crate::retry::RetryConfig::default() - } - } fn max_videos_per_call(&self) -> Option { Some(1) } @@ -950,7 +987,17 @@ mod tests { operation: &serde_json::Value, options: &VideoCallOptions, ) -> Result { - self.0.do_status(operation, options).await + self.0 + .do_status(operation, options) + .await + .map_err(|error| match error { + AiMuxError::ApiCall(mut detail) => { + detail.response_headers = + Some([("retry-after-ms".to_string(), "30000".to_string())].into()); + AiMuxError::ApiCall(detail) + } + other => other, + }) } } let model = SlowRetries(ScriptedVideoModel::new(vec![ @@ -974,6 +1021,70 @@ mod tests { /// Without a `poll` option the AI SDK defaults apply /// (generate-video.ts:543-544); the model has no poll settings. + #[tokio::test] + async fn url_videos_are_downloaded_with_the_injected_download() { + let model = ScriptedVideoModel::new(vec![StatusStep::Complete]); + let requests = Arc::new(Mutex::new(Vec::new())); + let mut options = fast_poll_options(); + let recorded = requests.clone(); + options.download = Some(Download::new( + move |request: crate::download::DownloadRequest| { + recorded.lock().unwrap().push(request.url); + async { + Ok(crate::download::DownloadedFile { + data: vec![7, 7], + media_type: Some("video/webm".into()), + }) + } + }, + )); + let result = generate_video(&model, options).await.unwrap(); + assert_eq!(*requests.lock().unwrap(), ["http://127.0.0.1:9/video.mp4"]); + // The video's own media type comes first. + assert!( + matches!(&result.videos[..], [VideoData::Binary { data, media_type }] + if *data == [7, 7] && media_type == "video/mp4"), + "{:?}", + result.videos + ); + } + + /// Without an injected download, the guarded download refuses a + /// loopback URL. + #[tokio::test] + async fn the_default_video_download_is_guarded() { + let model = ScriptedVideoModel::new(vec![StatusStep::Complete]); + let mut options = fast_poll_options(); + options.download = None; + let error = generate_video(&model, options).await.unwrap_err(); + assert!( + matches!(&error, AiMuxError::InvalidArgument(message) + if message == "download URL targets a non-public address: 127.0.0.1"), + "{error:?}" + ); + } + + #[test] + fn a_downloaded_video_takes_the_first_usable_media_type() { + assert_eq!( + downloaded_video_media_type("video/mp4", Some("video/webm")), + "video/mp4" + ); + assert_eq!( + downloaded_video_media_type("application/octet-stream", Some("video/webm")), + "video/webm" + ); + assert_eq!( + downloaded_video_media_type("", Some("video/webm")), + "video/webm" + ); + assert_eq!( + downloaded_video_media_type("", Some("application/octet-stream")), + "video/mp4" + ); + assert_eq!(downloaded_video_media_type("", None), "video/mp4"); + } + #[test] fn poll_pacing_defaults_to_the_ai_sdk_values() { let pacing = PollPacing::resolve(None); @@ -1038,7 +1149,7 @@ mod tests { assert_eq!(unique.len(), 3, "one key per batch: {keys:?}"); } - #[tokio::test] + #[tokio::test(start_paused = true)] async fn caller_key_is_distinct_per_batch_and_stable_across_retries_and_replays() { let model = ScriptedVideoModel::new(vec![]).with_start_failures(1); let mut options = fast_poll_options(); diff --git a/aimux-core/src/wrap_provider.rs b/aimux-core/src/wrap_provider.rs new file mode 100644 index 000000000..90a8cae84 --- /dev/null +++ b/aimux-core/src/wrap_provider.rs @@ -0,0 +1,276 @@ +//! Provider-level middleware: the AI SDK's `wrapProvider` +//! (ai@7.0.127 `middleware/wrap-provider.ts:20-53`). + +use std::sync::Arc; + +use crate::embedding_model::EmbeddingModel; +use crate::error::AiMuxError; +use crate::files_model::Files; +use crate::image_model::ImageModel; +use crate::image_model_middleware::{ImageModelMiddleware, wrap_image_model}; +use crate::language_model::LanguageModel; +use crate::language_model_middleware::{LanguageModelMiddleware, wrap_language_model}; +use crate::provider::{Provider, ProviderDiscovery}; +use crate::reranking_model::RerankingModel; +use crate::search_model::SearchModel; +use crate::skills_model::Skills; +use crate::speech_model::SpeechModel; +use crate::transcription_model::TranscriptionModel; +use crate::video_model::VideoModel; + +/// Wrap every language model of `provider` with `language_model_middleware` +/// and, when `image_model_middleware` is not empty, every image model with +/// it. The first middleware of each list is the outermost, as in +/// [`wrap_language_model`] and [`wrap_image_model`]. +/// +/// Every other member is the wrapped provider's own: embedding, +/// transcription, speech and reranking models, files and skills, as +/// upstream, and the aimux extensions (video, search and decision models, +/// discovery), which are members of [`Provider`] here. +#[must_use] +pub fn wrap_provider( + provider: Arc, + language_model_middleware: &[Arc], + image_model_middleware: &[Arc], +) -> Arc { + Arc::new(WrappedProvider { + provider, + language_model_middleware: language_model_middleware.to_vec(), + image_model_middleware: image_model_middleware.to_vec(), + }) +} + +struct WrappedProvider { + provider: Arc, + language_model_middleware: Vec>, + image_model_middleware: Vec>, +} + +impl Provider for WrappedProvider { + fn language_model(&self, id: &str) -> Result, AiMuxError> { + Ok(wrap_language_model( + self.provider.language_model(id)?, + &self.language_model_middleware, + None, + None, + )) + } + + fn embedding_model(&self, id: &str) -> Result, AiMuxError> { + self.provider.embedding_model(id) + } + + fn image_model(&self, id: &str) -> Result, AiMuxError> { + Ok(wrap_image_model( + self.provider.image_model(id)?, + &self.image_model_middleware, + None, + None, + )) + } + + fn transcription_model( + &self, + id: &str, + ) -> Option, AiMuxError>> { + self.provider.transcription_model(id) + } + + fn speech_model(&self, id: &str) -> Option, AiMuxError>> { + self.provider.speech_model(id) + } + + fn reranking_model(&self, id: &str) -> Option, AiMuxError>> { + self.provider.reranking_model(id) + } + + fn files(&self) -> Option> { + self.provider.files() + } + + fn skills(&self) -> Option> { + self.provider.skills() + } + + fn decision_model(&self, id: &str) -> Result, AiMuxError> { + self.provider.decision_model(id) + } + + fn video_model(&self, id: &str) -> Option, AiMuxError>> { + self.provider.video_model(id) + } + + fn search_model(&self, id: &str) -> Option, AiMuxError>> { + self.provider.search_model(id) + } + + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + self.provider.discovery() + } +} + +#[cfg(test)] +mod tests { + use async_trait::async_trait; + + use super::*; + use crate::embedding_model::{EmbeddingCallOptions, EmbeddingResult}; + use crate::image_model::{ImageCallOptions, ImageResult}; + use crate::options::CallOptions; + use crate::result::{GenerateResult, StreamResult}; + + struct Language; + + #[async_trait] + impl LanguageModel for Language { + fn provider(&self) -> &str { + "mock.chat" + } + + fn model_id(&self) -> &str { + "language" + } + + async fn do_generate(&self, _: &CallOptions) -> Result { + unreachable!("not called") + } + + async fn do_stream(&self, _: &CallOptions) -> Result { + unreachable!("not called") + } + } + + struct Image; + + #[async_trait] + impl ImageModel for Image { + fn provider(&self) -> &str { + "mock.image" + } + + fn model_id(&self) -> &str { + "image" + } + + fn max_images_per_call(&self) -> Option { + None + } + + async fn do_generate(&self, _: &ImageCallOptions) -> Result { + unreachable!("not called") + } + } + + struct Embedding; + + #[async_trait] + impl EmbeddingModel for Embedding { + fn provider(&self) -> &str { + "mock.embedding" + } + + fn model_id(&self) -> &str { + "embedding" + } + + fn max_embeddings_per_call(&self) -> Option { + None + } + + fn supports_parallel_calls(&self) -> bool { + true + } + + async fn do_embed(&self, _: &EmbeddingCallOptions) -> Result { + unreachable!("not called") + } + } + + struct Mock { + image: Arc, + embedding: Arc, + } + + impl Provider for Mock { + fn language_model(&self, _: &str) -> Result, AiMuxError> { + Ok(Arc::new(Language)) + } + + fn embedding_model(&self, _: &str) -> Result, AiMuxError> { + Ok(self.embedding.clone()) + } + + fn image_model(&self, _: &str) -> Result, AiMuxError> { + Ok(self.image.clone()) + } + } + + struct RenameLanguage; + + impl LanguageModelMiddleware for RenameLanguage { + fn override_model_id(&self, model: &dyn LanguageModel) -> Option { + Some(format!("wrapped-{}", model.model_id())) + } + } + + struct RenameImage; + + impl ImageModelMiddleware for RenameImage { + fn override_model_id(&self, model: &dyn ImageModel) -> Option { + Some(format!("wrapped-{}", model.model_id())) + } + } + + fn mock() -> ( + Arc, + Arc, + Arc, + ) { + let image: Arc = Arc::new(Image); + let embedding: Arc = Arc::new(Embedding); + let provider = Arc::new(Mock { + image: image.clone(), + embedding: embedding.clone(), + }); + (provider, image, embedding) + } + + /// `wrapProvider` (ai@7.0.127 `wrap-provider.ts:31-52`): language models + /// are always wrapped, image models only when image middleware is given, + /// and every other member is the provider's own. + #[test] + fn language_models_are_wrapped_and_other_members_pass_through() { + let (provider, image, embedding) = mock(); + let wrapped = wrap_provider(provider, &[Arc::new(RenameLanguage)], &[]); + + let language = wrapped.language_model("language").unwrap(); + assert_eq!(language.model_id(), "wrapped-language"); + assert_eq!(language.provider(), "mock.chat"); + assert!(Arc::ptr_eq(&wrapped.image_model("image").unwrap(), &image)); + assert!(Arc::ptr_eq( + &wrapped.embedding_model("embedding").unwrap(), + &embedding + )); + assert!(wrapped.transcription_model("t").is_none()); + assert!(wrapped.files().is_none()); + assert_eq!(wrapped.specification_version(), "v4"); + } + + #[test] + fn image_models_are_wrapped_when_image_middleware_is_given() { + let (provider, _, _) = mock(); + let wrapped = wrap_provider( + provider, + &[Arc::new(RenameLanguage)], + &[Arc::new(RenameImage)], + ); + assert_eq!( + wrapped.image_model("image").unwrap().model_id(), + "wrapped-image" + ); + assert_eq!( + wrapped.language_model("language").unwrap().model_id(), + "wrapped-language" + ); + } +} diff --git a/aimux-core/tests/error_value_golden_test.rs b/aimux-core/tests/error_value_golden_test.rs index ef1b8fa7a..d9e16e9f0 100644 --- a/aimux-core/tests/error_value_golden_test.rs +++ b/aimux-core/tests/error_value_golden_test.rs @@ -190,6 +190,20 @@ fn error_value_snapshots_plain_variants() { AiMuxError::UnsupportedFunctionality("no audio".into()), r#"{"UnsupportedFunctionality":"no audio"}"#, ), + ( + AiMuxError::LoadApiKey { + env_var: "OPENAI_API_KEY".into(), + description: "OpenAI".into(), + }, + r#"{"LoadApiKey":{"env_var":"OPENAI_API_KEY","description":"OpenAI"}}"#, + ), + ( + AiMuxError::LoadSetting { + env_var: "AWS_REGION".into(), + name: "region".into(), + }, + r#"{"LoadSetting":{"env_var":"AWS_REGION","name":"region"}}"#, + ), ( AiMuxError::NoSuchModel { model_id: "gpt-9".into(), @@ -200,8 +214,11 @@ fn error_value_snapshots_plain_variants() { ( AiMuxError::NoSuchProvider { provider_id: "acme".into(), + model_id: "acme".into(), + model_type: "languageModel".into(), + available_providers: vec!["other".into()], }, - r#"{"NoSuchProvider":{"provider_id":"acme"}}"#, + r#"{"NoSuchProvider":{"provider_id":"acme","model_id":"acme","model_type":"languageModel","available_providers":["other"]}}"#, ), ( AiMuxError::Timeout("total timeout".into()), @@ -293,6 +310,8 @@ fn variant_addition_breaks_this_match(error: &AiMuxError) { | AiMuxError::ToolCallRepair { .. } | AiMuxError::InvalidArgument(_) | AiMuxError::InvalidPrompt(_) + | AiMuxError::LoadApiKey { .. } + | AiMuxError::LoadSetting { .. } | AiMuxError::MissingToolResults { .. } | AiMuxError::ToolChoiceViolation(_) | AiMuxError::NoOutputGenerated(_) @@ -340,6 +359,14 @@ fn variant_set_is_exactly_the_wire_tag_list() { }, AiMuxError::InvalidArgument("x".into()), AiMuxError::InvalidPrompt("x".into()), + AiMuxError::LoadApiKey { + env_var: "X".into(), + description: "x".into(), + }, + AiMuxError::LoadSetting { + env_var: "X".into(), + name: "x".into(), + }, AiMuxError::MissingToolResults { tool_call_ids: vec!["x".into()], }, @@ -359,6 +386,9 @@ fn variant_set_is_exactly_the_wire_tag_list() { }, AiMuxError::NoSuchProvider { provider_id: "x".into(), + model_id: "x".into(), + model_type: String::new(), + available_providers: vec![], }, AiMuxError::Timeout("x".into()), AiMuxError::Aborted("request aborted".into()), @@ -384,6 +414,8 @@ fn variant_set_is_exactly_the_wire_tag_list() { "InvalidResponseData", "InvalidToolInput", "JsonParse", + "LoadApiKey", + "LoadSetting", "MissingToolResults", "NoOutputGenerated", "NoSuchModel", diff --git a/aimux-core/tests/fixtures/master_decision_recordings.jsonl b/aimux-core/tests/fixtures/master_decision_recordings.jsonl new file mode 100644 index 000000000..08c8490fa --- /dev/null +++ b/aimux-core/tests/fixtures/master_decision_recordings.jsonl @@ -0,0 +1,2 @@ +{"schema":2,"call_id":"call-1791293057374663321-0","recorded_at":"2026-10-06T13:24:17.375Z","input":{"operation":"decision","prompt":[],"options":{"state":"I was billed twice and payments are blocked. Please fix this urgently.","questions":[{"type":"boolean","id":"urgent","instructions":"Does the customer request urgent action?"},{"type":"choice","id":"department","instructions":"Choose the responsible department.","options":[{"label":"billing"},{"label":"technical"},{"label":"sales"}]},{"type":"score","id":"severity","instructions":"Rate the impact of the issue.","levels":["Minor inconvenience","Significant issue","Payments completely blocked"]}],"max_retries":0,"timeout":{"total_ms":30000,"step_ms":null,"first_chunk_ms":null,"chunk_ms":null}}},"provider":{"provider":"jev","model_id":"jev-latest","base_url":"https://api.typesafe.ai/v1/systemone","api_key_source":"explicit","profile":{"decision_capabilities":{"rounding":{"probability_decimals":2,"score_decimals":2},"probability_source":"native","supports_boolean":true,"supports_choice":true,"supports_score":true,"returns_distributions":true,"max_questions":null,"min_choices":1,"max_choices":255,"max_score_levels":10}},"provider_options":null},"exchanges":[{"attempt":1,"exchange_index":1,"request":{"method":"POST","url":"https://api.typesafe.ai/v1/systemone","headers":[["Authorization","[REDACTED]"]],"body":"{\"model\":\"jev-latest\",\"state\":\"I was billed twice and payments are blocked. Please fix this urgently.\",\"questions\":{\"urgent\":{\"type\":\"noul\",\"instructions\":\"Does the customer request urgent action?\"},\"department\":{\"type\":\"choice\",\"instructions\":\"Choose the responsible department.\",\"criteria\":{\"billing\":null,\"technical\":null,\"sales\":null}},\"severity\":{\"type\":\"score\",\"instructions\":\"Rate the impact of the issue.\",\"criteria\":[\"Minor inconvenience\",\"Significant issue\",\"Payments completely blocked\"]}}}"},"response":{"status":200,"headers":[["date","Tue, 06 Oct 2026 13:24:17 GMT"],["content-type","application/json"],["content-length","440"],["connection","keep-alive"],["server","cloudflare"],["x-typesafe-request-id","req_01a11162eaba71f5b6825f0466bca790"],["cf-cache-status","DYNAMIC"],["set-cookie","[REDACTED]"],["cf-ray","a464fc4aaee03760-LAX"],["alt-svc","h3=\":443\"; ma=86400"]],"body":"{\"model\":\"jev-1.13.0\",\"answers\":{\"urgent\":{\"type\":\"noul\",\"noul\":0.99},\"department\":{\"type\":\"choice\",\"choice\":\"billing\",\"confidence\":1.0,\"probabilities\":{\"billing\":1.0,\"technical\":0.0,\"sales\":0.0}},\"severity\":{\"type\":\"score\",\"score\":2.0,\"confidence\":1.0,\"legend\":{\"0\":\"Minor inconvenience\",\"1\":\"Significant issue\",\"2\":\"Payments completely blocked\"},\"probabilities\":{\"0\":0.0,\"1\":0.0,\"2\":1.0}}},\"usage\":{\"input_tokens\":365,\"output_tokens\":68}}","stream_chunks":1,"ttfb_ms":486},"timing":{"latency_ms":486,"ttfb_ms":486},"error":null,"finalized":true}],"outcome":{"decision_result":{"rounding":{"probability_decimals":2,"score_decimals":2},"answers":{"department":{"type":"choice","selected":"billing","probabilities":{"billing":1.0,"sales":0.0,"technical":0.0},"confidence":1.0},"severity":{"type":"score","expected_value":2.0,"levels":["Minor inconvenience","Significant issue","Payments completely blocked"],"probabilities":[0.0,0.0,1.0],"confidence":1.0},"urgent":{"type":"boolean","probability_true":0.99}},"provider":"jev","model":"jev-1.13.0","model_version":null,"probability_source":"native","usage":{"input_tokens":{"total":365},"output_tokens":{"total":68},"raw":{"input_tokens":365,"output_tokens":68}},"latency_ms":null,"provider_metadata":{"jev":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.99},"department":{"type":"choice","choice":"billing","confidence":1.0,"probabilities":{"billing":1.0,"technical":0.0,"sales":0.0}},"severity":{"type":"score","score":2.0,"confidence":1.0,"legend":{"0":"Minor inconvenience","1":"Significant issue","2":"Payments completely blocked"},"probabilities":{"0":0.0,"1":0.0,"2":1.0}}},"usage":{"input_tokens":365,"output_tokens":68}}},"response":{"headers":{"connection":"keep-alive","date":"Tue, 06 Oct 2026 13:24:17 GMT","set-cookie":"[REDACTED]","alt-svc":"h3=\":443\"; ma=86400","x-typesafe-request-id":"req_01a11162eaba71f5b6825f0466bca790","cf-cache-status":"DYNAMIC","server":"cloudflare","content-length":"440","content-type":"application/json","cf-ray":"a464fc4aaee03760-LAX"},"body":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.99},"department":{"type":"choice","choice":"billing","confidence":1.0,"probabilities":{"billing":1.0,"technical":0.0,"sales":0.0}},"severity":{"type":"score","score":2.0,"confidence":1.0,"legend":{"0":"Minor inconvenience","1":"Significant issue","2":"Payments completely blocked"},"probabilities":{"0":0.0,"1":0.0,"2":1.0}}},"usage":{"input_tokens":365,"output_tokens":68}}}},"status":"success","finish_reason":null,"error":null,"error_value":null,"usage":{"input_tokens":{"total":365},"output_tokens":{"total":68},"raw":{"input_tokens":365,"output_tokens":68}}},"complete":true,"transport_closed":true} +{"schema":2,"call_id":"call-1791293057868486266-1","recorded_at":"2026-10-06T13:24:17.868Z","input":{"operation":"decision","prompt":[],"options":{"state":{"ticket":"Payments blocked"},"questions":[{"type":"boolean","id":"urgent","instructions":{"question":"Is action urgent?","context":["payments"]},"criteria":{"true":{"includes":["blocked"]},"false":["routine","can wait"]}},{"type":"choice","id":"department","instructions":["Route the ticket",{"priority":"payments"}],"options":[{"label":"billing","description":{"includes":["payments"],"examples":["payout failure"]}},{"label":"technical"}]},{"type":"score","id":"severity","instructions":"Rate severity","levels":[{"label":"low","includes":["minor inconvenience"]},["blocking","No workaround"]]}],"max_retries":0,"timeout":{"total_ms":30000,"step_ms":null,"first_chunk_ms":null,"chunk_ms":null}}},"provider":{"provider":"jev","model_id":"jev-latest","base_url":"https://api.typesafe.ai/v1/systemone","api_key_source":"explicit","profile":{"decision_capabilities":{"rounding":{"probability_decimals":2,"score_decimals":2},"probability_source":"native","supports_boolean":true,"supports_choice":true,"supports_score":true,"returns_distributions":true,"max_questions":null,"min_choices":1,"max_choices":255,"max_score_levels":10}},"provider_options":null},"exchanges":[{"attempt":1,"exchange_index":1,"request":{"method":"POST","url":"https://api.typesafe.ai/v1/systemone","headers":[["Authorization","[REDACTED]"]],"body":"{\"model\":\"jev-latest\",\"state\":{\"ticket\":\"Payments blocked\"},\"questions\":{\"urgent\":{\"type\":\"noul\",\"instructions\":{\"question\":\"Is action urgent?\",\"context\":[\"payments\"]},\"criteria\":{\"true\":{\"includes\":[\"blocked\"]},\"false\":[\"routine\",\"can wait\"]}},\"department\":{\"type\":\"choice\",\"instructions\":[\"Route the ticket\",{\"priority\":\"payments\"}],\"criteria\":{\"billing\":{\"includes\":[\"payments\"],\"examples\":[\"payout failure\"]},\"technical\":null}},\"severity\":{\"type\":\"score\",\"instructions\":\"Rate severity\",\"criteria\":[{\"label\":\"low\",\"includes\":[\"minor inconvenience\"]},[\"blocking\",\"No workaround\"]]}}}"},"response":{"status":200,"headers":[["date","Tue, 06 Oct 2026 13:24:18 GMT"],["content-type","application/json"],["content-length","431"],["connection","keep-alive"],["server","cloudflare"],["set-cookie","[REDACTED]"],["x-typesafe-request-id","req_01a11162ebc076ee849d8e67ae47f547"],["cf-cache-status","DYNAMIC"],["cf-ray","a464fc4c4b1c3760-LAX"],["alt-svc","h3=\":443\"; ma=86400"]],"body":"{\"model\":\"jev-1.13.0\",\"answers\":{\"urgent\":{\"type\":\"noul\",\"noul\":0.92},\"department\":{\"type\":\"choice\",\"choice\":\"billing\",\"confidence\":0.98,\"probabilities\":{\"technical\":0.01,\"billing\":0.99}},\"severity\":{\"type\":\"score\",\"score\":0.99,\"confidence\":0.98,\"legend\":{\"0\":{\"label\":\"low\",\"includes\":[\"minor inconvenience\"]},\"1\":[\"blocking\",\"No workaround\"]},\"probabilities\":{\"0\":0.01,\"1\":0.99}}},\"usage\":{\"input_tokens\":448,\"output_tokens\":61}}","stream_chunks":1,"ttfb_ms":257},"timing":{"latency_ms":257,"ttfb_ms":257},"error":null,"finalized":true}],"outcome":{"decision_result":{"rounding":{"probability_decimals":2,"score_decimals":2},"answers":{"department":{"type":"choice","selected":"billing","probabilities":{"billing":0.99,"technical":0.01},"confidence":0.98},"severity":{"type":"score","expected_value":0.99,"levels":[{"label":"low","includes":["minor inconvenience"]},["blocking","No workaround"]],"probabilities":[0.01,0.99],"confidence":0.98},"urgent":{"type":"boolean","probability_true":0.92}},"provider":"jev","model":"jev-1.13.0","model_version":null,"probability_source":"native","usage":{"input_tokens":{"total":448},"output_tokens":{"total":61},"raw":{"input_tokens":448,"output_tokens":61}},"latency_ms":null,"provider_metadata":{"jev":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.92},"department":{"type":"choice","choice":"billing","confidence":0.98,"probabilities":{"technical":0.01,"billing":0.99}},"severity":{"type":"score","score":0.99,"confidence":0.98,"legend":{"0":{"label":"low","includes":["minor inconvenience"]},"1":["blocking","No workaround"]},"probabilities":{"0":0.01,"1":0.99}}},"usage":{"input_tokens":448,"output_tokens":61}}},"response":{"headers":{"date":"Tue, 06 Oct 2026 13:24:18 GMT","content-length":"431","set-cookie":"[REDACTED]","connection":"keep-alive","server":"cloudflare","cf-ray":"a464fc4c4b1c3760-LAX","alt-svc":"h3=\":443\"; ma=86400","x-typesafe-request-id":"req_01a11162ebc076ee849d8e67ae47f547","cf-cache-status":"DYNAMIC","content-type":"application/json"},"body":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.92},"department":{"type":"choice","choice":"billing","confidence":0.98,"probabilities":{"technical":0.01,"billing":0.99}},"severity":{"type":"score","score":0.99,"confidence":0.98,"legend":{"0":{"label":"low","includes":["minor inconvenience"]},"1":["blocking","No workaround"]},"probabilities":{"0":0.01,"1":0.99}}},"usage":{"input_tokens":448,"output_tokens":61}}}},"status":"success","finish_reason":null,"error":null,"error_value":null,"usage":{"input_tokens":{"total":448},"output_tokens":{"total":61},"raw":{"input_tokens":448,"output_tokens":61}}},"complete":true,"transport_closed":true} diff --git a/aimux-core/tests/fixtures/master_language_recordings.jsonl b/aimux-core/tests/fixtures/master_language_recordings.jsonl new file mode 100644 index 000000000..88d6bf85a --- /dev/null +++ b/aimux-core/tests/fixtures/master_language_recordings.jsonl @@ -0,0 +1,2 @@ +{"schema":2,"call_id":"call-1791920170631056022-0","recorded_at":"2026-10-13T19:36:10.631Z","input":{"prompt":[{"role":"user","content":[{"type":"text","text":"What is the capital of France?"}],"provider_options":null}],"options":{"prompt":[{"role":"user","content":[{"type":"text","text":"What is the capital of France?"}],"provider_options":null}],"max_output_tokens":null,"temperature":0.0,"stop_sequences":null,"top_p":null,"top_k":null,"presence_penalty":null,"frequency_penalty":null,"response_format":null,"seed":null,"tools":null,"tool_choice":"auto","headers":null,"provider_options":null,"reasoning":null,"body_overrides":null,"max_retries":0,"timeout":null,"include_raw_chunks":null}},"provider":{"provider":"openai","model_id":"gpt-4o","base_url":"http://127.0.0.1:37129","api_key_source":"explicit","profile":{"supports_top_k":true,"supports_tools":true,"supports_response_format":true,"stream_usage_key":null,"max_tokens_key":null},"provider_options":{"base_url":null,"headers":null,"organization":null,"project":null,"max_retries":2,"body_overrides":null}},"exchanges":[{"attempt":1,"exchange_index":1,"request":{"method":"POST","url":"http://127.0.0.1:37129/chat/completions","headers":[["Authorization","[REDACTED]"],["Content-Type","application/json"]],"body":"{\"model\":\"gpt-4o\",\"messages\":[{\"role\":\"user\",\"content\":\"What is the capital of France?\"}],\"temperature\":0.0}"},"response":{"status":200,"headers":[["content-type","application/json"],["content-length","261"],["date","Tue, 13 Oct 2026 19:36:10 GMT"]],"body":"{\"id\":\"chatcmpl-1\",\"object\":\"chat.completion\",\"created\":1791300000,\"model\":\"gpt-4o-2024-08-06\",\"choices\":[{\"index\":0,\"message\":{\"role\":\"assistant\",\"content\":\"Paris.\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":14,\"completion_tokens\":2,\"total_tokens\":16}}","stream_chunks":1,"ttfb_ms":3},"timing":{"latency_ms":3,"ttfb_ms":3},"error":null,"finalized":true}],"outcome":{"status":"success","finish_reason":"stop","error":null,"error_value":null,"usage":{"input_tokens":{"total":14,"no_cache":14,"cache_read":0},"output_tokens":{"total":2,"text":2,"reasoning":0},"raw":{"prompt_tokens":14,"completion_tokens":2,"total_tokens":16}}},"complete":true,"transport_closed":true} +{"schema":2,"call_id":"call-1791920170635884356-1","recorded_at":"2026-10-13T19:36:10.636Z","input":{"prompt":[{"role":"system","content":[{"type":"text","text":"Answer in one word."}],"provider_options":null},{"role":"user","content":[{"type":"text","text":"Capital of Italy?"}],"provider_options":null}],"options":{"prompt":[{"role":"system","content":[{"type":"text","text":"Answer in one word."}],"provider_options":null},{"role":"user","content":[{"type":"text","text":"Capital of Italy?"}],"provider_options":null}],"max_output_tokens":null,"temperature":null,"stop_sequences":null,"top_p":null,"top_k":null,"presence_penalty":null,"frequency_penalty":null,"response_format":null,"seed":null,"tools":null,"tool_choice":"auto","headers":null,"provider_options":null,"reasoning":null,"body_overrides":null,"max_retries":0,"timeout":null,"include_raw_chunks":null}},"provider":{"provider":"openai","model_id":"gpt-4o","base_url":"http://127.0.0.1:37129","api_key_source":"explicit","profile":{"supports_top_k":true,"supports_tools":true,"supports_response_format":true,"stream_usage_key":null,"max_tokens_key":null},"provider_options":{"base_url":null,"headers":null,"organization":null,"project":null,"max_retries":2,"body_overrides":null}},"exchanges":[{"attempt":1,"exchange_index":1,"request":{"method":"POST","url":"http://127.0.0.1:37129/chat/completions","headers":[["Authorization","[REDACTED]"],["Content-Type","application/json"]],"body":"{\"model\":\"gpt-4o\",\"messages\":[{\"role\":\"system\",\"content\":\"Answer in one word.\"},{\"role\":\"user\",\"content\":\"Capital of Italy?\"}]}"},"response":{"status":200,"headers":[["content-type","application/json"],["content-length","259"],["date","Tue, 13 Oct 2026 19:36:10 GMT"]],"body":"{\"id\":\"chatcmpl-2\",\"object\":\"chat.completion\",\"created\":1791300000,\"model\":\"gpt-4o-2024-08-06\",\"choices\":[{\"index\":0,\"message\":{\"role\":\"assistant\",\"content\":\"Rome\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":14,\"completion_tokens\":2,\"total_tokens\":16}}","stream_chunks":1,"ttfb_ms":1},"timing":{"latency_ms":1,"ttfb_ms":1},"error":null,"finalized":true}],"outcome":{"status":"success","finish_reason":"stop","error":null,"error_value":null,"usage":{"input_tokens":{"total":14,"no_cache":14,"cache_read":0},"output_tokens":{"total":2,"text":2,"reasoning":0},"raw":{"prompt_tokens":14,"completion_tokens":2,"total_tokens":16}}},"complete":true,"transport_closed":true} diff --git a/aimux-core/tests/master_recording_test.rs b/aimux-core/tests/master_recording_test.rs new file mode 100644 index 000000000..927f7fc62 --- /dev/null +++ b/aimux-core/tests/master_recording_test.rs @@ -0,0 +1,154 @@ +//! Recordings written by master (aimux 0.5, recording schema 2) load in this +//! version; what it cannot replay from them fails at replay time, saying why. +//! +//! Fixtures, both written by master's own `JsonlRecorder`: +//! - `fixtures/master_language_recordings.jsonl`: two OpenAI chat calls +//! recorded on 2026-10-13 by master a8e3a996 (`OpenAIProvider` chat model +//! `gpt-4o` against a local mock server, explicit key): "What is the capital +//! of France?" with temperature 0, then "Capital of Italy?" with the +//! instructions "Answer in one word.", which master records as a system +//! message whose content is a part array. Written by a scratch test that is +//! not part of the repository; the file is unchanged. +//! - `fixtures/master_decision_recordings.jsonl`: master's +//! `aimux-providers/tests/fixtures/jev_systemone_live.jsonl` (two live Jev +//! decisions, 2026-10-06), copied byte for byte. + +use aimux_core::error::AiMuxError; +use aimux_core::generate::{GenerateTextOptions, generate_text}; +use aimux_core::recording::{OutcomeStatus, RECORDING_SCHEMA, Recording, RecordingOperation}; +use aimux_core::replay::{MockDecisionReplayModel, MockReplayModel, replay_with_model}; + +const LANGUAGE: &str = include_str!("fixtures/master_language_recordings.jsonl"); +const DECISION: &str = include_str!("fixtures/master_decision_recordings.jsonl"); + +fn fixture(name: &str) -> String { + format!("{}/tests/fixtures/{name}", env!("CARGO_MANIFEST_DIR")) +} + +fn parse(lines: &str) -> Vec { + lines + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect() +} + +/// Every field master writes loads: its configuration snapshot is ignored, +/// the identity this version adds (`provider_id`) is empty, and the wire +/// exchange and outcome are kept. +#[test] +fn master_language_recordings_load() { + assert_eq!(RECORDING_SCHEMA, 2); + let recordings = parse(LANGUAGE); + assert_eq!(recordings.len(), 2); + for rec in &recordings { + assert_eq!(rec.schema, 2); + assert_eq!(rec.input.operation, RecordingOperation::LanguageModel); + assert_eq!(rec.provider.provider_id, ""); + assert_eq!(rec.provider.provider, "openai"); + assert_eq!(rec.provider.model_id, "gpt-4o"); + assert!(rec.complete && rec.transport_closed); + assert_eq!(rec.outcome.status, OutcomeStatus::Success); + assert_eq!(rec.exchanges.len(), 1); + let request = &rec.exchanges[0].request; + assert!(request.url.ends_with("/chat/completions")); + assert!( + request + .headers + .iter() + .any(|(name, value)| name == "Authorization" && value == "[REDACTED]") + ); + } + + // The first prompt (one user text message) reads; the second has a + // system message in master's part-array form, so it loads as empty. + assert_eq!(recordings[0].input.prompt.len(), 1); + assert!(recordings[0].input.unreadable_prompt().is_none()); + assert!(recordings[1].input.prompt.is_empty()); + assert!(recordings[1].input.unreadable_prompt().is_some()); + + // Round trip: written back as schema 2 with this version's fields. + let written = serde_json::to_value(&recordings[0]).unwrap(); + assert_eq!(written["schema"], 2); + assert!(written["provider"].get("base_url").is_none()); + serde_json::from_value::(written).unwrap(); + + let mock = MockReplayModel::from_jsonl(&fixture("master_language_recordings.jsonl")).unwrap(); + assert_eq!(mock.recordings().len(), 2); +} + +/// Offline replay of a master recording whose prompt this version reads +/// returns the recorded response. +#[tokio::test] +async fn a_master_recording_replays_offline() { + let mock = MockReplayModel::from_jsonl(&fixture("master_language_recordings.jsonl")).unwrap(); + let result = generate_text( + &mock, + "What is the capital of France?", + GenerateTextOptions { + temperature: Some(0.0), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(result.text(), "Paris."); +} + +/// What this version cannot replay from a master recording fails when it is +/// replayed, not when it is loaded, and the error says the recording is from +/// aimux 0.5 or earlier. +#[tokio::test] +async fn what_cannot_be_replayed_fails_at_replay_time() { + let recordings = parse(LANGUAGE); + let mock = MockReplayModel::from_jsonl(&fixture("master_language_recordings.jsonl")).unwrap(); + + // The prompt with a system message never matches offline. + let err = generate_text( + &mock, + "Capital of Italy?", + GenerateTextOptions { + instructions: Some("Answer in one word.".into()), + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("aimux 0.5 or earlier"), "{err}"); + + // Request replay refuses the unreadable prompt ... + let err = replay_with_model(&recordings[1], &mock, None) + .await + .unwrap_err(); + assert!(matches!(err, AiMuxError::InvalidArgument(_)), "{err}"); + assert!(err.to_string().contains("recorded prompt"), "{err}"); + + // ... and master's call options (`tool_choice: "auto"`), naming the writer. + let err = replay_with_model(&recordings[0], &mock, None) + .await + .unwrap_err(); + assert!(matches!(err, AiMuxError::JsonParse(_)), "{err}"); + assert!(err.to_string().contains("aimux 0.5 or earlier"), "{err}"); +} + +/// Master's decision recordings load; offline decision replay needs the +/// capabilities this version records under `input`, which master keeps in +/// the ignored `provider.profile`, and says so. +#[test] +fn master_decision_recordings_load_and_explain_what_replay_lacks() { + let recordings = parse(DECISION); + assert_eq!(recordings.len(), 2); + for rec in &recordings { + assert_eq!(rec.schema, 2); + assert_eq!(rec.input.operation, RecordingOperation::Decision); + assert_eq!(rec.provider.provider_id, ""); + assert_eq!(rec.provider.provider, "jev"); + assert!(rec.input.decision_capabilities.is_none()); + assert!(rec.outcome.decision_result.is_some()); + assert_eq!(rec.exchanges[0].response.as_ref().unwrap().status, 200); + } + + let err = MockDecisionReplayModel::new("jev", "jev-latest", recordings) + .err() + .unwrap(); + assert!(err.to_string().contains("provider.profile"), "{err}"); +} diff --git a/aimux-ffi/aimux-error.h b/aimux-ffi/aimux-error.h index bfe50be7c..20c14c130 100644 --- a/aimux-ffi/aimux-error.h +++ b/aimux-ffi/aimux-error.h @@ -6,7 +6,7 @@ * trailing out-parameter, which remains at its documented sentinel on failure. * * Every non-NULL error has one non-zero `aimux_error_code_t` and one message. - * Codes 1..18 come from `AiMuxError`, 100..105 from + * Codes 1..20 come from `AiMuxError`, 100..105 from * `RecordingError`, and 200..206 identify failures detected while crossing * the C ABI. Higher-level bindings reconstruct their native error types from * that code; they map all 200..206 codes to the language's existing @@ -40,7 +40,7 @@ typedef struct aimux_error aimux_error_t; typedef enum aimux_error_code { AIMUX_OK = 0, - /* AiMuxError: 1..18, except retired code 4. */ + /* AiMuxError: 1..20, except retired code 4. */ AIMUX_E_OTHER = 1, AIMUX_E_JSON_PARSE = 2, AIMUX_E_INVALID_RESPONSE_DATA = 3, @@ -61,6 +61,12 @@ typedef enum aimux_error_code { AIMUX_E_INVALID_TOOL_INPUT = 16, AIMUX_E_TOOL_CALL_REPAIR = 17, AIMUX_E_NO_OUTPUT_GENERATED = 18, + /* No API key was passed and the fallback environment variable is unset + * (the AI SDK's LoadAPIKeyError). Pre-HTTP: no request was made. */ + AIMUX_E_LOAD_API_KEY = 19, + /* A required provider setting was not passed and its fallback environment + * variable is unset (the AI SDK's LoadSettingError). */ + AIMUX_E_LOAD_SETTING = 20, /* RecordingError: 100..105. */ AIMUX_E_RECORDING_INIT = 100, @@ -108,9 +114,13 @@ int32_t aimux_error_status(const aimux_error_t *error); /** Retry hint in milliseconds (0 = retry now), or -1 when absent. */ int64_t aimux_error_retry_ms(const aimux_error_t *error); -/** Provider's own error code, e.g. "insufficient_quota". */ +/** Provider's own error code, e.g. "insufficient_quota". Under + * AIMUX_E_LOAD_API_KEY / AIMUX_E_LOAD_SETTING: the environment variable that + * was consulted, e.g. "OPENAI_API_KEY". */ char *aimux_error_provider_code(const aimux_error_t *error); -/** Failure text without Aimux's composed prefix. */ +/** Failure text without Aimux's composed prefix. Under AIMUX_E_LOAD_API_KEY: + * what the key is for ("OpenAI"); under AIMUX_E_LOAD_SETTING: the setting's + * parameter name ("region"). */ char *aimux_error_provider_message(const aimux_error_t *error); /** Raw provider response body. */ char *aimux_error_response_body(const aimux_error_t *error); diff --git a/aimux-ffi/aimux-ffi.h b/aimux-ffi/aimux-ffi.h index 90187b6b5..5d4790d14 100644 --- a/aimux-ffi/aimux-ffi.h +++ b/aimux-ffi/aimux-ffi.h @@ -15,7 +15,7 @@ * Getter strings are caller-owned (`aimux_free_string`). See aimux-error.h. * * Each prototype below identifies its expected high-level error: - * [AiMuxError] codes 1..17 + * [AiMuxError] codes 1..20 * [RecordingError] codes 100..105 * [C ABI] no expected high-level code * Every fallible call can additionally return a C ABI failure (200..206). @@ -177,9 +177,10 @@ aimux_error_t *aimux_azure_new_with_base(const char *api_key, const char *base_u * a C ABI failure. * @param model_id NUL-terminated model ID. * @param config_json Optional JSON object of ProviderOptions - * ({"base_url": "...", "headers": {...}, "max_retries": 0, - * "body_overrides": {...}}); - * NULL / empty / "null" for defaults. + * ({"base_url": "...", "headers": {...}, "organization": "...", + * "project": "..."}); "max_retries" (call-level) and + * "body_overrides" (removed) are rejected as invalid + * arguments. NULL / empty / "null" for defaults. * @param out_handle Receives the model handle. AiMuxError: unknown * provider, bad config, missing env key, invalid model id. */ @@ -735,10 +736,10 @@ aimux_error_t *aimux_recording_try_flush(void); rejects (bad schema, unknown protocol) is AiMuxError AIMUX_E_INVALID_ARGUMENT. */ aimux_error_t *aimux_register_providers(const char *config_json); -/* [AiMuxError] Set the global proxy configuration (M6, RFC-0016). Must be called - before the first generate_text / stream_text call; a no-op if the shared - HTTP client is already initialised. config_json is a serialized - ProxyConfig: { "http_url", "https_url", "all_url", "no_proxy" } (all +/* [AiMuxError] Replace the default transport with one routed by this proxy + configuration (M6, RFC-0016). Requests started afterwards (HTTP, built-in + WebSocket connects, download hops) use it; calling it again replaces the + previous configuration. config_json is a serialized ProxyConfig: { "http_url", "https_url", "all_url", "no_proxy" } (all optional). Malformed text is a C ABI failure; wrong shape is AiMuxError AIMUX_E_INVALID_ARGUMENT. */ aimux_error_t *aimux_init_proxy(const char *config_json); diff --git a/aimux-ffi/src/lib.rs b/aimux-ffi/src/lib.rs index da9bbac4d..7bc7a6a86 100644 --- a/aimux-ffi/src/lib.rs +++ b/aimux-ffi/src/lib.rs @@ -12,7 +12,7 @@ //! value on failure (the out-parameter is left at its sentinel: handle 0, //! pointer NULL). Every non-NULL error has one code from [`aimux_error_code`] //! and one message from [`aimux_error_message`], and is released exactly once -//! with [`aimux_error_free`]. Codes 1..17 come from `AiMuxError`, 100..105 +//! with [`aimux_error_free`]. Codes 1..20 come from `AiMuxError`, 100..105 //! from `RecordingError`, and 200..206 identify failures detected while //! crossing the C ABI. //! @@ -72,18 +72,20 @@ use aimux_core::provider::Provider; use aimux_core::recording::RecordingError; use aimux_core::tool::ToolCall; use aimux_core::trace::{RingTraceStore, TraceFilter, TraceLayer}; -use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider}; -use aimux_providers::anthropic_aws::{AnthropicAwsProvider, AnthropicAwsProviderConfig}; -use aimux_providers::azure::{AzureConfig, AzureProvider}; -use aimux_providers::bedrock::{BedrockProvider, BedrockProviderConfig}; -use aimux_providers::cohere::{CohereConfig, CohereProvider}; -use aimux_providers::google::{GoogleConfig, GoogleProvider}; -use aimux_providers::mistral::{MistralConfig, MistralProvider}; -use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; -use aimux_providers::tavily::{TavilyConfig, TavilyProvider}; -use aimux_providers::vertex::{VertexProvider, VertexProviderConfig}; -use aimux_providers::xai::{XAIConfig, XAIProvider}; -use aimux_providers::{ProviderOptions, provider, provider_handle}; +use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; +use aimux_providers::anthropic_aws::{ + AnthropicAwsAuth, AnthropicAwsProviderSettings, create_anthropic_aws, +}; +use aimux_providers::azure::{AzureOpenAIProviderSettings, create_azure}; +use aimux_providers::bedrock::{AmazonBedrockProviderSettings, create_amazon_bedrock}; +use aimux_providers::cohere::{CohereProviderSettings, create_cohere}; +use aimux_providers::google::{GoogleProviderSettings, create_google}; +use aimux_providers::mistral::{MistralProviderSettings, create_mistral}; +use aimux_providers::openai::{OpenAIProviderSettings, create_openai}; +use aimux_providers::provider::{ProviderOptions, provider, provider_handle}; +use aimux_providers::tavily::{TavilyProvider, TavilyProviderSettings, create_tavily}; +use aimux_providers::vertex::{VertexProviderSettings, create_google_vertex}; +use aimux_providers::xai::{XAIProviderSettings, create_xai}; use futures::StreamExt; use tokio::runtime::Runtime; @@ -97,7 +99,7 @@ use tokio::runtime::Runtime; #[derive(Clone)] enum HandleEntry { Language(Arc), - Provider(Arc), + Provider(ProviderEntry), Embedding(Arc), Speech(Arc), Image(Arc), @@ -112,6 +114,14 @@ enum HandleEntry { Abort(AbortSignal), } +/// A provider handle: the model factory plus its runtime-discovery side. +/// +/// Discovery is accessed through `Provider::discovery()`. +#[derive(Clone)] +struct ProviderEntry { + provider: Arc, +} + type Registry = HashMap; static REGISTRY: OnceLock> = OnceLock::new(); @@ -441,6 +451,12 @@ pub const AIMUX_E_NO_SUCH_TOOL: i32 = 15; pub const AIMUX_E_INVALID_TOOL_INPUT: i32 = 16; pub const AIMUX_E_TOOL_CALL_REPAIR: i32 = 17; pub const AIMUX_E_NO_OUTPUT_GENERATED: i32 = 18; +// Appended after NoOutputGenerated (4 stays retired): the AI SDK's +// `LoadAPIKeyError` and `LoadSettingError`. Their facts travel on the +// provider-code / provider-message getters: `env_var` and the description +// (key) or setting name. +pub const AIMUX_E_LOAD_API_KEY: i32 = 19; +pub const AIMUX_E_LOAD_SETTING: i32 = 20; // 100..105 preserve `RecordingError` as a separate high-level type while C // uses one code space for every returned error. @@ -485,6 +501,8 @@ fn aimux_error_code_of(err: &AiMuxError) -> i32 { AiMuxError::InvalidPrompt(_) | AiMuxError::MissingToolResults { .. } => { AIMUX_E_INVALID_PROMPT } + AiMuxError::LoadApiKey { .. } => AIMUX_E_LOAD_API_KEY, + AiMuxError::LoadSetting { .. } => AIMUX_E_LOAD_SETTING, AiMuxError::TokenExpired(_) => AIMUX_E_TOKEN_EXPIRED, AiMuxError::UnsupportedFunctionality(_) => AIMUX_E_UNSUPPORTED_FUNCTIONALITY, AiMuxError::NoSuchModel { .. } => AIMUX_E_NO_SUCH_MODEL, @@ -607,19 +625,35 @@ pub extern "C" fn aimux_error_retryable(err: *const aimux_error_t) -> i32 { } /// `AIMUX_E_API_CALL`: the provider's own error code, e.g. "insufficient_quota". +/// `AIMUX_E_LOAD_API_KEY` / `AIMUX_E_LOAD_SETTING`: the environment variable +/// that was consulted (`env_var`), e.g. "OPENAI_API_KEY". #[unsafe(no_mangle)] pub extern "C" fn aimux_error_provider_code(err: *const aimux_error_t) -> *mut c_char { - opt_cstring(map_aimux_error(err, |e| api_call(e)?.provider_code.clone()).flatten()) + opt_cstring( + map_aimux_error(err, |e| match e { + AiMuxError::LoadApiKey { env_var, .. } | AiMuxError::LoadSetting { env_var, .. } => { + Some(env_var.clone()) + } + _ => api_call(e)?.provider_code.clone(), + }) + .flatten(), + ) } /// `AIMUX_E_API_CALL`: the failure's own text ("slow down"), without the /// composed prefix `message` carries ("API call error: HTTP 429: slow down"). +/// `AIMUX_E_LOAD_API_KEY`: what the key is for ("OpenAI"); +/// `AIMUX_E_LOAD_SETTING`: the setting's parameter name ("region"). #[unsafe(no_mangle)] pub extern "C" fn aimux_error_provider_message(err: *const aimux_error_t) -> *mut c_char { opt_cstring( - map_aimux_error(err, |e| Some(api_call(e)?.message.clone())) - .flatten() - .filter(|m| !m.is_empty()), + map_aimux_error(err, |e| match e { + AiMuxError::LoadApiKey { description, .. } => Some(description.clone()), + AiMuxError::LoadSetting { name, .. } => Some(name.clone()), + _ => Some(api_call(e)?.message.clone()), + }) + .flatten() + .filter(|m| !m.is_empty()), ) } @@ -754,7 +788,7 @@ pub extern "C" fn aimux_error_model_type(err: *const aimux_error_t) -> *mut c_ch pub extern "C" fn aimux_error_provider_id(err: *const aimux_error_t) -> *mut c_char { opt_cstring( map_aimux_error(err, |e| match e { - AiMuxError::NoSuchProvider { provider_id } => Some(provider_id.clone()), + AiMuxError::NoSuchProvider { provider_id, .. } => Some(provider_id.clone()), _ => None, }) .flatten(), @@ -1050,6 +1084,20 @@ fn invoke_stream_callback(callback_name: &str, f: impl FnOnce()) -> Result<(), F // C ABI: provider constructors // ───────────────────────────────────────────────────────────────────────────── +/// An OpenAI provider for a key handed over the C ABI. The key is an explicit +/// value, so an empty string is sent as given and never falls back to +/// `OPENAI_API_KEY`. +fn openai_provider( + api_key: String, + base_url: Option, +) -> Result { + create_openai(OpenAIProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) +} + /// Create an OpenAI model instance. AiMuxError: invalid model id. #[unsafe(no_mangle)] pub extern "C" fn aimux_openai_new( @@ -1059,8 +1107,9 @@ pub extern "C" fn aimux_openai_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let m = OpenAIProvider::new(OpenAIConfig::new(api_key)).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + Ok(intern_model( + openai_provider(api_key, None)?.language_model(&model_id)?, + )) }) } @@ -1076,12 +1125,38 @@ pub extern "C" fn aimux_openai_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let m = OpenAIProvider::new(config).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let provider = openai_provider(api_key, parse_base_url(base_url)?)?; + Ok(intern_model(provider.language_model(&model_id)?)) + }) +} + +/// An Anthropic provider for an explicit key handed over by the host, with an +/// optional base URL. The key is an explicit value (`""` included), so it +/// never falls back to `ANTHROPIC_API_KEY`. +fn anthropic_provider( + api_key: String, + base_url: Option, +) -> Result { + create_anthropic(AnthropicProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) +} + +/// An Anthropic-on-AWS provider for an explicit API key and region. +fn anthropic_aws_provider( + api_key: String, + region: String, + base_url: Option, +) -> Result { + create_anthropic_aws(AnthropicAwsProviderSettings { + region: Some(region), + auth: Some(AnthropicAwsAuth::ApiKey( + aimux_provider_utils::Resolvable::Value(api_key), + )), + base_url, + ..Default::default() }) } @@ -1094,8 +1169,8 @@ pub extern "C" fn aimux_anthropic_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let m = AnthropicProvider::new(AnthropicConfig::new(api_key)).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = anthropic_provider(api_key, None)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1111,12 +1186,9 @@ pub extern "C" fn aimux_anthropic_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = AnthropicConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let m = AnthropicProvider::new(config).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = + anthropic_provider(api_key, parse_base_url(base_url)?)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1132,10 +1204,8 @@ pub extern "C" fn aimux_anthropic_aws_new( let api_key = str_arg(api_key, "api_key")?; let region = str_arg(region, "region")?; let model_id = str_arg(model_id, "model_id")?; - let m = - AnthropicAwsProvider::new(AnthropicAwsProviderConfig::with_api_key(api_key, region)) - .language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = anthropic_aws_provider(api_key, region, None)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1152,12 +1222,9 @@ pub extern "C" fn aimux_anthropic_aws_new_with_base( let api_key = str_arg(api_key, "api_key")?; let region = str_arg(region, "region")?; let model_id = str_arg(model_id, "model_id")?; - let mut config = AnthropicAwsProviderConfig::with_api_key(api_key, region); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let m = AnthropicAwsProvider::new(config).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = anthropic_aws_provider(api_key, region, parse_base_url(base_url)?)? + .language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1177,14 +1244,10 @@ pub extern "C" fn aimux_azure_new( let api_key = str_arg(api_key, "api_key")?; let resource_name = str_arg(resource_name, "resource_name")?; let deployment = str_arg(deployment, "deployment")?; - let mut config = AzureConfig::new() - .with_api_key(api_key) - .with_resource_name(resource_name); - if let Some(v) = opt_str_arg(api_version, "api_version")?.filter(|v| !v.is_empty()) { - config = config.with_api_version(v); - } - let m = AzureProvider::new(config)?.language_model(&deployment)?; - Ok(intern_model(Arc::from(m))) + let api_version = opt_str_arg(api_version, "api_version")?.filter(|v| !v.is_empty()); + let m = azure_provider(api_key, Some(resource_name), None, api_version)? + .language_model(&deployment)?; + Ok(intern_model(m)) }) } @@ -1202,17 +1265,129 @@ pub extern "C" fn aimux_azure_new_with_base( let api_key = str_arg(api_key, "api_key")?; let base_url = str_arg(base_url, "base_url")?; let deployment = str_arg(deployment, "deployment")?; - let mut config = AzureConfig::new() - .with_api_key(api_key) - .with_base_url(base_url); - if let Some(v) = opt_str_arg(api_version, "api_version")?.filter(|v| !v.is_empty()) { - config = config.with_api_version(v); - } - let m = AzureProvider::new(config)?.language_model(&deployment)?; - Ok(intern_model(Arc::from(m))) + let api_version = opt_str_arg(api_version, "api_version")?.filter(|v| !v.is_empty()); + let m = azure_provider(api_key, None, Some(base_url), api_version)? + .language_model(&deployment)?; + Ok(intern_model(m)) + }) +} + +/// A Bedrock provider for explicit SigV4 credentials handed over by the host, +/// with an optional base URL. The credentials are explicit values, so they +/// never fall back to the environment. +fn bedrock_provider( + access_key_id: String, + secret_access_key: String, + region: String, + base_url: Option, +) -> Result { + create_amazon_bedrock(AmazonBedrockProviderSettings { + credential_provider: Some(aimux_provider_utils::Resolvable::Value( + aimux_provider_utils::AwsCredentials { + access_key_id, + secret_access_key, + session_token: None, + region: region.clone(), + }, + )), + region: Some(region), + base_url, + ..Default::default() + }) +} + +/// A Vertex AI provider for an explicit access token, project and location, +/// with an optional base URL. +fn vertex_provider( + access_token: String, + project: String, + location: String, + base_url: Option, +) -> Result { + create_google_vertex(VertexProviderSettings { + access_token: Some(aimux_provider_utils::Resolvable::Value(access_token)), + project: Some(project), + location: Some(location), + base_url, + ..Default::default() }) } +/// A Google provider for an explicit key handed over by the host, with an +/// optional base URL. +fn google_provider( + api_key: String, + base_url: Option, +) -> Result { + create_google(GoogleProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) +} + +/// A Cohere provider for an explicit key handed over by the host, with an +/// optional base URL. +fn cohere_provider( + api_key: String, + base_url: Option, +) -> Result { + create_cohere(CohereProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) +} + +/// A Mistral provider for an explicit key handed over by the host, with an +/// optional base URL. +fn mistral_provider( + api_key: String, + base_url: Option, +) -> Result { + create_mistral(MistralProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) +} + +/// An xAI provider for an explicit key handed over by the host, with an +/// optional base URL. +fn xai_provider( + api_key: String, + base_url: Option, +) -> Result { + create_xai(XAIProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) +} + +/// An Azure OpenAI provider for an explicit key, with either a resource name +/// or a base URL. A dated `api_version` belongs to the deployment URL form +/// (`{base}/deployments/{deployment}{path}?api-version=`), so passing one +/// selects it; without one the AI SDK's v1 form is used. +fn azure_provider( + api_key: String, + resource_name: Option, + base_url: Option, + api_version: Option, +) -> Result { + let mut settings = AzureOpenAIProviderSettings { + api_key: Some(aimux_provider_utils::Resolvable::Value(api_key)), + resource_name, + base_url, + ..Default::default() + }; + if let Some(version) = api_version { + settings.api_version = Some(version); + settings.use_deployment_based_urls = true; + } + create_azure(settings) +} + /// Create a Bedrock model instance (AWS SigV4 credentials). #[unsafe(no_mangle)] pub extern "C" fn aimux_bedrock_new( @@ -1233,13 +1408,9 @@ pub extern "C" fn aimux_bedrock_new( model_id, "model_id", )?; - let m = BedrockProvider::new(BedrockProviderConfig::new( - access_key_id, - secret_access_key, - region, - )) - .language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = bedrock_provider(access_key_id, secret_access_key, region, None)? + .language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1264,12 +1435,14 @@ pub extern "C" fn aimux_bedrock_new_with_base( model_id, "model_id", )?; - let mut config = BedrockProviderConfig::new(access_key_id, secret_access_key, region); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let m = BedrockProvider::new(config).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = bedrock_provider( + access_key_id, + secret_access_key, + region, + parse_base_url(base_url)?, + )? + .language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1293,9 +1466,9 @@ pub extern "C" fn aimux_vertex_new( model_id, "model_id", )?; - let m = VertexProvider::new(VertexProviderConfig::new(access_token, project, location)) - .language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = + vertex_provider(access_token, project, location, None)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1320,12 +1493,9 @@ pub extern "C" fn aimux_vertex_new_with_base( model_id, "model_id", )?; - let mut config = VertexProviderConfig::new(access_token, project, location); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let m = VertexProvider::new(config).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = vertex_provider(access_token, project, location, parse_base_url(base_url)?)? + .language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1338,8 +1508,8 @@ pub extern "C" fn aimux_cohere_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let m = CohereProvider::new(CohereConfig::new(api_key)).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = cohere_provider(api_key, None)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1353,12 +1523,8 @@ pub extern "C" fn aimux_cohere_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = CohereConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let m = CohereProvider::new(config).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = cohere_provider(api_key, parse_base_url(base_url)?)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1371,8 +1537,8 @@ pub extern "C" fn aimux_mistral_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let m = MistralProvider::new(MistralConfig::new(api_key)).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = mistral_provider(api_key, None)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1386,12 +1552,8 @@ pub extern "C" fn aimux_mistral_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = MistralConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let m = MistralProvider::new(config).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = mistral_provider(api_key, parse_base_url(base_url)?)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1404,8 +1566,8 @@ pub extern "C" fn aimux_xai_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let m = XAIProvider::new(XAIConfig::new(api_key)).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = xai_provider(api_key, None)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1419,12 +1581,8 @@ pub extern "C" fn aimux_xai_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = XAIConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let m = XAIProvider::new(config).language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = xai_provider(api_key, parse_base_url(base_url)?)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1435,11 +1593,11 @@ pub extern "C" fn aimux_xai_new_with_base( /// entry (replaces the retired `aimux_deepseek_new` etc.). /// - `model_id` — model id string. /// - `config_json` — optional JSON object of `ProviderOptions` -/// (`{"base_url": "...", "headers": {...}, "max_retries": 0, "body_overrides": {...}}`); -/// NULL / empty / "null" for defaults. +/// (`{"base_url": "...", "headers": {...}, "organization": "...", "project": "..."}`); +/// NULL / empty / "null" for defaults. `max_retries` (call-level) and +/// `body_overrides` (removed) are rejected as invalid arguments. /// -/// AiMuxError: unknown provider, bad config shape, missing env key, or -/// invalid model id. +/// AiMuxError: unknown provider, bad config shape, or invalid model id. #[unsafe(no_mangle)] pub extern "C" fn aimux_provider_new( name: *const c_char, @@ -1455,12 +1613,12 @@ pub extern "C" fn aimux_provider_new( let key = opt_str_arg(api_key, "api_key")?; let opts = parse_provider_options(config_json)?; let m = provider(&name, key, &model_id, opts)?; - Ok(intern_model(Arc::from(m))) + Ok(intern_model(m)) }) } /// Convenience: create a language model by provider name, reading the API key -/// from the provider's env var. +/// from the provider's env var at request time. #[unsafe(no_mangle)] pub extern "C" fn aimux_provider_from_env( name: *const c_char, @@ -1470,8 +1628,8 @@ pub extern "C" fn aimux_provider_from_env( with_out_handle(out_handle, || { let name = str_arg(name, "name")?; let model_id = str_arg(model_id, "model_id")?; - let m = provider(&name, None, &model_id, None)?; - Ok(intern_model(Arc::from(m))) + let m = provider_handle(&name, None, None)?.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -1499,8 +1657,10 @@ pub extern "C" fn aimux_provider_handle_new( let name = str_arg(name, "name")?; let key = opt_str_arg(api_key, "api_key")?; let opts = parse_provider_options(config_json)?; - let p = provider_handle(&name, key, opts)?; - Ok(intern_handle(HandleEntry::Provider(Arc::from(p)))) + let provider = provider_handle(&name, key, opts)?; + Ok(intern_handle(HandleEntry::Provider(ProviderEntry { + provider, + }))) }) } @@ -1522,7 +1682,10 @@ pub extern "C" fn aimux_provider_list_models( } .into()); }; - run_json(p.list_models()) + let discovery = p.provider.discovery().ok_or_else(|| { + AiMuxError::InvalidArgument("provider does not support model discovery".into()) + })?; + run_json(discovery.list_models()) }) } @@ -1545,8 +1708,8 @@ pub extern "C" fn aimux_provider_model( .into()); }; let model_id = str_arg(model_id, "model_id")?; - let m = p.language_model(&model_id)?; - Ok(intern_model(Arc::from(m))) + let m = p.provider.language_model(&model_id)?; + Ok(intern_model(m)) }) } @@ -2196,7 +2359,7 @@ pub extern "C" fn aimux_openai_embedding_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let model = OpenAIProvider::new(OpenAIConfig::new(api_key)).embedding_model(&model_id); + let model = openai_provider(api_key, None)?.embedding(&model_id); Ok(intern_handle(HandleEntry::Embedding(Arc::new(model)))) }) } @@ -2210,11 +2373,8 @@ pub extern "C" fn aimux_openai_embedding_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = OpenAIProvider::new(config).embedding_model(&model_id); + let provider = openai_provider(api_key, parse_base_url(base_url)?)?; + let model = provider.embedding(&model_id); Ok(intern_handle(HandleEntry::Embedding(Arc::new(model)))) }) } @@ -2227,7 +2387,7 @@ pub extern "C" fn aimux_cohere_embedding_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let model = CohereProvider::new(CohereConfig::new(api_key)).embedding_model(&model_id); + let model = cohere_provider(api_key, None)?.embedding(&model_id); Ok(intern_handle(HandleEntry::Embedding(Arc::new(model)))) }) } @@ -2241,11 +2401,7 @@ pub extern "C" fn aimux_cohere_embedding_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = CohereConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = CohereProvider::new(config).embedding_model(&model_id); + let model = cohere_provider(api_key, parse_base_url(base_url)?)?.embedding(&model_id); Ok(intern_handle(HandleEntry::Embedding(Arc::new(model)))) }) } @@ -2258,7 +2414,7 @@ pub extern "C" fn aimux_google_embedding_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let model = GoogleProvider::new(GoogleConfig::new(api_key)).embedding_model(&model_id); + let model = google_provider(api_key, None)?.embedding(&model_id); Ok(intern_handle(HandleEntry::Embedding(Arc::new(model)))) }) } @@ -2272,11 +2428,7 @@ pub extern "C" fn aimux_google_embedding_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = GoogleProvider::new(config).embedding_model(&model_id); + let model = google_provider(api_key, parse_base_url(base_url)?)?.embedding(&model_id); Ok(intern_handle(HandleEntry::Embedding(Arc::new(model)))) }) } @@ -2322,7 +2474,7 @@ pub extern "C" fn aimux_openai_speech_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let model = OpenAIProvider::new(OpenAIConfig::new(api_key)).speech(&model_id); + let model = openai_provider(api_key, None)?.speech(&model_id); Ok(intern_handle(HandleEntry::Speech(Arc::new(model)))) }) } @@ -2336,11 +2488,8 @@ pub extern "C" fn aimux_openai_speech_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = OpenAIProvider::new(config).speech(&model_id); + let provider = openai_provider(api_key, parse_base_url(base_url)?)?; + let model = provider.speech(&model_id); Ok(intern_handle(HandleEntry::Speech(Arc::new(model)))) }) } @@ -2375,7 +2524,7 @@ pub extern "C" fn aimux_openai_image_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let model = OpenAIProvider::new(OpenAIConfig::new(api_key)).image(&model_id); + let model = openai_provider(api_key, None)?.image(&model_id); Ok(intern_handle(HandleEntry::Image(Arc::new(model)))) }) } @@ -2389,11 +2538,8 @@ pub extern "C" fn aimux_openai_image_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = OpenAIProvider::new(config).image(&model_id); + let provider = openai_provider(api_key, parse_base_url(base_url)?)?; + let model = provider.image(&model_id); Ok(intern_handle(HandleEntry::Image(Arc::new(model)))) }) } @@ -2406,7 +2552,7 @@ pub extern "C" fn aimux_google_image_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let model = GoogleProvider::new(GoogleConfig::new(api_key)).image(&model_id); + let model = google_provider(api_key, None)?.image(&model_id); Ok(intern_handle(HandleEntry::Image(Arc::new(model)))) }) } @@ -2420,11 +2566,7 @@ pub extern "C" fn aimux_google_image_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = GoogleProvider::new(config).image(&model_id); + let model = google_provider(api_key, parse_base_url(base_url)?)?.image(&model_id); Ok(intern_handle(HandleEntry::Image(Arc::new(model)))) }) } @@ -2457,7 +2599,7 @@ pub extern "C" fn aimux_openai_transcription_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let model = OpenAIProvider::new(OpenAIConfig::new(api_key)).transcription(&model_id); + let model = openai_provider(api_key, None)?.transcription(&model_id); Ok(intern_handle(HandleEntry::Transcription(Arc::new(model)))) }) } @@ -2471,11 +2613,8 @@ pub extern "C" fn aimux_openai_transcription_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = OpenAIProvider::new(config).transcription(&model_id); + let provider = openai_provider(api_key, parse_base_url(base_url)?)?; + let model = provider.transcription(&model_id); Ok(intern_handle(HandleEntry::Transcription(Arc::new(model)))) }) } @@ -2715,7 +2854,7 @@ pub extern "C" fn aimux_openai_files_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let api_key = str_arg(api_key, "api_key")?; - let files = OpenAIProvider::new(OpenAIConfig::new(api_key)).files(); + let files = openai_provider(api_key, None)?.files(); Ok(intern_handle(HandleEntry::Files(Arc::new(files)))) }) } @@ -2728,11 +2867,8 @@ pub extern "C" fn aimux_openai_files_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let api_key = str_arg(api_key, "api_key")?; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let files = OpenAIProvider::new(config).files(); + let provider = openai_provider(api_key, parse_base_url(base_url)?)?; + let files = provider.files(); Ok(intern_handle(HandleEntry::Files(Arc::new(files)))) }) } @@ -2774,7 +2910,7 @@ pub extern "C" fn aimux_cohere_reranking_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let model = CohereProvider::new(CohereConfig::new(api_key)).reranking_model(&model_id); + let model = cohere_provider(api_key, None)?.reranking(&model_id); Ok(intern_handle(HandleEntry::Reranking(Arc::new(model)))) }) } @@ -2788,11 +2924,7 @@ pub extern "C" fn aimux_cohere_reranking_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = CohereConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = CohereProvider::new(config).reranking_model(&model_id); + let model = cohere_provider(api_key, parse_base_url(base_url)?)?.reranking(&model_id); Ok(intern_handle(HandleEntry::Reranking(Arc::new(model)))) }) } @@ -2831,7 +2963,7 @@ pub extern "C" fn aimux_google_video_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let model = GoogleProvider::new(GoogleConfig::new(api_key)).video(&model_id); + let model = google_provider(api_key, None)?.video(&model_id); Ok(intern_handle(HandleEntry::Video(Arc::new(model)))) }) } @@ -2845,11 +2977,7 @@ pub extern "C" fn aimux_google_video_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let (api_key, model_id) = parse_two_args(api_key, "api_key", model_id, "model_id")?; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = GoogleProvider::new(config).video(&model_id); + let model = google_provider(api_key, parse_base_url(base_url)?)?.video(&model_id); Ok(intern_handle(HandleEntry::Video(Arc::new(model)))) }) } @@ -2891,7 +3019,7 @@ pub extern "C" fn aimux_provider_decision_model( .into()); }; let model_id = str_arg(model_id, "model_id")?; - let model = provider.decision_model(&model_id)?; + let model = provider.provider.decision_model(&model_id)?; Ok(intern_handle(HandleEntry::Decision(Arc::from(model)))) }) } @@ -2993,6 +3121,19 @@ pub extern "C" fn aimux_decide_with_abort( // C ABI: Search // ───────────────────────────────────────────────────────────────────────────── +/// A Tavily provider for an explicit key handed over by the host, with an +/// optional base URL. +fn tavily_provider( + api_key: String, + base_url: Option, +) -> Result { + create_tavily(TavilyProviderSettings { + api_key: Some(aimux_provider_utils::Resolvable::Value(api_key)), + base_url, + ..Default::default() + }) +} + /// Create a Tavily search model instance. `model_id` is accepted for API /// symmetry but ignored (Tavily uses a fixed endpoint). #[unsafe(no_mangle)] @@ -3003,7 +3144,7 @@ pub extern "C" fn aimux_tavily_search_new( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let api_key = str_arg(api_key, "api_key")?; - let model = TavilyProvider::new(TavilyConfig::new(api_key)).search_model(); + let model = tavily_provider(api_key, None)?.search_model(); Ok(intern_handle(HandleEntry::Search(Arc::new(model)))) }) } @@ -3017,11 +3158,7 @@ pub extern "C" fn aimux_tavily_search_new_with_base( ) -> *mut aimux_error_t { with_out_handle(out_handle, || { let api_key = str_arg(api_key, "api_key")?; - let mut config = TavilyConfig::new(api_key); - if let Some(url) = parse_base_url(base_url)? { - config = config.with_base_url(url); - } - let model = TavilyProvider::new(config).search_model(); + let model = tavily_provider(api_key, parse_base_url(base_url)?)?.search_model(); Ok(intern_handle(HandleEntry::Search(Arc::new(model)))) }) } @@ -3379,7 +3516,7 @@ pub extern "C" fn aimux_register_providers(config_json: *const c_char) -> *mut a // Malformed JSON text is this layer's finding; a well-formed document // that the registry rejects (bad schema, unknown protocol) is an AiMuxError. serde_json::from_str::(&json).map_err(|e| wire_err("config_json", e))?; - aimux_providers::load_providers_from_json(&json).map_err(|e| match e { + aimux_providers::provider::load_providers_from_json(&json).map_err(|e| match e { // The registry reports a schema mismatch as `JsonParse`; the text // already parsed above, so what it rejected is the shape — // `AiMuxError::InvalidArgument`, not a provider-response parse failure. @@ -3390,9 +3527,11 @@ pub extern "C" fn aimux_register_providers(config_json: *const c_char) -> *mut a }) } -/// Set the global proxy configuration (M6, RFC-0016). Must be called before -/// the first `aimux_generate_text` / `aimux_stream_text` call; a no-op if the -/// shared HTTP client is already initialised. +/// Replace the default Rust transport with one routed by this proxy +/// configuration (M6, RFC-0016). Every later request of a provider without +/// its own transport, every built-in WebSocket connect and every validated +/// download hop uses it; requests already in flight keep theirs. Calling it +/// again replaces the previous configuration. /// /// `config_json` is a serialized `ProxyConfig`: /// `{ "http_url": "...", "https_url": "...", "all_url": "...", "no_proxy": @@ -3402,10 +3541,9 @@ pub extern "C" fn aimux_register_providers(config_json: *const c_char) -> *mut a pub extern "C" fn aimux_init_proxy(config_json: *const c_char) -> *mut aimux_error_t { no_result(|| { let config: aimux_provider_utils::ProxyConfig = parse_json_arg(config_json, "config_json")?; - // `init_proxy` returns false when the shared client is already up; treat - // that as success (idempotent) so callers don't need to reason about - // ordering races. - let _ = aimux_provider_utils::init_proxy(config); + aimux_provider_utils::set_default_transport(aimux_provider_utils::ReqwestFetch::new( + config, + )); Ok(()) }) } @@ -3614,7 +3752,7 @@ mod tests { fn expect_aimux_error(e: *mut aimux_error_t) -> (i32, String) { assert!(!e.is_null(), "expected a returned error"); let code = aimux_error_code(e); - if !(AIMUX_E_OTHER..=AIMUX_E_NO_OUTPUT_GENERATED).contains(&code) { + if !(AIMUX_E_OTHER..=AIMUX_E_LOAD_SETTING).contains(&code) { panic!("expected an AiMuxError code, got {code}: {}", msg(e)); } let out = (code, take(aimux_error_message(e)).unwrap()); @@ -3792,8 +3930,57 @@ mod tests { assert_eq!(aimux_error_retryable(h), 0); aimux_error_free(owner); + // LoadApiKey / LoadSetting carry the consulted environment variable on + // the provider-code channel and the description / setting name on the + // provider-message channel; every other getter answers its sentinel. + let owner = boxed(AiMuxError::LoadApiKey { + env_var: "OPENAI_API_KEY".into(), + description: "OpenAI".into(), + }); + let h = owner; + assert_eq!(aimux_error_code(h), AIMUX_E_LOAD_API_KEY); + assert_eq!( + take(aimux_error_provider_code(h)).as_deref(), + Some("OPENAI_API_KEY") + ); + assert_eq!( + take(aimux_error_provider_message(h)).as_deref(), + Some("OpenAI") + ); + assert!( + take(aimux_error_message(h)) + .unwrap() + .contains("OPENAI_API_KEY") + ); + assert!(aimux_error_response_body(h).is_null()); + assert!(aimux_error_model_id(h).is_null()); + assert_eq!(aimux_error_status(h), -1); + assert_eq!(aimux_error_retryable(h), 0); + aimux_error_free(owner); + + let owner = boxed(AiMuxError::LoadSetting { + env_var: "AWS_REGION".into(), + name: "region".into(), + }); + let h = owner; + assert_eq!(aimux_error_code(h), AIMUX_E_LOAD_SETTING); + assert_eq!( + take(aimux_error_provider_code(h)).as_deref(), + Some("AWS_REGION") + ); + assert_eq!( + take(aimux_error_provider_message(h)).as_deref(), + Some("region") + ); + assert!(aimux_error_url(h).is_null()); + assert_eq!(aimux_error_status(h), -1); + aimux_error_free(owner); + let owner = boxed(AiMuxError::NoSuchProvider { provider_id: "p".into(), + model_id: String::new(), + model_type: String::new(), + available_providers: Vec::new(), }); let h = owner; assert_eq!(aimux_error_code(h), AIMUX_E_NO_SUCH_PROVIDER); @@ -3972,9 +4159,12 @@ mod tests { ); } - /// Pin the full 16-variant → code mapping. + /// Pin the full 19-variant → code mapping. #[test] fn error_code_mapping_covers_all_variants() { + assert_eq!(AIMUX_E_NO_OUTPUT_GENERATED, 18); + assert_eq!(AIMUX_E_LOAD_API_KEY, 19); + assert_eq!(AIMUX_E_LOAD_SETTING, 20); let s = |t: &str| t.to_string(); let cases: Vec<(AiMuxError, i32)> = vec![ ( @@ -3993,6 +4183,10 @@ mod tests { AIMUX_E_RETRY, ), (AiMuxError::JsonParse(s("x")), AIMUX_E_JSON_PARSE), + ( + AiMuxError::NoOutputGenerated(s("x")), + AIMUX_E_NO_OUTPUT_GENERATED, + ), ( AiMuxError::InvalidResponseData(s("x")), AIMUX_E_INVALID_RESPONSE_DATA, @@ -4029,6 +4223,20 @@ mod tests { AIMUX_E_INVALID_ARGUMENT, ), (AiMuxError::InvalidPrompt(s("x")), AIMUX_E_INVALID_PROMPT), + ( + AiMuxError::LoadApiKey { + env_var: s("OPENAI_API_KEY"), + description: s("OpenAI"), + }, + AIMUX_E_LOAD_API_KEY, + ), + ( + AiMuxError::LoadSetting { + env_var: s("AWS_REGION"), + name: s("region"), + }, + AIMUX_E_LOAD_SETTING, + ), (AiMuxError::TokenExpired(s("x")), AIMUX_E_TOKEN_EXPIRED), ( AiMuxError::UnsupportedFunctionality(s("x")), @@ -4044,6 +4252,9 @@ mod tests { ( AiMuxError::NoSuchProvider { provider_id: s("x"), + model_id: String::new(), + model_type: String::new(), + available_providers: Vec::new(), }, AIMUX_E_NO_SUCH_PROVIDER, ), diff --git a/aimux-ffi/tests/common/mod.rs b/aimux-ffi/tests/common/mod.rs index e0f76a45e..81777a24d 100644 --- a/aimux-ffi/tests/common/mod.rs +++ b/aimux-ffi/tests/common/mod.rs @@ -6,8 +6,8 @@ use std::ffi::{CStr, CString}; use std::os::raw::c_char; use aimux_ffi::{ - AIMUX_E_FFI_CALLBACK_FAILURE, AIMUX_E_FFI_NULL_POINTER, AIMUX_E_OTHER, AIMUX_E_RECORDING_INIT, - AIMUX_E_RECORDING_WRITE, AIMUX_E_RETRY, aimux_error_code, aimux_error_free, + AIMUX_E_FFI_CALLBACK_FAILURE, AIMUX_E_FFI_NULL_POINTER, AIMUX_E_LOAD_SETTING, AIMUX_E_OTHER, + AIMUX_E_RECORDING_INIT, AIMUX_E_RECORDING_WRITE, aimux_error_code, aimux_error_free, aimux_error_message, aimux_error_t, aimux_free_string, }; @@ -48,7 +48,7 @@ pub fn ok(e: *mut aimux_error_t, name: &str) { pub fn expect_aimux_error(e: *mut aimux_error_t, name: &str) -> (i32, String) { assert!(!e.is_null(), "{name}: expected a returned error"); let code = aimux_error_code(e); - if !(AIMUX_E_OTHER..=AIMUX_E_RETRY).contains(&code) { + if !(AIMUX_E_OTHER..=AIMUX_E_LOAD_SETTING).contains(&code) { panic!( "{name}: expected an AiMuxError code, got {code}: {}", msg(e) diff --git a/aimux-ffi/tests/exports_smoke_test.rs b/aimux-ffi/tests/exports_smoke_test.rs index 33d630eeb..d0cd7d41d 100644 --- a/aimux-ffi/tests/exports_smoke_test.rs +++ b/aimux-ffi/tests/exports_smoke_test.rs @@ -24,7 +24,7 @@ //! - Every returned error and every returned JSON string is released //! (`aimux_error_free` / `aimux_free_string`). //! -//! Note: multimodal session smoke calls use the default RetryConfig +//! Note: multimodal session smoke calls use the default retry settings //! (2 retries) against a refused port; under `--test-threads=1` this adds //! roughly a minute of backoff sleep — run the default parallel profile. //! @@ -495,7 +495,7 @@ fn constructor_exports_build_and_release_handles() { ); // Registry-backed constructors (RFC-0017 / RFC-0027). - let provider_opts = c(r#"{"base_url":"http://127.0.0.1:1","max_retries":0}"#); + let provider_opts = c(r#"{"base_url":"http://127.0.0.1:1"}"#); ctor!( "provider_new", aimux_provider_new( @@ -571,23 +571,17 @@ fn constructor_exports_build_and_release_handles() { } } -/// `provider_from_env` reads the provider's env var, which the harness cannot -/// control: a live handle or a clean missing-key error are both correct. +/// `provider_from_env` defers reading the key until request time. #[test] -fn provider_from_env_handle_or_clean_failure() { +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, ); - if e.is_null() { - assert_ne!(h, 0); - aimux_drop_handle(h); - } else { - assert_eq!(h, 0); - expect_aimux_error(e, "provider_from_env"); - } + expect_handle(e, h, "provider_from_env"); + aimux_drop_handle(h); } // ── session class: text generation (4 exports) ────────────────────────────── @@ -902,12 +896,12 @@ fn transcription_session_reaches_clean_terminal_error() { #[test] fn provider_discovery_and_catalogue_exports_fail_cleanly() { - // Provider handle bound to the unreachable host with retries disabled. + // Provider handle bound to the unreachable host. let mut provider = 0; let e = aimux_provider_handle_new( c("groq").as_ptr(), c(FAKE_KEY).as_ptr(), - c(r#"{"base_url":"http://127.0.0.1:1","max_retries":0}"#).as_ptr(), + c(r#"{"base_url":"http://127.0.0.1:1"}"#).as_ptr(), &mut provider, ); expect_handle(e, provider, "provider_handle_new"); diff --git a/aimux-provider-utils/Cargo.toml b/aimux-provider-utils/Cargo.toml index bf427a093..fb2f67265 100644 --- a/aimux-provider-utils/Cargo.toml +++ b/aimux-provider-utils/Cargo.toml @@ -12,6 +12,7 @@ documentation = "https://docs.rs/aimux-provider-utils" aimux-core = { workspace = true } aimux-stream = { workspace = true } +async-trait = { workspace = true } reqwest = { workspace = true } http = "1" serde = { workspace = true } @@ -25,6 +26,9 @@ futures = { workspace = true } async-stream = { workspace = true } url = { workspace = true } chrono = "0.4" +sha2 = "0.10" +hmac = "0.12" +hex = "0.4" # WebSocket client (RFC-0028 realtime transcription). Optional: enabled via # the `ws` feature so consumers without realtime needs compile it out. diff --git a/aimux-provider-utils/src/api_key.rs b/aimux-provider-utils/src/api_key.rs index 6257efa36..5ac0488f5 100644 --- a/aimux-provider-utils/src/api_key.rs +++ b/aimux-provider-utils/src/api_key.rs @@ -1,32 +1,62 @@ -//! API key loading (from parameter or environment variable). +//! API key and setting loading (from a parameter or an environment variable). use aimux_core::AiMuxError; /// Load an API key from the given value or environment variable. /// -/// If `api_key` is `Some`, returns it directly. -/// Otherwise reads from `environment_variable_name`. +/// A supplied value is returned exactly as given — including the empty +/// string. An explicit key never falls back to the environment: a caller that +/// deliberately passes `""` (a keyless local server) must not have a +/// different credential substituted from the process environment. Only +/// `None` reads `environment_variable_name`. /// /// # Errors /// -/// Returns `AiMuxError::InvalidArgument` when neither the `api_key` value nor -/// the environment variable yields a key. +/// Returns `AiMuxError::LoadApiKey` when `api_key` is `None` and the +/// environment variable is unset. pub fn load_api_key( api_key: Option<&str>, environment_variable_name: &str, description: &str, ) -> Result { - if let Some(key) = api_key - && !key.is_empty() - { + if let Some(key) = api_key { return Ok(key.to_string()); } - std::env::var(environment_variable_name).map_err(|_| { - AiMuxError::InvalidArgument(format!( - "No API key found for {description}. \ - Please provide it via the `api_key` parameter \ - or set the `{environment_variable_name}` environment variable." - )) + std::env::var(environment_variable_name).map_err(|_| AiMuxError::LoadApiKey { + env_var: environment_variable_name.to_string(), + description: description.to_string(), }) } + +/// Load a required setting from the given value or environment variable, +/// with the same no-fallback rule for an explicit value as [`load_api_key`]. +/// +/// # Errors +/// +/// Returns `AiMuxError::LoadSetting` when `value` is `None` and the +/// environment variable is unset. +pub fn load_setting( + value: Option<&str>, + environment_variable_name: &str, + setting_name: &str, +) -> Result { + load_optional_setting(value, environment_variable_name).ok_or_else(|| AiMuxError::LoadSetting { + env_var: environment_variable_name.to_string(), + name: setting_name.to_string(), + }) +} + +/// Load an optional setting from the given value or environment variable. +/// An explicit value (even `""`) wins; `None` reads the environment; an unset +/// variable yields `None`. +#[must_use] +pub fn load_optional_setting( + value: Option<&str>, + environment_variable_name: &str, +) -> Option { + match value { + Some(value) => Some(value.to_string()), + None => std::env::var(environment_variable_name).ok(), + } +} diff --git a/aimux-provider-utils/src/extract_response_headers.rs b/aimux-provider-utils/src/extract_response_headers.rs index c82435746..2b559fd90 100644 --- a/aimux-provider-utils/src/extract_response_headers.rs +++ b/aimux-provider-utils/src/extract_response_headers.rs @@ -5,9 +5,7 @@ use std::collections::HashMap; /// Extract response headers as ordered pairs while redacting sensitive values. /// Recording and public error context share this one policy. #[must_use] -pub fn extract_response_header_pairs( - headers: &reqwest::header::HeaderMap, -) -> Vec<(String, String)> { +pub(crate) fn extract_response_header_pairs(headers: &http::HeaderMap) -> Vec<(String, String)> { headers .iter() .map(|(name, value)| { @@ -23,6 +21,6 @@ pub fn extract_response_header_pairs( /// Extract response headers while redacting sensitive values. #[must_use] -pub fn extract_response_headers(headers: &reqwest::header::HeaderMap) -> HashMap { +pub fn extract_response_headers(headers: &http::HeaderMap) -> HashMap { extract_response_header_pairs(headers).into_iter().collect() } diff --git a/aimux-provider-utils/src/fetch.rs b/aimux-provider-utils/src/fetch.rs new file mode 100644 index 000000000..7d706bcd7 --- /dev/null +++ b/aimux-provider-utils/src/fetch.rs @@ -0,0 +1,567 @@ +//! Pluggable HTTP transport: the Rust shape of the AI SDK's `FetchFunction`. +//! +//! Every provider exchange ends in one [`Fetch::fetch`] call. A provider (or +//! its settings) may carry a [`FetchFunction`], as the AI SDK's `fetch` +//! setting: a mock, a signing decorator such as +//! [`crate::sigv4_fetch::SigV4Fetch`], or a [`ReqwestFetch`] with its own +//! [`ProxyConfig`]. When none is set the exchange uses [`default_fetch`], +//! resolved *per request* (never frozen when a provider is created), so the +//! process default can change after a provider exists, as replacing +//! `globalThis.fetch` does in JavaScript. +//! +//! This module owns the real network leaf: [`ReqwestFetch`] (its own +//! connection pool per tokio runtime, its own proxy routing) and +//! [`PinnedFetch`] (a single non-redirecting hop connected only to +//! pre-validated addresses, used by the SSRF download guard). Recording, +//! response handlers and retry stay above this layer, in `http.rs` and +//! `aimux-core`. + +use std::collections::HashMap; +use std::fmt; +use std::net::IpAddr; +use std::sync::{Arc, Mutex, OnceLock, RwLock}; +use std::time::Duration; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::StreamExt; +use futures::stream::BoxStream; +use http::{HeaderMap, Method, StatusCode}; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use url::Url; + +use aimux_core::AbortSignal; + +/// Proxy routing of one transport: a [`ReqwestFetch`], a [`PinnedFetch`] +/// hop, or the WebSocket connector. With every field empty, HTTP follows the +/// `HTTP_PROXY` / `HTTPS_PROXY` / `ALL_PROXY` / `NO_PROXY` environment, as +/// reqwest does, and the WebSocket connector connects directly. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct ProxyConfig { + pub http_url: Option, + pub https_url: Option, + pub all_url: Option, + pub no_proxy: Option, +} + +/// One outbound HTTP request, fully materialised: body bytes are final +/// (multipart already encoded), so a decorator can sign or inspect exactly +/// what goes on the wire. +#[derive(Debug, Clone)] +pub struct FetchRequest { + pub method: Method, + pub url: Url, + pub headers: HeaderMap, + pub body: Bytes, + /// Caller cancellation. The helper layer already drops the in-flight + /// future when this fires; transports may additionally use it to cancel + /// work they spawned. + pub signal: Option, + /// Whole-exchange bound, when the caller wants the transport itself to + /// enforce one. The helper layer enforces its own response deadline and + /// leaves this `None`. + pub timeout: Option, + /// Whether credential decorators may authenticate this request. + pub credentials: RequestCredentials, +} + +/// Whether a credential decorator may authenticate a request, after the +/// `credentials` member of the WHATWG fetch `RequestInit`. +/// +/// The helper layer sets [`RequestCredentials::Omit`] once a redirect leaves +/// the origin the credentials belong to. Caller headers are already reduced +/// to `User-Agent` by then; a decorator that adds credentials itself, such as +/// [`crate::sigv4_fetch::SigV4Fetch`], must forward the request untouched. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum RequestCredentials { + /// Decorators authenticate the request. + #[default] + Include, + /// Decorators add no credentials. + Omit, +} + +impl FetchRequest { + /// A request with no headers, no body and no bounds. + #[must_use] + pub fn new(method: Method, url: Url) -> Self { + Self { + method, + url, + headers: HeaderMap::new(), + body: Bytes::new(), + signal: None, + timeout: None, + credentials: RequestCredentials::Include, + } + } +} + +/// A response whose body has not been read yet. +pub struct FetchResponse { + pub status: StatusCode, + pub headers: HeaderMap, + /// The URL this response came from. + pub url: Url, + pub body: BoxStream<'static, Result>, +} + +impl FetchResponse { + /// A response whose whole body is already in memory (mocks, tests). + #[must_use] + pub fn from_bytes(status: StatusCode, headers: HeaderMap, url: Url, body: Bytes) -> Self { + Self { + status, + headers, + url, + body: futures::stream::once(async move { Ok(body) }).boxed(), + } + } + + /// The HTTP status. + #[must_use] + pub fn status(&self) -> StatusCode { + self.status + } + + /// The response headers. + #[must_use] + pub fn headers(&self) -> &HeaderMap { + &self.headers + } + + /// The declared `Content-Length`, when present and well-formed. + #[must_use] + pub fn content_length(&self) -> Option { + self.headers + .get(http::header::CONTENT_LENGTH)? + .to_str() + .ok()? + .parse() + .ok() + } + + /// Consume the response into its body stream. + pub fn bytes_stream(self) -> BoxStream<'static, Result> { + self.body + } +} + +impl fmt::Debug for FetchResponse { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FetchResponse") + .field("status", &self.status) + .field("headers", &self.headers) + .field("url", &self.url.as_str()) + .finish_non_exhaustive() + } +} + +/// A transport-level failure: no HTTP response was produced, or the body +/// stream broke. HTTP error statuses are *not* errors here — they come back +/// as a [`FetchResponse`] and are interpreted by the response handlers. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum FetchError { + /// The connection could not be established. + #[error("{0}")] + Connect(String), + /// The transport's own timeout fired. + #[error("request timed out")] + Timeout, + /// The caller's [`AbortSignal`] fired. + #[error("request aborted")] + Aborted, + /// The connection failed while sending the request or reading the body. + #[error("{0}")] + Io(String), + /// Anything else, including requests the transport refuses to build. + #[error("{0}")] + Other(String), +} + +impl FetchError { + /// Whether an operation retry may help. Connection, timeout and I/O + /// failures are transient; an abort is the caller's decision and `Other` + /// (a malformed request, a broken transport) will fail the same way again. + #[must_use] + pub fn is_retryable(&self) -> bool { + matches!(self, Self::Connect(_) | Self::Timeout | Self::Io(_)) + } +} + +/// The transport contract: one request in, one response out. A transport does +/// not follow redirects: a `3xx` is returned as it is and the helper layer +/// (`http.rs`) decides, hop by hop, whether and with which headers to follow. +#[async_trait] +pub trait Fetch: Send + Sync + 'static { + /// Perform one HTTP exchange. + /// + /// # Errors + /// + /// Returns a [`FetchError`] when no response could be produced. + async fn fetch(&self, request: FetchRequest) -> Result; +} + +impl fmt::Debug for dyn Fetch { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("dyn Fetch") + } +} + +/// A shareable transport, as carried by provider settings and requests. +pub type FetchFunction = Arc; + +static DEFAULT_TRANSPORT: OnceLock>> = OnceLock::new(); + +fn default_transport_slot() -> &'static RwLock> { + DEFAULT_TRANSPORT.get_or_init(|| RwLock::new(Arc::new(ReqwestFetch::default()))) +} + +/// The process default transport, read at request time by every exchange +/// whose provider has no `fetch`. +#[must_use] +pub fn default_fetch() -> FetchFunction { + default_transport() +} + +/// The process default transport as its concrete type: the built-in +/// WebSocket connector and the download hops share its proxy settings. +pub(crate) fn default_transport() -> Arc { + default_transport_slot() + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() +} + +/// Replace the process default transport. Requests started afterwards by +/// providers without their own `fetch`, built-in WebSocket connects and +/// validated download hops use it; requests in flight keep theirs. A +/// provider given `fetch` is unaffected. +pub fn set_default_transport(transport: ReqwestFetch) { + *default_transport_slot() + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::new(transport); +} + +/// The real network leaf: reqwest over rustls, routed by its own +/// [`ProxyConfig`], with its own connection pool per tokio runtime (clones +/// share the pool). +/// +/// The pool has a connect timeout but no client-wide response timeout: it +/// also serves streaming exchanges, so the non-streaming response bound lives +/// per exchange in the helper layer and `aimux-core` owns the operation +/// deadline. +#[derive(Debug, Clone, Default)] +pub struct ReqwestFetch { + proxy: ProxyConfig, + clients: Arc, Client>>>, +} + +impl ReqwestFetch { + /// A transport routed by `proxy`. + #[must_use] + pub fn new(proxy: ProxyConfig) -> Self { + Self { + proxy, + clients: Arc::default(), + } + } + + /// The proxy routing of this transport. + #[must_use] + pub fn proxy(&self) -> &ProxyConfig { + &self.proxy + } + + // One client (and so one connection pool) PER RUNTIME, not per + // transport: pooled connections are driven by tasks spawned onto the + // runtime that made the request, so when that runtime shuts down its + // pooled connections become unusable while staying checked in — and the + // OS can recycle their ports to fresh servers, handing later requests a + // dead connection. Production processes run a single runtime and still + // get exactly one client. Runtimes leave no drop signal, so dead entries + // are undetectable; the map is instead bounded by `CLIENTS_PER_TRANSPORT`. + fn client(&self) -> Result { + let key = tokio::runtime::Handle::try_current().ok().map(|h| h.id()); + let mut clients = self + .clients + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(client) = clients.get(&key) { + return Ok(client.clone()); + } + let client = build_client(&self.proxy).map_err(|error| { + FetchError::Other(format!("HTTP client initialization failed: {error}")) + })?; + if clients.len() >= CLIENTS_PER_TRANSPORT { + // Which entries are dead is unknowable, so evict them all; + // in-flight requests hold their own `Client` clone and are + // unaffected. + clients.clear(); + } + clients.insert(key, client.clone()); + Ok(client) + } +} + +#[async_trait] +impl Fetch for ReqwestFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + let client = self.client()?; + send(&client, request).await + } +} + +/// One non-redirecting hop of a validated download. The connection is pinned +/// to exactly the DNS answers that passed the SSRF guard (resolve overrides), +/// defeating TTL-0 rebinding; an empty pin list is the hop on the trusted +/// origin, which connects normally. The hop's proxy configuration is applied so reqwest makes the per-URL +/// routing decision itself: when a proxy carries the request the proxy +/// resolves the target (a trusted transport, the override is unused), and any +/// request the proxy rules send DIRECT still connects only through the +/// validated, pinned addresses. +/// +/// A client is built per hop: downloads are rare and each hop is one +/// request, so there is no pool to share — and no stale-runtime pool to +/// reuse. +#[derive(Debug, Clone, Default)] +pub struct PinnedFetch { + addresses: Vec, + proxy: ProxyConfig, +} + +impl PinnedFetch { + /// Pin the connection to these validated addresses, routed by `proxy`. + #[must_use] + pub fn new(addresses: Vec, proxy: ProxyConfig) -> Self { + Self { addresses, proxy } + } + + fn client(&self, url: &Url) -> Result { + let builder = reqwest::Client::builder() + .connect_timeout(Duration::from_millis(10_000)) + .redirect(reqwest::redirect::Policy::none()); + let builder = apply_proxy(builder, &self.proxy); + if self.addresses.is_empty() { + return builder.build().map_err(|error| { + FetchError::Other(format!("download client initialization failed: {error}")) + }); + } + let host = url + .host_str() + .ok_or_else(|| FetchError::Other("download url has no host".to_string()))?; + let port = url.port_or_known_default().unwrap_or(80); + let socket_addresses: Vec<_> = self + .addresses + .iter() + .map(|address| std::net::SocketAddr::new(*address, port)) + .collect(); + builder + .resolve_to_addrs(host, &socket_addresses) + .build() + .map_err(|error| { + FetchError::Other(format!( + "pinned download client initialization failed: {error}" + )) + }) + } +} + +#[async_trait] +impl Fetch for PinnedFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + let client = self.client(&request.url)?; + send(&client, request).await + } +} + +async fn send(client: &Client, request: FetchRequest) -> Result { + let mut builder = client + .request(request.method, request.url) + .headers(request.headers); + if !request.body.is_empty() { + builder = builder.body(request.body); + } + if let Some(timeout) = request.timeout { + builder = builder.timeout(timeout); + } + let response = builder.send().await.map_err(map_reqwest_error)?; + let status = response.status(); + let headers = response.headers().clone(); + let url = response.url().clone(); + Ok(FetchResponse { + status, + headers, + url, + body: response + .bytes_stream() + .map(|chunk| chunk.map_err(map_reqwest_error)) + .boxed(), + }) +} + +fn map_reqwest_error(error: reqwest::Error) -> FetchError { + if error.is_builder() { + FetchError::Other(error.to_string()) + } else if error.is_connect() { + FetchError::Connect(error.to_string()) + } else if error.is_timeout() { + FetchError::Timeout + } else { + FetchError::Io(error.to_string()) + } +} + +// Hosts that churn short-lived runtimes (a test binary, an FFI embedder +// creating a runtime per call) would otherwise grow a transport's client map +// without bound. Eviction is safe: correctness only requires never *reusing* +// a dead runtime's pool, and an evicted live runtime simply rebuilds a fresh +// client on its next request. +const CLIENTS_PER_TRANSPORT: usize = 8; + +fn build_client(proxy: &ProxyConfig) -> Result { + let builder = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .connect_timeout(Duration::from_millis(10_000)) + .pool_max_idle_per_host(10) + .pool_idle_timeout(Some(Duration::from_secs(30))) + .tcp_keepalive(Some(Duration::from_secs(20))); + apply_proxy(builder, proxy) + .build() + .map_err(|error| error.to_string()) +} + +fn apply_proxy(mut builder: reqwest::ClientBuilder, proxy: &ProxyConfig) -> reqwest::ClientBuilder { + let http = proxy.http_url.as_deref().or(proxy.all_url.as_deref()); + let https = proxy.https_url.as_deref().or(proxy.all_url.as_deref()); + if let Some(url) = http + && let Ok(reqwest_proxy) = reqwest::Proxy::http(url) + { + builder = apply_no_proxy(builder, reqwest_proxy, &proxy.no_proxy); + } + if let Some(url) = https + && let Ok(reqwest_proxy) = reqwest::Proxy::https(url) + { + builder = apply_no_proxy(builder, reqwest_proxy, &proxy.no_proxy); + } + builder +} + +// Kept separate so both proxy schemes use precisely the same no-proxy rule. +fn apply_no_proxy( + builder: reqwest::ClientBuilder, + proxy: reqwest::Proxy, + no_proxy: &Option, +) -> reqwest::ClientBuilder { + match no_proxy.as_deref().map(reqwest::NoProxy::from_string) { + Some(Some(no_proxy)) => builder.proxy(proxy.no_proxy(Some(no_proxy))), + _ => builder.proxy(proxy), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + /// A fake HTTP proxy answering every request with `ok`; returns its URL + /// and the request lines it saw. + async fn fake_proxy() -> (String, Arc>>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let seen = Arc::new(Mutex::new(Vec::new())); + let log = seen.clone(); + tokio::spawn(async move { + while let Ok((mut socket, _)) = listener.accept().await { + let log = log.clone(); + tokio::spawn(async move { + let mut buffer = vec![0u8; 4096]; + let read = socket.read(&mut buffer).await.unwrap_or(0); + let text = String::from_utf8_lossy(&buffer[..read]).to_string(); + log.lock() + .unwrap() + .push(text.lines().next().unwrap_or_default().to_string()); + let _ = socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-length: 2\r\nconnection: close\r\n\r\nok", + ) + .await; + }); + } + }); + (format!("http://{address}"), seen) + } + + fn get(url: &str) -> FetchRequest { + FetchRequest::new(Method::GET, Url::parse(url).unwrap()) + } + + /// The proxy belongs to the transport, as the AI SDK leaves routing to + /// the `fetch` a provider is given: a `ReqwestFetch` with a proxy sends + /// through it, whatever the process default is. + #[tokio::test] + async fn a_transport_routes_through_its_own_proxy() { + let (proxy, seen) = fake_proxy().await; + let transport = ReqwestFetch::new(ProxyConfig { + http_url: Some(proxy), + ..ProxyConfig::default() + }); + let response = transport + .fetch(get("http://upstream.invalid/v1/models")) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + seen.lock().unwrap().as_slice(), + ["GET http://upstream.invalid/v1/models HTTP/1.1"] + ); + } + + /// `default_fetch` is read per request: replacing the default transport + /// reroutes a transport handle taken later, and the WebSocket connector + /// and download hops read the same settings. + #[tokio::test] + #[serial_test::serial] + async fn the_default_transport_can_be_replaced() { + let (proxy, seen) = fake_proxy().await; + let previous = default_transport(); + let config = ProxyConfig { + http_url: Some(proxy), + // Tests running alongside talk to mock servers on 127.0.0.1; leave them + // on a direct route while the default is swapped. + no_proxy: Some("localhost,127.0.0.1".to_string()), + ..ProxyConfig::default() + }; + set_default_transport(ReqwestFetch::new(config.clone())); + assert_eq!(default_transport().proxy(), &config); + let response = default_fetch() + .fetch(get("http://upstream.invalid/v1")) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(seen.lock().unwrap().len(), 1); + set_default_transport((*previous).clone()); + assert_eq!(default_transport().proxy(), previous.proxy()); + } + + #[tokio::test] + async fn from_bytes_serves_the_body_once() { + let url = Url::parse("https://example.test/").unwrap(); + let mut headers = HeaderMap::new(); + headers.insert(http::header::CONTENT_LENGTH, "5".parse().unwrap()); + let response = + FetchResponse::from_bytes(StatusCode::OK, headers, url, Bytes::from("hello")); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.content_length(), Some(5)); + let chunks: Vec<_> = response.bytes_stream().collect().await; + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].as_ref().unwrap(), &Bytes::from("hello")); + } + + #[test] + fn default_fetch_is_a_single_process_wide_leaf() { + assert!(Arc::ptr_eq(&default_fetch(), &default_fetch())); + } +} diff --git a/aimux-provider-utils/src/handle_fetch_error.rs b/aimux-provider-utils/src/handle_fetch_error.rs index 4a935e327..76b7a1160 100644 --- a/aimux-provider-utils/src/handle_fetch_error.rs +++ b/aimux-provider-utils/src/handle_fetch_error.rs @@ -1,6 +1,8 @@ //! Fetch error normalization. -use aimux_core::AiMuxError; +use aimux_core::{AiMuxError, ApiCallError}; + +use crate::fetch::FetchError; /// Attach request context to transport failures without changing other errors. #[must_use] @@ -19,3 +21,21 @@ pub fn handle_fetch_error( other => other, } } + +/// Convert a transport failure into the public error model: an abort stays an +/// abort; everything else is an `ApiCall` error with no status (no response +/// arrived), retryable exactly when [`FetchError::is_retryable`] says so. +#[must_use] +pub(crate) fn fetch_error_to_ai_mux_error( + error: FetchError, + url: &str, + request_body_values: &serde_json::Value, +) -> AiMuxError { + match error { + FetchError::Aborted => AiMuxError::Aborted("request aborted".into()), + other => AiMuxError::ApiCall(Box::new(ApiCallError { + is_retryable: other.is_retryable(), + ..ApiCallError::new(other.to_string(), url, request_body_values.clone()) + })), + } +} diff --git a/aimux-provider-utils/src/headers.rs b/aimux-provider-utils/src/headers.rs index 96822c73b..eb00f6282 100644 --- a/aimux-provider-utils/src/headers.rs +++ b/aimux-provider-utils/src/headers.rs @@ -1,14 +1,92 @@ //! HTTP header utilities. +//! +//! Provider headers are layered — provider settings, then per-call headers, +//! then authentication — and later layers win. A layer entry of `None` means +//! "remove this header" (the AI SDK's `undefined` header value). use std::collections::HashMap; -/// The SDK version (injected at build time or hardcoded for now). -const SDK_VERSION: &str = env!("CARGO_PKG_VERSION"); +use crate::resolvable::Resolvable; -/// Add a `User-Agent` suffix to the headers. -/// -/// Pattern: `ai-sdk//` -pub fn with_user_agent_suffix(headers: &mut HashMap, provider_name: &str) { - let ua = format!("ai-sdk/{provider_name}/{SDK_VERSION}"); - headers.insert("User-Agent".to_string(), ua); +/// A header layer: `None` removes the header from earlier layers. +pub type HeaderMapOpt = HashMap>; + +/// Headers supplied directly or computed per request (token refresh, …). +pub type HeadersFn = Resolvable; + +/// Append a suffix to the case-insensitive `user-agent` header. +pub fn with_user_agent_suffix(headers: &mut HashMap, suffix: &str) { + let layer = headers + .drain() + .map(|(name, value)| (name, Some(value))) + .collect(); + *headers = normalize_headers(layer).into_iter().collect(); + let value = user_agent_value(headers.get("user-agent").map_or("", String::as_str), suffix); + headers.insert("user-agent".to_string(), value); +} + +pub(crate) fn user_agent_value(current: &str, suffix: &str) -> String { + [current, suffix] + .into_iter() + .filter(|value| !value.is_empty()) + .collect::>() + .join(" ") +} + +/// Resolve provider headers and append the pinned package's suffix per request. +#[must_use] +pub fn with_user_agent_suffix_fn( + headers: HeadersFn, + package: &'static str, + version: &'static str, +) -> HeadersFn { + Resolvable::from_async_fn(move || { + let headers = headers.clone(); + async move { + let mut headers = normalize_headers(headers.resolve().await?) + .into_iter() + .collect(); + with_user_agent_suffix(&mut headers, &format!("ai-sdk-{package}/{version}")); + Ok(headers + .into_iter() + .map(|(name, value)| (name, Some(value))) + .collect()) + } + }) +} + +/// Merge header layers, earliest first. Names are lowercased; a later layer +/// overrides an earlier one (never appends), including with `None`, which +/// removes the header. Entries of one layer whose names differ only by case +/// are applied in byte order of the original name, so the result does not +/// depend on hash-map iteration order. +#[must_use] +pub fn combine_headers(layers: &[&HeaderMapOpt]) -> HeaderMapOpt { + let mut combined = HeaderMapOpt::new(); + for layer in layers { + let mut entries: Vec<_> = layer.iter().collect(); + entries.sort_by(|a, b| a.0.cmp(b.0)); + for (name, value) in entries { + combined.insert(name.to_ascii_lowercase(), value.clone()); + } + } + combined +} + +/// Turn a header layer into the list that goes on the wire: names lowercased, +/// `None` entries dropped, one entry per name (when names differ only by case +/// the one that sorts last wins), sorted by name. +#[must_use] +pub fn normalize_headers(headers: HeaderMapOpt) -> Vec<(String, String)> { + let mut entries: Vec<_> = headers.into_iter().collect(); + entries.sort_by(|a, b| a.0.cmp(&b.0)); + let mut unique: std::collections::BTreeMap> = + std::collections::BTreeMap::new(); + for (name, value) in entries { + unique.insert(name.to_ascii_lowercase(), value); + } + unique + .into_iter() + .filter_map(|(name, value)| value.map(|value| (name, value))) + .collect() } diff --git a/aimux-provider-utils/src/http.rs b/aimux-provider-utils/src/http.rs index 5294190c7..501e7017b 100644 --- a/aimux-provider-utils/src/http.rs +++ b/aimux-provider-utils/src/http.rs @@ -1,192 +1,32 @@ -//! One shared HTTP client and one-exchange request transport. +//! One-exchange request transport. //! //! Retry and operation/stream timeouts belong to `aimux-core`. This module -//! sends exactly one request and leaves successful/failed body interpretation -//! to the response handlers selected by the provider. +//! sends exactly one request through a [`Fetch`] (the request's own, or the +//! process default) and leaves successful/failed body interpretation to the +//! response handlers selected by the provider. -use std::collections::HashMap; use std::pin::Pin; -use std::sync::{Mutex, OnceLock}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; use bytes::Bytes; use futures::{Stream, StreamExt, stream::BoxStream}; -use reqwest::Client; -use serde::{Deserialize, Serialize}; use aimux_core::recording::{ HttpExchange, HttpRecord, RecordingContext, ResponseRecord, TimingRecord, }; use aimux_core::{AiMuxError, ApiCallError}; +use crate::fetch::{ + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, PinnedFetch, RequestCredentials, + default_fetch, +}; +use crate::handle_fetch_error::fetch_error_to_ai_mux_error; use crate::logging::{body_logging_enabled, redact_body}; -/// Process-wide proxy configuration, fixed before the shared client is built. -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct ProxyConfig { - pub http_url: Option, - pub https_url: Option, - pub all_url: Option, - pub no_proxy: Option, -} - -static GLOBAL_PROXY: OnceLock = OnceLock::new(); -// One client (and so one connection pool) PER RUNTIME, not per process: -// pooled connections are driven by tasks spawned onto the runtime that -// made the request, so when that runtime shuts down its pooled -// connections become unusable while staying checked in — and the OS can -// recycle their ports to fresh servers, handing later requests a dead -// connection. Production processes run a single runtime and still get -// exactly one client. Runtimes leave no drop signal, so dead entries are -// undetectable; the map is instead bounded by `SHARED_CLIENT_CAP` — see -// `shared_client`. -static SHARED: OnceLock, Client>>> = OnceLock::new(); - -// Hosts that churn short-lived runtimes (a test binary, an FFI embedder -// creating a runtime per call) would otherwise grow the map without bound. -// Eviction is safe: correctness only requires never *reusing* a dead -// runtime's pool, and an evicted live runtime simply rebuilds a fresh -// client on its next request. -const SHARED_CLIENT_CAP: usize = 8; - -/// Set proxy configuration before the first HTTP operation. -/// -/// The HTTP client snapshots this config when it is first built; the -/// WebSocket connect path (RFC-0034 §2) re-reads it live on every connect, -/// so a late `init_proxy` can affect WS while HTTP keeps its first-built -/// client. Set it before any traffic either way. -pub fn init_proxy(config: ProxyConfig) -> bool { - GLOBAL_PROXY.set(config).is_ok() -} - -/// Return the process-wide proxy configuration (empty default when -/// `init_proxy` was never called). Read by both the HTTP client builder and -/// the WebSocket connect path (RFC-0034 §2), so `ws://`/`wss://` traffic -/// honors the same proxy/no_proxy settings as HTTP. -pub(crate) fn global_proxy() -> ProxyConfig { - GLOBAL_PROXY.get().cloned().unwrap_or_default() -} - -/// Return the single shared client. It has a connect timeout but no -/// client-wide response timeout: the shared client also serves streaming -/// exchanges, so the 30s non-streaming response bound lives per-exchange in -/// `call_to_api`, and Core owns the operation deadline. -/// -/// # Errors -/// -/// Returns an initialization error if the shared client cannot be built. -pub fn shared_client() -> Result { - let key = tokio::runtime::Handle::try_current().ok().map(|h| h.id()); - let mut clients = SHARED - .get_or_init(|| Mutex::new(HashMap::new())) - .lock() - .expect("shared HTTP client mutex poisoned"); - if let Some(client) = clients.get(&key) { - return Ok(client.clone()); - } - let client = build_client(global_proxy()).map_err(|error| { - AiMuxError::Other(format!("shared HTTP client initialization failed: {error}")) - })?; - if clients.len() >= SHARED_CLIENT_CAP { - // Which entries are dead is unknowable, so evict them all; in-flight - // requests hold their own `Client` clone and are unaffected. - clients.clear(); - } - clients.insert(key, client.clone()); - Ok(client) -} - -fn build_client(proxy: ProxyConfig) -> Result { - let builder = Client::builder() - .connect_timeout(Duration::from_millis(10_000)) - .pool_max_idle_per_host(10) - .pool_idle_timeout(Some(Duration::from_secs(30))) - .tcp_keepalive(Some(Duration::from_secs(20))); - apply_proxy(builder, &proxy) - .build() - .map_err(|error| error.to_string()) -} - /// Redirect ceiling for the manual validated-download loop. const MAX_REDIRECTS: usize = 10; -/// Non-redirecting client for a validated hop on the trusted origin. -/// Redirects are followed manually so every hop is re-validated; built per -/// exchange (downloads are rare and each hop is one request, so there is no -/// pool to share — and no stale-runtime pool to reuse). -fn download_client() -> Result { - let builder = Client::builder() - .connect_timeout(Duration::from_millis(10_000)) - .redirect(reqwest::redirect::Policy::none()); - apply_proxy(builder, &global_proxy()) - .build() - .map_err(|e| AiMuxError::Other(format!("download client initialization failed: {e}"))) -} - -/// Client for one validated hop off the trusted origin: the connection is -/// pinned to exactly the DNS answers that passed the guard (resolve -/// overrides), defeating TTL-0 rebinding. The proxy configuration is applied -/// so reqwest makes the per-URL routing decision itself: when a proxy -/// carries the request the proxy resolves the target (a trusted transport, -/// the override below is unused), and any request the proxy rules send -/// DIRECT still connects only through the validated, pinned addresses. -fn pinned_client(url: &str, addresses: &[std::net::IpAddr]) -> Result { - if addresses.is_empty() { - return Err(AiMuxError::Other( - "cannot build a pinned download client without an address".into(), - )); - } - let parsed = reqwest::Url::parse(url) - .map_err(|e| AiMuxError::Other(format!("invalid download url: {e}")))?; - let host = parsed - .host_str() - .ok_or_else(|| AiMuxError::Other("download url has no host".to_string()))? - .to_string(); - let port = parsed.port_or_known_default().unwrap_or(80); - let socket_addresses: Vec<_> = addresses - .iter() - .map(|address| std::net::SocketAddr::new(*address, port)) - .collect(); - let builder = Client::builder() - .connect_timeout(Duration::from_millis(10_000)) - .redirect(reqwest::redirect::Policy::none()); - apply_proxy(builder, &global_proxy()) - .resolve_to_addrs(&host, &socket_addresses) - .build() - .map_err(|e| { - AiMuxError::Other(format!("pinned download client initialization failed: {e}")) - }) -} - -fn apply_proxy(mut builder: reqwest::ClientBuilder, proxy: &ProxyConfig) -> reqwest::ClientBuilder { - let http = proxy.http_url.as_deref().or(proxy.all_url.as_deref()); - let https = proxy.https_url.as_deref().or(proxy.all_url.as_deref()); - if let Some(url) = http - && let Ok(reqwest_proxy) = reqwest::Proxy::http(url) - { - builder = apply_no_proxy(builder, reqwest_proxy, &proxy.no_proxy); - } - if let Some(url) = https - && let Ok(reqwest_proxy) = reqwest::Proxy::https(url) - { - builder = apply_no_proxy(builder, reqwest_proxy, &proxy.no_proxy); - } - builder -} - -// Kept separate so both proxy schemes use precisely the same no-proxy rule. -fn apply_no_proxy( - builder: reqwest::ClientBuilder, - proxy: reqwest::Proxy, - no_proxy: &Option, -) -> reqwest::ClientBuilder { - match no_proxy.as_deref().map(reqwest::NoProxy::from_string) { - Some(Some(no_proxy)) => builder.proxy(proxy.no_proxy(Some(no_proxy))), - _ => builder.proxy(proxy), - } -} - /// Explicit raw request body used only by `post_to_api`. #[derive(Debug, Clone)] pub enum HttpBody { @@ -225,8 +65,7 @@ pub struct HttpRequest { /// (e.g. Replicate `prefer: wait`) declares its own bound here. Streaming /// exchanges are exempt regardless. pub response_timeout: Option, - /// Per-request override of the successful-JSON-body size cap (default - /// [`crate::read_response_with_size_limit::DEFAULT_MAX_JSON_RESPONSE_SIZE`]). + /// Per-request override of the successful-JSON-body size cap (default 64 MiB). /// Only consulted by [`crate::response_handler::create_json_response_handler`]; /// binary downloads keep the separate, larger /// `DEFAULT_MAX_DOWNLOAD_SIZE` bound regardless of this field. @@ -248,6 +87,11 @@ pub struct HttpRequest { /// this value — from the first request and on every redirect hop. `None` /// means the caller gates its own headers (e.g. BFL's host allowlist). pub credentialed_origin: Option, + /// Transport for this exchange. `None` resolves [`default_fetch`] when the + /// request is sent. Ignored when `validate_url` is set: downloads of + /// provider-supplied URLs always use the SSRF guard's pinned transport, + /// never a provider-injected one. + pub fetch: Option, } /// The per-operation context of an exchange: the call options' cancellation, @@ -351,6 +195,10 @@ pub(crate) struct PreparedRequest { pub(crate) response_timeout: Option, pub(crate) max_json_response_bytes: Option, pub(crate) validation: Option, + pub(crate) fetch: Option, + /// [`RequestCredentials::Omit`] once the target is off the credentialed + /// origin; never switched back. + credentials: RequestCredentials, } impl HttpRequest { @@ -359,10 +207,12 @@ impl HttpRequest { // when the target is off the credentialed origin they never leave, // whichever transport the exchange takes. let mut headers = self.headers; + let mut credentials = RequestCredentials::Include; if let Some(origin) = self.credentialed_origin.as_deref() - && !crate::download_guard::same_origin(&self.url, origin) + && !aimux_core::download_guard::same_origin(&self.url, origin) { - crate::download_guard::retain_user_agent(&mut headers); + aimux_core::download_guard::retain_user_agent(&mut headers); + credentials = RequestCredentials::Omit; } PreparedRequest { method, @@ -378,6 +228,8 @@ impl HttpRequest { trusted_origin: self.trusted_origin, credentialed_origin: self.credentialed_origin, }), + fetch: self.fetch, + credentials, } } } @@ -442,7 +294,7 @@ impl Drop for ExchangeGuard<'_> { /// Execute exactly one HTTP exchange. Core owns retry and operation timeout. pub(crate) async fn send_request_once( request: &PreparedRequest, -) -> Result { +) -> Result { let started = Instant::now(); let (attempt, exchange_index) = request .recording_context @@ -470,14 +322,7 @@ pub(crate) async fn send_request_once( ); } - let sent = if request.validation.is_some() { - send_validated_redirects(request).await - } else { - match shared_client() { - Ok(client) => send_one_request(&client, request).await, - Err(error) => Err(error), - } - }; + let sent = send_following_redirects(request).await; let response = match sent { Ok(response) => response, Err(error) => { @@ -489,7 +334,7 @@ pub(crate) async fn send_request_once( let latency_ms = started.elapsed().as_millis() as u64; tracing::debug!( target: "aimux_provider_utils::http", - status = response.status().as_u16(), + status = response.status.as_u16(), latency_ms, "response" ); @@ -504,51 +349,48 @@ pub(crate) async fn send_request_once( } async fn send_one_request( - client: &Client, + fetch: &dyn Fetch, request: &PreparedRequest, -) -> Result { - if let Some(signal) = &request.abort_signal { - if signal.is_aborted() { - return Err(AiMuxError::from_abort_signal(signal)); - } - tokio::select! { - biased; - () = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)), - response = build_request_builder(client, request)?.send() => { - response.map_err(|error| AiMuxError::ApiCall(Box::new(ApiCallError { - is_retryable: true, - ..api_call_error(request, error.to_string()) - }))) +) -> Result { + let mut fetch_request = build_fetch_request(request)?; + fetch_request.signal = request.abort_signal.clone(); + let sent = match &request.abort_signal { + Some(signal) => { + if signal.is_aborted() { + return Err(AiMuxError::from_abort_signal(signal)); + } + tokio::select! { + biased; + () = signal.cancelled() => return Err(AiMuxError::from_abort_signal(signal)), + sent = fetch.fetch(fetch_request) => sent, } } - } else { - build_request_builder(client, request)? - .send() - .await - .map_err(|error| { - AiMuxError::ApiCall(Box::new(ApiCallError { - is_retryable: true, - ..api_call_error(request, error.to_string()) - })) - }) - } + None => fetch.fetch(fetch_request).await, + }; + sent.map_err(|error| { + fetch_error_to_ai_mux_error( + error, + &sanitized_request_url(&request.url), + &crate::logging::redact_request_values(&request.body), + ) + }) } -fn is_redirect_status(status: reqwest::StatusCode) -> bool { +fn is_redirect_status(status: http::StatusCode) -> bool { matches!( status, - reqwest::StatusCode::MOVED_PERMANENTLY - | reqwest::StatusCode::FOUND - | reqwest::StatusCode::SEE_OTHER - | reqwest::StatusCode::TEMPORARY_REDIRECT - | reqwest::StatusCode::PERMANENT_REDIRECT + http::StatusCode::MOVED_PERMANENTLY + | http::StatusCode::FOUND + | http::StatusCode::SEE_OTHER + | http::StatusCode::TEMPORARY_REDIRECT + | http::StatusCode::PERMANENT_REDIRECT ) } fn redirect_error( request: &PreparedRequest, message: impl Into, - status: reqwest::StatusCode, + status: http::StatusCode, ) -> AiMuxError { AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(status.as_u16()), @@ -556,53 +398,73 @@ fn redirect_error( })) } -/// Send one validated-download exchange, following redirects manually so -/// every hop is re-validated and its connection pinned to the DNS answers -/// that passed the guard (see `download_guard`). The whole chain is one -/// exchange to the caller, matching how an auto-following client behaves. -async fn send_validated_redirects( - request: &PreparedRequest, -) -> Result { - let validation = request - .validation - .clone() - .expect("send_validated_redirects requires a validation config"); - let trusted_origin = validation.trusted_origin.as_deref(); - let credentialed_origin = validation.credentialed_origin.as_deref(); +/// Send one exchange, following redirects hop by hop in this layer (a +/// transport never follows one itself). The whole chain is one exchange to +/// the caller, matching how an auto-following client behaves. +/// +/// - An ordinary API call follows every redirect, as the AI SDK's `postToApi` +/// does under fetch's default `redirect: "follow"` (provider-utils@5.0.53 +/// `post-to-api.ts:98-107`). Once a hop leaves the request's origin, every +/// header except `User-Agent` is dropped and the transport is told +/// `RequestCredentials::Omit`, as `fetchWithValidatedRedirects` does +/// (`fetch-with-validated-redirects.ts:195-206`), so neither a provider key +/// in a custom header nor anything a transport decorator adds (a +/// signature, a bearer token) reaches another origin. Every followed hop +/// goes through the request's transport again, so a signing transport +/// signs the URL it actually sends. +/// - A validated download (`validation` set) re-validates every hop and pins +/// its connection to the DNS answers that passed the guard (see +/// `download_guard`); it follows cross-origin hops with the caller headers +/// stripped. D26: its transport is never provider-injected. +async fn send_following_redirects(request: &PreparedRequest) -> Result { + let validation = request.validation.as_ref(); + let trusted_origin = validation.and_then(|v| v.trusted_origin.as_deref()); let mut current = request.clone(); - // First-request credential gating already ran in `prepare`; this strips - // hop-by-hop, forwarding, and metadata-service headers for every hop. - crate::download_guard::sanitize_download_headers(&mut current.headers); // Headers never travel past this origin on a redirect; stripping is // one-way, so a hop back onto it cannot restore them. - let credential_anchor = credentialed_origin.unwrap_or(&request.url); - let mut pinned = - crate::download_guard::validate_download_target(¤t.url, trusted_origin).await?; + let credential_anchor = validation + .and_then(|v| v.credentialed_origin.as_deref()) + .unwrap_or(&request.url); + let mut pinned = match validation { + Some(_) => { + // First-request credential gating already ran in `prepare`; this + // strips hop-by-hop, forwarding, and metadata-service headers. + aimux_core::download_guard::sanitize_download_headers(&mut current.headers); + Some( + aimux_core::download_guard::validate_download_target(¤t.url, trusted_origin) + .await?, + ) + } + None => None, + }; + // Resolved per request, so a late change of the process default is + // honored by providers created before it. + let api_fetch = request.fetch.clone().unwrap_or_else(default_fetch); for redirect_count in 0..=MAX_REDIRECTS { - // Empty pins mean the hop is on the trusted origin. - let client = if pinned.is_empty() { - download_client()? - } else { - pinned_client(¤t.url, &pinned)? + let response = match &pinned { + // Empty pins mean the hop is on the trusted origin. + // D26: the download hop is never provider-injected; it routes + // through the default transport's proxy settings. + Some(pins) => { + let hop = PinnedFetch::new( + pins.clone(), + crate::fetch::default_transport().proxy().clone(), + ); + send_one_request(&hop, ¤t).await? + } + None => send_one_request(api_fetch.as_ref(), ¤t).await?, }; - let response = send_one_request(&client, ¤t).await?; - let status = response.status(); + let status = response.status; if !is_redirect_status(status) { return Ok(response); } - let Some(location) = response.headers().get(reqwest::header::LOCATION) else { + let Some(location) = response.headers.get(http::header::LOCATION) else { return Ok(response); }; - if redirect_count == MAX_REDIRECTS { - return Err(redirect_error(request, "too many redirects", status)); - } let location = location .to_str() .map(str::to_owned) .map_err(|_| redirect_error(request, "redirect location is not valid UTF-8", status))?; - // Dropping the unconsumed 3xx body releases its connection before a - // potentially slow DNS check for the next hop. - drop(response); let base = url::Url::parse(¤t.url) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid request URL: {e}")))?; let next = base @@ -619,17 +481,34 @@ async fn send_validated_redirects( )); } let next = next.to_string(); - let hop_trusted = crate::download_guard::hop_trusted_origin(trusted_origin, ¤t.url); - pinned = crate::download_guard::validate_download_target(&next, hop_trusted).await?; + if redirect_count == MAX_REDIRECTS { + return Err(redirect_error(request, "too many redirects", status)); + } + // Dropping the unconsumed 3xx body releases its connection before a + // potentially slow DNS check for the next hop. + drop(response); + if pinned.is_some() { + let hop_trusted = + aimux_core::download_guard::hop_trusted_origin(trusted_origin, ¤t.url); + pinned = Some( + aimux_core::download_guard::validate_download_target(&next, hop_trusted).await?, + ); + } // Credentials are scoped to the credential anchor, not the previous - // hop: once a redirect leaves it, headers cannot come back. - if !crate::download_guard::same_origin(&next, credential_anchor) { - crate::download_guard::retain_user_agent(&mut current.headers); + // hop: once a redirect leaves it, headers cannot come back. An API + // request follows every redirect, like the AI SDK's `postToApi` under + // fetch's default `redirect: "follow"`, but off the anchor it keeps + // only `User-Agent`, as `fetchWithValidatedRedirects` does: the fetch + // spec strips only `Authorization`, and provider keys also travel in + // custom headers. Signing decorators see `RequestCredentials::Omit`. + if !aimux_core::download_guard::same_origin(&next, credential_anchor) { + aimux_core::download_guard::retain_user_agent(&mut current.headers); + current.credentials = RequestCredentials::Omit; } - if status == reqwest::StatusCode::SEE_OTHER + if status == http::StatusCode::SEE_OTHER || (matches!( status, - reqwest::StatusCode::MOVED_PERMANENTLY | reqwest::StatusCode::FOUND + http::StatusCode::MOVED_PERMANENTLY | http::StatusCode::FOUND ) && matches!(current.method, HttpMethod::Post)) { current.method = HttpMethod::Get; @@ -640,28 +519,71 @@ async fn send_validated_redirects( unreachable!("redirect loop returns on response or error") } -fn build_request_builder( - client: &Client, - request: &PreparedRequest, -) -> Result { - let mut builder = match request.method { - HttpMethod::Get => client.get(&request.url), - HttpMethod::Post => client.post(&request.url), - }; +/// Materialise one prepared request for a [`Fetch`]: validated URL, final body +/// bytes and an ordered-by-name header map. Headers are *inserted*, so a +/// later entry for the same (case-insensitive) name replaces an earlier one — +/// a request never goes out with duplicate `Authorization` headers. +fn build_fetch_request(request: &PreparedRequest) -> Result { + let url = url::Url::parse(&request.url) + .map_err(|error| AiMuxError::InvalidArgument(format!("invalid request URL: {error}")))?; + let mut headers = http::HeaderMap::new(); for (name, value) in &request.headers { - let header_name = reqwest::header::HeaderName::try_from(name) + let header_name = http::HeaderName::try_from(name) .map_err(|_| AiMuxError::InvalidArgument(format!("invalid header name: {name}")))?; - let header_value = reqwest::header::HeaderValue::try_from(value).map_err(|_| { - AiMuxError::InvalidArgument(format!("invalid header value for {name}: {value}")) - })?; - builder = builder.header(header_name, header_value); + // The value is never echoed: it is commonly a credential. + let header_value = http::HeaderValue::try_from(value) + .map_err(|_| AiMuxError::InvalidArgument(format!("invalid header value for {name}")))?; + headers.insert(header_name, header_value); } - Ok(match &request.body { - HttpBody::Json(value) => builder.json(value), - HttpBody::Bytes(bytes, content_type) => builder - .header(reqwest::header::CONTENT_TYPE, content_type) - .body(bytes.clone()), - HttpBody::Empty => builder, + let suffix = "ai-sdk-provider-utils/5.0.51"; + let current = request + .headers + .iter() + .rev() + .find_map(|(name, value)| { + name.eq_ignore_ascii_case("user-agent") + .then_some(value.as_str()) + }) + .unwrap_or_default(); + let user_agent = crate::headers::user_agent_value(current, suffix); + headers.insert( + http::header::USER_AGENT, + http::HeaderValue::try_from(user_agent).map_err(|_| { + AiMuxError::InvalidArgument("invalid header value for user-agent".to_string()) + })?, + ); + let body = match &request.body { + HttpBody::Json(value) => { + if !headers.contains_key(http::header::CONTENT_TYPE) { + headers.insert( + http::header::CONTENT_TYPE, + http::HeaderValue::from_static("application/json"), + ); + } + Bytes::from(serde_json::to_vec(value).map_err(|error| { + AiMuxError::InvalidArgument(format!("request body is not serializable: {error}")) + })?) + } + HttpBody::Bytes(bytes, content_type) => { + let content_type = http::HeaderValue::try_from(content_type).map_err(|_| { + AiMuxError::InvalidArgument(format!("invalid content type: {content_type}")) + })?; + headers.insert(http::header::CONTENT_TYPE, content_type); + Bytes::copy_from_slice(bytes) + } + HttpBody::Empty => Bytes::new(), + }; + Ok(FetchRequest { + method: match request.method { + HttpMethod::Get => http::Method::GET, + HttpMethod::Post => http::Method::POST, + }, + url, + headers, + body, + signal: None, + timeout: None, + credentials: request.credentials, }) } @@ -689,16 +611,15 @@ pub async fn sleep_or_abort( } fn observe_response( - response: reqwest::Response, + response: FetchResponse, request: &PreparedRequest, started: Instant, latency_ms: u64, attempt: u32, exchange_index: u32, -) -> reqwest::Response { - let status = response.status(); - let version = response.version(); - let headers = response.headers().clone(); +) -> FetchResponse { + let status = response.status; + let headers = response.headers.clone(); let recording = request.recording_context.as_ref().map(|context| { let redacted_headers = redacted_response_headers(&headers); let response = ResponseRecord { @@ -740,23 +661,22 @@ fn observe_response( } }); let observed = ObservedBody { - inner: response.bytes_stream().boxed(), + inner: response.body, started, chunks: 0, done: false, recording, }; - let mut rebuilt = http::Response::builder() - .status(status) - .version(version) - .body(reqwest::Body::wrap_stream(observed)) - .expect("status and version came from a valid reqwest response"); - *rebuilt.headers_mut() = headers; - rebuilt.into() + FetchResponse { + status, + headers, + url: response.url, + body: observed.boxed(), + } } struct ObservedBody { - inner: BoxStream<'static, Result>, + inner: BoxStream<'static, Result>, started: Instant, chunks: usize, done: bool, @@ -764,7 +684,7 @@ struct ObservedBody { } impl Stream for ObservedBody { - type Item = Result; + type Item = Result; fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { if self.done { @@ -965,7 +885,7 @@ fn record_failed_exchange( ); } -fn redacted_response_headers(headers: &reqwest::header::HeaderMap) -> Vec<(String, String)> { +fn redacted_response_headers(headers: &http::HeaderMap) -> Vec<(String, String)> { crate::extract_response_headers::extract_response_header_pairs(headers) } diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index e7d5e5a77..9687b13e0 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -2,26 +2,29 @@ //! //! Shared utilities for provider implementations. //! -//! Provides one-exchange HTTP helpers, response handlers, API key loading, -//! header management, URL utilities and the streamed tool-call tracker for -//! the OpenAI chat-completions wire format — the Rust equivalents of -//! `@ai-sdk/provider-utils`. Operation retry and timeout live in `aimux-core`. +//! Provides the pluggable [`Fetch`] transport, one-exchange HTTP helpers, +//! response handlers, API key and setting loading, [`Resolvable`] values and +//! layered header management, URL utilities and the streamed tool-call +//! tracker for the OpenAI chat-completions wire format — the Rust equivalents +//! of `@ai-sdk/provider-utils`. Operation retry and timeout live in +//! `aimux-core`. pub mod api_key; -#[doc(hidden)] -pub use aimux_core::download_guard; pub mod extract_response_headers; +pub mod fetch; pub mod generate_id; pub mod get_from_api; pub mod handle_fetch_error; pub mod headers; pub mod http; pub mod logging; +pub mod media_type; pub mod multipart; pub mod post_to_api; pub mod read_response_with_size_limit; +pub mod resolvable; pub mod response_handler; -pub mod retry; +pub mod sigv4_fetch; mod streaming_tool_call_argument_state; mod streaming_tool_call_tracker; pub mod url; @@ -30,25 +33,33 @@ pub mod url; #[cfg(feature = "ws")] pub mod ws; -pub use api_key::load_api_key; -pub use download_guard::same_origin; +pub use aimux_core::download_guard::same_origin; +pub use api_key::{load_api_key, load_optional_setting, load_setting}; +pub use fetch::{ + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, PinnedFetch, ProxyConfig, + RequestCredentials, ReqwestFetch, default_fetch, set_default_transport, +}; pub use generate_id::generate_id; pub use get_from_api::get_from_api; -pub use headers::with_user_agent_suffix; -pub use http::{ - ExchangeContext, HttpBody, HttpRequest, ProxyConfig, init_proxy, shared_client, sleep_or_abort, +pub use headers::{ + HeaderMapOpt, HeadersFn, combine_headers, normalize_headers, with_user_agent_suffix, +}; +pub use http::{ExchangeContext, HttpBody, HttpRequest, sleep_or_abort}; +pub use logging::{init_logging, redact_error_context}; +pub use media_type::{ + MAX_ID3_TAG_BYTES, MediaTypeData, detect_media_type, get_top_level_media_type, + is_full_media_type, resolve_full_media_type, }; -pub use logging::{body_logging_enabled, init_logging, redact_body, redact_error_context}; pub use multipart::{MultipartForm, media_type_to_extension}; pub use post_to_api::{post_form_data_to_api, post_json_to_api, post_to_api}; +pub use resolvable::Resolvable; pub use response_handler::{ ProviderErrorParts, ResponseHandler, ResponseHandlerInput, ResponseHandlerOutput, create_binary_response_handler, create_event_source_response_handler, create_json_error_response_handler, create_json_response_handler, - create_standard_json_error_response_handler, create_status_code_error_response_handler, - invalid_response_api_call, stream_error_api_call, + create_status_code_error_response_handler, invalid_response_api_call, stream_error_api_call, }; -pub use retry::RetryConfig; +pub use sigv4_fetch::{AwsCredentials, SigV4Fetch}; // As `@ai-sdk/provider-utils` (`index.ts:115-119`): the tracker is exported // from the crate root, its argument state stays internal. pub use streaming_tool_call_tracker::{ diff --git a/aimux-provider-utils/src/logging.rs b/aimux-provider-utils/src/logging.rs index 3b28790ed..8eaab9110 100644 --- a/aimux-provider-utils/src/logging.rs +++ b/aimux-provider-utils/src/logging.rs @@ -28,15 +28,15 @@ static INIT: Once = Once::new(); /// Default level: `warn` — only retries and failures are printed, normal /// traffic stays silent. -pub const DEFAULT_LEVEL: &str = "warn"; +const DEFAULT_LEVEL: &str = "warn"; /// Body log truncation limit (bytes), RFC-0014 §4.3. -pub const BODY_LOG_LIMIT: usize = 4096; +const BODY_LOG_LIMIT: usize = 4096; /// Initialize the global logger (idempotent). Every entry point funnels here. /// /// Precedence: `AIMUX_LOG` (RUST_LOG-style directives) > `AIMUX_LOG_LEVEL` -/// (simple level name) > `level` argument > [`DEFAULT_LEVEL`]. +/// (simple level name) > `level` argument > `DEFAULT_LEVEL`. /// /// No-ops (without registering anything) when a global subscriber is already /// installed by the consumer — we never override the host's own logger. @@ -66,7 +66,7 @@ pub fn init_logging(level: &str) { /// Cheap when already initialized (`Once::is_completed` is a single atomic /// load). Registers a subscriber only when an `AIMUX_LOG*` env var is present /// **and** no global subscriber exists yet. -pub fn auto_init_from_env() { +pub(crate) fn auto_init_from_env() { if INIT.is_completed() || tracing::dispatcher::has_been_set() { return; } @@ -79,7 +79,7 @@ pub fn auto_init_from_env() { /// Whether request/response body trace logging is enabled (`AIMUX_LOG_BODY=1`). #[must_use] -pub fn body_logging_enabled() -> bool { +pub(crate) fn body_logging_enabled() -> bool { std::env::var("AIMUX_LOG_BODY") .map(|v| v == "1") .unwrap_or(false) @@ -104,7 +104,7 @@ fn build_filter(level: Option<&str>) -> EnvFilter { EnvFilter::try_new(format!("aimux={level}")).unwrap_or_else(|_| EnvFilter::new(DEFAULT_LEVEL)) } -/// Truncate `body` to [`BODY_LOG_LIMIT`] bytes at a UTF-8 char boundary. +/// Truncate `body` to `BODY_LOG_LIMIT` bytes at a UTF-8 char boundary. fn truncate(body: &str) -> &str { if body.len() <= BODY_LOG_LIMIT { return body; @@ -119,9 +119,9 @@ fn truncate(body: &str) -> &str { /// Redact a body for trace logging: JSON object keys whose lowercased name /// match the shared recording sensitivity policy have their values replaced /// with `[REDACTED]`. Non-JSON bodies pass through unchanged. Always truncated -/// to [`BODY_LOG_LIMIT`]. +/// to `BODY_LOG_LIMIT`. #[must_use] -pub fn redact_body(body: &str) -> String { +pub(crate) fn redact_body(body: &str) -> String { let truncated = truncate(body); match serde_json::from_str::(truncated) { Ok(value) => serde_json::to_string(&redact_error_context(value)) @@ -184,7 +184,7 @@ fn redact_value(value: serde_json::Value, depth: usize) -> serde_json::Value { /// Sanitize request values before storing them in a public `ApiCallError`. #[must_use] -pub fn redact_request_values(body: &crate::http::HttpBody) -> serde_json::Value { +pub(crate) fn redact_request_values(body: &crate::http::HttpBody) -> serde_json::Value { match body { crate::http::HttpBody::Json(value) => redact_error_context(value.clone()), crate::http::HttpBody::Bytes(bytes, content_type) => serde_json::json!({ diff --git a/aimux-provider-utils/src/media_type.rs b/aimux-provider-utils/src/media_type.rs new file mode 100644 index 000000000..1127ac960 --- /dev/null +++ b/aimux-provider-utils/src/media_type.rs @@ -0,0 +1,210 @@ +//! Media type detection and resolution using the provider-utils signatures. + +use std::borrow::Cow; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model_message::FilePart; +use aimux_core::shared::{FileBytes, FileData}; + +/// Inline data accepted by the media type detector. +#[derive(Clone, Copy, Debug)] +pub enum MediaTypeData<'a> { + Bytes(&'a [u8]), + Base64(&'a str), +} + +/// Largest ID3v2 tag (including its ten-byte header) scanned for an audio frame. +pub const MAX_ID3_TAG_BYTES: usize = 128 * 1024; +const DEFAULT_SNIFF_BYTES: usize = 18; +const MAX_SIGNATURE_BYTES: usize = 12; + +// -1 marks a variable byte. Order is significant for overlapping containers. +const SIGNATURES: &[(&str, &[i16])] = &[ + ("image/gif", &[0x47, 0x49, 0x46, 0x38, 0x37, 0x61]), + ("image/gif", &[0x47, 0x49, 0x46, 0x38, 0x39, 0x61]), + ("image/png", &[0x89, 0x50, 0x4e, 0x47]), + ("image/jpeg", &[0xff, 0xd8]), + ( + "image/webp", + &[ + 0x52, 0x49, 0x46, 0x46, -1, -1, -1, -1, 0x57, 0x45, 0x42, 0x50, + ], + ), + ("image/bmp", &[0x42, 0x4d, -1, -1, -1, -1, 0, 0, 0, 0]), + ("image/tiff", &[0x49, 0x49, 0x2a, 0]), + ("image/tiff", &[0x4d, 0x4d, 0, 0x2a]), + ( + "image/avif", + &[0, 0, 0, -1, 0x66, 0x74, 0x79, 0x70, 0x61, 0x76, 0x69, 0x66], + ), + ( + "image/heic", + &[0, 0, 0, -1, 0x66, 0x74, 0x79, 0x70, 0x68, 0x65, 0x69, 0x63], + ), + ("application/pdf", &[0x25, 0x50, 0x44, 0x46]), + ("audio/aac", &[0xff, 0xf0]), + ("audio/aac", &[0xff, 0xf1]), + ("audio/aac", &[0xff, 0xf8]), + ("audio/aac", &[0xff, 0xf9]), + ("audio/mpeg", &[0xff, 0xfb]), + ("audio/mpeg", &[0xff, 0xfa]), + ("audio/mpeg", &[0xff, 0xf3]), + ("audio/mpeg", &[0xff, 0xf2]), + ("audio/mpeg", &[0xff, 0xe3]), + ("audio/mpeg", &[0xff, 0xe2]), + ( + "audio/wav", + &[ + 0x52, 0x49, 0x46, 0x46, -1, -1, -1, -1, 0x57, 0x41, 0x56, 0x45, + ], + ), + ("audio/ogg", &[0x4f, 0x67, 0x67, 0x53]), + ("audio/flac", &[0x66, 0x4c, 0x61, 0x43]), + ("audio/aac", &[0x40, 0x15, 0, 0]), + ("audio/webm", &[0x1a, 0x45, 0xdf, 0xa3]), + ("video/mp4", &[0, 0, 0, -1, 0x66, 0x74, 0x79, 0x70]), + ("video/webm", &[0x1a, 0x45, 0xdf, 0xa3]), + ( + "video/quicktime", + &[0, 0, 0, 0x14, 0x66, 0x74, 0x79, 0x70, 0x71, 0x74], + ), + ("video/x-msvideo", &[0x52, 0x49, 0x46, 0x46]), + ("audio/mp4", &[0, 0, 0, -1, 0x66, 0x74, 0x79, 0x70]), +]; + +fn decode_prefix(data: MediaTypeData<'_>, max_bytes: usize) -> Result, AiMuxError> { + let data = match data { + MediaTypeData::Bytes(bytes) => { + return Ok(Cow::Borrowed(&bytes[..bytes.len().min(max_bytes)])); + } + MediaTypeData::Base64(data) => data, + }; + // Match atob: URL alphabet, ASCII whitespace, optional padding, and + // noncanonical trailing bits are accepted; malformed input is an error. + let invalid = || AiMuxError::InvalidArgument("invalid base64 data".into()); + let mut chars = data + .encode_utf16() + .take(max_bytes.div_ceil(3) * 4) + .filter(|&c| !matches!(c, 0x09 | 0x0a | 0x0c | 0x0d | 0x20)) + .collect::>(); + if chars.len().is_multiple_of(4) { + for _ in 0..2 { + if chars.last() == Some(&0x3d) { + chars.pop(); + } + } + } + if chars.len() % 4 == 1 { + return Err(invalid()); + } + let mut bytes = Vec::with_capacity(chars.len() * 3 / 4); + let mut value = 0u32; + let mut bits = 0; + for c in chars { + let digit = match c { + 0x41..=0x5a => c - 0x41, + 0x61..=0x7a => c - 0x61 + 26, + 0x30..=0x39 => c - 0x30 + 52, + 0x2b | 0x2d => 62, + 0x2f | 0x5f => 63, + _ => return Err(invalid()), + }; + value = (value << 6) | u32::from(digit); + bits += 6; + if bits >= 8 { + bits -= 8; + bytes.push((value >> bits) as u8); + value &= (1 << bits) - 1; + } + } + bytes.truncate(max_bytes); + Ok(Cow::Owned(bytes)) +} + +/// Detect a known media type, optionally restricted to its top-level segment. +/// +/// # Errors +/// Returns an error if the inspected base64 prefix is malformed. +pub fn detect_media_type( + data: MediaTypeData<'_>, + top_level_type: Option<&str>, +) -> Result, AiMuxError> { + if top_level_type + .is_some_and(|kind| !matches!(kind, "image" | "audio" | "video" | "application")) + { + return Ok(None); + } + let mut bytes = decode_prefix(data, DEFAULT_SNIFF_BYTES)?; + let mut offset = 0; + if bytes.len() > 10 && bytes.starts_with(b"ID3") { + bytes = decode_prefix(data, MAX_ID3_TAG_BYTES + MAX_SIGNATURE_BYTES)?; + offset = 10 + + ((usize::from(bytes[6] & 0x7f) << 21) + | (usize::from(bytes[7] & 0x7f) << 14) + | (usize::from(bytes[8] & 0x7f) << 7) + | usize::from(bytes[9] & 0x7f)); + } + let bytes = &bytes[offset.min(bytes.len())..]; + Ok(SIGNATURES.iter().find_map(|&(media_type, prefix)| { + let eligible = match top_level_type { + Some(kind) => get_top_level_media_type(media_type) == kind, + None => media_type != "audio/mp4", + }; + (eligible + && bytes.len() >= prefix.len() + && prefix + .iter() + .zip(bytes) + .all(|(&expected, &actual)| expected == -1 || expected == i16::from(actual))) + .then_some(media_type) + })) +} + +/// Return the segment preceding the first slash. +#[must_use] +pub fn get_top_level_media_type(media_type: &str) -> &str { + media_type + .split_once('/') + .map_or(media_type, |(kind, _)| kind) +} + +/// Whether a media type has a nonempty subtype other than `*`. +#[must_use] +pub fn is_full_media_type(media_type: &str) -> bool { + media_type + .split_once('/') + .is_some_and(|(_, subtype)| !subtype.is_empty() && subtype != "*") +} + +/// Resolve an inline file's partial media type, or preserve its full type. +/// +/// # Errors +/// Returns an error for malformed base64, or partial media types +/// that cannot be detected or whose data is not inline. +pub fn resolve_full_media_type(part: &FilePart) -> Result { + let media_type = &part.media_type; + let data = match &part.data { + FileData::Data { + data: FileBytes::Binary(bytes), + } => Some(MediaTypeData::Bytes(bytes)), + FileData::Data { + data: FileBytes::Base64(base64), + } => Some(MediaTypeData::Base64(base64)), + _ => None, + }; + if is_full_media_type(media_type) { + return Ok(media_type.clone()); + } + 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-provider-utils/src/read_response_with_size_limit.rs b/aimux-provider-utils/src/read_response_with_size_limit.rs index bd5d1934a..794accf6f 100644 --- a/aimux-provider-utils/src/read_response_with_size_limit.rs +++ b/aimux-provider-utils/src/read_response_with_size_limit.rs @@ -5,6 +5,8 @@ use futures::StreamExt; use aimux_core::{AiMuxError, ApiCallError}; +use crate::fetch::FetchResponse; + /// Default maximum buffered response size for binary downloads (2 GiB, /// matching AI SDK). pub const DEFAULT_MAX_DOWNLOAD_SIZE: usize = 2 * 1024 * 1024 * 1024; @@ -15,7 +17,7 @@ pub const DEFAULT_MAX_DOWNLOAD_SIZE: usize = 2 * 1024 * 1024 * 1024; /// `serde_json::Value`, and a deserialized struct — a 2 GiB cap lets a single /// response balloon to several times that in resident memory. This bound is /// per-request configurable via `HttpRequest::max_json_response_bytes`. -pub const DEFAULT_MAX_JSON_RESPONSE_SIZE: usize = 64 * 1024 * 1024; +pub(crate) const DEFAULT_MAX_JSON_RESPONSE_SIZE: usize = 64 * 1024 * 1024; /// Read a response incrementally and fail before unbounded allocation. /// @@ -24,7 +26,7 @@ pub const DEFAULT_MAX_JSON_RESPONSE_SIZE: usize = 64 * 1024 * 1024; /// Returns an API-call error when the body exceeds `max_bytes` or cannot be /// read, and an abort error when the caller cancels. pub async fn read_response_with_size_limit( - response: reqwest::Response, + response: FetchResponse, url: &str, request_body_values: &serde_json::Value, max_bytes: usize, @@ -107,7 +109,7 @@ pub async fn read_response_with_size_limit( /// Returns an abort error when the caller cancels; never fails on read or /// size problems. pub(crate) async fn read_error_body_truncated( - response: reqwest::Response, + response: FetchResponse, max_bytes: usize, abort_signal: Option<&aimux_core::AbortSignal>, ) -> Result<(Vec, bool), AiMuxError> { diff --git a/aimux-provider-utils/src/resolvable.rs b/aimux-provider-utils/src/resolvable.rs new file mode 100644 index 000000000..8d2cee627 --- /dev/null +++ b/aimux-provider-utils/src/resolvable.rs @@ -0,0 +1,157 @@ +//! A value that may be supplied directly, computed on demand, or awaited. +//! +//! The AI SDK's provider settings accept `string | (() => string | Promise)` +//! for API keys and headers (`Resolvable` in `@ai-sdk/provider-utils`). +//! This is the Rust shape of that union. + +use std::fmt; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::{BoxFuture, Shared}; + +use aimux_core::AiMuxError; + +/// A value of `T`, or a way to obtain one when a request is made. +/// +/// `Future` and `AsyncFn` differ in *when* they compute: a `Future` is +/// evaluated at most once and every [`resolve`](Self::resolve) shares the +/// cached outcome (a promise value), while an `AsyncFn` / `Fn` runs again on +/// every resolve (a token refresher, a credential chain). +pub enum Resolvable { + /// A plain value. + Value(T), + /// A synchronous, fallible producer, called on every resolve. + Fn(Arc Result + Send + Sync>), + /// An asynchronous producer, called on every resolve. + AsyncFn(Arc BoxFuture<'static, Result> + Send + Sync>), + /// A one-shot future whose outcome (success or failure) is cached. + Future(Shared>>), +} + +impl Resolvable { + /// Obtain the value. + /// + /// # Errors + /// + /// Returns whatever error the producer yields. + pub async fn resolve(&self) -> Result { + match self { + Self::Value(value) => Ok(value.clone()), + Self::Fn(produce) => produce(), + Self::AsyncFn(produce) => produce().await, + Self::Future(future) => future.clone().await, + } + } + + /// A synchronous producer evaluated on every resolve. + #[must_use] + pub fn from_fn(produce: impl Fn() -> Result + Send + Sync + 'static) -> Self { + Self::Fn(Arc::new(produce)) + } + + /// An asynchronous producer evaluated on every resolve. + #[must_use] + pub fn from_async_fn(produce: impl Fn() -> Fut + Send + Sync + 'static) -> Self + where + Fut: std::future::Future> + Send + 'static, + { + Self::AsyncFn(Arc::new(move || produce().boxed())) + } + + /// A future evaluated at most once, its outcome shared by every resolve. + #[must_use] + pub fn from_future(future: Fut) -> Self + where + Fut: std::future::Future> + Send + 'static, + { + Self::Future(future.boxed().shared()) + } +} + +impl From for Resolvable { + fn from(value: T) -> Self { + Self::Value(value) + } +} + +impl Clone for Resolvable { + fn clone(&self) -> Self { + match self { + Self::Value(value) => Self::Value(value.clone()), + Self::Fn(produce) => Self::Fn(produce.clone()), + Self::AsyncFn(produce) => Self::AsyncFn(produce.clone()), + Self::Future(future) => Self::Future(future.clone()), + } + } +} + +impl fmt::Debug for Resolvable { + /// Never prints the value: a `Resolvable` commonly carries credentials. + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Value(_) => "Resolvable::Value(..)", + Self::Fn(_) => "Resolvable::Fn(..)", + Self::AsyncFn(_) => "Resolvable::AsyncFn(..)", + Self::Future(_) => "Resolvable::Future(..)", + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + #[tokio::test] + async fn value_and_fn_resolve() { + assert_eq!( + Resolvable::from("key".to_string()).resolve().await.unwrap(), + "key" + ); + let value = Resolvable::from_fn(|| Ok::<_, AiMuxError>(7_u32)); + assert_eq!(value.resolve().await.unwrap(), 7); + } + + #[tokio::test] + async fn async_fn_recomputes_on_every_resolve() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let value = Resolvable::from_async_fn(move || { + let counter = counter.clone(); + async move { Ok::<_, AiMuxError>(counter.fetch_add(1, Ordering::SeqCst)) } + }); + assert_eq!(value.resolve().await.unwrap(), 0); + assert_eq!(value.resolve().await.unwrap(), 1); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn future_is_evaluated_once_and_shared_across_clones() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let value = Resolvable::from_future(async move { + Ok::<_, AiMuxError>(counter.fetch_add(1, Ordering::SeqCst)) + }); + let clone = value.clone(); + assert_eq!(value.resolve().await.unwrap(), 0); + assert_eq!(value.resolve().await.unwrap(), 0); + assert_eq!(clone.resolve().await.unwrap(), 0); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn errors_propagate_and_a_future_caches_its_failure() { + let value: Resolvable = + Resolvable::from_future(async { Err(AiMuxError::Other("boom".into())) }); + assert!(matches!(value.resolve().await, Err(AiMuxError::Other(_)))); + assert!(matches!(value.resolve().await, Err(AiMuxError::Other(_)))); + } + + #[test] + fn debug_never_prints_the_value() { + let value = Resolvable::from("sk-secret".to_string()); + assert!(!format!("{value:?}").contains("sk-secret")); + } +} diff --git a/aimux-provider-utils/src/response_handler.rs b/aimux-provider-utils/src/response_handler.rs index 82604e08b..a19022bb1 100644 --- a/aimux-provider-utils/src/response_handler.rs +++ b/aimux-provider-utils/src/response_handler.rs @@ -10,6 +10,7 @@ use serde::de::DeserializeOwned; use aimux_core::{AiMuxError, ApiCallError}; use crate::extract_response_headers::extract_response_headers; +use crate::fetch::FetchResponse; use crate::read_response_with_size_limit::{ DEFAULT_MAX_DOWNLOAD_SIZE, DEFAULT_MAX_JSON_RESPONSE_SIZE, read_response_with_size_limit, }; @@ -18,7 +19,7 @@ use crate::read_response_with_size_limit::{ pub struct ResponseHandlerInput { pub url: String, pub request_body_values: serde_json::Value, - pub response: reqwest::Response, + pub response: FetchResponse, pub abort_signal: Option, /// Per-request override of the successful-JSON-body size cap, forwarded /// from `HttpRequest::max_json_response_bytes`. Only @@ -96,7 +97,7 @@ impl ResponseHandler { /// Parse a successful JSON response using the endpoint's response type. /// -/// The body is size-limited to [`DEFAULT_MAX_JSON_RESPONSE_SIZE`] (or +/// The body is size-limited to 64 MiB (or /// `HttpRequest::max_json_response_bytes`, when the caller overrides it) — /// deliberately smaller than the binary-download bound, since a JSON success /// body is deserialized straight into `T` and held alongside the raw bytes. @@ -214,28 +215,6 @@ pub fn invalid_response_api_call( })) } -/// Handler for the `{ "error": { "message", "type" | "code" } }` error shape, -/// which several providers share verbatim. Providers whose error JSON differs -/// keep their own mapping via [`create_json_error_response_handler`]. -#[must_use] -pub fn create_standard_json_error_response_handler() -> ResponseHandler { - create_json_error_response_handler(|data| { - let error = data.get("error").unwrap_or(data); - ProviderErrorParts { - message: error - .get("message") - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_owned(), - provider_code: error - .get("type") - .or_else(|| error.get("code")) - .and_then(serde_json::Value::as_str) - .map(str::to_owned), - } - }) -} - /// Parse an error JSON and map its data to a provider message and code. #[must_use] pub fn create_json_error_response_handler( diff --git a/aimux-provider-utils/src/retry.rs b/aimux-provider-utils/src/retry.rs deleted file mode 100644 index aba3b9821..000000000 --- a/aimux-provider-utils/src/retry.rs +++ /dev/null @@ -1,3 +0,0 @@ -//! Compatibility exports for retry configuration. - -pub use aimux_core::retry::RetryConfig; diff --git a/aimux-provider-utils/src/sigv4_fetch.rs b/aimux-provider-utils/src/sigv4_fetch.rs new file mode 100644 index 000000000..b4037ad11 --- /dev/null +++ b/aimux-provider-utils/src/sigv4_fetch.rs @@ -0,0 +1,742 @@ +//! AWS Signature Version 4 request signing, as a [`Fetch`] decorator. +//! +//! [`SigV4Fetch`] wraps another transport and signs each request's *final* +//! body bytes just before they are sent, so AWS-hosted providers +//! can authenticate through ordinary provider settings (`fetch`) instead of +//! each model re-implementing header construction. +//! +//! The signer is a self-contained implementation using only `sha2` + `hmac` +//! (no AWS SDK dependency). It supports static credentials +//! (`access_key_id` + `secret_access_key` + optional `session_token`), any AWS +//! service name and region, and any method/body. +//! +//! Reference: + +use std::fmt; +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use hmac::{Hmac, Mac}; +use sha2::{Digest, Sha256}; + +use crate::fetch::{ + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, RequestCredentials, +}; +use crate::resolvable::Resolvable; + +/// A set of AWS credentials for SigV4 signing. +/// +/// `Debug` is hand-written so the secret key and session token never reach +/// logs or error output; the access key id and region stay visible because +/// they identify the credentials without granting access. +#[derive(Clone)] +pub struct AwsCredentials { + pub access_key_id: String, + pub secret_access_key: String, + /// Optional temporary session token (STS). + pub session_token: Option, + pub region: String, +} + +impl fmt::Debug for AwsCredentials { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AwsCredentials") + .field("access_key_id", &self.access_key_id) + .field("secret_access_key", &"[REDACTED]") + .field( + "session_token", + &self.session_token.as_ref().map(|_| "[REDACTED]"), + ) + .field("region", &self.region) + .finish() + } +} + +/// The result of signing a request. +#[derive(Debug, Clone)] +struct SignedRequest { + /// All headers to send with the request, including the `Authorization` + /// header and any `X-Amz-*` headers. + pub headers: Vec<(String, String)>, +} + +type HmacSha256 = Hmac; + +/// Sign an HTTP request using AWS SigV4, stamped with the current time. +/// +/// # Arguments +/// - `credentials` - The AWS credentials to sign with. +/// - `service` - The AWS service name (e.g. `"bedrock"`, `"aws-external-anthropic"`). +/// - `method` - HTTP method (e.g. `"POST"`). +/// - `url` - The full request URL. +/// - `body` - The request body as a string (e.g. JSON). +/// - `extra_headers` - Additional headers to include in the request (these are +/// also included in the canonical headers for signing). +#[cfg(test)] +#[must_use] +fn sign_request( + credentials: &AwsCredentials, + service: &str, + method: &str, + url: &str, + body: &str, + extra_headers: &[(String, String)], +) -> SignedRequest { + sign_request_at( + credentials, + service, + method, + url, + body.as_bytes(), + extra_headers, + Utc::now(), + ) +} + +/// Sign an HTTP request with an explicit signing time and a byte body. +#[must_use] +fn sign_request_at( + credentials: &AwsCredentials, + service: &str, + method: &str, + url: &str, + body: &[u8], + extra_headers: &[(String, String)], + now: DateTime, +) -> SignedRequest { + let parsed = + url::Url::parse(url).unwrap_or_else(|_| url::Url::parse("http://localhost").unwrap()); + // The signed host must match the wire host: keep the port whenever the + // URL carries a non-default one (Url::port() is Some only for non-default + // ports). Signing a bare host diverges from the Host header reqwest would + // otherwise so gateways and proxies on non-standard ports reject the + // signature (observed as 502 when an env proxy forwarded a signed + // loopback request whose host lacked the port). + let host = match parsed.port() { + Some(port) => format!("{}:{}", parsed.host_str().unwrap_or(""), port), + None => parsed.host_str().unwrap_or("").to_string(), + }; + let path = canonical_uri(parsed.path()); + let query = canonical_query(parsed.query().unwrap_or("")); + + // Timestamp:yyyyMMddTHHmmssZ and date:yyyyMMdd + let amz_date = now.format("%Y%m%dT%H%M%SZ").to_string(); + let date_stamp = now.format("%Y%m%d").to_string(); + + // Body hash + let body_hash = hex::encode(Sha256::digest(body)); + + // Build canonical headers (sorted, lowercased, trimmed). + // Always include host, x-amz-date, and x-amz-content-sha256. + let mut canonical_headers: Vec<(String, String)> = vec![ + ("host".to_string(), host.to_string()), + ("x-amz-content-sha256".to_string(), body_hash.clone()), + ("x-amz-date".to_string(), amz_date.clone()), + ]; + + if let Some(ref token) = credentials.session_token { + canonical_headers.push(("x-amz-security-token".to_string(), token.clone())); + } + + for (k, v) in extra_headers { + let lower = k.to_lowercase(); + // Don't override the required headers. + if lower != "host" + && lower != "x-amz-date" + && lower != "x-amz-content-sha256" + && lower != "x-amz-security-token" + { + canonical_headers.push((lower, canonical_header_value(v))); + } + } + + canonical_headers.sort_by(|a, b| a.0.cmp(&b.0)); + + let canonical_headers_str: String = canonical_headers + .iter() + .map(|(k, v)| format!("{k}:{v}\n")) + .collect(); + let signed_headers: String = canonical_headers + .iter() + .map(|(k, _)| k.as_str()) + .collect::>() + .join(";"); + + // Canonical request + let canonical_request = format!( + "{}\n{}\n{}\n{}\n{}\n{}", + method.to_uppercase(), + path, + query, + canonical_headers_str, + signed_headers, + body_hash + ); + + let canonical_request_hash = hex::encode(Sha256::digest(canonical_request.as_bytes())); + + // Credential scope + let credential_scope = format!( + "{}/{}/{}/aws4_request", + date_stamp, credentials.region, service + ); + + // String to sign + let string_to_sign = + format!("AWS4-HMAC-SHA256\n{amz_date}\n{credential_scope}\n{canonical_request_hash}"); + + // Signing key: derived through chained HMAC, each step keyed by the + // previous one and starting from "AWS4" + secret. + let k_date = hmac_sha256( + format!("AWS4{}", credentials.secret_access_key).as_bytes(), + date_stamp.as_bytes(), + ); + let k_region = hmac_sha256(&k_date, credentials.region.as_bytes()); + let k_service = hmac_sha256(&k_region, service.as_bytes()); + let k_signing = hmac_sha256(&k_service, b"aws4_request"); + + let signature = hex::encode(hmac_sha256(&k_signing, string_to_sign.as_bytes())); + + let authorization = format!( + "AWS4-HMAC-SHA256 Credential={}/{}, SignedHeaders={}, Signature={}", + credentials.access_key_id, credential_scope, signed_headers, signature + ); + + // Build the final header list. + let mut headers: Vec<(String, String)> = canonical_headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + headers.push(("Authorization".to_string(), authorization)); + + SignedRequest { headers } +} + +/// SigV4 `UriEncode`: every byte except the unreserved set +/// (`A-Z a-z 0-9 - _ . ~`) becomes `%XX` with uppercase hex. +fn uri_encode(input: &[u8]) -> String { + let mut out = String::with_capacity(input.len()); + for &byte in input { + if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b'~') { + out.push(byte as char); + } else { + out.push_str(&format!("%{byte:02X}")); + } + } + out +} + +/// Canonical URI for every service but S3: each segment of the path as sent +/// is URI-encoded again, so `:` signs as `%3A` and an already-escaped `%3A` +/// signs as `%253A`. +fn canonical_uri(path: &str) -> String { + if path.is_empty() { + return "/".to_string(); + } + path.split('/') + .map(|segment| uri_encode(segment.as_bytes())) + .collect::>() + .join("/") +} + +/// Canonical query string: each name and value percent-decoded, then +/// URI-encoded, and the pairs sorted by name and then value. +fn canonical_query(query: &str) -> String { + let mut pairs: Vec<(String, String)> = query + .split('&') + .filter(|pair| !pair.is_empty()) + .map(|pair| { + let (name, value) = pair.split_once('=').unwrap_or((pair, "")); + ( + uri_encode(&percent_decode(name)), + uri_encode(&percent_decode(value)), + ) + }) + .collect(); + pairs.sort(); + pairs + .iter() + .map(|(name, value)| format!("{name}={value}")) + .collect::>() + .join("&") +} + +fn percent_decode(input: &str) -> Vec { + let bytes = input.as_bytes(); + let mut out = Vec::with_capacity(bytes.len()); + let mut i = 0; + while i < bytes.len() { + let escaped = (bytes[i] == b'%') + .then(|| bytes.get(i + 1..i + 3)) + .flatten() + .and_then(|hex| std::str::from_utf8(hex).ok()) + .and_then(|hex| u8::from_str_radix(hex, 16).ok()); + match escaped { + Some(byte) => { + out.push(byte); + i += 3; + } + None => { + out.push(bytes[i]); + i += 1; + } + } + } + out +} + +/// Canonical header value: trimmed, with runs of spaces collapsed to one. +fn canonical_header_value(value: &str) -> String { + value.split_whitespace().collect::>().join(" ") +} + +fn hmac_sha256(key: &[u8], data: &[u8]) -> Vec { + let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length"); + mac.update(data); + mac.finalize().into_bytes().to_vec() +} + +type Clock = Arc DateTime + Send + Sync>; + +/// A [`Fetch`] decorator that signs every request with SigV4 and forwards it +/// to the wrapped transport. A request marked [`RequestCredentials::Omit`] is +/// forwarded unsigned. +/// +/// Credentials are resolved on each request, so a `Resolvable::AsyncFn` can +/// hand out rotating STS credentials. The signature covers the request's +/// final body bytes and every header present on the request (except +/// `Authorization`, which it replaces); `host`, `x-amz-date`, +/// `x-amz-content-sha256` and, for temporary credentials, +/// `x-amz-security-token` are set from the signed values. +pub struct SigV4Fetch { + inner: FetchFunction, + credentials: Resolvable, + service: String, + clock: Clock, +} + +impl SigV4Fetch { + /// Sign requests for `service` (e.g. `"bedrock"`) with `credentials` + /// (whose `region` scopes the signature) and forward them to `inner`. + #[must_use] + pub fn new( + inner: FetchFunction, + credentials: Resolvable, + service: impl Into, + ) -> Self { + Self { + inner, + credentials, + service: service.into(), + clock: Arc::new(Utc::now), + } + } + + /// Replace the signing clock (deterministic signatures in tests). + #[must_use] + pub fn with_clock(mut self, clock: impl Fn() -> DateTime + Send + Sync + 'static) -> Self { + self.clock = Arc::new(clock); + self + } +} + +impl fmt::Debug for SigV4Fetch { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SigV4Fetch") + .field("service", &self.service) + .finish_non_exhaustive() + } +} + +#[async_trait] +impl Fetch for SigV4Fetch { + async fn fetch(&self, mut request: FetchRequest) -> Result { + // A redirect hop off the credentialed origin: neither the signature + // nor `x-amz-security-token` may reach that host. + if request.credentials == RequestCredentials::Omit { + return self.inner.fetch(request).await; + } + let credentials = self.credentials.resolve().await.map_err(|error| { + FetchError::Other(format!("failed to resolve AWS credentials: {error}")) + })?; + let mut extra_headers = Vec::with_capacity(request.headers.len()); + for (name, value) in &request.headers { + if name == http::header::AUTHORIZATION { + continue; + } + let value = value.to_str().map_err(|_| { + FetchError::Other(format!("cannot sign non-ASCII header value for {name}")) + })?; + extra_headers.push((name.as_str().to_string(), value.to_string())); + } + let signed = sign_request_at( + &credentials, + &self.service, + request.method.as_str(), + request.url.as_str(), + &request.body, + &extra_headers, + (self.clock)(), + ); + for (name, value) in signed.headers { + let name = http::HeaderName::try_from(name.as_str()).map_err(|error| { + FetchError::Other(format!("invalid signed header name {name}: {error}")) + })?; + let value = http::HeaderValue::try_from(value.as_str()).map_err(|error| { + FetchError::Other(format!("invalid signed header value for {name}: {error}")) + })?; + request.headers.insert(name, value); + } + self.inner.fetch(request).await + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use bytes::Bytes; + use chrono::TimeZone; + use http::{HeaderMap, Method, StatusCode}; + + use super::*; + + fn creds() -> AwsCredentials { + AwsCredentials { + access_key_id: "AKIAIOSFODNN7EXAMPLE".to_string(), + secret_access_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(), + session_token: None, + region: "us-east-1".to_string(), + } + } + + #[test] + fn debug_redacts_secret_key_and_session_token() { + let mut credentials = creds(); + credentials.session_token = Some("FwoGZXIvYXdzEXAMPLETOKEN".to_string()); + let printed = format!("{credentials:?}"); + assert!(!printed.contains("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY")); + assert!(!printed.contains("FwoGZXIvYXdzEXAMPLETOKEN")); + assert!(printed.contains("[REDACTED]")); + assert!(printed.contains("AKIAIOSFODNN7EXAMPLE")); + assert!(printed.contains("us-east-1")); + } + + fn fixed_time() -> DateTime { + Utc.with_ymd_and_hms(2024, 1, 2, 3, 4, 5).unwrap() + } + + fn header(signed: &SignedRequest, name: &str) -> String { + signed + .headers + .iter() + .find(|(k, _)| k.eq_ignore_ascii_case(name)) + .map(|(_, v)| v.clone()) + .unwrap_or_else(|| panic!("signed request carries a {name} header")) + } + + /// Records the (already signed) request the decorator forwards. + #[derive(Default)] + struct Capture(Mutex>); + + #[async_trait] + impl Fetch for Capture { + async fn fetch(&self, request: FetchRequest) -> Result { + let url = request.url.clone(); + *self.0.lock().unwrap() = Some(request); + Ok(FetchResponse::from_bytes( + StatusCode::OK, + HeaderMap::new(), + url, + Bytes::new(), + )) + } + } + + fn decorated(credentials: AwsCredentials) -> (SigV4Fetch, Arc) { + let capture = Arc::new(Capture::default()); + let fetch = SigV4Fetch::new(capture.clone(), Resolvable::from(credentials), "bedrock") + .with_clock(fixed_time); + (fetch, capture) + } + + const URL: &str = "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-5-sonnet-20240620-v1:0/converse"; + const BODY: &str = r#"{"messages":[]}"#; + + #[test] + fn host_keeps_non_default_port() { + let signed = sign_request( + &creds(), + "bedrock", + "POST", + "http://127.0.0.1:54321/model/x/invoke", + "{}", + &[], + ); + assert_eq!(header(&signed, "host"), "127.0.0.1:54321"); + } + + #[test] + fn host_omits_default_port() { + let signed = sign_request( + &creds(), + "bedrock", + "POST", + "https://bedrock-runtime.us-east-1.amazonaws.com/model/x/invoke", + "{}", + &[], + ); + assert_eq!( + header(&signed, "host"), + "bedrock-runtime.us-east-1.amazonaws.com" + ); + } + + fn signature(signed: &SignedRequest) -> String { + header(signed, "authorization") + .rsplit_once("Signature=") + .expect("authorization carries a signature") + .1 + .to_string() + } + + /// Expected values come from botocore 1.43's `SigV4Auth` (AWS's own + /// Python signer) for the same credentials, headers, body and fixed + /// time, not from this implementation. + #[test] + fn signatures_match_botocore() { + let signed = sign_request_at( + &creds(), + "bedrock", + "POST", + URL, + BODY.as_bytes(), + &[("X-Custom".to_string(), " value ".to_string())], + fixed_time(), + ); + assert_eq!( + header(&signed, "authorization"), + "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/20240102/us-east-1/bedrock/aws4_request, \ + SignedHeaders=host;x-amz-content-sha256;x-amz-date;x-custom, \ + Signature=538444fb4b17f3cdd777bbd5f1cd4e3e1f53d6ac97167da96c81ec8d14d6468a" + ); + + let session = AwsCredentials { + session_token: Some("SESSION".to_string()), + ..creds() + }; + let signed = sign_request_at( + &session, + "bedrock", + "POST", + "http://127.0.0.1:54321/model/x/invoke?a=1", + b"{}", + &[], + fixed_time(), + ); + assert_eq!( + header(&signed, "authorization"), + "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/20240102/us-east-1/bedrock/aws4_request, \ + SignedHeaders=host;x-amz-content-sha256;x-amz-date;x-amz-security-token, \ + Signature=dd6a2384cb9c8f0db8ea8ee383a61bf06dd85da679cec54d5c944b1b0cb8d9f6" + ); + assert_eq!(header(&signed, "x-amz-security-token"), "SESSION"); + } + + /// An escaped path segment is encoded again, header whitespace collapses, + /// and query pairs are sorted (same botocore reference as above). + #[test] + fn canonicalization_matches_botocore() { + let signed = sign_request_at( + &creds(), + "bedrock", + "POST", + "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-3-5-sonnet-20240620-v1%3A0/converse-stream", + b"{}", + &[ + ("content-type".to_string(), "application/json".to_string()), + ("X-Multi".to_string(), "a b c".to_string()), + ], + fixed_time(), + ); + assert_eq!( + signature(&signed), + "a6abc7f26ad04c0922b03e05c9184a20a8fcf4b68b7dfce358e9a6436c8a9489" + ); + + let signed = sign_request_at( + &creds(), + "bedrock", + "POST", + "https://polly.us-east-1.amazonaws.com/v1/speech?b=2&a=x%20y&a=1", + b"{}", + &[], + fixed_time(), + ); + assert_eq!( + signature(&signed), + "70c6e5794c4f1456f33f2b41cebc483ee0d19c98b95660411f4b51ccac464006" + ); + + let signed = sign_request_at( + &creds(), + "bedrock", + "GET", + "https://h.example.com", + b"", + &[], + fixed_time(), + ); + assert_eq!( + signature(&signed), + "4810ddfb5f22377b8c39477909c5abf54ea3c7e03f9db0f8ccb9377ed39ef26d" + ); + } + + #[test] + fn canonical_uri_and_query_follow_sigv4() { + assert_eq!(canonical_uri(""), "/"); + assert_eq!( + canonical_uri("/model/a.b-v1:0/converse"), + "/model/a.b-v1%3A0/converse" + ); + assert_eq!(canonical_uri("/a+b/c%2Fd"), "/a%2Bb/c%252Fd"); + assert_eq!( + canonical_query("b=2&a=x%20y&a=1&e=&z"), + "a=1&a=x%20y&b=2&e=&z=" + ); + } + + #[tokio::test] + async fn decorator_authorization_is_byte_identical_to_sign_request() { + let (fetch, capture) = decorated(creds()); + let mut request = FetchRequest::new(Method::POST, URL.parse().unwrap()); + request + .headers + .insert("x-custom", " value ".parse().unwrap()); + request.body = Bytes::from(BODY); + + fetch.fetch(request).await.unwrap(); + + let sent = capture + .0 + .lock() + .unwrap() + .take() + .expect("inner fetch was called"); + let expected = sign_request_at( + &creds(), + "bedrock", + "POST", + URL, + BODY.as_bytes(), + &[("x-custom".to_string(), " value ".to_string())], + fixed_time(), + ); + assert_eq!( + sent.headers["authorization"].to_str().unwrap(), + header(&expected, "authorization") + ); + for name in ["host", "x-amz-date", "x-amz-content-sha256"] { + assert_eq!( + sent.headers[name].to_str().unwrap(), + header(&expected, name), + "{name}" + ); + } + assert_eq!(sent.body, Bytes::from(BODY)); + assert!(!sent.headers.contains_key("x-amz-security-token")); + } + + #[tokio::test] + async fn decorator_replaces_a_caller_authorization_and_signs_the_session_token() { + let session = AwsCredentials { + session_token: Some("SESSION".to_string()), + ..creds() + }; + let (fetch, capture) = decorated(session); + let mut request = FetchRequest::new(Method::POST, URL.parse().unwrap()); + request + .headers + .insert(http::header::AUTHORIZATION, "Bearer stale".parse().unwrap()); + request.body = Bytes::from(BODY); + + fetch.fetch(request).await.unwrap(); + + let sent = capture.0.lock().unwrap().take().unwrap(); + let authorization = sent.headers["authorization"].to_str().unwrap(); + assert!(authorization.starts_with("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/")); + assert!(authorization.contains("x-amz-security-token")); + assert_eq!(sent.headers.get_all("authorization").iter().count(), 1); + assert_eq!(sent.headers["x-amz-security-token"], "SESSION"); + } + + #[tokio::test] + async fn a_request_that_omits_credentials_is_forwarded_unsigned() { + let session = AwsCredentials { + session_token: Some("SESSION".to_string()), + ..creds() + }; + let (fetch, capture) = decorated(session); + let mut request = FetchRequest::new(Method::POST, URL.parse().unwrap()); + request.body = Bytes::from(BODY); + request.credentials = RequestCredentials::Omit; + + fetch.fetch(request).await.unwrap(); + + let sent = capture.0.lock().unwrap().take().unwrap(); + for name in [ + "authorization", + "x-amz-date", + "x-amz-content-sha256", + "x-amz-security-token", + ] { + assert!(!sent.headers.contains_key(name), "{name}"); + } + assert_eq!(sent.body, Bytes::from(BODY)); + } + + #[tokio::test] + async fn credentials_are_resolved_per_request() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let counter = calls.clone(); + let credentials = Resolvable::from_async_fn(move || { + let counter = counter.clone(); + async move { + counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(creds()) + } + }); + let capture = Arc::new(Capture::default()); + let fetch = SigV4Fetch::new(capture, credentials, "bedrock"); + for _ in 0..2 { + fetch + .fetch(FetchRequest::new(Method::GET, URL.parse().unwrap())) + .await + .unwrap(); + } + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn credential_failures_surface_without_calling_the_inner_fetch() { + let capture = Arc::new(Capture::default()); + let fetch = SigV4Fetch::new( + capture.clone(), + Resolvable::from_fn(|| { + Err(aimux_core::AiMuxError::LoadSetting { + env_var: "AWS_ACCESS_KEY_ID".into(), + name: "accessKeyId".into(), + }) + }), + "bedrock", + ); + let error = fetch + .fetch(FetchRequest::new(Method::GET, URL.parse().unwrap())) + .await + .unwrap_err(); + assert!(matches!(&error, FetchError::Other(m) if m.contains("AWS_ACCESS_KEY_ID"))); + assert!(capture.0.lock().unwrap().is_none()); + } +} diff --git a/aimux-provider-utils/src/ws.rs b/aimux-provider-utils/src/ws.rs index 2ad1c6322..8d4c11031 100644 --- a/aimux-provider-utils/src/ws.rs +++ b/aimux-provider-utils/src/ws.rs @@ -13,17 +13,20 @@ //! the loop alone does not cover the send path). //! - **Backpressure is socket-level**: tungstenite's `send().await` drives //! flush and pends while the socket write buffer is full. -//! - **Proxy support (RFC-0034 §2)**: the global `ProxyConfig` that governs -//! HTTP also governs `ws://`/`wss://` — `wss` uses `https_url` (or -//! `all_url`), `ws` uses `http_url` (or `all_url`), `no_proxy` entries are -//! honored, and tunneled connections are established with a manual HTTP -//! CONNECT before the WebSocket handshake. tokio-tungstenite has no proxy +//! - **Proxy support (RFC-0034 §2)**: a [`TungsteniteConnector`](crate::ws::TungsteniteConnector) carries its +//! own `ProxyConfig`; the built-in default shares the default HTTP +//! transport's. `wss` uses `https_url` (or `all_url`), `ws` uses `http_url` +//! (or `all_url`), `no_proxy` entries are honored, and tunneled +//! connections are established with a manual HTTP CONNECT before the +//! WebSocket handshake. tokio-tungstenite has no proxy //! support (and will not add it), so the tunnel is ours. SOCKS proxies are //! rejected loudly rather than silently bypassed. use std::future::pending; +use std::sync::Arc; -use futures::{SinkExt, StreamExt}; +use async_trait::async_trait; +use futures::{Sink, SinkExt, Stream, StreamExt}; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; @@ -34,6 +37,7 @@ use aimux_core::error::{AiMuxError, ApiCallError}; use aimux_core::options::TimeoutConfiguration; /// A request to open a WebSocket connection. +#[derive(Debug)] pub struct WebSocketRequest { /// `wss://` or `ws://` URL. pub url: String, @@ -48,6 +52,49 @@ pub struct WebSocketRequest { /// session establishment, `chunk_ms` bounds the gap between events, /// `total_ms` bounds the whole connection lifetime. pub timeout: Option, + /// Transport that performs the handshake. `None` uses a + /// [`TungsteniteConnector`] with the default HTTP transport's proxy + /// settings, read at connect time. Injection point for tests, for a + /// connector with its own proxy, and for providers that bring their own + /// socket. + pub connector: Option>, +} + +/// 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`]. +#[async_trait] +pub trait WsConnector: Send + Sync + 'static { + /// Establish the connection described by `request`. + /// + /// # Errors + /// + /// Returns the connect failure, abort or timeout. + async fn connect(&self, request: &WebSocketRequest) -> Result; +} + +impl std::fmt::Debug for dyn WsConnector { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("dyn WsConnector") + } +} + +/// The message transport behind a [`WsConnection`]: any tungstenite-message +/// stream/sink pair. +pub trait WsTransport: + Stream> + + Sink + + Send + + Unpin +{ +} + +impl WsTransport for T where + T: Stream> + + Sink + + Send + + Unpin +{ } /// An inbound WebSocket message. @@ -59,7 +106,7 @@ pub enum WsMessage { /// A connected WebSocket with abort/timeout enforcement built in. pub struct WsConnection { - stream: WebSocketStream>, + stream: Box, url: String, abort: Option, /// Deadline for the first event (connect + session ack). Cleared after @@ -160,13 +207,13 @@ pub struct ProxyTunnel { pub target_authority: String, } -/// Test-only exposure of [`resolve_proxy`] (integration tests construct -/// arbitrary configs without the set-once global). +/// Test-only exposure of [`resolve_proxy`] (integration tests check proxy +/// selection for arbitrary configs without connecting). #[doc(hidden)] #[allow(non_snake_case)] pub fn ws__proxy_decision_for( target: &url::Url, - config: &crate::http::ProxyConfig, + config: &crate::fetch::ProxyConfig, ) -> Result, AiMuxError> { resolve_proxy(target, config) } @@ -192,10 +239,10 @@ fn sanitized_proxy_url(raw: &str) -> String { raw.to_string() } -/// Pick direct vs. tunneled for a WS target under the global proxy config. +/// Pick direct vs. tunneled for a WS target under a proxy config. fn resolve_proxy( target: &url::Url, - config: &crate::http::ProxyConfig, + config: &crate::fetch::ProxyConfig, ) -> Result, AiMuxError> { let raw = match target.scheme() { "wss" => config.https_url.clone().or_else(|| config.all_url.clone()), @@ -398,8 +445,9 @@ async fn connect_through_proxy( Ok((websocket, response)) } -/// Open a WebSocket connection. The connect phase races abort and the -/// `first_chunk_ms` timeout (which doubles as the connect timeout). +/// Open a WebSocket connection through the request's [`WsConnector`], or the +/// built-in one when none is set. The built-in connect phase races abort and +/// the `first_chunk_ms` timeout (which doubles as the connect timeout). /// /// # Errors /// @@ -407,6 +455,42 @@ async fn connect_through_proxy( /// cancellation wins the connect race, and `Timeout` when `first_chunk_ms` /// expires while connecting. pub async fn ws_connect(req: &WebSocketRequest) -> Result { + match &req.connector { + Some(connector) => connector.connect(req).await, + None => { + let proxy = crate::fetch::default_transport().proxy().clone(); + TungsteniteConnector::new(proxy).connect(req).await + } + } +} + +/// The built-in connector: tungstenite, tunneled through its own proxy +/// settings (the WebSocket counterpart of a +/// [`ReqwestFetch`](crate::fetch::ReqwestFetch)). +#[derive(Debug, Clone, Default)] +pub struct TungsteniteConnector { + proxy: crate::fetch::ProxyConfig, +} + +impl TungsteniteConnector { + /// A connector routed by `proxy`. + #[must_use] + pub fn new(proxy: crate::fetch::ProxyConfig) -> Self { + Self { proxy } + } +} + +#[async_trait] +impl WsConnector for TungsteniteConnector { + async fn connect(&self, request: &WebSocketRequest) -> Result { + connect_tungstenite(request, &self.proxy).await + } +} + +async fn connect_tungstenite( + req: &WebSocketRequest, + config: &crate::fetch::ProxyConfig, +) -> Result { let mut http_req = req .url .as_str() @@ -441,12 +525,11 @@ pub async fn ws_connect(req: &WebSocketRequest) -> Result Result Result, req: &WebSocketRequest) -> Self { + let first_chunk_deadline = req + .timeout + .as_ref() + .and_then(|t| t.first_chunk_ms) + .map(|ms| tokio::time::Instant::now() + tokio::time::Duration::from_millis(ms)); + Self::with_deadline(transport, req, first_chunk_deadline) + } + + fn with_deadline( + stream: Box, + req: &WebSocketRequest, + first_chunk_deadline: Option, + ) -> Self { + Self { + stream, + url: req.url.clone(), + abort: req.abort_signal.clone(), + first_chunk_deadline, + chunk_timeout: req + .timeout + .as_ref() + .and_then(|t| t.chunk_ms) + .map(tokio::time::Duration::from_millis), + total_deadline: req + .timeout + .as_ref() + .and_then(|t| t.total_ms) + .map(|ms| tokio::time::Instant::now() + tokio::time::Duration::from_millis(ms)), + } + } + /// Send a text message. Aborted / timed out sends surface as errors; the /// pending-while-buffer-full behavior is the socket-level backpressure. /// @@ -640,12 +748,14 @@ impl WsConnection { /// a close handshake against a dead peer must not hang the caller (the /// send pends while the socket buffer can't drain). pub async fn close(&mut self) { - let close_fut = self.stream.close(Some( + // What `WebSocketStream::close(Some(frame))` does, expressed through + // the transport trait so injected connectors close the same way. + let close_fut = self.stream.send(Message::Close(Some( tokio_tungstenite::tungstenite::protocol::frame::CloseFrame { code: tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Normal, reason: std::borrow::Cow::Borrowed("finished"), }, - )); + ))); tokio::select! { res = close_fut => { let _ = res; diff --git a/aimux-provider-utils/tests/api_redirect_test.rs b/aimux-provider-utils/tests/api_redirect_test.rs new file mode 100644 index 000000000..5b2cccabc --- /dev/null +++ b/aimux-provider-utils/tests/api_redirect_test.rs @@ -0,0 +1,260 @@ +//! Redirects on an API request (`validate_url: false`): every hop is +//! followed, as the AI SDK's `postToApi` does under fetch's default +//! `redirect: "follow"`, and a hop off the request's origin carries only +//! `User-Agent`, as `fetchWithValidatedRedirects` does. + +use std::sync::Arc; + +use aimux_core::AiMuxError; +use aimux_provider_utils::{ + AwsCredentials, HttpRequest, ProviderErrorParts, Resolvable, SigV4Fetch, + create_json_error_response_handler, create_json_response_handler, default_fetch, + post_json_to_api, +}; +use chrono::{TimeZone, Utc}; +use serde::Deserialize; +use serde_json::json; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, Request, ResponseTemplate}; + +#[derive(Deserialize)] +struct Reply { + #[allow(dead_code)] + ok: bool, +} + +fn failed() -> aimux_provider_utils::ResponseHandler { + create_json_error_response_handler(|_| ProviderErrorParts { + message: "request failed".into(), + provider_code: None, + }) +} + +fn credentialed(url: String) -> HttpRequest { + HttpRequest { + url, + headers: vec![ + ("authorization".into(), "Bearer secret".into()), + ("x-api-key".into(), "secret-key".into()), + ("user-agent".into(), "caller-agent".into()), + ], + ..Default::default() + } +} + +async fn post(request: HttpRequest) { + post_json_to_api( + request, + json!({"prompt": "hi"}), + create_json_response_handler::(), + failed(), + ) + .await + .unwrap(); +} + +async fn only_request(server: &MockServer) -> Request { + let mut received = server.received_requests().await.unwrap(); + assert_eq!(received.len(), 1); + received.remove(0) +} + +fn has(request: &Request, name: &str) -> bool { + request.headers.contains_key(name) +} + +async fn redirect(server: &MockServer, from: &str, to: String) { + Mock::given(path(from.to_string())) + .respond_with(ResponseTemplate::new(307).insert_header("location", to.as_str())) + .mount(server) + .await; +} + +async fn ok(server: &MockServer, at: &str) { + Mock::given(path(at.to_string())) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true}))) + .mount(server) + .await; +} + +#[tokio::test] +async fn a_cross_origin_redirect_is_followed_without_caller_headers() { + let provider = MockServer::start().await; + let foreign = MockServer::start().await; + redirect(&provider, "/call", format!("{}/moved", foreign.uri())).await; + ok(&foreign, "/moved").await; + + post(credentialed(format!("{}/call", provider.uri()))).await; + + let first = only_request(&provider).await; + assert!(has(&first, "authorization") && has(&first, "x-api-key")); + let hop = only_request(&foreign).await; + assert_eq!(hop.method.as_str(), "POST", "307 keeps the method"); + assert!(!has(&hop, "authorization")); + assert!(!has(&hop, "x-api-key")); + let agent = hop.headers["user-agent"].to_str().unwrap(); + assert!(agent.starts_with("caller-agent"), "{agent}"); +} + +#[tokio::test] +async fn a_same_origin_redirect_keeps_caller_headers() { + let provider = MockServer::start().await; + redirect(&provider, "/call", "/moved".into()).await; + ok(&provider, "/moved").await; + + post(credentialed(format!("{}/call", provider.uri()))).await; + + let received = provider.received_requests().await.unwrap(); + assert_eq!(received.len(), 2); + assert!(has(&received[1], "authorization") && has(&received[1], "x-api-key")); +} + +#[tokio::test] +async fn headers_dropped_off_origin_do_not_return_on_a_hop_back() { + let provider = MockServer::start().await; + let foreign = MockServer::start().await; + redirect(&provider, "/call", format!("{}/bounce", foreign.uri())).await; + redirect(&foreign, "/bounce", format!("{}/back", provider.uri())).await; + ok(&provider, "/back").await; + + post(credentialed(format!("{}/call", provider.uri()))).await; + + let received = provider.received_requests().await.unwrap(); + assert_eq!(received.len(), 2); + assert_eq!(received[1].url.path(), "/back"); + assert!(!has(&received[1], "authorization")); + assert!(!has(&received[1], "x-api-key")); +} + +#[tokio::test] +async fn a_signing_transport_does_not_sign_a_cross_origin_hop() { + let provider = MockServer::start().await; + let foreign = MockServer::start().await; + redirect(&provider, "/call", format!("{}/moved", foreign.uri())).await; + ok(&foreign, "/moved").await; + let signing = SigV4Fetch::new( + default_fetch(), + Resolvable::from(AwsCredentials { + access_key_id: "AKIDEXAMPLE".into(), + secret_access_key: "secret".into(), + session_token: Some("SESSION".into()), + region: "us-east-1".into(), + }), + "bedrock", + ); + + post(HttpRequest { + url: format!("{}/call", provider.uri()), + fetch: Some(Arc::new(signing)), + ..Default::default() + }) + .await; + + let first = only_request(&provider).await; + assert!(has(&first, "authorization") && has(&first, "x-amz-security-token")); + let hop = only_request(&foreign).await; + for name in ["authorization", "x-amz-security-token", "x-amz-date"] { + assert!(!has(&hop, name), "{name}"); + } +} + +#[tokio::test] +async fn a_303_turns_the_hop_into_a_get() { + let provider = MockServer::start().await; + let foreign = MockServer::start().await; + Mock::given(path("/call")) + .respond_with( + ResponseTemplate::new(303) + .insert_header("location", format!("{}/moved", foreign.uri())), + ) + .mount(&provider) + .await; + Mock::given(method("GET")) + .and(path("/moved")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true}))) + .mount(&foreign) + .await; + + post(credentialed(format!("{}/call", provider.uri()))).await; + + assert_eq!(only_request(&foreign).await.method.as_str(), "GET"); +} + +/// `MAX_REDIRECTS` (http.rs:28): the request and ten followed hops go out; +/// the eleventh redirect fails the call with the 3xx status (http.rs:488-489). +/// The fetch spec allows 20; aimux stops at 10 like the AI SDK's +/// `MAX_DOWNLOAD_REDIRECTS` (provider-utils@5.0.55 +/// `fetch-with-validated-redirects.ts:10`). +#[tokio::test] +async fn an_eleventh_redirect_fails_the_call() { + let provider = MockServer::start().await; + Mock::given(path("/loop")) + .respond_with(ResponseTemplate::new(307).insert_header("location", "/loop")) + .mount(&provider) + .await; + + let result = post_json_to_api( + credentialed(format!("{}/loop", provider.uri())), + json!({"prompt": "hi"}), + create_json_response_handler::(), + failed(), + ) + .await; + + let Err(AiMuxError::ApiCall(error)) = result else { + panic!("a redirect loop must fail with an API call error"); + }; + assert!( + error.message.contains("too many redirects"), + "{}", + error.message + ); + assert_eq!(error.status_code, Some(307)); + assert_eq!(provider.received_requests().await.unwrap().len(), 11); +} + +/// A same-origin hop goes through the signing transport again and is signed +/// for the URL it is sent to (RFC-0040 §3.2: 同 origin 重定向也必须重新签名， +/// 不能复用旧签名). With the clock fixed, the two signatures can differ only +/// because the path differs. +#[tokio::test] +async fn a_same_origin_hop_is_signed_again_for_its_own_url() { + let provider = MockServer::start().await; + redirect(&provider, "/call", "/moved".into()).await; + ok(&provider, "/moved").await; + let signing = SigV4Fetch::new( + default_fetch(), + Resolvable::from(AwsCredentials { + access_key_id: "AKIDEXAMPLE".into(), + secret_access_key: "secret".into(), + session_token: Some("SESSION".into()), + region: "us-east-1".into(), + }), + "bedrock", + ) + .with_clock(|| Utc.with_ymd_and_hms(2026, 10, 15, 0, 0, 0).unwrap()); + + post(HttpRequest { + url: format!("{}/call", provider.uri()), + fetch: Some(Arc::new(signing)), + ..Default::default() + }) + .await; + + let received = provider.received_requests().await.unwrap(); + assert_eq!(received.len(), 2); + assert_eq!(received[1].url.path(), "/moved"); + let header = |request: &Request, name: &str| request.headers[name].to_str().unwrap().to_owned(); + for request in &received { + assert!(has(request, "authorization") && has(request, "x-amz-security-token")); + } + assert_eq!( + header(&received[0], "x-amz-date"), + header(&received[1], "x-amz-date") + ); + assert_ne!( + header(&received[0], "authorization"), + header(&received[1], "authorization"), + "the hop must carry a signature for /moved, not the one for /call" + ); +} diff --git a/aimux-provider-utils/tests/download_guard_test.rs b/aimux-provider-utils/tests/download_guard_test.rs index 52995b512..0dd42b9f0 100644 --- a/aimux-provider-utils/tests/download_guard_test.rs +++ b/aimux-provider-utils/tests/download_guard_test.rs @@ -34,6 +34,7 @@ fn request( validate_url: true, trusted_origin: trusted_origin.map(str::to_owned), credentialed_origin: credentialed_origin.map(str::to_owned), + fetch: None, } } diff --git a/aimux-provider-utils/tests/fetch_injection_test.rs b/aimux-provider-utils/tests/fetch_injection_test.rs new file mode 100644 index 000000000..bd360205f --- /dev/null +++ b/aimux-provider-utils/tests/fetch_injection_test.rs @@ -0,0 +1,108 @@ +//! Validated downloads use the guarded transport instead of an injected fetch. + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use bytes::Bytes; +use http::{HeaderMap, StatusCode}; +use serde::Deserialize; +use serde_json::json; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use aimux_core::AiMuxError; +use aimux_provider_utils::{ + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HttpRequest, ProviderErrorParts, + create_json_error_response_handler, create_json_response_handler, get_from_api, +}; + +#[derive(Debug, Deserialize, PartialEq)] +struct Reply { + answer: u32, +} + +/// Records every request and replies with a canned result. +struct MockFetch { + requests: Mutex>, + reply: Box Result<(StatusCode, String), FetchError> + Send + Sync>, +} + +impl MockFetch { + fn json(status: u16, body: serde_json::Value) -> Arc { + let body = body.to_string(); + Arc::new(Self { + requests: Mutex::new(Vec::new()), + reply: Box::new(move || Ok((StatusCode::from_u16(status).unwrap(), body.clone()))), + }) + } + + fn calls(&self) -> Vec { + self.requests.lock().unwrap().clone() + } +} + +#[async_trait] +impl Fetch for MockFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + let url = request.url.clone(); + self.requests.lock().unwrap().push(request); + let (status, body) = (self.reply)()?; + Ok(FetchResponse::from_bytes( + status, + HeaderMap::new(), + url, + Bytes::from(body), + )) + } +} + +fn failed() -> aimux_provider_utils::ResponseHandler { + create_json_error_response_handler(|value| ProviderErrorParts { + message: value["error"] + .as_str() + .unwrap_or("request failed") + .to_string(), + provider_code: None, + }) +} + +fn request(url: &str, fetch: &Arc) -> HttpRequest { + let fetch: FetchFunction = fetch.clone(); + HttpRequest { + url: url.to_string(), + headers: vec![("x-custom".into(), "visible".into())], + fetch: Some(fetch), + ..Default::default() + } +} + +/// D26: a validated download is never routed through a provider-injected +/// transport — it uses the guard's pinned transport instead. +#[tokio::test] +async fn validate_url_ignores_the_injected_fetch() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/asset")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9}))) + .mount(&server) + .await; + let mock = MockFetch::json(200, json!({"answer": 0})); + let mut http_request = request(&format!("{}/asset", server.uri()), &mock); + http_request.headers = vec![]; + http_request.validate_url = true; + http_request.trusted_origin = Some(server.uri()); + + let output = get_from_api( + http_request, + create_json_response_handler::(), + failed(), + ) + .await + .unwrap(); + + assert_eq!(output.value, Reply { answer: 9 }); + assert!( + mock.calls().is_empty(), + "the injected fetch must not see downloads" + ); +} diff --git a/aimux-provider-utils/tests/live_sigv4_test.rs b/aimux-provider-utils/tests/live_sigv4_test.rs new file mode 100644 index 000000000..32b4c0ef3 --- /dev/null +++ b/aimux-provider-utils/tests/live_sigv4_test.rs @@ -0,0 +1,101 @@ +//! Live SigV4 check against real AWS endpoints. Ignored by default; run with +//! real credentials in `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` (and +//! `AWS_SESSION_TOKEN` for temporary ones): +//! +//! ```sh +//! cargo test -p aimux-provider-utils --test live_sigv4_test -- --ignored --nocapture +//! ``` +//! +//! STS `GetCallerIdentity` needs no IAM permission, so it must succeed. The +//! Bedrock and Polly calls may be denied by IAM or model access, but AWS +//! checks the signature first, so any response other than a signature error +//! means the signature was accepted. + +use aimux_provider_utils::sigv4_fetch::{AwsCredentials, SigV4Fetch}; +use aimux_provider_utils::{Fetch, FetchRequest, Resolvable, default_fetch}; +use futures::StreamExt; +use http::Method; + +const SIGNATURE_ERRORS: &[&str] = &[ + "SignatureDoesNotMatch", + "IncompleteSignature", + "InvalidSignatureException", + "signature we calculated does not match", +]; + +async fn call( + service: &str, + method: Method, + url: &str, + body: &str, + content_type: Option<&str>, +) -> (u16, String) { + let credentials = AwsCredentials { + access_key_id: std::env::var("AWS_ACCESS_KEY_ID").expect("AWS_ACCESS_KEY_ID"), + secret_access_key: std::env::var("AWS_SECRET_ACCESS_KEY").expect("AWS_SECRET_ACCESS_KEY"), + session_token: std::env::var("AWS_SESSION_TOKEN").ok(), + region: "us-east-1".into(), + }; + let fetch = SigV4Fetch::new(default_fetch(), Resolvable::from(credentials), service); + let mut request = FetchRequest::new(method, url.parse().unwrap()); + if let Some(content_type) = content_type { + request + .headers + .insert("content-type", content_type.parse().unwrap()); + } + request.body = bytes::Bytes::from(body.to_string()); + let response = fetch.fetch(request).await.unwrap(); + let status = response.status().as_u16(); + let mut stream = response.bytes_stream(); + let mut buf = Vec::new(); + while let Some(chunk) = stream.next().await { + buf.extend_from_slice(&chunk.unwrap()); + } + let body = String::from_utf8_lossy(&buf).into_owned(); + println!("{service} {url} -> {status}"); + assert!( + !SIGNATURE_ERRORS.iter().any(|error| body.contains(error)), + "{service} rejected the signature: {body}" + ); + (status, body) +} + +#[tokio::test] +#[ignore = "calls AWS with real credentials"] +async fn aws_accepts_the_signature() { + let (status, body) = call( + "sts", + Method::POST, + "https://sts.us-east-1.amazonaws.com/", + "Action=GetCallerIdentity&Version=2011-06-15", + Some("application/x-www-form-urlencoded"), + ) + .await; + assert_eq!(status, 200, "sts: {body}"); + + call( + "bedrock", + Method::GET, + "https://bedrock.us-east-1.amazonaws.com/foundation-models?byProvider=anthropic", + "", + None, + ) + .await; + // `:` in the model id exercises path canonicalization. + call( + "bedrock", + Method::POST, + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1:0/converse", + r#"{"messages":[{"role":"user","content":[{"text":"hi"}]}],"inferenceConfig":{"maxTokens":5}}"#, + Some("application/json"), + ) + .await; + call( + "polly", + Method::GET, + "https://polly.us-east-1.amazonaws.com/v1/voices?LanguageCode=en-US", + "", + None, + ) + .await; +} diff --git a/aimux-provider-utils/tests/load_api_key.rs b/aimux-provider-utils/tests/load_api_key.rs index 1af14f5e7..9c5962c97 100644 --- a/aimux-provider-utils/tests/load_api_key.rs +++ b/aimux-provider-utils/tests/load_api_key.rs @@ -1,4 +1,4 @@ -//! Rust tests for `load_api_key`, the Rust equivalent of +//! Rust tests for `load_api_key`, the Rust equivalent of //! `packages/provider-utils/src/load-api-key.ts`. //! //! There is no `load-api-key.test.ts` in the upstream SDK, so these tests are @@ -8,6 +8,7 @@ //! they cannot run concurrently and race on the shared env var. use aimux_core::AiMuxError; +use aimux_provider_utils::api_key::{load_optional_setting, load_setting}; use aimux_provider_utils::load_api_key; use serial_test::serial; @@ -52,25 +53,31 @@ fn reads_api_key_from_environment_variable() { #[test] #[serial] -fn returns_invalid_argument_when_neither_provided() { +fn returns_load_api_key_error_when_neither_provided() { // TS: throws LoadAPIKeyError when apiKey is null and the env var is unset. cleanup(); remove_env(); let err = load_api_key(None, ENV_VAR, "Test API key").unwrap_err(); - assert!(matches!(err, AiMuxError::InvalidArgument(_))); + assert!(matches!(err, AiMuxError::LoadApiKey { .. })); } #[test] #[serial] -fn error_message_mentions_description_and_env_var() { - // The error message guides the user to both the parameter and the env var. +fn error_carries_description_and_env_var_and_mentions_both() { + // The error guides the user to both the parameter and the env var. cleanup(); remove_env(); let err = load_api_key(None, ENV_VAR, "Test API key").unwrap_err(); - let msg = match err { - AiMuxError::InvalidArgument(m) => m, - other => panic!("expected InvalidArgument error, got {other:?}"), + let AiMuxError::LoadApiKey { + env_var, + description, + } = &err + else { + panic!("expected LoadApiKey error, got {err:?}"); }; + assert_eq!(env_var, ENV_VAR); + assert_eq!(description, "Test API key"); + let msg = err.to_string(); assert!( msg.contains("Test API key"), "message should mention description: {msg}" @@ -83,23 +90,23 @@ fn error_message_mentions_description_and_env_var() { #[test] #[serial] -fn empty_string_api_key_falls_back_to_env_var() { - // Rust behaviour (stricter than TS): an empty `api_key` string is treated - // as "not provided" and falls through to the environment variable. +fn empty_string_api_key_is_returned_verbatim_without_env_fallback() { + // TS: `typeof apiKey === 'string'` returns the value, even "". An explicit + // empty key must never be replaced by a credential from the environment. cleanup(); set_env("env-key"); let key = load_api_key(Some(""), ENV_VAR, "Test").unwrap(); - assert_eq!(key, "env-key"); + assert_eq!(key, ""); cleanup(); } #[test] #[serial] -fn empty_string_api_key_errors_when_env_var_also_unset() { +fn empty_string_api_key_is_returned_when_env_var_is_unset_too() { cleanup(); remove_env(); - let err = load_api_key(Some(""), ENV_VAR, "Test").unwrap_err(); - assert!(matches!(err, AiMuxError::InvalidArgument(_))); + let key = load_api_key(Some(""), ENV_VAR, "Test").unwrap(); + assert_eq!(key, ""); } #[test] @@ -115,12 +122,59 @@ fn explicit_api_key_takes_precedence_over_env_var() { #[test] #[serial] -fn whitespace_only_api_key_falls_back_to_env_var() { - // A whitespace-only string is non-empty, so the Rust impl returns it as-is - // (it does not trim). This documents that behaviour. +fn whitespace_only_api_key_is_returned_as_is() { + // An explicit value is never trimmed or second-guessed. cleanup(); set_env("env-key"); let key = load_api_key(Some(" "), ENV_VAR, "Test").unwrap(); assert_eq!(key, " "); cleanup(); } + +#[test] +#[serial] +fn load_setting_prefers_the_value_then_the_environment() { + cleanup(); + set_env("env-value"); + assert_eq!( + load_setting(Some("given"), ENV_VAR, "region").unwrap(), + "given" + ); + assert_eq!(load_setting(Some(""), ENV_VAR, "region").unwrap(), ""); + assert_eq!(load_setting(None, ENV_VAR, "region").unwrap(), "env-value"); + cleanup(); +} + +#[test] +#[serial] +fn load_setting_reports_the_missing_setting() { + cleanup(); + let err = load_setting(None, ENV_VAR, "region").unwrap_err(); + let AiMuxError::LoadSetting { env_var, name } = &err else { + panic!("expected LoadSetting error, got {err:?}"); + }; + assert_eq!(env_var, ENV_VAR); + assert_eq!(name, "region"); + assert!(err.to_string().contains("region")); +} + +#[test] +#[serial] +fn load_optional_setting_never_fails() { + cleanup(); + assert_eq!(load_optional_setting(None, ENV_VAR), None); + assert_eq!( + load_optional_setting(Some(""), ENV_VAR), + Some(String::new()) + ); + set_env("env-value"); + assert_eq!( + load_optional_setting(None, ENV_VAR).as_deref(), + Some("env-value") + ); + assert_eq!( + load_optional_setting(Some("given"), ENV_VAR).as_deref(), + Some("given") + ); + cleanup(); +} diff --git a/aimux-provider-utils/tests/ws_proxy_test.rs b/aimux-provider-utils/tests/ws_proxy_test.rs index 90e015d49..6516f60c3 100644 --- a/aimux-provider-utils/tests/ws_proxy_test.rs +++ b/aimux-provider-utils/tests/ws_proxy_test.rs @@ -1,4 +1,5 @@ -//! RFC-0034 §2: the global `ProxyConfig` governs WebSocket connections. +//! RFC-0034 §2: a WebSocket connector's `ProxyConfig` governs its +//! connections; the built-in default shares the default HTTP transport's. //! //! Integration tests run a real local WS server and a fake CONNECT proxy on //! a dedicated runtime thread (the fixture must outlive every `#[tokio::test]` @@ -10,8 +11,8 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; use aimux_core::options::TimeoutConfiguration; -use aimux_provider_utils::http::ProxyConfig; -use aimux_provider_utils::ws::{WebSocketRequest, ws_connect}; +use aimux_provider_utils::ws::{TungsteniteConnector, WebSocketRequest, ws_connect}; +use aimux_provider_utils::{ProxyConfig, ReqwestFetch, set_default_transport}; use futures::{SinkExt, StreamExt}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; @@ -65,10 +66,11 @@ fn fixture() -> &'static Fixture { }); let (proxy_port, tunnel_port, direct_port) = ready_rx.recv().expect("fixture thread died"); - // One process-wide proxy config (init_proxy is a set-once global): - // everything tunnels through the fake proxy except the direct - // server's port, which is covered by a port-specific no_proxy entry. - let configured = aimux_provider_utils::http::init_proxy(ProxyConfig { + // The default transport's proxy config, which the built-in connector + // shares: everything tunnels through the fake proxy except the + // direct server's port, which is covered by a port-specific no_proxy + // entry. + set_default_transport(ReqwestFetch::new(ProxyConfig { http_url: Some(format!("http://127.0.0.1:{proxy_port}")), // wss targets tunnel too — the live smoke relies on it to drive // the rustls branch through the CONNECT proxy against the real @@ -77,8 +79,7 @@ fn fixture() -> &'static Fixture { all_url: None, // Port-specific entry: also exercises the port-aware matching. no_proxy: Some(format!("127.0.0.1:{direct_port}")), - }); - assert!(configured, "init_proxy must win (no earlier test set it)"); + })); Fixture { tunnel_port, @@ -205,6 +206,7 @@ fn ws_request(url: String) -> WebSocketRequest { step_ms: None, total_ms: Some(10_000), }), + connector: None, } } @@ -239,6 +241,25 @@ async fn ws_tunnels_through_connect_proxy() { assert!(fixture.proxy_connections.load(Ordering::SeqCst) >= 1); } +/// A connector with its own settings ignores the default transport's: one +/// without a proxy reaches the tunnel-only server directly. +#[tokio::test] +#[serial_test::serial] +async fn a_connector_uses_its_own_proxy_settings() { + let fixture = fixture(); + let before = fixture.proxy_connections.load(Ordering::SeqCst); + let mut request = ws_request(format!("ws://127.0.0.1:{}", fixture.tunnel_port)); + request.connector = Some(Arc::new(TungsteniteConnector::new(ProxyConfig::default()))); + let mut connection = ws_connect(&request).await.expect("direct connect"); + connection.send_text("own").await.expect("send direct"); + match connection.next().await { + Some(Ok(aimux_provider_utils::ws::WsMessage::Text(text))) => assert_eq!(text, "own"), + other => panic!("expected direct echo, got {other:?}"), + } + connection.close().await; + assert_eq!(fixture.proxy_connections.load(Ordering::SeqCst), before); +} + #[tokio::test] #[serial_test::serial] async fn no_proxy_entry_connects_directly() { @@ -563,6 +584,7 @@ async fn live_wss_handshake_through_connect_proxy() { step_ms: None, total_ms: Some(20_000), }), + connector: None, }; let _ = &mut request; diff --git a/aimux-providers/Cargo.toml b/aimux-providers/Cargo.toml index 942abdb5c..bfca8c5d6 100644 --- a/aimux-providers/Cargo.toml +++ b/aimux-providers/Cargo.toml @@ -27,12 +27,10 @@ bytes = { workspace = true } tracing = { workspace = true } async-stream = { workspace = true } base64 = "0.22" -sha2 = "0.10" -hmac = "0.12" -hex = "0.4" chrono = "0.4" url = { workspace = true } percent-encoding = "2" +regex = "1" [features] default = ["realtime"] @@ -50,6 +48,9 @@ tokio = { workspace = true } wiremock = "0.6" futures = { workspace = true } serial_test = { workspace = true } +# SigV4 payload-hash assertions in the Bedrock / Anthropic-on-AWS tests. +sha2 = "0.10" +hex = "0.4" [lints] workspace = true diff --git a/aimux-providers/src/anthropic/cache_control.rs b/aimux-providers/src/anthropic/cache_control.rs index 66e1d0ffb..31e82e210 100644 --- a/aimux-providers/src/anthropic/cache_control.rs +++ b/aimux-providers/src/anthropic/cache_control.rs @@ -7,21 +7,34 @@ //! set on a non-cacheable context such as a thinking block) are dropped with a //! `Warning::Unsupported`. -use aimux_core::shared::JsonObject; +use aimux_core::shared::SharedProviderOptions; use aimux_core::types::Warning; use serde_json::Value; +use super::options::{CANONICAL, anthropic_options}; + /// Anthropic allows a maximum of 4 cache breakpoints per request. const MAX_CACHE_BREAKPOINTS: u32 = 4; /// Extract the `cache_control` value from provider metadata, mirroring the TS /// `getCacheControl` helper. Accepts both the camelCase `cacheControl` and the -/// snake_case `cache_control` keys in the `providerOptions.anthropic` namespace. +/// snake_case `cache_control` keys under `providerOptions.anthropic`. /// /// The value is passed through unchanged (the Anthropic API validates it). #[must_use] -pub fn extract_cache_control(anthropic: Option<&JsonObject>) -> Option { - let anthropic = anthropic?; +pub fn extract_cache_control(provider_options: Option<&SharedProviderOptions>) -> Option { + extract_cache_control_for(provider_options, CANONICAL) +} + +/// [`extract_cache_control`] for a provider created with a custom `name`: the +/// options are read from `providerOptions.anthropic` merged with +/// `providerOptions[name]`, the custom key winning. +#[must_use] +pub(crate) fn extract_cache_control_for( + provider_options: Option<&SharedProviderOptions>, + options_name: &str, +) -> Option { + let anthropic = anthropic_options(provider_options, options_name)?; anthropic .get("cacheControl") .or_else(|| anthropic.get("cache_control")) @@ -31,23 +44,42 @@ pub fn extract_cache_control(anthropic: Option<&JsonObject>) -> Option { /// Tracks cache breakpoint usage across a single prompt conversion and emits /// warnings when cache_control is set in an unsupported context or when the /// breakpoint limit is exceeded. -#[derive(Debug, Default)] +#[derive(Debug)] pub struct CacheControlValidator { breakpoint_count: u32, warnings: Vec, + options_name: String, +} + +impl Default for CacheControlValidator { + fn default() -> Self { + Self::new() + } } impl CacheControlValidator { + /// A validator reading `providerOptions.anthropic`. #[must_use] pub fn new() -> Self { - Self::default() + Self::for_options_name(CANONICAL) + } + + /// A validator for a provider whose providerOptions key is `options_name` + /// (read in addition to `anthropic`). + #[must_use] + pub(crate) fn for_options_name(options_name: &str) -> Self { + Self { + breakpoint_count: 0, + warnings: Vec::new(), + options_name: options_name.to_string(), + } } /// Resolve the `cache_control` value (if any) for a content block, applying /// the validation rules. /// - /// - `anthropic` is the part- or message-level Anthropic namespace to - /// read `cacheControl` from. + /// - `provider_options` is the part- or message-level provider metadata to + /// read `anthropic.cacheControl` from. /// - `context_type` is a human-readable label for the block kind (e.g. /// `"user message part"`, `"thinking block"`) used in warning messages. /// - `can_cache` is `false` for contexts that cannot carry cache_control @@ -58,11 +90,11 @@ impl CacheControlValidator { /// or `None` when there is none or it was rejected. pub fn get_cache_control( &mut self, - anthropic: Option<&JsonObject>, + provider_options: Option<&SharedProviderOptions>, context_type: &str, can_cache: bool, ) -> Option { - let cache_control_value = extract_cache_control(anthropic); + let cache_control_value = extract_cache_control_for(provider_options, &self.options_name); let cc = cache_control_value?; diff --git a/aimux-providers/src/anthropic/config.rs b/aimux-providers/src/anthropic/config.rs new file mode 100644 index 000000000..d0b1a36d7 --- /dev/null +++ b/aimux-providers/src/anthropic/config.rs @@ -0,0 +1,217 @@ +//! The private per-model configuration every host of the Anthropic Messages +//! API shares. +//! +//! This is the Rust shape of the object `createAnthropic` hands +//! `AnthropicLanguageModel` (`{ provider, baseURL, headers, fetch, +//! supportedUrls, ... }`). The first-party provider, Claude Platform on AWS and +//! Anthropic on Vertex each build one; the model reads nothing else, so what +//! differs between the hosts (endpoint, credentials, a signing transport, a +//! body envelope) is data here, never a branch in the model. It has no getters +//! for the credential or the settings that produced it. + +use std::collections::{BTreeSet, HashMap}; +use std::sync::Arc; + +use futures::future::BoxFuture; +use serde_json::Value; + +use aimux_core::AiMuxError; +use aimux_core::language_model::SupportedUrls; +use aimux_provider_utils::{ + ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, ResponseHandler, + combine_headers, normalize_headers, +}; + +use super::convert::RequestProfile; +pub use crate::shared::TransformRequestBody; + +/// Maps an endpoint path (`"/messages"`) to the full request URL. +pub(crate) type UrlFn = Arc String + Send + Sync>; + +/// `(model id, streaming)` to the full URL of the messages call, for hosts +/// that put the model and the mode in the path. +pub(crate) type RequestUrlFn = Arc String + Send + Sync>; + +/// Rewrites the request body a host-specific way. +pub(crate) type BodyFn = Arc Value + Send + Sync>; + +/// Builds the configuration of one request, for hosts whose endpoint or +/// credentials (a Vertex project and location, say) can only be known when a +/// call is made. +pub(crate) type ResolveFn = + Arc BoxFuture<'static, Result> + Send + Sync>; + +/// Builds the host's error handler for non-2xx responses. +pub(crate) type ErrorHandlerFn = fn() -> ResponseHandler; + +/// How a host differs from the first-party Messages API (the upstream +/// `buildRequestUrl`, `transformRequestBody`, `supportsNativeStructuredOutput` +/// and `supportsStrictTools` config members, plus its error shape). +#[derive(Clone)] +pub(crate) struct AnthropicModelHooks { + /// The messages URL when it depends on the model and the mode. `None`: + /// `url("/messages")`. + pub(crate) request_url: Option, + /// Applied to the finished request body before the provider's own + /// `transform_request_body`. + pub(crate) prepare_body: Option, + /// Parses a failed response. + pub(crate) failed_response_handler: ErrorHandlerFn, + /// Structured outputs and their beta header. + pub(crate) supports_native_structured_output: bool, + /// `strict` on tool definitions. + pub(crate) supports_strict_tools: bool, +} + +impl Default for AnthropicModelHooks { + /// The first-party API. + fn default() -> Self { + Self { + request_url: None, + prepare_body: None, + failed_response_handler: super::anthropic_failed_response_handler, + supports_native_structured_output: true, + supports_strict_tools: true, + } + } +} + +/// What a model needs to talk to the API. Built by the provider that owns the +/// model; read by [`AnthropicMessagesModel`](super::model::AnthropicMessagesModel) +/// and the files interface. +#[derive(Clone)] +pub(crate) struct AnthropicModelConfig { + /// The identity the model reports from `provider()`. + pub(crate) provider: String, + /// Endpoint path to full URL. + pub(crate) url: UrlFn, + /// Provider headers (credential, `anthropic-version`, user headers), + /// resolved on every request. + pub(crate) headers: HeadersFn, + /// Transport; `None` uses the process default, resolved per request. + pub(crate) fetch: Option, + /// URLs the model fetches itself. Empty: the caller downloads them. + pub(crate) supported_urls: SupportedUrls, + /// Provider-level request-body rewrite. + pub(crate) transform_request_body: Option, + /// The origin credentialed headers may be sent to. Never read for URLs. + pub(crate) base_url: String, + /// The providerOptions key read in addition to `anthropic` and written + /// on response metadata. + pub(crate) provider_options_name: String, + /// Host differences. + pub(crate) hooks: AnthropicModelHooks, + /// Late binding: when set, every request uses the configuration this + /// returns (its own `resolve` is `None`) instead of the fields above, + /// which then only describe the model's static identity. + pub(crate) resolve: Option, +} + +impl AnthropicModelConfig { + /// The configuration to use for one request: itself, or what the host's + /// `resolve` builds now. + /// + /// # Errors + /// + /// Returns the error of the host's setting or credential loading, for + /// instance `AiMuxError::LoadSetting` for an unset Vertex project. + pub(crate) async fn resolved(&self) -> Result { + match &self.resolve { + Some(resolve) => resolve().await, + None => Ok(self.clone()), + } + } + + /// The full URL of an endpoint path. + pub(crate) fn url(&self, path: &str) -> String { + (self.url)(path) + } + + /// The URL of the messages call for a model. + pub(crate) fn messages_url(&self, model_id: &str, stream: bool) -> String { + match &self.hooks.request_url { + Some(request_url) => request_url(model_id, stream), + None => self.url("/messages"), + } + } + + /// What the request body builder needs to know about this host. + pub(crate) fn request_profile(&self) -> RequestProfile { + RequestProfile { + options_name: self.provider_options_name.clone(), + supports_native_structured_output: self.hooks.supports_native_structured_output, + supports_strict_tools: self.hooks.supports_strict_tools, + } + } + + /// The host's error handler. + pub(crate) fn failed_response_handler(&self) -> ResponseHandler { + (self.hooks.failed_response_handler)() + } + + /// 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. + /// + /// # Errors + /// + /// Returns the error of the credential or header producer, for instance + /// `AiMuxError::LoadApiKey` when no key was given and the environment + /// variable is unset. + pub(crate) async fn request_headers( + &self, + call_headers: Option<&HashMap>, + 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 + .iter() + .map(|(name, value)| (name.clone(), Some(value.clone()))) + .collect() + }) + .unwrap_or_default(); + Ok(normalize_headers(combine_headers(&[ + &provider, + &beta_layer, + &call, + ]))) + } + + /// An exchange that inherits the operation's cancellation and recording + /// context and carries this config's transport and credential origin. + pub(crate) fn http_request( + &self, + url: String, + headers: Vec<(String, String)>, + options: &impl ExchangeContext, + ) -> HttpRequest { + self.with_transport(HttpRequest::new(url, headers, options)) + } + + /// Attach the transport and the credential origin to a request. + pub(crate) fn with_transport(&self, mut request: HttpRequest) -> HttpRequest { + request.fetch = self.fetch.clone(); + request.credentialed_origin = Some(self.base_url.clone()); + request + } + + /// Run the host's body preparation, then the provider-level rewrite. + pub(crate) fn transform_body(&self, mut body: Value) -> Value { + if let Some(prepare) = &self.hooks.prepare_body { + body = prepare(body); + } + match &self.transform_request_body { + Some(transform) => transform(body), + None => body, + } + } +} diff --git a/aimux-providers/src/anthropic/convert.rs b/aimux-providers/src/anthropic/convert.rs index 4d2f4111d..040ab56aa 100644 --- a/aimux-providers/src/anthropic/convert.rs +++ b/aimux-providers/src/anthropic/convert.rs @@ -6,17 +6,13 @@ //! `build_request_body` / `parse_stop_reason` helpers used by the language //! models. //! -//! Two public prompt-conversion entry points are exposed: -//! - [`convert_prompt_to_anthropic_full`] — the complete conversion, returning -//! the Anthropic `system` + `messages` shape alongside the beta headers and -//! warnings produced during conversion. Supports mid-conversation system -//! messages, trailing-whitespace trimming on the final assistant message, -//! reasoning/thinking parts, URL & base64 file parts (with PDF beta and -//! top-level media-type sniffing), and provider-referenced files. -//! - [`convert_prompt_to_anthropic`] — the legacy two-tuple `(system, messages)` -//! return form used by [`build_request_body`]. It is equivalent to the full -//! conversion with `send_reasoning = false`, discarding the betas and -//! warnings. +//! The prompt-conversion entry point is +//! [`convert_prompt_to_anthropic_full_fallible`]: the complete conversion, +//! returning the Anthropic `system` + `messages` shape alongside the beta +//! headers and warnings produced during conversion. It supports +//! mid-conversation system messages, trailing-whitespace trimming on the final +//! assistant message, reasoning/thinking parts, URL & base64 file parts (with +//! PDF beta and top-level media-type sniffing), and provider-referenced files. use std::collections::{BTreeSet, HashSet}; @@ -28,10 +24,14 @@ use aimux_core::language_model_message::{ use aimux_core::options::{CallOptions, Tool}; 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 serde_json::{Map, Value, json}; use crate::anthropic::cache_control::CacheControlValidator; -use crate::anthropic::prepare_tools::{AnthropicTool, prepare_tools_with_provider}; +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::tool_name_mapping::ToolNameMapping; /// Beta header emitted when a PDF file part is present. @@ -41,7 +41,7 @@ const BETA_MID_CONVERSATION_SYSTEM: &str = "mid-conversation-system-2026-04-07"; /// Beta header emitted when a provider-referenced file part is present. const BETA_FILES_API: &str = "files-api-2025-04-14"; -/// Result of [`convert_prompt_to_anthropic_full`]. +/// Result of [`convert_prompt_to_anthropic_full_fallible`]. #[derive(Debug, Clone)] pub struct AnthropicPromptConversion { /// Anthropic `system` blocks, or `None` when there is no system prompt. @@ -98,18 +98,35 @@ pub fn convert_prompt_to_anthropic_full_with_tools( prompt: &LanguageModelPrompt, send_reasoning: bool, tool_names: &ToolNameMapping, +) -> Result { + convert_prompt_for(prompt, send_reasoning, tool_names, CANONICAL) +} + +/// [`convert_prompt_to_anthropic_full_with_tools`] for a provider whose +/// providerOptions key is `options_name`: the part and message options are +/// read from `anthropic` merged with `options_name` (the custom key wins) and +/// the metadata the response side wrote under `options_name` is read back. +/// +/// # Errors +/// +/// Same as [`convert_prompt_to_anthropic_full_with_tools`]. +pub(crate) fn convert_prompt_for( + prompt: &LanguageModelPrompt, + send_reasoning: bool, + tool_names: &ToolNameMapping, + options_name: &str, ) -> 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 // scopes this set to one merged assistant block; scanning the whole prompt // is a superset of that and is safe because tool call ids are unique. - let mcp_tool_use_ids = collect_mcp_tool_use_ids(prompt); + let mcp_tool_use_ids = collect_mcp_tool_use_ids(prompt, options_name); let mut system: Vec = Vec::new(); let mut messages: Vec = Vec::new(); let mut betas: BTreeSet = BTreeSet::new(); let mut warnings: Vec = Vec::new(); - let mut validator = CacheControlValidator::new(); + let mut validator = CacheControlValidator::for_options_name(options_name); #[derive(Clone, Copy, PartialEq, Eq)] enum Eff { @@ -134,31 +151,14 @@ pub fn convert_prompt_to_anthropic_full_with_tools( } for msg in prompt { - let role = match msg { - LanguageModelMessage::System { .. } => None, - LanguageModelMessage::User { .. } | LanguageModelMessage::Tool { .. } => { - Some(Eff::User) - } - LanguageModelMessage::Assistant { .. } => Some(Eff::Assistant), - }; - if let Some(role) = role { - seen_non_system = true; - if last != Some(role) { - flush(&mut messages, &mut last, &mut acc); - } - last = Some(role); - } - match msg { + let eff = match msg { LanguageModelMessage::System { content, provider_options, } => { flush(&mut messages, &mut last, &mut acc); - let cc = validator.get_cache_control( - provider_options.as_ref().and_then(|o| o.get("anthropic")), - "system message", - true, - ); + 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); @@ -166,13 +166,24 @@ pub fn convert_prompt_to_anthropic_full_with_tools( messages.push(json!({ "role": "system", "content": blocks })); betas.insert(BETA_MID_CONVERSATION_SYSTEM.to_string()); } + continue; } + LanguageModelMessage::Assistant { .. } => Eff::Assistant, + LanguageModelMessage::User { .. } | LanguageModelMessage::Tool { .. } => Eff::User, + }; + seen_non_system = true; + if last != Some(eff) { + flush(&mut messages, &mut last, &mut acc); + } + match msg { + LanguageModelMessage::System { .. } => unreachable!(), LanguageModelMessage::User { content, provider_options, } => { for (idx, part) in content.iter().enumerate() { - acc.push(match part { + let is_last = idx + 1 == content.len(); + let block = match part { UserPart::Text(TextPart { text, provider_options: part_options, @@ -181,7 +192,7 @@ pub fn convert_prompt_to_anthropic_full_with_tools( &mut validator, part_options.as_ref(), provider_options.as_ref(), - idx + 1 == content.len(), + is_last, "user message part", "user message", ); @@ -192,11 +203,13 @@ pub fn convert_prompt_to_anthropic_full_with_tools( &mut betas, &mut validator, provider_options.as_ref(), - idx + 1 == content.len(), + is_last, "user message part", "user message", + options_name, )?, - }); + }; + acc.push(block); } } LanguageModelMessage::Tool { @@ -222,22 +235,37 @@ pub fn convert_prompt_to_anthropic_full_with_tools( provider_options, } => { for (idx, part) in content.iter().enumerate() { - if let Some(block) = convert_assistant_part( - part, - send_reasoning, - tool_names, - &mcp_tool_use_ids, - &mut betas, - &mut warnings, - &mut validator, - idx + 1 == content.len(), - provider_options.as_ref(), - )? { + let is_last = idx + 1 == content.len(); + let block = if let AssistantPart::ToolResult(part) = part { + convert_assistant_tool_result( + part, + tool_names, + &mcp_tool_use_ids, + &mut warnings, + &mut validator, + is_last, + provider_options.as_ref(), + ) + } else { + convert_assistant_part( + part, + send_reasoning, + tool_names, + &mut betas, + &mut warnings, + &mut validator, + is_last, + provider_options.as_ref(), + options_name, + )? + }; + if let Some(block) = block { acc.push(block); } } } } + last = Some(eff); } flush(&mut messages, &mut last, &mut acc); @@ -271,56 +299,6 @@ pub fn convert_prompt_to_anthropic_full_with_tools( }) } -/// Convert a prompt into the full Anthropic shape. -/// -/// Panics on conversion failure. Production paths use the fallible variant -/// [`convert_prompt_to_anthropic_full_fallible`]; this panic wrapper exists -/// only for integration tests under `tests/`. It is `#[doc(hidden)]` and -/// `#[deprecated]` so it neither appears on the public API surface nor can be -/// pulled in by accident (release uses `panic = "abort"`, so reaching a panic -/// here via FFI would kill the host process). -#[doc(hidden)] -#[deprecated( - since = "0.2.1", - note = "panics on failure; use convert_prompt_to_anthropic_full_fallible instead (issue #90 R1)" -)] -#[must_use] -pub fn convert_prompt_to_anthropic_full( - prompt: &LanguageModelPrompt, - send_reasoning: bool, -) -> AnthropicPromptConversion { - convert_prompt_to_anthropic_full_fallible(prompt, send_reasoning) - .expect("convert_prompt_to_anthropic_full: conversion failed") -} - -/// Convert a prompt into the Anthropic `system` + `messages` shape. -/// -/// This is the legacy two-tuple return form used by [`build_request_body`]. It -/// is equivalent to [`convert_prompt_to_anthropic_full`] with -/// `send_reasoning = false`, discarding the betas and warnings. Consecutive -/// messages that map to the same effective Anthropic role (`user`/`tool` → -/// `user`, `assistant` → `assistant`) are merged into a single message, matching -/// the SDK behaviour. -#[doc(hidden)] -#[deprecated( - since = "0.2.1", - note = "panics on failure; use convert_prompt_to_anthropic_full_fallible instead (issue #90 R1)" -)] -#[must_use] -pub fn convert_prompt_to_anthropic( - prompt: &LanguageModelPrompt, -) -> (Option>, Vec) { - match convert_prompt_to_anthropic_full_fallible(prompt, false) { - Ok(result) => (result.system, result.messages), - Err(e) => panic!("{}", e), - } -} - -/// Return the top-level media type (the segment before the first `/`). -fn top_level_media_type(media_type: &str) -> &str { - media_type.split('/').next().unwrap_or("") -} - fn resolve_cache_control( validator: &mut CacheControlValidator, part_options: Option<&SharedProviderOptions>, @@ -330,18 +308,10 @@ fn resolve_cache_control( message_context: &str, ) -> Option { validator - .get_cache_control( - part_options.and_then(|o| o.get("anthropic")), - part_context, - true, - ) + .get_cache_control(part_options, part_context, true) .or_else(|| { if is_last_part { - validator.get_cache_control( - message_options.and_then(|o| o.get("anthropic")), - message_context, - true, - ) + validator.get_cache_control(message_options, message_context, true) } else { None } @@ -355,6 +325,7 @@ fn apply_cc(mut block: Value, cc: Option) -> Value { block } +#[allow(clippy::too_many_arguments)] fn convert_file_part( file: &FilePart, betas: &mut BTreeSet, @@ -363,6 +334,7 @@ fn convert_file_part( is_last_part: bool, part_context: &str, message_context: &str, + options_name: &str, ) -> Result { let FilePart { data, @@ -374,7 +346,14 @@ fn convert_file_part( FileData::Data { data: FileBytes::Binary(bytes), } => { - let full = resolve_full_media_type(media_type, bytes)?; + let full = if matches!( + get_top_level_media_type(media_type), + "image" | "application" + ) { + resolve_full_media_type(file)? + } else { + media_type.clone() + }; route_file_bytes(&full, bytes, filename.as_deref(), betas)? } FileData::Data { @@ -384,29 +363,36 @@ fn convert_file_part( let bytes = base64::engine::general_purpose::STANDARD .decode(data) .unwrap_or_default(); - let full = resolve_full_media_type(media_type, &bytes)?; + let full = if matches!( + get_top_level_media_type(media_type), + "image" | "application" + ) { + resolve_full_media_type(file)? + } else { + media_type.clone() + }; route_file_base64(&full, data, &bytes, filename.as_deref(), betas)? } FileData::Url { url, .. } => route_file_url(media_type, url, betas)?, FileData::Reference { reference } => { let file_id = resolve_anthropic_reference(reference)?; betas.insert(BETA_FILES_API.to_string()); - let container_upload = provider_options + let container_upload = anthropic_options(provider_options.as_ref(), options_name) .as_ref() - .and_then(|o| o.get("anthropic")) .and_then(|a| a.get("containerUpload")) .and_then(Value::as_bool) .unwrap_or(false); if container_upload { return Ok(json!({ "type": "container_upload", "file_id": file_id })); - } else if top_level_media_type(media_type) == "image" { + } else if get_top_level_media_type(media_type) == "image" { json!({ "type": "image", "source": { "type": "file", "file_id": file_id } }) } else { json!({ "type": "document", "source": { "type": "file", "file_id": file_id } }) } } FileData::Text { text } => { - let options = provider_options.as_ref().and_then(|o| o.get("anthropic")); + let document_options = anthropic_options(provider_options.as_ref(), options_name); + let options = document_options.as_ref(); let mut block = json!({ "type": "document", "source": { "type": "text", "media_type": "text/plain", "data": text, }}); @@ -464,16 +450,9 @@ fn convert_tool_result( if let ToolResultOutput::Content { value } = output { for item in value { if let ToolResultContent::File(file) = item - && let FileData::Data { data } = &file.data + && let FileData::Data { .. } = &file.data { - use base64::Engine; - let bytes = match data { - FileBytes::Binary(bytes) => bytes.clone(), - FileBytes::Base64(data) => base64::engine::general_purpose::STANDARD - .decode(data) - .unwrap_or_default(), - }; - resolve_full_media_type(&file.media_type, &bytes)?; + resolve_full_media_type(file)?; } } } @@ -487,45 +466,40 @@ fn convert_tool_result( block["is_error"] = json!(true); } // cache_control: part ?? output ?? (is_last_part ? message). - let cc = match validator.get_cache_control( - provider_options.as_ref().and_then(|o| o.get("anthropic")), - "tool result part", - true, - ) { - Some(v) => Some(v), - None => match extract_tool_result_output_provider_options(output) { - Some(out_opts) => { - validator.get_cache_control(out_opts.get("anthropic"), "tool result output", true) - } - None => { - if is_last_part { + let cc = validator + .get_cache_control(provider_options.as_ref(), "tool result 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.and_then(|o| o.get("anthropic")), + message_provider_options, "tool result message", true, ) - } else { - None - } - } - }, - }; + }) + .flatten() + }); Ok(apply_cc(block, cc)) } -/// Convert an assistant part, preserving reasoning and provider-executed tool results. -/// Part-level cache control falls back to message options on the final part. #[allow(clippy::too_many_arguments)] fn convert_assistant_part( part: &AssistantPart, send_reasoning: bool, tool_names: &ToolNameMapping, - mcp_tool_use_ids: &HashSet<&str>, betas: &mut BTreeSet, warnings: &mut Vec, validator: &mut CacheControlValidator, is_last_part: bool, message_provider_options: Option<&SharedProviderOptions>, + options_name: &str, ) -> Result, AiMuxError> { let resolve_cc = |validator: &mut CacheControlValidator, part_opts| { resolve_cache_control( @@ -555,6 +529,7 @@ fn convert_assistant_part( is_last_part, "assistant message part", "assistant message", + options_name, )?, AssistantPart::Reasoning(ReasoningPart { @@ -567,6 +542,7 @@ fn convert_assistant_part( send_reasoning, warnings, validator, + options_name, )); } @@ -582,16 +558,16 @@ fn convert_assistant_part( if *provider_executed == Some(true) { let provider_name = tool_names.to_provider_tool_name(tool_name); - let anthropic_options = provider_options - .as_ref() - .and_then(|options| options.get("anthropic")); + let own_options = anthropic_options(provider_options.as_ref(), options_name); - if anthropic_options + if own_options + .as_ref() .and_then(|options| options.get("type")) .and_then(Value::as_str) == Some("mcp-tool-use") { - let Some(server_name) = anthropic_options + let Some(server_name) = own_options + .as_ref() .and_then(|options| options.get("serverName")) .and_then(Value::as_str) else { @@ -667,9 +643,8 @@ fn convert_assistant_part( } else { json!({ "rawInvalidInput": input }) }; - let caller = provider_options + let caller = anthropic_options(provider_options.as_ref(), options_name) .as_ref() - .and_then(|options| options.get("anthropic")) .and_then(|anthropic| anthropic.get("caller")) .and_then(|caller| { let caller_type = caller.get("type")?.as_str()?; @@ -694,18 +669,8 @@ fn convert_assistant_part( apply_cc(block, cc) } - AssistantPart::ReasoningFile(_) | AssistantPart::Custom(_) => return Ok(None), - AssistantPart::ToolResult(part) => { - return Ok(convert_assistant_tool_result( - part, - tool_names, - mcp_tool_use_ids, - warnings, - validator, - is_last_part, - message_provider_options, - )); - } + AssistantPart::Custom(_) | AssistantPart::ReasoningFile(_) => return Ok(None), + AssistantPart::ToolResult(_) => unreachable!(), })) } @@ -717,7 +682,10 @@ fn convert_assistant_part( /// before any message is converted. The marker is the one the response side /// writes on an `mcp_tool_use` block (`stream.rs`): /// `providerOptions.anthropic.type == "mcp-tool-use"`. -fn collect_mcp_tool_use_ids(prompt: &LanguageModelPrompt) -> HashSet<&str> { +fn collect_mcp_tool_use_ids<'a>( + prompt: &'a LanguageModelPrompt, + options_name: &str, +) -> HashSet<&'a str> { let mut ids = HashSet::new(); for msg in prompt { let LanguageModelMessage::Assistant { content, .. } = msg else { @@ -729,8 +697,8 @@ fn collect_mcp_tool_use_ids(prompt: &LanguageModelPrompt) -> HashSet<&str> { provider_options: Some(opts), .. }) = part - && opts - .get("anthropic") + && anthropic_options(Some(opts), options_name) + .as_ref() .and_then(|a| a.get("type")) .and_then(|t| t.as_str()) == Some("mcp-tool-use") @@ -791,29 +759,22 @@ fn convert_assistant_tool_result( // cache_control: part ?? output ?? (is_last_part ? message) — the same // resolution order the bare `tool_result` path uses. - let cache_control = match validator.get_cache_control( - provider_options.as_ref().and_then(|o| o.get("anthropic")), - "assistant message part", - true, - ) { - Some(v) => Some(v), - None => match extract_tool_result_output_provider_options(output) { - Some(out_opts) => { - validator.get_cache_control(out_opts.get("anthropic"), "tool result output", true) - } - None => { - if is_last_part { - validator.get_cache_control( - message_provider_options.and_then(|o| o.get("anthropic")), - "assistant message", - true, - ) - } else { - None - } - } - }, - }; + 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 raw_value = match output { ToolResultOutput::Json { value, .. } | ToolResultOutput::ErrorJson { value, .. } => { @@ -1144,7 +1105,7 @@ fn route_file_bytes( betas: &mut BTreeSet, ) -> Result { use base64::Engine; - match top_level_media_type(full_media_type) { + match get_top_level_media_type(full_media_type) { "image" => { let b64 = base64::engine::general_purpose::STANDARD.encode(bytes); Ok(json!({ @@ -1192,7 +1153,7 @@ fn route_file_base64( title: Option<&str>, betas: &mut BTreeSet, ) -> Result { - match top_level_media_type(full_media_type) { + match get_top_level_media_type(full_media_type) { "image" => Ok(json!({ "type": "image", "source": { "type": "base64", "media_type": full_media_type, "data": b64 } @@ -1234,7 +1195,7 @@ fn route_file_url( url: &str, betas: &mut BTreeSet, ) -> Result { - match top_level_media_type(media_type) { + match get_top_level_media_type(media_type) { "image" => Ok(json!({ "type": "image", "source": { "type": "url", "url": url } })), "application" if media_type == "application/pdf" => { betas.insert(BETA_PDFS.to_string()); @@ -1249,75 +1210,6 @@ fn route_file_url( } } -/// Returns `true` only when the media type has a non-empty, non-wildcard -/// subtype (i.e. `type/subtype` where `subtype` is not `*`). -fn is_full_media_type(media_type: &str) -> bool { - match media_type.find('/') { - Some(i) => { - let subtype = &media_type[i + 1..]; - !subtype.is_empty() && subtype != "*" - } - None => false, - } -} - -/// Sniff the concrete media type from inline bytes for the given top-level -/// segment, mirroring the TS `detectMediaType` signature tables. -fn detect_media_type(bytes: &[u8], top_level: &str) -> Option<&'static str> { - match top_level { - "image" => detect_image_media_type(bytes), - "application" => detect_document_media_type(bytes), - _ => None, - } -} - -fn detect_image_media_type(bytes: &[u8]) -> Option<&'static str> { - const SIGS: &[(&[u8], &str)] = &[ - (&[0x47, 0x49, 0x46], "image/gif"), - (&[0x89, 0x50, 0x4e, 0x47], "image/png"), - (&[0xff, 0xd8], "image/jpeg"), - (&[0x42, 0x4d], "image/bmp"), - (&[0x49, 0x49, 0x2a, 0x00], "image/tiff"), - (&[0x4d, 0x4d, 0x00, 0x2a], "image/tiff"), - ]; - for (prefix, mt) in SIGS { - if bytes.starts_with(prefix) { - return Some(mt); - } - } - // WEBP: "RIFF" .... "WEBP" - if bytes.len() >= 12 && &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" { - return Some("image/webp"); - } - None -} - -fn detect_document_media_type(bytes: &[u8]) -> Option<&'static str> { - if bytes.starts_with(&[0x25, 0x50, 0x44, 0x46]) { - return Some("application/pdf"); - } - None -} - -/// Resolve a file part's media type to a full `type/subtype` form, mirroring -/// the TS `resolveFullMediaType`. When the media type is already a full type it -/// is returned as-is; otherwise the subtype is sniffed from the inline bytes. -pub(crate) fn resolve_full_media_type( - media_type: &str, - bytes: &[u8], -) -> Result { - if is_full_media_type(media_type) { - return Ok(media_type.to_string()); - } - let top = top_level_media_type(media_type); - match detect_media_type(bytes, top) { - Some(detected) => Ok(detected.to_string()), - None => Err(AiMuxError::UnsupportedFunctionality(format!( - "file of media type \"{media_type}\" must specify subtype since it could not be auto-detected" - ))), - } -} - /// Convert an assistant `Reasoning` part into a `thinking` or /// `redacted_thinking` block, or drop it with a warning. Returns `None` when /// the part is dropped. @@ -1331,6 +1223,7 @@ fn convert_reasoning_part( send_reasoning: bool, warnings: &mut Vec, validator: &mut CacheControlValidator, + options_name: &str, ) -> Option { if !send_reasoning { warnings.push(Warning::Other { @@ -1341,13 +1234,14 @@ fn convert_reasoning_part( // `redactedData` is read from `providerOptions.anthropic.redactedData` // (mirroring the TS `anthropicReasoningMetadataSchema`). - let redacted_data = provider_options - .and_then(|o| o.get("anthropic")) + let reasoning_options = anthropic_options(provider_options, options_name); + let redacted_data = reasoning_options + .as_ref() .and_then(|a| a.get("redactedData")) .and_then(|v| v.as_str()); - let effective_signature = provider_options - .and_then(|o| o.get("anthropic")) + let effective_signature = reasoning_options + .as_ref() .and_then(|a| a.get("signature")) .and_then(Value::as_str); @@ -1355,11 +1249,7 @@ fn convert_reasoning_part( // Thinking blocks cannot carry cache_control directly — they are cached // implicitly when in previous assistant turns. Validate to emit a // helpful warning if a value was set. - validator.get_cache_control( - provider_options.and_then(|o| o.get("anthropic")), - "thinking block", - false, - ); + validator.get_cache_control(provider_options, "thinking block", false); Some(json!({ "type": "thinking", "thinking": text, @@ -1367,11 +1257,7 @@ fn convert_reasoning_part( })) } else if let Some(data) = redacted_data { // Redacted thinking blocks likewise cannot carry cache_control. - validator.get_cache_control( - provider_options.and_then(|o| o.get("anthropic")), - "redacted thinking block", - false, - ); + validator.get_cache_control(provider_options, "redacted thinking block", false); Some(json!({ "type": "redacted_thinking", "data": data, @@ -1385,17 +1271,19 @@ fn convert_reasoning_part( } /// Resolve the Anthropic file id from a provider-reference object, mirroring the -/// TS `resolveProviderReference`. Returns an error when the key is missing. +/// TS `resolveProviderReference`: the reference is always keyed by the +/// canonical name, whatever the provider is called. Errors when that key is +/// absent, matching the TS `UnsupportedFunctionalityError`. fn resolve_anthropic_reference( reference: &aimux_core::shared::SharedProviderReference, ) -> Result { - if let Some(id) = reference.get("anthropic") { + 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(); Err(AiMuxError::InvalidArgument(format!( - "No provider reference found for provider 'anthropic'. Available providers: {}", + "No provider reference found for provider '{CANONICAL}'. Available providers: {}", providers.join(", ") ))) } @@ -1428,11 +1316,11 @@ fn resolve_tool_result_output( let source = match &part.data { FileData::Url { url, .. } => json!({ "type": "url", "url": url }), FileData::Data { data } => { - let (data, bytes) = match data { - FileBytes::Binary(bytes) => (base64::engine::general_purpose::STANDARD.encode(bytes), bytes.clone()), - FileBytes::Base64(data) => (data.clone(), base64::engine::general_purpose::STANDARD.decode(data).unwrap_or_default()), + let data = match data { + FileBytes::Binary(bytes) => base64::engine::general_purpose::STANDARD.encode(bytes), + FileBytes::Base64(data) => data.clone(), }; - let full = resolve_full_media_type(&part.media_type, &bytes).ok()?; + 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 }) @@ -1442,7 +1330,7 @@ fn resolve_tool_result_output( 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("anthropic")) + 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 } @@ -1698,7 +1586,7 @@ pub struct RequestBodyResult { fn parse_anthropic_option_object(shape: &str, value: &Value) -> Result { let object = value.as_object().ok_or(())?; - let parsed = crate::openai::convert::parse_option_fields(object, "anthropic", |key, value| { + let parsed = crate::openai::convert::parse_option_fields(object, CANONICAL, |key, value| { let kind = match (shape, key) { ("options", "sendReasoning" | "disableParallelToolUse" | "toolStreaming") => "bool", ("options", "structuredOutputMode") => "outputFormat|jsonTool|auto", @@ -1749,7 +1637,7 @@ fn parse_anthropic_option_object(shape: &str, value: &Value) -> Result "[]skill", ("skill", "type") => "anthropic|custom", ("skill", "skillId") - if object.get("type").and_then(Value::as_str) == Some("anthropic") => + if object.get("type").and_then(Value::as_str) == Some(PROVIDER_SKILL_TYPE) => { "string" } @@ -1913,13 +1801,39 @@ fn parse_anthropic_option_value(kind: &str, value: &Value) -> Result if valid { Ok(value.clone()) } else { Err(()) } } -/// Read a value from `provider_options["anthropic"][key]`. -fn anthropic_option(options: &Option, key: &str) -> Option { - options - .as_ref() - .and_then(|m| m.get("anthropic")) - .and_then(|o| o.get(key)) - .cloned() +/// What differs between the hosts of the Messages API as far as the request +/// body goes: which providerOptions key the caller's options live under (read +/// in addition to the canonical one) and which tool/output features the host +/// accepts. +#[derive(Debug, Clone)] +pub(crate) struct RequestProfile { + /// The custom providerOptions key; `anthropic` is always read too. + pub(crate) options_name: String, + /// Structured outputs and their beta header (Vertex has neither). + pub(crate) supports_native_structured_output: bool, + /// `strict` on tool definitions (Vertex rejects it). + pub(crate) supports_strict_tools: bool, +} + +impl Default for RequestProfile { + /// The first-party API under its canonical name. + fn default() -> Self { + Self { + options_name: CANONICAL.to_string(), + supports_native_structured_output: true, + supports_strict_tools: true, + } + } +} + +/// Read a value from the Anthropic options of a call (`anthropic` merged with +/// the profile's custom key). +fn anthropic_option( + options: &Option, + profile: &RequestProfile, + key: &str, +) -> Option { + anthropic_options_in(options.as_ref(), &profile.options_name).and_then(|o| o.get(key).cloned()) } /// Recursively remove `null`-valued fields from JSON objects (mirroring the @@ -2016,12 +1930,13 @@ fn strip_anthropic_sampling_params( fn resolve_anthropic_thinking( model_id: &str, options: &CallOptions, + profile: &RequestProfile, caps: &ModelCapabilities, warnings: &mut Vec, ) -> (Option, Option) { let mut thinking_config: Option = - anthropic_option(&options.provider_options, "thinking"); - let mut effort: Option = anthropic_option(&options.provider_options, "effort") + anthropic_option(&options.provider_options, profile, "thinking"); + let mut effort: Option = anthropic_option(&options.provider_options, profile, "effort") .and_then(|v| v.as_str().map(std::string::ToString::to_string)); if let Some(reasoning) = options.reasoning @@ -2231,17 +2146,18 @@ fn insert_anthropic_sampling( fn apply_anthropic_tools( body: &mut Value, options: &CallOptions, + profile: &RequestProfile, stream: bool, caps: &ModelCapabilities, betas: &mut BTreeSet, warnings: &mut Vec, ) { let disable_parallel_tool_use = - anthropic_option(&options.provider_options, "disableParallelToolUse") + anthropic_option(&options.provider_options, profile, "disableParallelToolUse") .and_then(|v| v.as_bool()) .unwrap_or(false); let default_eager_input_streaming = stream - && anthropic_option(&options.provider_options, "toolStreaming") + && anthropic_option(&options.provider_options, profile, "toolStreaming") .and_then(|v| v.as_bool()) .unwrap_or(true); @@ -2260,7 +2176,7 @@ fn apply_anthropic_tools( None => Vec::new(), }; - let prepared = prepare_tools_with_provider( + let prepared = prepare_tools_for( if options.tools.is_some() { Some(&anthropic_tools) } else { @@ -2268,9 +2184,10 @@ fn apply_anthropic_tools( }, options.tool_choice.as_ref(), disable_parallel_tool_use, - caps.supports_structured_output, - caps.supports_structured_output, + caps.supports_structured_output && profile.supports_native_structured_output, + caps.supports_structured_output && profile.supports_strict_tools, default_eager_input_streaming, + &profile.options_name, ); warnings.extend(prepared.tool_warnings); @@ -2296,9 +2213,10 @@ fn apply_anthropic_tools( fn append_anthropic_mcp_servers( body: &mut Value, options: &CallOptions, + profile: &RequestProfile, betas: &mut BTreeSet, ) { - let Some(mcp) = anthropic_option(&options.provider_options, "mcpServers") else { + let Some(mcp) = anthropic_option(&options.provider_options, profile, "mcpServers") else { return; }; let Some(arr) = mcp.as_array() else { @@ -2351,10 +2269,11 @@ fn append_anthropic_mcp_servers( fn append_anthropic_container( body: &mut Value, options: &CallOptions, + profile: &RequestProfile, betas: &mut BTreeSet, warnings: &mut Vec, ) -> Result<(), AiMuxError> { - let Some(container) = anthropic_option(&options.provider_options, "container") else { + let Some(container) = anthropic_option(&options.provider_options, profile, "container") else { return Ok(()); }; let skills = container.get("skills").and_then(|s| s.as_array()); @@ -2371,12 +2290,12 @@ fn append_anthropic_container( let skill_id = if stype == "custom" { match skill .get("providerReference") - .and_then(|r| r.get("anthropic")) + .and_then(|r| r.get(CANONICAL)) { Some(id) => id.clone(), None => { return Err(AiMuxError::UnsupportedFunctionality(format!( - "skill provider reference is missing the 'anthropic' key: {skill}" + "skill provider reference is missing the '{CANONICAL}' key: {skill}" ))); } } @@ -2444,16 +2363,38 @@ pub fn build_request_body_with_warnings( model_id: &str, options: &CallOptions, stream: bool, +) -> Result { + build_request_body_for(model_id, options, stream, &RequestProfile::default()) +} + +/// [`build_request_body_with_warnings`] for a host of the Messages API other +/// than the first-party endpoint under its canonical name. +/// +/// # Errors +/// +/// Same as [`build_request_body_with_warnings`]. +pub(crate) fn build_request_body_for( + model_id: &str, + options: &CallOptions, + stream: bool, + profile: &RequestProfile, ) -> Result { let mut parsed_options = options.clone(); - if let Some(namespaces) = parsed_options.provider_options.as_mut() - && let Some(raw) = namespaces.get("anthropic") + if let Some(raw) = + anthropic_options_in(options.provider_options.as_ref(), &profile.options_name) { - let parsed = parse_anthropic_option_object("options", &Value::Object(raw.clone())) - .map_err(|()| { + let parsed = + parse_anthropic_option_object("options", &Value::Object(raw)).map_err(|()| { AiMuxError::InvalidArgument("invalid anthropic provider options".to_string()) })?; - namespaces.insert("anthropic".to_string(), parsed.as_object().unwrap().clone()); + let namespaces = parsed_options + .provider_options + .get_or_insert_with(Default::default); + namespaces.shift_remove(CANONICAL); + namespaces.insert( + profile.options_name.clone(), + parsed.as_object().unwrap().clone(), + ); } let options = &parsed_options; let mut warnings: Vec = Vec::new(); @@ -2479,7 +2420,7 @@ pub fn build_request_body_with_warnings( // providerOptions.anthropic.thinking / .effort + top-level `reasoning`. let (thinking_config, thinking_effort) = - resolve_anthropic_thinking(model_id, options, &caps, &mut warnings); + resolve_anthropic_thinking(model_id, options, profile, &caps, &mut warnings); let (thinking_type, send_thinking, mut thinking_budget, thinking_display) = derive_anthropic_thinking(&thinking_config); @@ -2493,10 +2434,11 @@ pub fn build_request_body_with_warnings( // 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_to_anthropic_full_with_tools( + let conversion = convert_prompt_for( &options.prompt, send_input_reasoning, &ToolNameMapping::new(options.tools.as_deref()), + &profile.options_name, )?; let system = conversion.system; let messages = conversion.messages; @@ -2507,8 +2449,11 @@ pub fn build_request_body_with_warnings( "model": model_id, "messages": messages, "max_tokens": max_tokens, - "stream": stream, }); + // `stream` is sent only when streaming, as the AI SDK does. + if stream { + body["stream"] = json!(true); + } if let Some(sys) = system { body["system"] = json!(sys); @@ -2541,25 +2486,35 @@ pub fn build_request_body_with_warnings( insert_anthropic_sampling(&mut body, temperature, top_p, top_k, options); + // 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()) + .filter(|user_id| !user_id.is_null()) + { + body["metadata"] = json!({ "user_id": user_id }); + } + // Tools — provider-defined tools alongside function tools, plus the // required beta headers / tool warnings. - apply_anthropic_tools(&mut body, options, stream, &caps, &mut betas, &mut warnings); + apply_anthropic_tools( + &mut body, + options, + profile, + stream, + &caps, + &mut betas, + &mut warnings, + ); // providerOptions.anthropic.mcpServers → mcp_servers + beta header. - append_anthropic_mcp_servers(&mut body, options, &mut betas); + append_anthropic_mcp_servers(&mut body, options, profile, &mut betas); // providerOptions.anthropic.container — programmatic tool calling / skills. - append_anthropic_container(&mut body, options, &mut betas, &mut warnings)?; + append_anthropic_container(&mut body, options, profile, &mut betas, &mut warnings)?; // providerOptions.anthropic.contextManagement → context_management is not // yet implemented (build_context_management pending). - // Per-call request body overrides (RFC-0017): deep-merge user-supplied - // JSON into the built body. `null` values delete the corresponding key. - if let Some(ref overrides) = options.body_overrides { - crate::openai::convert::deep_merge_json(&mut body, overrides); - } - Ok(RequestBodyResult { body, warnings, diff --git a/aimux-providers/src/anthropic/files.rs b/aimux-providers/src/anthropic/files.rs index a06821719..361cfd6b2 100644 --- a/aimux-providers/src/anthropic/files.rs +++ b/aimux-providers/src/anthropic/files.rs @@ -4,8 +4,6 @@ //! Aligned with Vercel AI SDK `AnthropicFiles` //! (`reference/ai/packages/anthropic/src/anthropic-files.ts`). -use std::collections::HashMap; - use async_trait::async_trait; use serde::Deserialize; use serde_json::json; @@ -14,9 +12,10 @@ use aimux_core::error::AiMuxError; use aimux_core::files_model::{Files, UploadFileCallOptions, UploadFileData, UploadFileResult}; use aimux_core::shared::FileBytes; -use aimux_provider_utils::{HttpBody, HttpRequest}; +use aimux_provider_utils::HttpBody; -use super::AnthropicConfig; +use super::config::AnthropicModelConfig; +use super::options::CANONICAL; /// The beta header value for the Anthropic Files API. const FILES_BETA_HEADER: &str = "files-api-2025-04-14"; @@ -58,45 +57,20 @@ struct AnthropicFilesResponse { /// Aligned with TS `AnthropicFiles`. 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 AnthropicFiles { - config: AnthropicConfig, + config: AnthropicModelConfig, } impl AnthropicFiles { - #[must_use] - pub fn new(config: AnthropicConfig) -> Self { + pub(crate) fn from_config(config: AnthropicModelConfig) -> Self { Self { config } } - - fn build_headers(&self) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "anthropic-version".to_string(), - self.config.api_version.clone(), - ); - // Auth: prefer bearer token, fall back to x-api-key. - if let Some(token) = &self.config.auth_token { - headers.insert("authorization".to_string(), format!("Bearer {token}")); - } else { - headers.insert("x-api-key".to_string(), self.config.api_key.clone()); - } - // Extra config-level headers. - if let Some(cfg_headers) = &self.config.headers { - for (k, v) in cfg_headers { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/v1/files", self.config.base_url) - } } #[async_trait] impl Files for AnthropicFiles { + /// `"{name}.files"` with the `.messages` suffix removed. fn provider(&self) -> &str { - "anthropic.files" + &self.config.provider } async fn upload_file( @@ -118,12 +92,9 @@ impl Files for AnthropicFiles { None, ); - let mut headers = self.build_headers(); - headers.insert("anthropic-beta".to_string(), FILES_BETA_HEADER.to_string()); - let header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + 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?; // `send()` returns Ok only for 2xx; non-2xx responses are mapped to an // error internally using the shared error structure. The multipart body @@ -131,25 +102,16 @@ impl Files for AnthropicFiles { // sets `Content-Type` from it, so it is intentionally not added to the // header list above. // - // Nothing above `upload_file` retries it (there is no Core - // `do_upload_file`), so the retry lives here — safe for the whole - // exchange because an upload is not billable and a failed - // create-a-file request returns no id to replay against (§9.4). - let retries = aimux_core::retry::prepare_retries( - None, - self.config.retry_config, - options.abort_signal.clone(), - ); - let resp = retries - .retry(|| { - aimux_provider_utils::post_to_api( - HttpRequest::new(self.endpoint(), header_list.clone(), options), - HttpBody::Bytes(body.clone(), content_type.clone()), - aimux_provider_utils::create_json_response_handler(), - super::anthropic_failed_response_handler(), - ) - }) - .await?; + // One attempt: like the AI SDK's `uploadFile`, an upload is never + // replayed implicitly. A caller that wants a retry makes it around + // `upload_file`. + let resp = aimux_provider_utils::post_to_api( + config.http_request(config.url("/files"), header_list, options), + HttpBody::Bytes(body, content_type), + aimux_provider_utils::create_json_response_handler(), + self.config.failed_response_handler(), + ) + .await?; let data: AnthropicFilesResponse = resp.value; @@ -172,7 +134,7 @@ impl Files for AnthropicFiles { } let mut provider_ref = aimux_core::shared::SharedProviderReference::new(); - provider_ref.insert("anthropic".to_string(), data.id); + provider_ref.insert(CANONICAL.to_string(), data.id); let result_media_type = data .mime_type @@ -184,10 +146,9 @@ impl Files for AnthropicFiles { provider_reference: provider_ref, media_type: Some(result_media_type), filename: result_filename, - provider_metadata: Some(aimux_core::shared::NamespaceMap::from([( - "anthropic".to_string(), - metadata, - )])), + provider_metadata: Some( + std::iter::once((self.config.provider_options_name.clone(), metadata)).collect(), + ), warnings: Vec::new(), }) } diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs index 404c74750..b807fcdb4 100644 --- a/aimux-providers/src/anthropic/mod.rs +++ b/aimux-providers/src/anthropic/mod.rs @@ -1,25 +1,53 @@ -//! Anthropic (Claude) provider. +//! Anthropic provider. +//! +//! [`create_anthropic`] is the Rust form of the AI SDK's `createAnthropic`: it +//! takes [`AnthropicProviderSettings`], validates the base URL, rejects an +//! `api_key` given together with an `auth_token`, fixes the provider name and +//! returns an [`AnthropicProvider`]. Credentials are not read there; they are +//! loaded in the request headers of every call, from the settings or (for the +//! key) from `ANTHROPIC_API_KEY`. [`anthropic()`] is the default instance. +//! +//! The Messages model is shared with Anthropic on AWS +//! ([`crate::anthropic_aws`]) and Anthropic on Vertex +//! ([`crate::vertex`]): they build the same private model configuration and differ +//! only in what it holds. pub mod cache_control; +pub(crate) mod config; pub mod convert; pub mod files; pub mod model; +pub(crate) mod options; pub mod prepare_tools; pub mod sanitize_json_schema; +pub mod skills; pub mod stream; pub mod tool_name_mapping; pub mod types; pub mod usage; -use std::collections::HashMap; -use std::sync::Arc; +pub use files::AnthropicFiles; +pub use model::AnthropicMessagesModel; +pub use skills::AnthropicSkills; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::{RetryConfig, load_api_key}; +use std::sync::{Arc, OnceLock}; + +use futures::future::BoxFuture; use serde_json::Value; +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::error::AiMuxError; +use aimux_core::files_model::Files; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{AuthScheme, Credential, credential_headers}; + +use config::{AnthropicModelConfig, AnthropicModelHooks}; + pub(crate) fn anthropic_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -35,423 +63,322 @@ pub(crate) fn anthropic_failed_response_handler() }) } -/// The bare (unversioned) Anthropic API URL. -const ANTHROPIC_API_URL: &str = "https://api.anthropic.com"; +const API_ORIGIN: &str = "https://api.anthropic.com"; +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.anthropic.com/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "ANTHROPIC_API_KEY"; +const API_VERSION: &str = "2023-06-01"; +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 { + let url = validate_base_url(url)?; + Ok(if url == API_ORIGIN { + DEFAULT_BASE_URL.to_string() + } else { + url + }) +} -/// Configuration for the Anthropic provider. -#[derive(Clone)] -pub struct AnthropicConfig { - /// API key sent via the `x-api-key` header. Empty when authenticating with - /// an `auth_token` instead. - pub api_key: String, - /// Bearer token sent via the `Authorization` header. When set, `api_key` - /// is ignored. +/// The URL patterns the API fetches itself (`supportedUrls` in `createAnthropic`). +fn supported_urls() -> SupportedUrls { + let https = regex::Regex::new(r"^https?://.*$").expect("static pattern"); + SupportedUrls( + [ + ("image/*".to_string(), vec![https.clone()]), + ("application/pdf".to_string(), vec![https]), + ] + .into_iter() + .collect(), + ) +} + +/// Settings of [`create_anthropic`] (the AI SDK's `AnthropicProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url`, `name` and the `api_key`/`auth_token` conflict; +/// credentials and `headers` are evaluated on every request. +#[derive(Clone, Default)] +pub struct AnthropicProviderSettings { + /// Base URL for the API calls, version segment included. `None` reads + /// `ANTHROPIC_BASE_URL` when the provider is created, else + /// `https://api.anthropic.com/v1`; a trailing slash is removed. The bare + /// `https://api.anthropic.com` means the default. An invalid explicit + /// value fails [`create_anthropic`]; an invalid `ANTHROPIC_BASE_URL` + /// fails every request instead. + pub base_url: Option, + /// The API key, sent as `x-api-key`. `None` (with no `auth_token`) loads + /// `ANTHROPIC_API_KEY` when a request is made and fails that request with + /// `AiMuxError::LoadApiKey` if it is unset. An explicit value is used as + /// given, `""` included: it never falls back to the environment. + pub api_key: Option, + /// 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. pub auth_token: Option, - /// Unversioned API root (no trailing slash, no `/v1` suffix). The endpoint - /// is built as `{base_url}/v1/messages`. - pub base_url: String, - /// `anthropic-version` header value. - pub api_version: String, - /// Provider name surfaced by `LanguageModel::provider` (TS default - /// `anthropic.messages`). - pub name: String, - /// Extra headers merged into every request. - pub headers: Option>, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, - /// Provider 级请求体覆盖（RFC-0017）。在标准请求体之后 deep-merge。 - pub body_overrides: Option, - /// api_key 来源(RFC-0023 `ProviderRecord.api_key_source`):`None` = 显式 - /// (config_snapshot 记为 "explicit");`Some("env:VAR")` = 来自环境变量; - /// `Some("none")` = 本地/匿名占位。不存明文之外的信息。 - pub api_key_source: Option, - /// Generator for source IDs. Defaults to a 16-character alphanumeric ID. + /// Extra headers on every request. A `None` value removes the header, + /// including one of the fixed ones. Per-call headers win over these. + pub headers: Option, + /// The provider name, the `provider()` string of the messages model. + /// Default `"anthropic.messages"`. Its first dot-separated segment is + /// also the providerOptions key read in addition to `anthropic` (and the + /// key response metadata is written under). + 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, + /// Generator for source IDs. Defaults to the provider utility generator. pub generate_id: Option String + Send + Sync>>, } -impl std::fmt::Debug for AnthropicConfig { +impl std::fmt::Debug for AnthropicProviderSettings { + /// Never prints credentials or header values. fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AnthropicConfig") - .field("api_key", &"") - .field("auth_token", &self.auth_token.is_some()) + f.debug_struct("AnthropicProviderSettings") .field("base_url", &self.base_url) - .field("api_version", &self.api_version) + .field("api_key", &self.api_key.is_some()) + .field("auth_token", &self.auth_token.is_some()) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) .field("name", &self.name) - .field("headers", &self.headers) - .field("retry_config", &self.retry_config) - .field("body_overrides", &self.body_overrides) - .field("api_key_source", &self.api_key_source) + .field("fetch", &self.fetch.is_some()) .field("generate_id", &self.generate_id.is_some()) .finish() } } -impl AnthropicConfig { - /// Create a config authenticated with an API key, using the default - /// Anthropic base URL. - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - auth_token: None, - base_url: ANTHROPIC_API_URL.to_string(), - api_version: "2023-06-01".to_string(), - name: "anthropic.messages".to_string(), - headers: None, - retry_config: RetryConfig::default(), - body_overrides: None, - api_key_source: None, - generate_id: None, - } - } - - /// Set the generator used for source IDs. - #[must_use] - pub fn with_generate_id(mut self, generate_id: F) -> Self - where - F: Fn() -> String + Send + Sync + 'static, +/// Create an Anthropic provider. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when the explicit `base_url` is not +/// an `http(s)` URL with a host, or when both `api_key` and `auth_token` are +/// given. Those are the only ways this fails: credentials are loaded per +/// request, and an invalid `ANTHROPIC_BASE_URL` is reported by each request, +/// not here. +pub fn create_anthropic( + settings: AnthropicProviderSettings, +) -> Result { + let base_url = crate::shared::base_url_setting( + settings.base_url.as_deref(), + "ANTHROPIC_BASE_URL", + DEFAULT_BASE_URL, + normalize_base_url, + )?; + let is_truthy = |value: &String| !value.is_empty(); + if settings.api_key.as_ref().is_some_and(is_truthy) + && settings.auth_token.as_ref().is_some_and(is_truthy) { - self.generate_id = Some(Arc::new(generate_id)); - self - } - - /// Start a builder for more involved configurations (e.g. `auth_token`, - /// custom `name`, extra `headers`). The builder validates that `api_key` - /// and `auth_token` are not both provided. - #[must_use] - pub fn builder() -> AnthropicConfigBuilder { - AnthropicConfigBuilder::default() - } - - /// Override the base URL. Normalizes the value (strips a trailing slash and - /// a trailing `/v1` segment so the endpoint formula `{base_url}/v1/messages` - /// never doubles the version) and rejects empty strings. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = normalize_base_url(&url.into()); - self - } - - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self - } - - /// Set a custom provider name. - #[must_use] - pub fn with_name(mut self, name: impl Into) -> Self { - self.name = name.into(); - self - } - - /// Authenticate with a bearer token instead of an API key. - #[must_use] - pub fn with_auth_token(mut self, token: impl Into) -> Self { - self.auth_token = Some(token.into()); - self - } - - /// Attach extra headers merged into every request. - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } - - /// Set provider-level request body overrides (RFC-0017). - #[must_use] - pub fn with_body_overrides(mut self, overrides: Value) -> Self { - self.body_overrides = Some(overrides); - self - } - - /// 标注 api_key 来源(RFC-0023 回放重建用)。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self + return Err(AiMuxError::InvalidArgument( + "Both apiKey and authToken were provided. Please use only one authentication method." + .to_string(), + )); } + let (credential, scheme) = match settings.auth_token.filter(is_truthy) { + Some(token) => ( + Credential::Explicit(Resolvable::Value(token)), + AuthScheme::Bearer, + ), + None => ( + Credential::explicit_or_env( + settings.api_key.map(Resolvable::Value), + API_KEY_ENV_VAR, + "Anthropic", + ), + AuthScheme::Header("x-api-key"), + ), + }; + Ok(AnthropicProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( + credential_headers( + credential, + scheme, + vec![("anthropic-version".to_string(), API_VERSION.to_string())], + settings.headers, + ), + crate::anthropic::options::CANONICAL, + "4.0.68", + ), + fetch: settings.fetch, + generate_id: settings.generate_id, + supported_urls: supported_urls(), + }) +} - /// Load the configuration from the environment. - /// - /// Reads `ANTHROPIC_API_KEY` (required) and the optional - /// `ANTHROPIC_BASE_URL`. `ANTHROPIC_AUTH_TOKEN` is not auto-loaded here; - /// use [`AnthropicConfig::builder`]`.auth_token` for token auth. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `ANTHROPIC_API_KEY` is not set - /// and no key was supplied. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "ANTHROPIC_API_KEY", "Anthropic")?; - let mut config = Self::new(api_key).with_api_key_source(Some("env:ANTHROPIC_API_KEY")); - if let Ok(base) = std::env::var("ANTHROPIC_BASE_URL") - && !base.is_empty() - { - config = config.with_base_url(base); - } - Ok(config) - } +/// The default provider: `create_anthropic` with default settings, created on +/// first use. `ANTHROPIC_BASE_URL` is read at creation (an invalid value fails +/// each request, not the creation); the key is loaded per request. +pub fn anthropic() -> &'static AnthropicProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_anthropic(AnthropicProviderSettings::default()) + .expect("default settings have no explicit base URL or credentials to reject") + }) } -/// Builder for [`AnthropicConfig`] supporting both `api_key` and `auth_token` -/// authentication, with conflict validation mirroring the TS `createAnthropic`. -#[derive(Default, Clone)] -pub struct AnthropicConfigBuilder { - api_key: Option, - auth_token: Option, - base_url: Option, - name: Option, - headers: Option>, - body_overrides: Option, +/// An Anthropic provider (the AI SDK's `AnthropicProvider`). Cheap to clone +/// the models out of; it holds no HTTP client. +pub struct AnthropicProvider { + name: String, + /// The normalized base URL, or the error of an invalid + /// `ANTHROPIC_BASE_URL`, which every request reports. + base_url: Result, + headers: HeadersFn, + fetch: Option, generate_id: Option String + Send + Sync>>, + supported_urls: SupportedUrls, } -impl std::fmt::Debug for AnthropicConfigBuilder { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AnthropicConfigBuilder") - .field("api_key", &self.api_key.is_some()) - .field("auth_token", &self.auth_token.is_some()) - .field("base_url", &self.base_url) - .field("name", &self.name) - .field("headers", &self.headers) - .field("body_overrides", &self.body_overrides) - .field("generate_id", &self.generate_id.is_some()) - .finish() +impl AnthropicProvider { + /// The model configuration reporting `provider` as its identity. + fn model_config(&self, provider: String) -> AnthropicModelConfig { + let base_url = self.base_url.clone().unwrap_or_default(); + let base = base_url.clone(); + AnthropicModelConfig { + provider, + url: Arc::new(move |path| format!("{base}{path}")), + headers: self.headers.clone(), + fetch: self.fetch.clone(), + supported_urls: self.supported_urls.clone(), + transform_request_body: None, + base_url, + provider_options_name: options::options_name_of(&self.name), + hooks: AnthropicModelHooks::default(), + // Every request resolves its configuration first; an invalid + // `ANTHROPIC_BASE_URL` fails it there. + resolve: self.base_url.as_ref().err().map(|error| { + let error = error.clone(); + let resolve: config::ResolveFn = + Arc::new(move || Box::pin(std::future::ready(Err(error.clone())))); + resolve + }), + } } -} -impl AnthropicConfigBuilder { - #[must_use] - pub fn api_key(mut self, api_key: impl Into) -> Self { - self.api_key = Some(api_key.into()); - self + /// The name with the `.messages` suffix removed: the prefix of the + /// provider strings of the other modalities. + fn bare_name(&self) -> &str { + self.name.strip_suffix(".messages").unwrap_or(&self.name) } + /// A Messages model; `provider()` is the provider name + /// (`"anthropic.messages"` by default). #[must_use] - pub fn auth_token(mut self, auth_token: impl Into) -> Self { - self.auth_token = Some(auth_token.into()); - self + pub fn messages(&self, model_id: &str) -> AnthropicMessagesModel { + AnthropicMessagesModel::with_config( + model_id.to_string(), + self.model_config(self.name.clone()), + ) + .with_generate_id(self.generate_id.clone()) } + /// The files interface; `provider()` is the provider name. #[must_use] - pub fn base_url(mut self, base_url: impl Into) -> Self { - self.base_url = Some(base_url.into()); - self + pub fn files(&self) -> AnthropicFiles { + AnthropicFiles::from_config(self.model_config(self.name.clone())) } + /// The skills upload interface. #[must_use] - pub fn name(mut self, name: impl Into) -> Self { - self.name = Some(name.into()); - self + pub fn skills(&self) -> AnthropicSkills { + AnthropicSkills::from_config( + self.model_config(format!("{}.skills", self.name.replacen(".messages", "", 1))), + ) } + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; the same model as + /// [`messages`](Self::messages) and + /// [`language_model`](Provider::language_model). #[must_use] - pub fn headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.messages(model_id)) } +} - #[must_use] - pub fn body_overrides(mut self, overrides: Value) -> Self { - self.body_overrides = Some(overrides); - self +impl Provider for AnthropicProvider { + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) } - #[must_use] - pub fn generate_id(mut self, generate_id: F) -> Self - where - F: Fn() -> String + Send + Sync + 'static, - { - self.generate_id = Some(Arc::new(generate_id)); - self + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - /// Build the config, validating that `api_key` and `auth_token` are not - /// both set and that a provided `base_url` is non-empty. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when both `api_key` and `auth_token` - /// are set, or when a provided `base_url` is empty. - pub fn build(self) -> Result { - if self.api_key.is_some() && self.auth_token.is_some() { - return Err(AiMuxError::InvalidArgument( - "Both apiKey and authToken were provided. Please use only one authentication method." - .to_string(), - )); - } - let base_url = match &self.base_url { - Some(url) => normalize_base_url(url), - None => ANTHROPIC_API_URL.to_string(), - }; - Ok(AnthropicConfig { - api_key: self.api_key.unwrap_or_default(), - auth_token: self.auth_token, - base_url, - api_version: "2023-06-01".to_string(), - name: self - .name - .unwrap_or_else(|| "anthropic.messages".to_string()), - headers: self.headers, - retry_config: RetryConfig::default(), - body_overrides: self.body_overrides, - api_key_source: None, - generate_id: self.generate_id, - }) + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) } -} -/// Normalize an Anthropic base URL: reject empty strings, strip a trailing -/// slash, then strip a trailing `/v1` segment (so the endpoint formula -/// `{base_url}/v1/messages` produces a single version segment regardless of -/// whether the caller supplied `https://api.anthropic.com/`, -/// `https://api.anthropic.com/v1`, or a proxy like -/// `https://proxy.example/v1/`). -fn normalize_base_url(url: &str) -> String { - if url.is_empty() { - return ANTHROPIC_API_URL.to_string(); - } - let mut s = url.trim_end_matches('/').to_string(); - if s.ends_with("/v1") { - s = s[..s.len() - "/v1".len()].trim_end_matches('/').to_string(); + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) } - s -} -/// Anthropic provider — creates `AnthropicModel` instances. -pub struct AnthropicProvider { - config: AnthropicConfig, -} - -impl AnthropicProvider { - #[must_use] - pub fn new(config: AnthropicConfig) -> Self { - Self { config } + fn files(&self) -> Option> { + Some(Arc::new(self.files())) } - #[must_use] - pub fn model(&self, model_id: &str) -> model::AnthropicModel { - model::AnthropicModel::new(model_id.to_string(), self.config.clone()) + fn skills(&self) -> Option> { + Some(Arc::new(self.skills())) } +} - /// Create a Files interface for uploading files to Anthropic. - #[must_use] - pub fn files(&self) -> files::AnthropicFiles { - files::AnthropicFiles::new(self.config.clone()) +impl ProviderDiscovery for AnthropicProvider { + /// `GET {base_url}/models`: one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config(format!("{}.models", self.bare_name())); + Box::pin(async move { list_models_once(&config).await }) } } -impl Provider for AnthropicProvider { - fn name(&self) -> &str { - "anthropic" +/// One `GET {base_url}/models` exchange against an Anthropic-shaped API: no +/// retry, no recording. Discovery is not a Core operation and the AI SDK has +/// no equivalent, so a failure is reported to the caller as it happened. +/// +/// # Errors +/// +/// Returns the header-resolution error (a missing key is `LoadApiKey`), +/// `ApiCall` for HTTP/transport failures and `JsonParse` when the body does +/// not deserialize into the models list. +pub(crate) async fn list_models_once( + config: &AnthropicModelConfig, +) -> Result, AiMuxError> { + let config = &config.resolved().await?; + #[derive(serde::Deserialize)] + struct Resp { + #[serde(default)] + data: Vec, } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + #[derive(serde::Deserialize)] + struct Entry { + id: String, + #[serde(default)] + display_name: Option, } - /// List models via `GET {base_url}/v1/models` (Anthropic native), enriched - /// with the community catalogue portrait when available (RFC-0027). - fn list_models( - &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let config = self.config.clone(); - Box::pin(async move { - let base = config - .base_url - .trim_end_matches("/v1") - .trim_end_matches('/'); - let url = format!("{base}/v1/models"); - let mut headers = std::collections::HashMap::new(); - headers.insert("anthropic-version".to_string(), config.api_version.clone()); - if let Some(token) = &config.auth_token { - headers.insert("authorization".to_string(), format!("Bearer {token}")); - } else { - headers.insert("x-api-key".to_string(), config.api_key.clone()); - } - if let Some(extra) = &config.headers { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - let mut header_list: Vec<(String, String)> = headers.into_iter().collect(); - header_list.push(("Content-Type".to_string(), "application/json".to_string())); - - use aimux_provider_utils::HttpRequest; - // Retry rationale: see `openai::model::execute_list_models`. - let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: header_list.clone(), - abort_signal: None, - call_id: None, - recording_context: None, - ..Default::default() - }, - aimux_provider_utils::create_json_response_handler(), - anthropic_failed_response_handler(), - ) - }) - .await?; - - #[derive(serde::Deserialize)] - struct Resp { - #[serde(default)] - data: Vec, - } - #[derive(serde::Deserialize)] - struct Entry { - id: String, - #[serde(default)] - display_name: Option, - } - let parsed: Resp = resp.value; - let runtime: Vec = parsed - .data - .into_iter() - .map(|e| aimux_core::model_catalogue::RuntimeModel { - id: e.id, - owned_by: e.display_name, - created: None, - }) - .collect(); - Ok(runtime) + let headers = config + .request_headers(None, &std::collections::BTreeSet::new()) + .await?; + let resp = aimux_provider_utils::get_from_api( + config.with_transport(aimux_provider_utils::HttpRequest { + url: config.url("/models"), + headers, + ..Default::default() + }), + aimux_provider_utils::create_json_response_handler(), + config.failed_response_handler(), + ) + .await?; + let parsed: Resp = resp.value; + Ok(parsed + .data + .into_iter() + .map(|entry| RuntimeModel { + id: entry.id, + owned_by: entry.display_name, + created: None, }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn normalize_strips_slash_and_v1() { - assert_eq!( - normalize_base_url("https://api.anthropic.com/"), - "https://api.anthropic.com" - ); - assert_eq!( - normalize_base_url("https://api.anthropic.com/v1"), - "https://api.anthropic.com" - ); - assert_eq!( - normalize_base_url("https://proxy.example/v1/"), - "https://proxy.example" - ); - assert_eq!( - normalize_base_url("https://proxy.example"), - "https://proxy.example" - ); - } + .collect()) } diff --git a/aimux-providers/src/anthropic/model.rs b/aimux-providers/src/anthropic/model.rs index b40bb84f9..c8e6b7d56 100644 --- a/aimux-providers/src/anthropic/model.rs +++ b/aimux-providers/src/anthropic/model.rs @@ -1,145 +1,115 @@ -//! Anthropic language model — implements `LanguageModel` trait. +//! The Anthropic Messages language model — implements `LanguageModel`. //! -//! The request body building, HTTP send and response/SSE parsing live in the -//! shared [`super::stream`] core; this module only wires the standard Anthropic -//! endpoint, Bearer/x-api-key auth and `Json` body encoding. - -use std::collections::{BTreeSet, HashMap}; +//! One model serves every host of the Messages API (`api.anthropic.com`, +//! Claude Platform on AWS, Anthropic on Vertex): it builds the request body, +//! asks its model configuration for the URL, the headers and the transport +//! and hands the result to the shared [`super::stream`] core, which sends it +//! and parses the response or the SSE stream. use async_trait::async_trait; +use serde_json::Value; 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::{GenerateResult, StreamResult}; +use aimux_core::types::Warning; +use aimux_provider_utils::HttpRequest; -use super::AnthropicConfig; -use super::convert::build_request_body_with_warnings; -use super::stream::{BodyEncoding, anthropic_generate_core, anthropic_stream_core}; +use super::config::AnthropicModelConfig; +use super::convert::build_request_body_for; +use super::stream::{anthropic_generate_core, anthropic_stream_core}; use super::tool_name_mapping::ToolNameMapping; -/// An Anthropic language model (e.g. `claude-sonnet-4-20250514`). -pub struct AnthropicModel { +/// One call, ready to send. +struct PreparedCall { + http: HttpRequest, + body: Value, + warnings: Vec, +} + +/// An Anthropic Messages model (e.g. `claude-sonnet-4-20250514`), created by +/// `provider.messages(id)`. +pub struct AnthropicMessagesModel { model_id: String, - config: AnthropicConfig, + config: AnthropicModelConfig, generate_id: std::sync::Arc String + Send + Sync>, } -impl AnthropicModel { - #[must_use] - pub fn new(model_id: String, config: AnthropicConfig) -> Self { - let generate_id = config - .generate_id - .clone() - .unwrap_or_else(|| std::sync::Arc::new(super::stream::generate_id)); +impl AnthropicMessagesModel { + /// A model of the host the config describes. + pub(crate) fn with_config(model_id: String, config: AnthropicModelConfig) -> Self { Self { model_id, config, - generate_id, + generate_id: std::sync::Arc::new(aimux_provider_utils::generate_id), } } - fn endpoint(&self) -> String { - format!("{}/v1/messages", self.config.base_url) + 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 } - /// Build the auth-header closure for the standard (Bearer/x-api-key) path. - /// - /// The body bytes are ignored (no request signing). Config-level and - /// per-call headers and the `anthropic-beta` header are merged last-wins, - /// matching the original `HashMap` semantics. - fn make_header_builder( + /// 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. + async fn prepare( &self, - extra: Option<&HashMap>, - betas: BTreeSet, - ) -> impl Fn(&[u8], &str) -> Result, AiMuxError> { - let api_version = self.config.api_version.clone(); - let auth_token = self.config.auth_token.clone(); - let api_key = self.config.api_key.clone(); - let cfg_headers = self.config.headers.clone(); - let extra = extra.cloned(); - move |_body: &[u8], _url: &str| -> Result, AiMuxError> { - let mut headers: HashMap = HashMap::new(); - headers.insert("anthropic-version".to_string(), api_version.clone()); - // Auth: prefer bearer token, fall back to x-api-key. - if let Some(token) = &auth_token { - headers.insert("authorization".to_string(), format!("Bearer {token}")); - } else { - headers.insert("x-api-key".to_string(), api_key.clone()); - } - // Extra config-level headers (lower precedence than per-call headers). - if let Some(cfg_headers) = &cfg_headers { - for (k, v) in cfg_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = &extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - if !betas.is_empty() { - headers.insert( - "anthropic-beta".to_string(), - betas.iter().cloned().collect::>().join(","), - ); - } - Ok(headers.into_iter().collect()) - } + options: &CallOptions, + stream: bool, + ) -> Result { + let config = self.config.resolved().await?; + let built = + build_request_body_for(&self.model_id, options, stream, &config.request_profile())?; + let body = config.transform_body(built.body); + let headers = config + .request_headers(options.headers.as_ref(), &built.betas) + .await?; + let http = config.http_request( + config.messages_url(&self.model_id, stream), + headers, + options, + ); + Ok(PreparedCall { + http, + body, + warnings: built.warnings, + }) } } #[async_trait] -impl LanguageModel for AnthropicModel { +impl LanguageModel for AnthropicMessagesModel { + /// `"{name}"`: `"anthropic.messages"` unless the provider was named. fn provider(&self) -> &str { - &self.config.name + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.base_url.clone()), - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - // RFC-0023 C6: ProviderOptions-shaped object (headers/max_retries/ - // body_overrides are the shared ProviderOptions fields) plus the - // Anthropic-specific `api_version`. Sensitive headers are redacted - // at the recording boundary (recording.rs redacts provider_options). - provider_options: Some(serde_json::json!({ - "headers": self.config.headers, - "api_version": self.config.api_version, - "max_retries": self.config.retry_config.max_retries, - "body_overrides": self.config.body_overrides, - })), - } + async fn supported_urls(&self) -> SupportedUrls { + self.config.supported_urls.clone() } async fn do_generate(&self, options: &CallOptions) -> Result { - let options = merge_anthropic_body_overrides(options, &self.config.body_overrides); - let req = build_request_body_with_warnings(&self.model_id, &options, false)?; - let endpoint = self.endpoint(); - let build_headers = self.make_header_builder(options.headers.as_ref(), req.betas); + let call = self.prepare(options, false).await?; anthropic_generate_core( - &endpoint, - req.body, - req.warnings, - build_headers, - BodyEncoding::Json, - options.abort_signal.clone(), + call.http, + 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(), ) @@ -147,58 +117,15 @@ impl LanguageModel for AnthropicModel { } async fn do_stream(&self, options: &CallOptions) -> Result { - let options = merge_anthropic_body_overrides(options, &self.config.body_overrides); - let req = build_request_body_with_warnings(&self.model_id, &options, true)?; - let endpoint = self.endpoint(); - let build_headers = self.make_header_builder(options.headers.as_ref(), req.betas); + let call = self.prepare(options, true).await?; anthropic_stream_core( - &endpoint, - req.body, - req.warnings, - build_headers, - BodyEncoding::Json, - options.abort_signal.clone(), - ToolNameMapping::new(options.tools.as_deref()), + call.http, + call.body, + call.warnings, + &self.config, + options, self.generate_id.clone(), ) .await } } - -/// Merge provider-level body_overrides into per-call options (RFC-0017). -fn merge_anthropic_body_overrides( - options: &CallOptions, - provider_overrides: &Option, -) -> CallOptions { - match (provider_overrides, &options.body_overrides) { - (Some(provider), Some(call)) => { - let mut merged = provider.clone(); - crate::openai::convert::deep_merge_json(&mut merged, call); - let mut opts = options.clone(); - opts.body_overrides = Some(merged); - opts - } - (Some(provider), None) => { - let mut opts = options.clone(); - opts.body_overrides = Some(provider.clone()); - opts - } - (None, _) => options.clone(), - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn config_snapshot_matches_model_identity() { - let config = AnthropicConfig::new("sk-test"); - let model = AnthropicModel::new("claude-sonnet-4-20250514".to_string(), config); - let snap = model.config_snapshot(); - assert_eq!(snap.provider, model.provider()); - assert_eq!(snap.model_id, model.model_id()); - assert!(snap.base_url.is_some()); - assert_eq!(snap.api_key_source, "explicit"); - } -} diff --git a/aimux-providers/src/anthropic/options.rs b/aimux-providers/src/anthropic/options.rs new file mode 100644 index 000000000..bd7adb84a --- /dev/null +++ b/aimux-providers/src/anthropic/options.rs @@ -0,0 +1,64 @@ +//! The providerOptions namespace of the Anthropic Messages protocol. +//! +//! `@ai-sdk/anthropic` always reads `providerOptions.anthropic` and, when the +//! provider was created with a custom `name`, also `providerOptions[]`; the custom key wins where both set the same option. +//! What the model writes back (reasoning signatures, MCP server names, ...) +//! is keyed by that custom name too, so a transcript round-trips under the +//! name the caller chose. This module is the one place that knows the +//! canonical key and how the two are merged. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The canonical providerOptions key, read for every model of the protocol +/// whatever its name. +pub(crate) const CANONICAL: &str = "anthropic"; + +/// The type discriminator for provider-managed container skills. +pub(crate) const PROVIDER_SKILL_TYPE: &str = CANONICAL; + +/// The providerOptions key of a provider name: its first dot-separated +/// segment, trimmed (`"proxy.messages"` -> `"proxy"`, `"proxy"` -> `"proxy"`). +pub(crate) fn options_name_of(provider: &str) -> String { + provider + .split('.') + .next() + .unwrap_or_default() + .trim() + .to_string() +} + +/// Merge the canonical options with the custom-name options (custom wins, +/// shallowly). `None` when neither is set. +fn merge(canonical: Option<&JsonObject>, custom: Option<&JsonObject>) -> Option { + match (canonical, custom) { + (None, None) => None, + (Some(only), None) | (None, Some(only)) => Some(only.clone()), + (Some(canonical), Some(custom)) => { + let mut merged = canonical.clone(); + merged.extend(custom.iter().map(|(k, v)| (k.clone(), v.clone()))); + Some(merged) + } + } +} + +/// The Anthropic options in a providerOptions object (a part's, a message's): +/// `anthropic` merged with `name` when `name` is a different key. +pub(crate) fn anthropic_options( + provider_options: Option<&SharedProviderOptions>, + name: &str, +) -> Option { + let canonical = provider_options.and_then(|options| options.get(CANONICAL)); + let custom = (name != CANONICAL) + .then(|| provider_options.and_then(|options| options.get(name))) + .flatten(); + merge(canonical, custom) +} + +/// [`anthropic_options`] for `CallOptions::provider_options` or a tool's options. +pub(crate) fn anthropic_options_in( + provider_options: Option<&SharedProviderOptions>, + name: &str, +) -> Option { + anthropic_options(provider_options, name) +} diff --git a/aimux-providers/src/anthropic/prepare_tools.rs b/aimux-providers/src/anthropic/prepare_tools.rs index e51249552..69be1f706 100644 --- a/aimux-providers/src/anthropic/prepare_tools.rs +++ b/aimux-providers/src/anthropic/prepare_tools.rs @@ -15,6 +15,8 @@ use aimux_core::tool::LanguageModelV4ToolChoice; use aimux_core::types::Warning; use serde_json::{Map, Value, json}; +use super::options::{CANONICAL, anthropic_options_in}; + /// Result of [`prepare_tools`] / [`prepare_tools_with_provider`]. #[derive(Debug, Clone)] pub struct PreparedTools { @@ -62,6 +64,7 @@ pub fn prepare_tools( supports_structured_output, supports_strict_tools, default_eager_input_streaming, + CANONICAL, &mut betas, &mut tool_warnings, )); @@ -142,6 +145,30 @@ pub fn prepare_tools_with_provider( supports_structured_output: bool, supports_strict_tools: bool, default_eager_input_streaming: bool, +) -> PreparedTools { + prepare_tools_for( + tools, + tool_choice, + disable_parallel_tool_use, + supports_structured_output, + supports_strict_tools, + default_eager_input_streaming, + CANONICAL, + ) +} + +/// [`prepare_tools_with_provider`] for a provider whose providerOptions key is +/// `options_name`: a tool's options are read from `anthropic` merged with that +/// key, the custom one winning. +#[must_use] +pub(crate) fn prepare_tools_for( + tools: Option<&[AnthropicTool]>, + tool_choice: Option<&LanguageModelV4ToolChoice>, + disable_parallel_tool_use: bool, + supports_structured_output: bool, + supports_strict_tools: bool, + default_eager_input_streaming: bool, + options_name: &str, ) -> PreparedTools { // Empty arrays are coerced to "no tools" to match the TS behaviour. let non_empty = tools.filter(|&t| !t.is_empty()); @@ -161,6 +188,7 @@ pub fn prepare_tools_with_provider( supports_structured_output, supports_strict_tools, default_eager_input_streaming, + options_name, &mut betas, &mut tool_warnings, )); @@ -252,13 +280,12 @@ fn prepare_function_tool( supports_structured_output: bool, supports_strict_tools: bool, default_eager_input_streaming: bool, + options_name: &str, betas: &mut BTreeSet, tool_warnings: &mut Vec, ) -> Value { - let anthropic_options = tool - .provider_options - .as_ref() - .and_then(|po| po.get("anthropic")); + let tool_options = anthropic_options_in(tool.provider_options.as_ref(), options_name); + let anthropic_options = tool_options.as_ref(); let eager_input_streaming = anthropic_options .and_then(|o| o.get("eagerInputStreaming")) diff --git a/aimux-providers/src/anthropic/skills.rs b/aimux-providers/src/anthropic/skills.rs new file mode 100644 index 000000000..2a72cb4a4 --- /dev/null +++ b/aimux-providers/src/anthropic/skills.rs @@ -0,0 +1,186 @@ +//! Skill uploads and version metadata from the Anthropic Skills API. + +use std::collections::BTreeSet; + +use async_trait::async_trait; +use serde::Deserialize; +use serde_json::json; + +use aimux_core::AiMuxError; +use aimux_core::files_model::UploadFileData; +use aimux_core::shared::FileBytes; +use aimux_core::skills_model::{Skills, UploadSkillCallOptions, UploadSkillResult}; +use aimux_provider_utils::{HttpBody, HttpRequest}; + +use super::config::AnthropicModelConfig; +use super::options::CANONICAL; + +#[derive(Deserialize)] +struct SkillResponse { + id: String, + display_title: Option, + name: Option, + description: Option, + latest_version: Option, + source: String, + created_at: String, + updated_at: String, +} + +#[derive(Deserialize)] +struct VersionResponse { + #[serde(rename = "type")] + _kind: String, + #[serde(rename = "skill_id")] + _skill_id: String, + name: Option, + description: Option, +} + +/// The provider's skill upload interface. +pub struct AnthropicSkills { + config: AnthropicModelConfig, +} + +impl AnthropicSkills { + pub(crate) fn from_config(config: AnthropicModelConfig) -> Self { + Self { config } + } +} + +// Dot segments need double encoding because URL parsing normalizes even +// percent-encoded dots. Otherwise this is encodeURIComponent's byte encoding. +fn encode_path_segment(value: &str) -> String { + if value == "." { + return "%252E".to_string(); + } + if value == ".." { + return "%252E%252E".to_string(); + } + let mut encoded = String::new(); + for byte in value.bytes() { + if byte.is_ascii_alphanumeric() || b"-_.!~*'()".contains(&byte) { + encoded.push(char::from(byte)); + } else { + use std::fmt::Write; + write!(encoded, "%{byte:02X}").expect("writing to a string cannot fail"); + } + } + encoded +} + +#[async_trait] +impl Skills for AnthropicSkills { + fn provider(&self) -> &str { + &self.config.provider + } + + async fn upload_skill( + &self, + options: &UploadSkillCallOptions, + ) -> Result { + let boundary = format!("----{}", aimux_provider_utils::generate_id()); + let mut body = Vec::new(); + if let Some(title) = &options.display_title { + body.extend_from_slice(format!("--{boundary}\r\nContent-Disposition: form-data; name=\"display_title\"\r\n\r\n{title}\r\n").as_bytes()); + } + for file in &options.files { + // FormData encodes these characters in multipart filenames. + let path = file + .path + .replace('\r', "%0D") + .replace('\n', "%0A") + .replace('"', "%22"); + body.extend_from_slice(format!("--{boundary}\r\nContent-Disposition: form-data; name=\"files[]\"; filename=\"{path}\"\r\nContent-Type: application/octet-stream\r\n\r\n").as_bytes()); + match &file.data { + UploadFileData::Text { text } => body.extend_from_slice(text.as_bytes()), + UploadFileData::Data { + data: FileBytes::Binary(bytes), + } => body.extend_from_slice(bytes), + UploadFileData::Data { + data: FileBytes::Base64(data), + } => { + use base64::Engine; + let normalized = data + .replace('-', "+") + .replace('_', "/") + .chars() + .filter(|character| { + !matches!(character, '\t' | '\n' | '\u{000C}' | '\r' | ' ') + }) + .collect::(); + let engine = base64::engine::GeneralPurpose::new( + &base64::alphabet::STANDARD, + base64::engine::GeneralPurposeConfig::new().with_decode_padding_mode( + base64::engine::DecodePaddingMode::Indifferent, + ), + ); + body.extend_from_slice(&engine.decode(normalized).map_err(|error| { + AiMuxError::InvalidArgument(format!("invalid base64: {error}")) + })?); + } + } + body.extend_from_slice(b"\r\n"); + } + body.extend_from_slice(format!("--{boundary}--\r\n").as_bytes()); + let config = self.config.resolved().await?; + let headers = config + .request_headers(None, &BTreeSet::from(["skills-2025-10-02".to_string()])) + .await?; + let response = aimux_provider_utils::post_to_api( + config.with_transport(HttpRequest { + url: config.url("/skills"), + headers: headers.clone(), + ..HttpRequest::default() + }), + HttpBody::Bytes(body, format!("multipart/form-data; boundary={boundary}")), + aimux_provider_utils::create_json_response_handler(), + config.failed_response_handler(), + ) + .await?; + let response: SkillResponse = response.value; + let version: Option = if let Some(version) = &response.latest_version { + let result = aimux_provider_utils::get_from_api( + config.with_transport(HttpRequest { + url: config.url(&format!( + "/skills/{}/versions/{}", + encode_path_segment(&response.id), + encode_path_segment(version) + )), + headers, + ..HttpRequest::default() + }), + aimux_provider_utils::create_json_response_handler(), + config.failed_response_handler(), + ) + .await?; + Some(result.value) + } else { + None + }; + Ok(UploadSkillResult { + provider_reference: aimux_core::shared::SharedProviderReference::from([( + CANONICAL.to_string(), + response.id, + )]), + display_title: response.display_title, + name: version + .as_ref() + .and_then(|version| version.name.clone()) + .or(response.name), + description: version + .and_then(|version| version.description) + .or(response.description), + latest_version: response.latest_version, + provider_metadata: Some(aimux_core::shared::provider_namespace( + CANONICAL, + json!({ + "source": response.source, + "createdAt": response.created_at, + "updatedAt": response.updated_at, + }), + )?), + warnings: Vec::new(), + }) + } +} diff --git a/aimux-providers/src/anthropic/stream.rs b/aimux-providers/src/anthropic/stream.rs index ec59bcc24..0e919a4ec 100644 --- a/aimux-providers/src/anthropic/stream.rs +++ b/aimux-providers/src/anthropic/stream.rs @@ -1,15 +1,12 @@ //! Shared Anthropic request/streaming core. //! -//! The standard Anthropic provider ([`crate::anthropic::model`]) and the -//! Anthropic-AWS provider (`crate::anthropic_aws::model`) speak the exact -//! same Messages API. The only differences are endpoint, authentication -//! (Bearer/x-api-key vs AWS SigV4) and the wire body encoding (`Json` vs -//! `Bytes`, the latter preventing re-serialization from invalidating the SigV4 -//! signature). +//! The standard Anthropic provider ([`crate::anthropic::model`]), the +//! Anthropic-AWS provider and Anthropic on Vertex all speak the same Messages +//! API through one [`AnthropicMessagesModel`](super::model::AnthropicMessagesModel); +//! what differs (endpoint, credentials, a SigV4 transport decorator, a body +//! envelope) lives in the model's configuration, not here. //! -//! This module factors out the parts that are identical across both providers: -//! - `build_anthropic_request` — serialize the body once, build auth headers -//! via a closure, and choose the wire encoding. +//! This module holds the parts that are identical across all of them: //! - `anthropic_generate_core` — non-streaming send + response parsing. //! - `anthropic_stream_core` — streaming send + the Anthropic SSE event loop. //! - `parse_anthropic_content` — shared content-block → `GenerateContent` @@ -18,21 +15,20 @@ use aimux_core::tool::RawToolCall; use aimux_core::tool::ToolResult; use std::collections::HashMap; -use std::hash::{BuildHasher, RandomState}; -use std::sync::atomic::{AtomicU64, Ordering}; use futures::StreamExt; -use aimux_core::AbortSignal; use aimux_core::error::AiMuxError; +use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; use aimux_core::types::{ProviderMetadata, Warning}; -use aimux_provider_utils::{HttpBody, HttpRequest}; +use aimux_provider_utils::HttpRequest; use serde_json::{Value, json}; +use super::config::AnthropicModelConfig; use super::convert::parse_stop_reason; use super::tool_name_mapping::ToolNameMapping; use super::types::{AnthropicResponse, ContentBlock, StreamErrorData, StreamEvent, ToolCallCaller}; @@ -70,79 +66,11 @@ pub(crate) fn anthropic_stream_error( ) } -/// How the request body is sent over the wire. -#[derive(Debug, Clone, Copy)] -pub(crate) enum BodyEncoding { - /// Send as `HttpBody::Json` — the HTTP layer re-serializes the body value. - /// Standard Anthropic path. - Json, - /// Send as `HttpBody::Bytes` using the *exact* bytes the headers were built - /// over. AWS SigV4 path — prevents re-serialization from breaking the - /// signature (the signature is computed over `body_bytes`, and - /// `body_bytes` is what is sent). - Bytes, -} - -/// Build an Anthropic `HttpRequest`. -/// -/// The body is serialized to bytes once so the same bytes are used both for -/// header construction (the `build_headers` closure — needed by SigV4, which -/// signs the request body) and for the wire body in the [`BodyEncoding::Bytes`] -/// path. For [`BodyEncoding::Json`] the closure receives the serialized bytes -/// but the body is sent as a re-serializable `Value` (the signature is not -/// involved, so re-serialization is harmless). -fn build_anthropic_request( - endpoint: &str, - body: &serde_json::Value, - build_headers: &impl Fn(&[u8], &str) -> Result, AiMuxError>, - body_encoding: BodyEncoding, - abort_signal: Option, -) -> Result<(HttpRequest, HttpBody), AiMuxError> { - // Serialize once; the Bytes path sends these exact bytes and the closure - // signs over them, guaranteeing signature/body agreement. - let body_bytes = serde_json::to_vec(body).map_err(|e| AiMuxError::JsonParse(e.to_string()))?; - let headers = build_headers(&body_bytes, endpoint)?; - - let http_body = match body_encoding { - BodyEncoding::Json => HttpBody::Json(body.clone()), - BodyEncoding::Bytes => HttpBody::Bytes(body_bytes, "application/json".to_string()), - }; - - // The recording context of the operation the call runs in. - let recording_context = aimux_core::recording::RecordingContext::current(); - Ok(( - HttpRequest { - url: endpoint.to_string(), - headers, - abort_signal, - call_id: recording_context.as_ref().map(|c| c.call_id.clone()), - recording_context, - ..Default::default() - }, - http_body, - )) -} - /// Read a string field, dropping absent / non-string values. fn str_field(v: &Value, key: &str) -> Option { v.get(key).and_then(|x| x.as_str()).map(str::to_string) } -/// Generate a non-cryptographic, 16-character alphanumeric ID, like `generateId`. -#[must_use] -pub(crate) fn generate_id() -> String { - const ALPHABET: &[u8] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; - static COUNTER: AtomicU64 = AtomicU64::new(0); - let counter = COUNTER.fetch_add(1, Ordering::Relaxed); - let random = RandomState::new(); - (0..16) - .map(|index| { - let value = random.hash_one((counter, index)); - char::from(ALPHABET[(value % ALPHABET.len() as u64) as usize]) - }) - .collect() -} - /// `web_search_result` → the camel-cased shape upstream exposes (:1243-1249). fn map_web_search_result(result: &Value) -> Value { json!({ @@ -355,6 +283,7 @@ pub(crate) fn initial_tool_input(input: &Value) -> String { pub(crate) fn tool_call_caller_metadata( caller: Option<&ToolCallCaller>, + options_name: &str, ) -> Option { let caller = match caller? { ToolCallCaller::CodeExecution20250825 { tool_id } => json!({ @@ -368,7 +297,7 @@ pub(crate) fn tool_call_caller_metadata( ToolCallCaller::Direct => json!({ "type": "direct" }), }; Some( - provider_namespace("anthropic", json!({ "caller": caller })) + provider_namespace(options_name, json!({ "caller": caller })) .expect("provider metadata must be an object"), ) } @@ -405,6 +334,7 @@ fn tool_search_provider_name<'a>( /// mapping is shared with [`parse_anthropic_content`]; only the part type /// differs. Returns an empty vec for blocks that are not results. pub(crate) fn stream_parts_for_result_block( + options_name: &str, block: &ContentBlock, names: &ToolNameMapping, mcp_tool_calls: &HashMap, @@ -461,7 +391,7 @@ pub(crate) fn stream_parts_for_result_block( title: str_field(result, "title"), provider_metadata: Some( provider_namespace( - "anthropic", + options_name, json!({ "pageAge": result.get("page_age").cloned().unwrap_or(Value::Null), }), @@ -536,7 +466,7 @@ pub(crate) fn stream_parts_for_result_block( dynamic: Some(true), provider_metadata: call.map(|(_, server)| { provider_namespace( - "anthropic", + options_name, json!({ "type": "mcp-tool-use", "serverName": server }), ) .expect("provider metadata must be an object") @@ -557,6 +487,7 @@ 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, blocks: &[ContentBlock], names: &ToolNameMapping, generate_id: &(dyn Fn() -> String + Send + Sync), @@ -603,7 +534,7 @@ pub(crate) fn parse_anthropic_content( input: input.to_string(), provider_executed: None, dynamic: None, - provider_metadata: tool_call_caller_metadata(caller.as_ref()), + provider_metadata: tool_call_caller_metadata(caller.as_ref(), options_name), })); } ContentBlock::Thinking { @@ -613,7 +544,7 @@ pub(crate) fn parse_anthropic_content( content.push(GenerateContent::Reasoning(ReasoningOutput { text: thinking.clone(), provider_metadata: Some( - provider_namespace("anthropic", json!({ "signature": signature })) + provider_namespace(options_name, json!({ "signature": signature })) .expect("provider metadata must be an object"), ), })); @@ -648,7 +579,7 @@ pub(crate) fn parse_anthropic_content( dynamic: Some(true), provider_metadata: Some( provider_namespace( - "anthropic", + options_name, json!({ "type": "mcp-tool-use", "serverName": server_name }), ) .expect("provider metadata must be an object"), @@ -660,7 +591,7 @@ pub(crate) fn parse_anthropic_content( content.push(GenerateContent::Reasoning(ReasoningOutput { text: String::new(), provider_metadata: Some( - provider_namespace("anthropic", json!({ "redactedData": data })) + provider_namespace(options_name, json!({ "redactedData": data })) .expect("provider metadata must be an object"), ), })); @@ -698,7 +629,7 @@ pub(crate) fn parse_anthropic_content( title: str_field(result, "title"), provider_metadata: Some( provider_namespace( - "anthropic", + options_name, json!({ "pageAge": result.get("page_age").cloned() .unwrap_or(Value::Null), @@ -827,7 +758,7 @@ pub(crate) fn parse_anthropic_content( dynamic: Some(true), provider_metadata: call.map(|(_, server)| { provider_namespace( - "anthropic", + options_name, json!({ "type": "mcp-tool-use", "serverName": server }), ) .expect("provider metadata must be an object") @@ -842,29 +773,26 @@ pub(crate) fn parse_anthropic_content( /// Shared non-streaming Anthropic core. /// -/// Builds and sends the request (using `build_headers` for auth and -/// `body_encoding` for the wire body), then parses the `AnthropicResponse` into -/// a `GenerateResult`. The usage breakdown (reasoning / text token split) is the -/// full version, so both the standard and the AWS provider report the same -/// detailed token accounting. -#[allow(clippy::too_many_arguments)] // core plumbing: endpoint/retry/body/warnings/auth/encoding/abort/timeout +/// Sends `body` through `request` (URL, headers and transport already +/// resolved by the model), then parses the `AnthropicResponse` into a +/// `GenerateResult`. The usage breakdown (reasoning / text token split) is the +/// 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. pub(crate) async fn anthropic_generate_core( - endpoint: &str, + request: HttpRequest, body: serde_json::Value, warnings: Vec, - build_headers: impl Fn(&[u8], &str) -> Result, AiMuxError>, - body_encoding: BodyEncoding, - abort_signal: Option, + config: &AnthropicModelConfig, + used_custom_provider_key: bool, tool_names: &ToolNameMapping, generate_id: &(dyn Fn() -> String + Send + Sync), ) -> Result { - let (request, request_body) = - build_anthropic_request(endpoint, &body, &build_headers, body_encoding, abort_signal)?; - let resp = aimux_provider_utils::post_to_api( + let resp = aimux_provider_utils::post_json_to_api( request, - request_body, + body.clone(), aimux_provider_utils::create_json_response_handler(), - super::anthropic_failed_response_handler(), + config.failed_response_handler(), ) .await?; @@ -873,7 +801,12 @@ pub(crate) async fn anthropic_generate_core( let data: AnthropicResponse = resp.value; - let content = parse_anthropic_content(&data.content, tool_names, generate_id); + let content = parse_anthropic_content( + &config.provider_options_name, + &data.content, + tool_names, + generate_id, + ); let finish_reason = data .stop_reason @@ -889,12 +822,21 @@ pub(crate) async fn anthropic_generate_core( // 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( + &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(), + ); + Ok(GenerateResult { content, finish_reason, usage, warnings, - provider_metadata: None, + provider_metadata: Some(provider_metadata), response: Some(aimux_core::shared::ResponseInfo { id: Some(data.id), timestamp: None, @@ -936,49 +878,62 @@ enum BlockState { /// Shared Anthropic streaming core. /// -/// Builds and sends the request (using `build_headers` for auth and -/// `body_encoding` for the wire body), then runs the Anthropic SSE event loop -/// to produce a `StreamResult`. -/// -/// `build_headers` receives the serialized body bytes and the endpoint URL — -/// the standard path returns a Bearer/x-api-key header set (ignoring the body), -/// the AWS path returns a SigV4-signed header set (signing over the body). -#[allow(clippy::too_many_arguments)] // core plumbing: endpoint/retry/body/warnings/auth/encoding/abort/timeout +/// Sends `body` through `request` (URL, headers and transport already +/// 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. pub(crate) async fn anthropic_stream_core( - endpoint: &str, + request: HttpRequest, body: serde_json::Value, warnings: Vec, - build_headers: impl Fn(&[u8], &str) -> Result, AiMuxError>, - body_encoding: BodyEncoding, - abort_signal: Option, - tool_names: ToolNameMapping, + config: &AnthropicModelConfig, + options: &CallOptions, generate_id: std::sync::Arc String + Send + Sync>, ) -> Result { - let (request, request_body) = - build_anthropic_request(endpoint, &body, &build_headers, body_encoding, abort_signal)?; - let resp = aimux_provider_utils::post_to_api( + let used_custom_provider_key = options + .provider_options + .as_ref() + .is_some_and(|options| options.contains_key(&config.provider_options_name)); + let tool_names = ToolNameMapping::new(options.tools.as_deref()); + let include_raw_chunks = options.include_raw_chunks == Some(true); + let endpoint = request.url.clone(); + let options_name = config.provider_options_name.clone(); + let resp = aimux_provider_utils::post_json_to_api( request, - request_body, - aimux_provider_utils::create_event_source_response_handler::(), - super::anthropic_failed_response_handler(), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + config.failed_response_handler(), ) .await?; let response_headers = resp.response_headers; - let mut sse_stream = resp.value; + // Each event as the AI SDK's `ParseResult`: the raw JSON when the data + // parsed as JSON, and the chunk schema's result. + let mut sse_stream = resp.value.map(move |item| match item { + Ok(raw) if include_raw_chunks => { + let event = + serde_json::from_value::(raw.clone()).map_err(AiMuxError::from); + (Some(raw), event) + } + Ok(raw) => ( + None, + serde_json::from_value::(raw).map_err(AiMuxError::from), + ), + Err(error) => (None, Err(error)), + }); let first_event = match sse_stream.next().await { - Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + Some((_, Err(error @ AiMuxError::ApiCall(_)))) => return Err(error), first_event => first_event, }; - if let Some(Ok(StreamEvent::Error { error })) = first_event.as_ref() { + if let Some((_, Ok(StreamEvent::Error { error }))) = first_event.as_ref() { return Err(anthropic_stream_error( error, - endpoint, + &endpoint, body.clone(), response_headers, )); } - let stream_error_url = endpoint.to_string(); + let stream_error_url = endpoint; let stream_request_body = body.clone(); let stream_response_headers = response_headers.clone(); @@ -990,6 +945,13 @@ pub(crate) async fn anthropic_stream_core( let mut blocks: HashMap = HashMap::new(); 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 stop_sequence: Option = None; + let mut container: Option = None; + let mut context_management: Option = None; let mut response_meta_emitted = false; let mut stream_errored = false; // id → (tool name, server name), so `mcp_tool_result` can inherit them @@ -999,7 +961,18 @@ pub(crate) async fn anthropic_stream_core( // one result block type, so the id is required to disambiguate aliases. let mut server_tool_calls: HashMap = HashMap::new(); - while let Some(event) = sse.next().await { + while let Some((raw, event)) = sse.next().await { + // Every chunk, valid or not, is forwarded first when raw chunks + // are requested (anthropic@4.0.71 `anthropic-language-model.ts:1956-1958`); + // a chunk that is not JSON has no raw value. + if include_raw_chunks + && (raw.is_some() + || event.as_ref().is_err_and(AiMuxError::is_recoverable_stream_error)) + { + yield Ok(StreamPart::Raw { + raw_value: raw.unwrap_or(Value::Null), + }); + } match event { Ok(stream_event) => { match stream_event { @@ -1009,6 +982,7 @@ pub(crate) async fn anthropic_stream_core( // 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); } if !response_meta_emitted { yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { @@ -1060,7 +1034,7 @@ 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()), + provider_metadata: tool_call_caller_metadata(caller.as_ref(), options_name.as_str()), }); } // Server-side tool use follows the same input @@ -1116,10 +1090,10 @@ pub(crate) async fn anthropic_stream_core( input: input.to_string(), provider_executed: Some(true), dynamic: Some(true), - provider_metadata: Some(provider_namespace("anthropic", json!({ - "type": "mcp-tool-use", - "serverName": server_name, - })).expect("provider metadata must be an object")), + provider_metadata: Some(provider_namespace(options_name.as_str(), json!({ + "type": "mcp-tool-use", + "serverName": server_name, + })).expect("provider metadata must be an object")), })); } // Redacted thinking — emit as ReasoningStart. @@ -1127,7 +1101,7 @@ pub(crate) async fn anthropic_stream_core( let id = index.to_string(); yield Ok(StreamPart::ReasoningStart { id: id.clone(), - provider_metadata: Some(provider_namespace("anthropic", json!({ "redactedData": data })).expect("provider metadata must be an object")), + provider_metadata: Some(provider_namespace(options_name.as_str(), json!({ "redactedData": data })).expect("provider metadata must be an object")), }); blocks.insert( index, @@ -1144,6 +1118,7 @@ pub(crate) async fn anthropic_stream_core( // (anthropic-language-model.ts:1901-2178). other => { for part in stream_parts_for_result_block( + options_name.as_str(), &other, &tool_names, &mcp_tool_calls, @@ -1293,7 +1268,7 @@ pub(crate) async fn anthropic_stream_core( yield Ok(StreamPart::ReasoningEnd { id: index.to_string(), provider_metadata: signature.map(|s| { - provider_namespace("anthropic", json!({ "signature": s })).expect("provider metadata must be an object") + provider_namespace(options_name.as_str(), json!({ "signature": s })).expect("provider metadata must be an object") }), }); } @@ -1315,7 +1290,8 @@ pub(crate) async fn anthropic_stream_core( id: id.clone(), provider_metadata: None, }); - let input = finalize_streamed_tool_input( + let input = + finalize_streamed_tool_input( accumulated_json, provider_tool_name.as_deref(), provider_tool_input_type.as_deref(), @@ -1332,11 +1308,20 @@ pub(crate) async fn anthropic_stream_core( } } } - StreamEvent::MessageDelta { delta, usage } => { + StreamEvent::MessageDelta { delta, usage, context_management: edits } => { 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); + context_management = edits.or(context_management); if let Some(u) = usage { + if let (Value::Object(raw), Ok(Value::Object(update))) = + (&mut raw_usage, serde_json::to_value(&u)) + { + raw.extend(update); + } + final_usage.raw = raw_usage.as_object().cloned(); let reasoning_tokens = u .output_tokens_details .as_ref() @@ -1404,7 +1389,16 @@ pub(crate) async fn anthropic_stream_core( } else { final_usage }, - provider_metadata: None, + provider_metadata: (!stream_errored).then(|| { + super::usage::result_provider_metadata( + options_name.as_str(), + used_custom_provider_key, + &raw_usage, + stop_sequence.as_deref(), + container.as_ref(), + context_management.as_ref(), + ) + }), }); }; diff --git a/aimux-providers/src/anthropic/types.rs b/aimux-providers/src/anthropic/types.rs index 50b018ea8..066f09519 100644 --- a/aimux-providers/src/anthropic/types.rs +++ b/aimux-providers/src/anthropic/types.rs @@ -15,6 +15,15 @@ pub struct AnthropicResponse { pub content: Vec, pub stop_reason: Option, pub usage: AnthropicUsage, + /// The matched custom stop sequence, when generation ended on one + /// (`providerMetadata.anthropic.stopSequence`). + #[serde(default)] + pub stop_sequence: Option, + /// The code-execution container of the turn, if any + /// (`providerMetadata.anthropic.container`). Left opaque; the model layer + /// maps it into the result-level provider metadata. + #[serde(default)] + pub container: Option, /// Optional `context_management` object echoed by the API, carrying /// `applied_edits`. Left opaque here; the model layer maps it into /// `providerMetadata.anthropic.contextManagement`. @@ -169,22 +178,30 @@ pub enum ContentBlock { /// the count of thinking/reasoning tokens. #[derive(Debug, Default, Serialize, Deserialize)] pub struct OutputTokenDetails { - #[serde(default)] + #[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. #[derive(Debug, Default, Serialize, Deserialize)] pub struct AnthropicUsage { - #[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 output_tokens_details: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_creation_input_tokens: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_read_input_tokens: Option, + #[serde(flatten)] + pub extra: serde_json::Map, } // ── Streaming events ── @@ -208,6 +225,8 @@ pub enum StreamEvent { delta: MessageDeltaBody, #[serde(default)] usage: Option, + #[serde(default)] + context_management: Option, }, #[serde(rename = "message_stop")] MessageStop, @@ -255,4 +274,8 @@ pub struct DeltaBlock { pub struct MessageDeltaBody { #[serde(default)] pub stop_reason: Option, + #[serde(default)] + pub stop_sequence: Option, + #[serde(default)] + pub container: Option, } diff --git a/aimux-providers/src/anthropic/usage.rs b/aimux-providers/src/anthropic/usage.rs index a132f06d0..9568f71cc 100644 --- a/aimux-providers/src/anthropic/usage.rs +++ b/aimux-providers/src/anthropic/usage.rs @@ -5,10 +5,12 @@ //! `packages/anthropic/src/convert-anthropic-usage.ts`. use serde::Deserialize; -use serde_json::Value; +use serde_json::{Value, json}; -use aimux_core::types::Usage; +use aimux_core::shared::provider_namespace; +use aimux_core::types::{ProviderMetadata, Usage}; +use crate::anthropic::options::CANONICAL as CANONICAL_KEY; use crate::anthropic::types::AnthropicUsage; /// A single iteration entry inside an `AnthropicUsage` object. @@ -136,6 +138,126 @@ pub fn convert_anthropic_usage(usage: &Value, raw_usage: Option<&Value>) -> Anth } } +/// The result-level `providerMetadata` of a response, as `@ai-sdk/anthropic` +/// builds it: under the canonical `anthropic` key and, for a provider created +/// with another `name`, under that key too when the request used custom options. +/// +/// `usage` is the raw usage object (for a stream: `message_start`'s, updated +/// by `message_delta`'s). `iterations`, `container` and `contextManagement` +/// 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>, +) -> ProviderMetadata { + let field = |value: &Value, key: &str| value.get(key).cloned().unwrap_or(Value::Null); + let iterations = + usage + .get("iterations") + .and_then(Value::as_array) + .map_or(Value::Null, |iterations| { + Value::Array( + iterations + .iter() + .map(|i| { + let mut iteration = json!({ + "type": field(i, "type"), + "inputTokens": field(i, "input_tokens"), + "outputTokens": field(i, "output_tokens"), + }); + if let Some(model) = i.get("model").filter(|model| !model.is_null()) { + iteration["model"] = model.clone(); + } + for (source, target) 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(); + } + } + iteration + }) + .collect(), + ) + }); + let container = container.map_or(Value::Null, |c| { + let skills = c + .get("skills") + .and_then(Value::as_array) + .map_or(Value::Null, |skills| { + Value::Array( + skills + .iter() + .map(|s| { + json!({ + "type": field(s, "type"), + "skillId": field(s, "skill_id"), + "version": field(s, "version"), + }) + }) + .collect(), + ) + }); + json!({ "expiresAt": field(c, "expires_at"), "id": field(c, "id"), "skills": skills }) + }); + let context_management = context_management.map_or(Value::Null, |cm| { + let edits = cm + .get("applied_edits") + .and_then(Value::as_array) + .map(|edits| { + edits + .iter() + .map(|edit| match edit.as_object() { + Some(edit) => Value::Object( + edit.iter() + .map(|(k, v)| (snake_to_camel(k), v.clone())) + .collect(), + ), + None => edit.clone(), + }) + .collect::>() + }) + .unwrap_or_default(); + json!({ "appliedEdits": edits }) + }); + let metadata = json!({ + "usage": usage, + "stopSequence": stop_sequence, + "iterations": iterations, + "container": container, + "contextManagement": context_management, + }); + let mut result = + provider_namespace(CANONICAL_KEY, metadata).expect("provider metadata must be an object"); + if used_custom_provider_key && options_name != CANONICAL_KEY { + result.insert(options_name.to_string(), result[CANONICAL_KEY].clone()); + } + result +} + +/// `cleared_input_tokens` -> `clearedInputTokens`. +fn snake_to_camel(key: &str) -> String { + let mut out = String::with_capacity(key.len()); + let mut upper = false; + for c in key.chars() { + if c == '_' { + upper = true; + } else if upper { + out.extend(c.to_uppercase()); + upper = false; + } else { + out.push(c); + } + } + out +} + /// Convert a typed `AnthropicUsage` (response/`message_start`) into the /// unified core `Usage`, filling cache fields and the raw payload. /// diff --git a/aimux-providers/src/anthropic_aws/mod.rs b/aimux-providers/src/anthropic_aws/mod.rs index 52699f9a9..12bca8694 100644 --- a/aimux-providers/src/anthropic_aws/mod.rs +++ b/aimux-providers/src/anthropic_aws/mod.rs @@ -1,293 +1,261 @@ //! Anthropic-AWS (Claude Platform on AWS) provider. //! -//! Implements the `LanguageModel` trait against the Claude Platform on AWS API -//! (`aws-external-anthropic.{region}.api.aws/v1/messages`). This is the -//! Anthropic Messages API hosted in AWS, authenticated with either AWS SigV4 -//! or an AWS-provisioned API key. -//! -//! The request/response format is identical to the standard Anthropic API, so -//! this provider reuses the shared [`crate::anthropic::convert`] message -//! conversion logic. The differences are: -//! - **Endpoint**: `aws-external-anthropic.{region}.api.aws/v1/messages` -//! instead of `api.anthropic.com/v1/messages`. -//! - **Authentication**: AWS SigV4 signing (service name -//! `aws-external-anthropic`) or `x-api-key` header. +//! The Claude Platform on AWS API (`aws-external-anthropic.{region}.api.aws/v1`) +//! is the Anthropic Messages API hosted in AWS, so [`create_anthropic_aws`] +//! returns a provider whose models are the shared +//! [`AnthropicMessagesModel`]. Only +//! the endpoint and the authentication differ: an AWS-provisioned API key, or +//! AWS SigV4 applied by a [`SigV4Fetch`](aimux_provider_utils::SigV4Fetch) +//! transport decorator over the final request. //! //! Reference: -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::{RetryConfig, without_trailing_slash}; - mod model; -pub use model::{AnthropicAwsConfig, AnthropicAwsModel}; +pub use model::AnthropicAwsModel; -use crate::bedrock::sigv4::AwsCredentials; +use std::sync::{Arc, OnceLock}; -/// Authentication method for the Anthropic-AWS provider. -#[derive(Debug, Clone)] -pub enum AnthropicAwsAuth { - /// `x-api-key` header authentication (takes precedence when provided). - ApiKey(String), - /// AWS SigV4 signing with static credentials (service: - /// `aws-external-anthropic`). - SigV4(AwsCredentials), -} +use futures::future::BoxFuture; -/// Configuration for the Anthropic-AWS provider. -#[derive(Debug, Clone)] -pub struct AnthropicAwsProviderConfig { - pub base_url: String, - pub auth: AnthropicAwsAuth, - pub api_version: String, - /// Optional workspace ID sent as `anthropic-workspace-id` header. - pub workspace_id: Option, - /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 - pub api_key_source: Option, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, -} +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_provider_utils::{ + AwsCredentials, FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url, +}; -impl AnthropicAwsProviderConfig { - /// Create a config using AWS SigV4 credentials. - pub fn new( - access_key_id: impl Into, - secret_access_key: impl Into, - region: impl Into, - ) -> Self { - let region = region.into(); - let base_url = format!("https://aws-external-anthropic.{region}.api.aws/v1"); - Self { - base_url, - auth: AnthropicAwsAuth::SigV4(AwsCredentials { - access_key_id: access_key_id.into(), - secret_access_key: secret_access_key.into(), - session_token: None, - region, - }), - api_version: "2023-06-01".to_string(), - workspace_id: None, - api_key_source: None, - retry_config: RetryConfig::default(), - } - } +use crate::anthropic::AnthropicMessagesModel; +use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks}; +use crate::anthropic::options::{CANONICAL, options_name_of}; - /// Create a config using an API key (`x-api-key` auth). - pub fn with_api_key(api_key: impl Into, region: impl Into) -> Self { - let region = region.into(); - let base_url = format!("https://aws-external-anthropic.{region}.api.aws/v1"); - Self { - base_url, - auth: AnthropicAwsAuth::ApiKey(api_key.into()), - api_version: "2023-06-01".to_string(), - workspace_id: None, - api_key_source: None, - retry_config: RetryConfig::default(), - } - } +pub use crate::shared::TransformRequestBody; - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// The region used when the settings name none and SigV4 credentials do not +/// carry one. +const DEFAULT_REGION: &str = "us-east-1"; - /// Set the Anthropic API version header. - #[must_use] - pub fn with_api_version(mut self, version: impl Into) -> Self { - self.api_version = version.into(); - self - } +/// The provider name when the settings give none. +const DEFAULT_NAME: &str = "anthropic-aws"; - /// Set the workspace ID. - #[must_use] - pub fn with_workspace_id(mut self, workspace_id: impl Into) -> Self { - self.workspace_id = Some(workspace_id.into()); - self - } +/// Authentication method for the Anthropic-AWS provider. +#[derive(Clone, Debug)] +pub enum AnthropicAwsAuth { + /// `x-api-key` header authentication with an AWS-provisioned key. A + /// [`Resolvable::Future`] is awaited once, an [`Resolvable::AsyncFn`] on + /// every request. + ApiKey(Resolvable), + /// AWS SigV4 signing (service `aws-external-anthropic`). The credentials + /// are resolved on every request, so a refreshing producer can hand out + /// rotating STS credentials. + SigV4(Resolvable), +} - /// 标注凭证来源(RFC-0023 回放重建用)。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self - } +/// Settings of [`create_anthropic_aws`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `region`, `base_url` and `name`; credentials and `headers` +/// are evaluated on every request. +#[derive(Clone, Default)] +pub struct AnthropicAwsProviderSettings { + /// Base URL for the API calls, version segment included. Default + /// `https://aws-external-anthropic.{region}.api.aws/v1`; a trailing slash + /// is removed. + pub base_url: Option, + /// The AWS region of the default base URL. When unset, the region of + /// `Resolvable::Value` SigV4 credentials, else `us-east-1`. Never read + /// from the environment. + pub region: Option, + /// How requests are authenticated. `None` sends `x-api-key` from + /// `ANTHROPIC_AWS_API_KEY`, loaded when a request is made (a missing + /// variable fails that request with `AiMuxError::LoadApiKey`). + pub auth: Option, + /// Sent as `anthropic-workspace-id`. + pub workspace_id: Option, + /// Extra headers on every request. A `None` value removes the header, + /// including one of the fixed ones. Per-call headers win over these. + pub headers: Option, + /// The provider name, the `provider()` string of the messages model. + /// Default `"anthropic-aws"`. The providerOptions key is always + /// `anthropic`, plus the first dot-separated segment of a name that was + /// given here. + pub name: Option, + /// The transport SigV4 signing (when used) wraps. `None` uses the process + /// default, resolved per request. + pub fetch: Option, + /// Rewrites every JSON request body once, after it is serialized and + /// before it is sent (and before it is signed). + pub transform_request_body: Option, +} - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self +impl std::fmt::Debug for AnthropicAwsProviderSettings { + /// Never prints credentials or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AnthropicAwsProviderSettings") + .field("base_url", &self.base_url) + .field("region", &self.region) + .field("auth", &self.auth) + .field("workspace_id", &self.workspace_id) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ) + .finish() } +} - /// Add a session token for temporary STS credentials. - #[must_use] - pub fn with_session_token(mut self, token: impl Into) -> Self { - if let AnthropicAwsAuth::SigV4(ref mut creds) = self.auth { - creds.session_token = Some(token.into()); - } - self +/// A region is a hostname label: refuse anything that could change the host +/// of the default base URL. +fn validate_region(region: &str) -> Result<(), AiMuxError> { + let valid = !region.is_empty() + && region + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'-'); + if valid { + Ok(()) + } else { + Err(AiMuxError::InvalidArgument(format!( + "invalid AWS region {region:?}: expected letters, digits and hyphens" + ))) } +} - /// Create from environment variables. - /// - /// Checks for `ANTHROPIC_AWS_API_KEY` first (API key auth), then falls - /// back to `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` + `AWS_REGION`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when neither `ANTHROPIC_AWS_API_KEY` - /// nor the `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` pair is present. - pub fn from_env() -> Result { - if let Ok(key) = std::env::var("ANTHROPIC_AWS_API_KEY") - && !key.trim().is_empty() - { - let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string()); - return Ok(Self::with_api_key(key, region) - .with_api_key_source(Some("env:ANTHROPIC_AWS_API_KEY"))); +/// Create an Anthropic-AWS provider. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` +/// URL with a host or `region` is not a valid hostname label. Those are the +/// only ways this fails: credentials are loaded per request, not here. +pub fn create_anthropic_aws( + settings: AnthropicAwsProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => { + let from_credentials = match &settings.auth { + Some(AnthropicAwsAuth::SigV4(Resolvable::Value(credentials))) => { + Some(credentials.region.clone()) + } + _ => None, + }; + let region = settings + .region + .clone() + .or(from_credentials) + .unwrap_or_else(|| DEFAULT_REGION.to_string()); + validate_region(®ion)?; + format!("https://aws-external-anthropic.{region}.api.aws/v1") } + }; + // The providerOptions key stays `anthropic` unless the caller named the + // provider. + let provider_options_name = settings + .name + .as_deref() + .map_or_else(|| CANONICAL.to_string(), options_name_of); + let name = settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()); + let (headers, fetch) = model::authenticated( + settings.auth, + settings.workspace_id, + settings.headers, + settings.fetch, + ); + Ok(AnthropicAwsProvider { + name, + provider_options_name, + base_url, + headers, + fetch, + transform_request_body: settings.transform_request_body, + }) +} - let access_key_id = std::env::var("AWS_ACCESS_KEY_ID").map_err(|_| { - AiMuxError::InvalidArgument( - "AWS_ACCESS_KEY_ID (or ANTHROPIC_AWS_API_KEY) environment variable is required for Anthropic-AWS" - .to_string(), - ) - })?; - let secret_access_key = std::env::var("AWS_SECRET_ACCESS_KEY").map_err(|_| { - AiMuxError::InvalidArgument( - "AWS_SECRET_ACCESS_KEY environment variable is required for Anthropic-AWS" - .to_string(), - ) - })?; - let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string()); - - let mut config = Self::new(access_key_id, secret_access_key, region) - .with_api_key_source(Some("env:AWS_ACCESS_KEY_ID")); - if let Ok(token) = std::env::var("AWS_SESSION_TOKEN") { - config = config.with_session_token(token); - } - if let Ok(ws) = std::env::var("ANTHROPIC_AWS_WORKSPACE_ID") { - config = config.with_workspace_id(ws); - } - Ok(config) - } +/// The default provider: `create_anthropic_aws` with default settings, created +/// on first use. Creating it reads nothing from the environment and cannot +/// fail; a missing key surfaces from the first request instead. +pub fn anthropic_aws() -> &'static AnthropicAwsProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_anthropic_aws(AnthropicAwsProviderSettings::default()) + .expect("default Anthropic-AWS settings are always valid") + }) } -/// Anthropic-AWS provider — creates [`AnthropicAwsModel`] instances. +/// An Anthropic-AWS provider. Cheap to clone the models out of; it holds no +/// HTTP client. pub struct AnthropicAwsProvider { - config: AnthropicAwsProviderConfig, + name: String, + provider_options_name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, + transform_request_body: Option, } impl AnthropicAwsProvider { + fn model_config(&self) -> AnthropicModelConfig { + let base = self.base_url.clone(); + AnthropicModelConfig { + provider: self.name.clone(), + url: Arc::new(move |path| format!("{base}{path}")), + headers: self.headers.clone(), + fetch: self.fetch.clone(), + supported_urls: SupportedUrls::default(), + transform_request_body: self.transform_request_body.clone(), + base_url: self.base_url.clone(), + provider_options_name: self.provider_options_name.clone(), + hooks: AnthropicModelHooks::default(), + resolve: None, + } + } + + /// A Messages model; `provider()` is the provider name + /// (`"anthropic-aws"` by default). #[must_use] - pub fn new(config: AnthropicAwsProviderConfig) -> Self { - Self { config } + pub fn messages(&self, model_id: &str) -> AnthropicAwsModel { + AnthropicMessagesModel::with_config(model_id.to_string(), self.model_config()) } - /// Create a model instance for the given Anthropic model id - /// (e.g. `"claude-sonnet-4-20250514"`). + /// The provider as a function: the default language model for an id. #[must_use] - pub fn model(&self, model_id: &str) -> AnthropicAwsModel { - AnthropicAwsModel::new( - model_id.to_string(), - AnthropicAwsConfig { - base_url: self.config.base_url.clone(), - auth: self.config.auth.clone(), - api_version: self.config.api_version.clone(), - workspace_id: self.config.workspace_id.clone(), - api_key_source: self.config.api_key_source.clone(), - retry_config: self.config.retry_config, - }, - ) + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.messages(model_id)) } } impl Provider for AnthropicAwsProvider { - fn name(&self) -> &str { - "anthropic-aws" + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - /// List models via `GET {base_url}/models` (Anthropic-compatible, RFC-0027). - /// API-key auth uses a simple header; SigV4 auth is not supported for - /// listing (returns `Unsupported`). - fn list_models( - &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let config = self.config.clone(); - Box::pin(async move { - let base = config.base_url.trim_end_matches('/'); - let url = format!("{base}/models"); - let mut headers = vec![ - ("anthropic-version".to_string(), config.api_version.clone()), - ("Content-Type".to_string(), "application/json".to_string()), - ]; - match &config.auth { - AnthropicAwsAuth::ApiKey(key) => { - headers.push(("x-api-key".to_string(), key.clone())); - } - AnthropicAwsAuth::SigV4(_) => { - return Err(AiMuxError::UnsupportedFunctionality( - "list_models via SigV4 not supported for Anthropic-AWS; use API key auth" - .into(), - )); - } - } + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } - use aimux_provider_utils::HttpRequest; - // Retry rationale: see `openai::model::execute_list_models`. - let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: headers.clone(), - abort_signal: None, - call_id: None, - recording_context: None, - ..Default::default() - }, - aimux_provider_utils::create_json_response_handler(), - crate::anthropic::anthropic_failed_response_handler(), - ) - }) - .await?; + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } +} - #[derive(serde::Deserialize)] - struct Resp { - #[serde(default)] - data: Vec, - } - #[derive(serde::Deserialize)] - struct Entry { - id: String, - #[serde(default)] - display_name: Option, - } - let parsed: Resp = resp.value; - let runtime: Vec = parsed - .data - .into_iter() - .map(|e| aimux_core::model_catalogue::RuntimeModel { - id: e.id, - owned_by: e.display_name, - created: None, - }) - .collect(); - Ok(runtime) - }) +impl ProviderDiscovery for AnthropicAwsProvider { + /// `GET {base_url}/models`: one exchange, no retry. Signed like any other + /// request when the provider uses SigV4. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config(); + Box::pin(async move { crate::anthropic::list_models_once(&config).await }) } } diff --git a/aimux-providers/src/anthropic_aws/model.rs b/aimux-providers/src/anthropic_aws/model.rs index 2bc277ee1..524870370 100644 --- a/aimux-providers/src/anthropic_aws/model.rs +++ b/aimux-providers/src/anthropic_aws/model.rs @@ -1,196 +1,74 @@ -//! Anthropic-AWS language model — implements `LanguageModel`. +//! Anthropic-AWS model wiring. //! -//! Reuses the shared [`crate::anthropic::convert`] message conversion logic, -//! [`crate::anthropic::types`] response types, and the shared streaming / -//! non-streaming core in [`crate::anthropic::stream`]. Only the endpoint, -//! authentication (AWS SigV4 / x-api-key) and `Bytes` body encoding differ from -//! the standard Anthropic provider — those are supplied via the header-builder -//! closure and [`BodyEncoding::Bytes`]. +//! Claude Platform on AWS serves the Anthropic Messages API unchanged, so its +//! model *is* the shared [`AnthropicMessagesModel`]. What is specific to the +//! host is only how a request is authenticated, and that is data on the +//! model's configuration: an `x-api-key` header, or a [`SigV4Fetch`] transport +//! decorator that signs the final method, URL, headers and body bytes just +//! before they are sent. -use std::collections::HashMap; +use std::sync::Arc; -use async_trait::async_trait; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, SigV4Fetch, default_fetch, +}; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateResult, StreamResult}; -use aimux_provider_utils::RetryConfig; - -use crate::anthropic::convert::build_request_body_with_warnings; -use crate::anthropic::stream::{BodyEncoding, anthropic_generate_core, anthropic_stream_core}; -use crate::anthropic::tool_name_mapping::ToolNameMapping; -use crate::bedrock::sigv4::sign_request; +use crate::anthropic::AnthropicMessagesModel; +use crate::shared::{AuthScheme, Credential, credential_headers}; use super::AnthropicAwsAuth; -/// Configuration for an Anthropic-AWS model instance. -#[derive(Debug, Clone)] -pub struct AnthropicAwsConfig { - pub base_url: String, - pub auth: AnthropicAwsAuth, - pub api_version: String, - pub workspace_id: Option, - /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 - pub api_key_source: Option, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, -} - -/// An Anthropic-AWS language model (e.g. `claude-sonnet-4-20250514`). -pub struct AnthropicAwsModel { - model_id: String, - config: AnthropicAwsConfig, -} - -impl AnthropicAwsModel { - #[must_use] - pub fn new(model_id: String, config: AnthropicAwsConfig) -> Self { - Self { model_id, config } - } - - fn endpoint(&self) -> String { - format!("{}/messages", self.config.base_url) +/// An Anthropic-AWS language model: the shared Messages model configured for +/// the AWS host. +pub type AnthropicAwsModel = AnthropicMessagesModel; + +/// The SigV4 service name of Claude Platform on AWS. +const SIGV4_SERVICE: &str = "aws-external-anthropic"; + +const API_KEY_ENV_VAR: &str = "ANTHROPIC_AWS_API_KEY"; +const API_VERSION: &str = "2023-06-01"; + +/// The provider headers and the transport that authenticate requests with +/// `auth`. +/// +/// An API key is sent as `x-api-key` (`ANTHROPIC_AWS_API_KEY` when none was +/// given). SigV4 sends no credential header at all: the returned transport +/// wraps `fetch` (the process default when `None`) and signs every request it +/// forwards. +pub(super) fn authenticated( + auth: Option, + workspace_id: Option, + user_headers: Option, + fetch: Option, +) -> (HeadersFn, Option) { + let mut fixed = vec![("anthropic-version".to_string(), API_VERSION.to_string())]; + if let Some(workspace_id) = workspace_id { + fixed.push(("anthropic-workspace-id".to_string(), workspace_id)); } - - /// Build the auth-header closure for the Anthropic-AWS path. - /// - /// The serialized body bytes are passed to AWS SigV4 signing (so the - /// signature is computed over the exact bytes that will be sent, preventing - /// re-serialization from invalidating it). For `x-api-key` auth the body is - /// unused. - fn make_header_builder( - &self, - extra: Option<&HashMap>, - ) -> impl Fn(&[u8], &str) -> Result, AiMuxError> { - let api_version = self.config.api_version.clone(); - let workspace_id = self.config.workspace_id.clone(); - let auth = self.config.auth.clone(); - let extra = extra.cloned(); - move |body: &[u8], url: &str| -> Result, AiMuxError> { - let body_str = std::str::from_utf8(body).unwrap_or_default(); - let mut base_headers = vec![ - ("Content-Type".to_string(), "application/json".to_string()), - ("anthropic-version".to_string(), api_version.clone()), - ]; - - if let Some(ref ws) = workspace_id { - base_headers.push(("anthropic-workspace-id".to_string(), ws.clone())); - } - - if let Some(extra) = &extra { - for (k, v) in extra { - base_headers.push((k.clone(), v.clone())); - } - } - - match &auth { - AnthropicAwsAuth::ApiKey(key) => { - base_headers.push(("x-api-key".to_string(), key.clone())); - Ok(base_headers) - } - AnthropicAwsAuth::SigV4(creds) => { - // Sign the request with AWS SigV4 (service: - // aws-external-anthropic). - let extra_for_signing: Vec<(String, String)> = base_headers - .iter() - .filter(|(k, _)| k != "Content-Type") - .cloned() - .collect(); - - let signed = sign_request( - creds, - "aws-external-anthropic", - "POST", - url, - body_str, - &extra_for_signing, - ); - - let mut headers = - vec![("Content-Type".to_string(), "application/json".to_string())]; - for (k, v) in &signed.headers { - headers.push((k.clone(), v.clone())); - } - Ok(headers) - } - } + match auth { + Some(AnthropicAwsAuth::SigV4(credentials)) => { + let inner = fetch.unwrap_or_else(default_fetch); + let signing: FetchFunction = + Arc::new(SigV4Fetch::new(inner, credentials, SIGV4_SERVICE)); + ( + credential_headers(Credential::None, AuthScheme::Bearer, fixed, user_headers), + Some(signing), + ) } - } -} - -#[async_trait] -impl LanguageModel for AnthropicAwsModel { - fn provider(&self) -> &str { - "anthropic-aws" - } - - fn model_id(&self) -> &str { - &self.model_id - } - - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - // M2b: record identity + credential source + endpoint config. Only the - // auth *kind* is serialized — never the SigV4 credentials or api-key - // plaintext. - let auth_kind = match &self.config.auth { - AnthropicAwsAuth::ApiKey(_) => "api_key", - AnthropicAwsAuth::SigV4(_) => "sigv4", - }; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.base_url.clone()), - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: Some(serde_json::json!({ - "auth_kind": auth_kind, - "api_version": self.config.api_version, - "workspace_id": self.config.workspace_id, - })), + key => { + let key: Option> = match key { + Some(AnthropicAwsAuth::ApiKey(key)) => Some(key), + _ => None, + }; + ( + credential_headers( + Credential::explicit_or_env(key, API_KEY_ENV_VAR, "Anthropic AWS"), + AuthScheme::Header("x-api-key"), + fixed, + user_headers, + ), + fetch, + ) } } - - async fn do_generate(&self, options: &CallOptions) -> Result { - let req = build_request_body_with_warnings(&self.model_id, options, false)?; - let endpoint = self.endpoint(); - let build_headers = self.make_header_builder(options.headers.as_ref()); - anthropic_generate_core( - &endpoint, - req.body, - req.warnings, - build_headers, - BodyEncoding::Bytes, - options.abort_signal.clone(), - &ToolNameMapping::new(options.tools.as_deref()), - &crate::anthropic::stream::generate_id, - ) - .await - } - - async fn do_stream(&self, options: &CallOptions) -> Result { - let req = build_request_body_with_warnings(&self.model_id, options, true)?; - let endpoint = self.endpoint(); - let build_headers = self.make_header_builder(options.headers.as_ref()); - anthropic_stream_core( - &endpoint, - req.body, - req.warnings, - build_headers, - BodyEncoding::Bytes, - options.abort_signal.clone(), - ToolNameMapping::new(options.tools.as_deref()), - std::sync::Arc::new(crate::anthropic::stream::generate_id), - ) - .await - } } diff --git a/aimux-providers/src/assemblyai.rs b/aimux-providers/src/assemblyai.rs index fd9b749c5..852942b02 100644 --- a/aimux-providers/src/assemblyai.rs +++ b/aimux-providers/src/assemblyai.rs @@ -7,8 +7,14 @@ //! 1. POST `/v2/upload` with raw audio bytes → returns `upload_url` //! 2. POST `/v2/transcript` with JSON body (audio_url + options) → returns transcript `id` //! 3. GET `/v2/transcript/{id}` polling until status is `completed` +//! +//! [`create_assemblyai`] takes [`AssemblyAIProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`AssemblyAIProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `ASSEMBLYAI_API_KEY`. +//! [`assemblyai()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::Duration; use async_trait::async_trait; use serde::Deserialize; @@ -16,14 +22,17 @@ use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; use aimux_core::error::ApiCallError; -use aimux_core::retry; use aimux_core::shared::{SharedProviderMetadata, Warning}; use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::{ - HttpBody, HttpRequest, load_api_key, sleep_or_abort, without_trailing_slash, +use aimux_provider_utils::HttpBody; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{ + AuthScheme, Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, credential_headers, + is_poll_control_key, poll_interval_ms, poll_until, }; /// AssemblyAI errors: `{"error": "..."}` where `error` is a plain string @@ -59,62 +68,133 @@ fn assemblyai_failed_response_handler() -> aimux_provider_utils::ResponseHandler // ── Config ────────────────────────────────────────────────────────────────── -#[derive(Debug, Clone)] -pub struct AssemblyAIConfig { - pub api_key: String, - pub base_url: String, - pub headers: Option>, +/// Milliseconds between two polls of a transcript +/// (`providerOptions.assemblyai.pollIntervalMs` overrides it for one call). +const POLL_INTERVAL_MS: u64 = 100; +/// How many times a transcript is polled before the call gives up (ten minutes +/// at the default interval). +const MAX_POLL_ATTEMPTS: u32 = 6_000; + +pub(crate) mod options; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.assemblyai.com"; +pub(crate) const API_KEY_ENV_VAR: &str = "ASSEMBLYAI_API_KEY"; +const DEFAULT_NAME: &str = "assemblyai"; + +/// Settings of [`create_assemblyai`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct AssemblyAIProviderSettings { + /// Base URL for the API calls. Default `https://api.assemblyai.com`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `ASSEMBLYAI_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.transcription"`). + /// Default `"assemblyai"`. The providerOptions key stays `assemblyai`. + 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, } -impl AssemblyAIConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.assemblyai.com".to_string(), - headers: None, - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self +impl std::fmt::Debug for AssemblyAIProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AssemblyAIProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } +/// Create a AssemblyAI provider. +/// +/// # 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. +pub fn create_assemblyai( + settings: AssemblyAIProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(AssemblyAIProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "AssemblyAI"), + AuthScheme::Header("Authorization"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `ASSEMBLYAI_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "ASSEMBLYAI_API_KEY", "AssemblyAI")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_assemblyai` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn assemblyai() -> &'static AssemblyAIProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_assemblyai(AssemblyAIProviderSettings::default()) + .expect("default AssemblyAI settings are always valid") + }) } +/// A AssemblyAI provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct AssemblyAIProvider { - config: AssemblyAIConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl AssemblyAIProvider { - #[must_use] - pub fn new(config: AssemblyAIConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// A transcription (STT) model (e.g. `"universal-3-5-pro"`); `provider()` is `"{name}.transcription"`. #[must_use] pub fn transcription(&self, model_id: &str) -> AssemblyAITranscriptionModel { - AssemblyAITranscriptionModel::new(model_id.to_string(), self.config.clone()) + AssemblyAITranscriptionModel::from_config( + model_id.to_string(), + self.model_config("transcription"), + ) } } +crate::impl_single_modality_provider!(AssemblyAIProvider, transcription_model, |p, id| p + .transcription(id)); + // ── Response schema ───────────────────────────────────────────────────────── #[derive(Debug, Deserialize)] @@ -172,53 +252,25 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid base64: {e}"))) } + AudioInput::Url(_) => Err(AudioInput::url_error()), } } pub struct AssemblyAITranscriptionModel { model_id: String, - config: AssemblyAIConfig, + config: EndpointConfig, } impl AssemblyAITranscriptionModel { - #[must_use] - pub fn new(model_id: String, config: AssemblyAIConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("Authorization".to_string(), self.config.api_key.clone()); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn upload_url(&self) -> String { - format!("{}/v2/upload", self.config.base_url) - } - - fn transcript_url(&self) -> String { - format!("{}/v2/transcript", self.config.base_url) - } - - fn transcript_status_url(&self, id: &str) -> String { - format!("{}/v2/transcript/{}", self.config.base_url, id) - } } #[async_trait] impl TranscriptionModel for AssemblyAITranscriptionModel { fn provider(&self) -> &str { - "assemblyai" + &self.config.provider } fn model_id(&self) -> &str { @@ -233,30 +285,18 @@ impl TranscriptionModel for AssemblyAITranscriptionModel { let audio_bytes = audio_input_to_bytes(&options.audio)?; - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); - - // Per-stage retry, not one retry around `do_generate`: Core's outer - // retry would re-upload the audio to retry a later stage. An - // exhausted inner retry returns `AiMuxError::Retry`, which the outer - // retry passes through, so `do_generate` is never replayed (§6.2). - let retries = retry::prepare_retries( - options.max_retries, - self.retry_config(), - options.abort_signal.clone(), - ); - - // Step 1: Upload audio. - let resp = retries - .retry(|| { - aimux_provider_utils::post_to_api( - HttpRequest::new(self.upload_url(), header_list.clone(), options), - HttpBody::Bytes(audio_bytes.clone(), "application/octet-stream".to_string()), - aimux_provider_utils::create_json_response_handler(), - assemblyai_failed_response_handler(), - ) - }) - .await?; + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let aai_options = options::assemblyai_options(options.provider_options.as_ref()); + + // Step 1: Upload audio. One attempt: an upload is never replayed + // implicitly. + let resp = aimux_provider_utils::post_to_api( + exchange.request(exchange.url("/v2/upload"), options), + HttpBody::Bytes(audio_bytes, "application/octet-stream".to_string()), + aimux_provider_utils::create_json_response_handler(), + assemblyai_failed_response_handler(), + ) + .await?; let upload: AssemblyAIUploadResponse = resp.value; @@ -279,176 +319,178 @@ impl TranscriptionModel for AssemblyAITranscriptionModel { body.insert("audio_url".to_string(), json!(upload.upload_url)); - // Parse and forward provider options. - if let Some(ref po) = options.provider_options - && let Some(obj) = po.get("assemblyai") - { + // Parse and forward provider options, except the poll interval, which + // is read here. + if let Some(obj) = aai_options { for (k, v) in obj { - body.insert(k.clone(), v.clone()); + if !is_poll_control_key(k) { + body.insert(k.clone(), v.clone()); + } } } - let body = Value::Object(body); - let resp = retries - .retry(|| { - aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.transcript_url(), header_list.clone(), options), - body.clone(), - aimux_provider_utils::create_json_response_handler(), - assemblyai_failed_response_handler(), - ) - }) - .await?; + // This creates the transcript and is sent exactly once; only the status + // poll below repeats. + let resp = aimux_provider_utils::post_json_to_api( + exchange.request(exchange.url("/v2/transcript"), options), + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + assemblyai_failed_response_handler(), + ) + .await?; let submit: AssemblyAISubmitResponse = resp.value; // Step 3: Poll for completion. - let mut raw_body: Value; - let mut response_headers: HashMap; - loop { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - - let poll_url = self.transcript_status_url(&submit.id); - let resp = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest::new(poll_url.clone(), header_list.clone(), options), - aimux_provider_utils::create_json_response_handler::< - AssemblyAITranscriptResponse, - >(), - assemblyai_failed_response_handler(), - ) - }) + let poll_url = exchange.url(&format!("/v2/transcript/{}", submit.id)); + let resp = poll_until( + &format!("assemblyai transcript {}", submit.id), + options.abort_signal.as_ref(), + Duration::from_millis(poll_interval_ms( + aai_options, + POLL_INTERVAL_MS_KEY, + POLL_INTERVAL_MS, + )), + MAX_POLL_ATTEMPTS, + || async { + let resp = aimux_provider_utils::get_from_api( + exchange.request(poll_url.clone(), options), + aimux_provider_utils::create_json_response_handler::< + AssemblyAITranscriptResponse, + >(), + assemblyai_failed_response_handler(), + ) .await?; - - response_headers = resp.response_headers; - - raw_body = resp.raw_value.unwrap_or(Value::Null); - let parsed = resp.value; - - if parsed.status == "completed" { - // Build segments from words (timestamps are in milliseconds). - let segments: Vec = parsed - .words - .as_ref() - .map(|words| { - words - .iter() - .map(|w| TranscriptionSegment { - text: w.text.clone(), - start_second: w.start / 1000.0, - end_second: w.end / 1000.0, - }) - .collect() - }) - .unwrap_or_default(); - - let language = parsed.language_code.clone(); - let duration_in_seconds = parsed.audio_duration.or_else(|| { - parsed - .words - .as_ref() - .and_then(|w| w.last()) - .map(|w| w.end / 1000.0) - }); - - let text = parsed.text.unwrap_or_default(); - - // Build provider metadata from extra fields. - let mut provider_metadata: Option = None; - let mut md = SharedProviderMetadata::new(); - let mut aai_meta = Map::new(); - if parsed.utterances.is_some() { - aai_meta.insert( - "utterances".to_string(), - raw_body.get("utterances").cloned().unwrap_or(Value::Null), - ); - } - if parsed.sentiment_analysis_results.is_some() { - aai_meta.insert( - "sentimentAnalysisResults".to_string(), - raw_body - .get("sentiment_analysis_results") - .cloned() - .unwrap_or(Value::Null), - ); - } - if parsed.entities.is_some() { - aai_meta.insert( - "entities".to_string(), - raw_body.get("entities").cloned().unwrap_or(Value::Null), - ); - } - if parsed.content_safety_labels.is_some() { - aai_meta.insert( - "contentSafetyLabels".to_string(), - raw_body - .get("content_safety_labels") - .cloned() - .unwrap_or(Value::Null), - ); - } - if parsed.iab_categories_result.is_some() { - aai_meta.insert( - "iabCategoriesResult".to_string(), - raw_body - .get("iab_categories_result") - .cloned() - .unwrap_or(Value::Null), - ); - } - if parsed.auto_highlights_result.is_some() { - aai_meta.insert( - "autoHighlightsResult".to_string(), - raw_body - .get("auto_highlights_result") - .cloned() - .unwrap_or(Value::Null), - ); - } - if !aai_meta.is_empty() { - md.insert("assemblyai".to_string(), aai_meta); - provider_metadata = Some(md); - } - - let timestamp = chrono::Utc::now().to_rfc3339(); - - return Ok(TranscriptionResult { - text, - segments, - language, - duration_in_seconds, - warnings, - request: Some(TranscriptionRequest { body: None }), - response: TranscriptionResponse { - timestamp: Some(timestamp), - model_id: Some(self.model_id.clone()), - headers: Some(response_headers), - body: Some(raw_body), - }, - provider_metadata, - }); - } - - if parsed.status == "error" { - return Err(AiMuxError::ApiCall(Box::new(ApiCallError { - status_code: Some(200), - provider_code: Some("error".to_string()), - response_body: Some(raw_body.to_string()), - ..ApiCallError::new( - format!( - "Transcription failed: {}", - parsed.error.unwrap_or_else(|| "Unknown error".to_string()) + match resp.value.status.clone().as_str() { + "completed" => Ok(PollStep::Ready(resp)), + "error" => Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some("error".to_string()), + response_body: Some( + resp.raw_value.clone().unwrap_or(Value::Null).to_string(), ), - poll_url, - serde_json::json!({}), - ) - }))); - } + ..ApiCallError::new( + format!( + "Transcription failed: {}", + resp.value + .error + .clone() + .unwrap_or_else(|| "Unknown error".to_string()) + ), + poll_url.clone(), + serde_json::json!({}), + ) + }))), + _ => Ok(PollStep::Pending), + } + }, + ) + .await?; + + let response_headers = resp.response_headers; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; + + // Build segments from words (timestamps are in milliseconds). + let segments: Vec = parsed + .words + .as_ref() + .map(|words| { + words + .iter() + .map(|w| TranscriptionSegment { + text: w.text.clone(), + start_second: w.start / 1000.0, + end_second: w.end / 1000.0, + }) + .collect() + }) + .unwrap_or_default(); + + let language = parsed.language_code.clone(); + let duration_in_seconds = parsed.audio_duration.or_else(|| { + parsed + .words + .as_ref() + .and_then(|w| w.last()) + .map(|w| w.end / 1000.0) + }); + + let text = parsed.text.unwrap_or_default(); + + // Build provider metadata from extra fields. + let mut provider_metadata: Option = None; + let mut md = SharedProviderMetadata::new(); + let mut aai_meta = Map::new(); + if parsed.utterances.is_some() { + aai_meta.insert( + "utterances".to_string(), + raw_body.get("utterances").cloned().unwrap_or(Value::Null), + ); + } + if parsed.sentiment_analysis_results.is_some() { + aai_meta.insert( + "sentimentAnalysisResults".to_string(), + raw_body + .get("sentiment_analysis_results") + .cloned() + .unwrap_or(Value::Null), + ); } + if parsed.entities.is_some() { + aai_meta.insert( + "entities".to_string(), + raw_body.get("entities").cloned().unwrap_or(Value::Null), + ); + } + if parsed.content_safety_labels.is_some() { + aai_meta.insert( + "contentSafetyLabels".to_string(), + raw_body + .get("content_safety_labels") + .cloned() + .unwrap_or(Value::Null), + ); + } + if parsed.iab_categories_result.is_some() { + aai_meta.insert( + "iabCategoriesResult".to_string(), + raw_body + .get("iab_categories_result") + .cloned() + .unwrap_or(Value::Null), + ); + } + if parsed.auto_highlights_result.is_some() { + aai_meta.insert( + "autoHighlightsResult".to_string(), + raw_body + .get("auto_highlights_result") + .cloned() + .unwrap_or(Value::Null), + ); + } + if !aai_meta.is_empty() { + md.insert(options::NAMESPACE.to_string(), aai_meta); + provider_metadata = Some(md); + } + + let timestamp = chrono::Utc::now().to_rfc3339(); + + Ok(TranscriptionResult { + text, + segments, + language, + duration_in_seconds, + warnings, + request: Some(TranscriptionRequest { body: None }), + response: TranscriptionResponse { + timestamp: Some(timestamp), + model_id: Some(self.model_id.clone()), + headers: Some(response_headers), + body: Some(raw_body), + }, + provider_metadata, + }) } } diff --git a/aimux-providers/src/assemblyai/options.rs b/aimux-providers/src/assemblyai/options.rs new file mode 100644 index 000000000..8fe772c00 --- /dev/null +++ b/aimux-providers/src/assemblyai/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of AssemblyAI. +//! +//! AssemblyAI reads `providerOptions.assemblyai`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "assemblyai"; + +/// The AssemblyAI options in a providerOptions map. +pub(crate) fn assemblyai_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/aws_polly.rs b/aimux-providers/src/aws_polly.rs index 63f9d4c67..39540ace3 100644 --- a/aimux-providers/src/aws_polly.rs +++ b/aimux-providers/src/aws_polly.rs @@ -3,9 +3,9 @@ //! Implements the `SpeechModel` trait against the Amazon Polly `SynthesizeSpeech` //! API (`POST https://polly.{region}.amazonaws.com/v1/speech`). //! -//! Authentication uses AWS Signature Version 4 (SigV4), reusing the -//! [`AwsCredentials`] / [`sign_request`] helpers shared with the Bedrock -//! provider. The service name is `"polly"`. +//! Authentication uses AWS Signature Version 4 (SigV4), applied by the shared +//! [`SigV4Fetch`] transport decorator over the final request bytes (the same +//! decorator Bedrock uses). The service name is `"polly"`. //! //! The Polly API accepts a JSON request body describing the synthesis job and //! returns the synthesized audio as a binary stream (the response body is *not* @@ -16,23 +16,30 @@ //! `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_REGION` (defaulting to //! `us-east-1`), and the optional `AWS_SESSION_TOKEN` for temporary STS //! credentials. +//! +//! [`create_aws_polly`] takes [`AwsPollyProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`AwsPollyProvider`]. The credentials are not +//! read there: they are loaded for every request, from the settings or from the +//! AWS environment variables. +//! [`aws_polly()`] is the default instance; it reads nothing and cannot fail. + +use std::sync::{Arc, OnceLock}; use async_trait::async_trait; use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::shared::{SharedProviderOptions, Warning}; use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; +use aimux_provider_utils::HttpBody; +use aimux_provider_utils::{ + AwsCredentials, FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, SigV4Fetch, default_fetch, + load_optional_setting, load_setting, validate_base_url, +}; -use aimux_provider_utils::{HttpBody, HttpRequest}; - -use crate::bedrock::sigv4::{AwsCredentials, sign_request}; - -/// Provider canonical name. -const PROVIDER_NAME: &str = "amazon-polly"; +use crate::shared::{Endpoint, EndpointConfig}; /// AWS service name used for SigV4 signing. const SERVICE_NAME: &str = "polly"; @@ -80,171 +87,281 @@ fn aws_error_parts(data: &Value) -> aimux_provider_utils::ProviderErrorParts { // ── Config ─────────────────────────────────────────────────────────────────── -/// Configuration for the Amazon Polly provider. +pub(crate) mod options; + +const DEFAULT_NAME: &str = "amazon-polly"; +const REGION_ENV_VAR: &str = "AWS_REGION"; +const ACCESS_KEY_ID_ENV_VAR: &str = "AWS_ACCESS_KEY_ID"; +const SECRET_ACCESS_KEY_ENV_VAR: &str = "AWS_SECRET_ACCESS_KEY"; +const SESSION_TOKEN_ENV_VAR: &str = "AWS_SESSION_TOKEN"; + +/// Settings of [`create_aws_polly`]. /// -/// Holds AWS SigV4 credentials. The `Debug` implementation redacts secrets so -/// credentials never appear in logs or error messages. -pub struct AwsPollyConfig { - access_key_id: String, - secret_access_key: String, - region: String, - session_token: Option, - base_url: String, +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; the region, the credentials and +/// `headers` are evaluated on every request. +#[derive(Clone, Default)] +pub struct AwsPollyProviderSettings { + /// The AWS region. `None` loads `AWS_REGION` when a request is made and + /// falls back to `us-east-1`. + pub region: Option, + /// The AWS access key id of SigV4 signing. `None` loads + /// `AWS_ACCESS_KEY_ID` when a request is made and fails that request with + /// `AiMuxError::LoadSetting` if it is unset. + pub access_key_id: Option, + /// The AWS secret access key of SigV4 signing. `None` loads + /// `AWS_SECRET_ACCESS_KEY`. + pub secret_access_key: Option, + /// The AWS session token of temporary credentials. When `access_key_id` + /// and `secret_access_key` are both given only this field is used; when + /// either comes from the environment the token also falls back to + /// `AWS_SESSION_TOKEN`. + pub session_token: Option, + /// Dynamic AWS credentials. When set they are used instead of + /// `access_key_id`, `secret_access_key` and `session_token`; it is + /// resolved on every request, so an [`Resolvable::AsyncFn`] can hand out + /// rotating STS credentials. Its `region` is used only when neither + /// `region` nor `AWS_REGION` names one. + pub credential_provider: Option>, + /// Base URL for the API calls. Default `https://polly.{region}.amazonaws.com`; + /// 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. + pub headers: Option, + /// The provider name, the prefix of the `provider()` string + /// (`"{name}.speech"`). Default `"amazon-polly"`. The providerOptions key + /// is `aws_polly`. + pub name: Option, + /// The transport SigV4 signing wraps. `None` uses the process default, + /// resolved per request. + pub fetch: Option, } -impl AwsPollyConfig { - /// Create a config with explicit static SigV4 credentials. - pub fn new( - access_key_id: impl Into, - secret_access_key: impl Into, - region: impl Into, - ) -> Self { - let region = region.into(); - let base_url = format!("https://polly.{region}.amazonaws.com"); - Self { - access_key_id: access_key_id.into(), - secret_access_key: secret_access_key.into(), - region, - session_token: None, - base_url, - } +impl std::fmt::Debug for AwsPollyProviderSettings { + /// Never prints keys, credentials or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AwsPollyProviderSettings") + .field("region", &self.region) + .field("access_key_id", &self.access_key_id.is_some()) + .field("secret_access_key", &self.secret_access_key.is_some()) + .field("session_token", &self.session_token.is_some()) + .field("credential_provider", &self.credential_provider.is_some()) + .field("base_url", &self.base_url) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Override the base URL (for testing or proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = url.into().trim_end_matches('/').to_string(); - self - } +/// Where the region and the AWS credentials of a request come from. +#[derive(Clone)] +struct Auth { + region: Option, + access_key_id: Option, + secret_access_key: Option, + session_token: Option, + credential_provider: Option>, +} - /// Add a session token for temporary STS credentials. - #[must_use] - pub fn with_session_token(mut self, token: impl Into) -> Self { - self.session_token = Some(token.into()); - self +impl Auth { + /// The region: the setting, `AWS_REGION`, the region of a plain-value + /// credential provider, then `us-east-1`. + fn region(&self) -> String { + if let Some(region) = load_optional_setting(self.region.as_deref(), REGION_ENV_VAR) { + return region; + } + if let Some(Resolvable::Value(credentials)) = &self.credential_provider + && !credentials.region.is_empty() + { + return credentials.region.clone(); + } + DEFAULT_REGION.to_string() } - /// Create from environment variables. - /// - /// Reads `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` + `AWS_REGION` - /// (defaulting to `us-east-1`), and the optional `AWS_SESSION_TOKEN`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `AWS_ACCESS_KEY_ID` or - /// `AWS_SECRET_ACCESS_KEY` is not set. - pub fn from_env() -> Result { - let access_key_id = std::env::var("AWS_ACCESS_KEY_ID").map_err(|_| { - AiMuxError::InvalidArgument( - "AWS_ACCESS_KEY_ID environment variable is required for Amazon Polly SigV4 auth" - .to_string(), - ) - })?; - let secret_access_key = std::env::var("AWS_SECRET_ACCESS_KEY").map_err(|_| { - AiMuxError::InvalidArgument( - "AWS_SECRET_ACCESS_KEY environment variable is required for Amazon Polly SigV4 auth" - .to_string(), - ) - })?; - let region = std::env::var("AWS_REGION").unwrap_or_else(|_| DEFAULT_REGION.to_string()); + /// The credentials of the settings and the environment (no provider). + fn static_credentials(&self) -> Result { + let access_key_id = load_setting( + self.access_key_id.as_deref(), + ACCESS_KEY_ID_ENV_VAR, + "access_key_id", + )?; + let secret_access_key = load_setting( + self.secret_access_key.as_deref(), + SECRET_ACCESS_KEY_ENV_VAR, + "secret_access_key", + )?; + let session_token = if self.access_key_id.is_some() && self.secret_access_key.is_some() { + self.session_token.clone() + } else { + load_optional_setting(self.session_token.as_deref(), SESSION_TOKEN_ENV_VAR) + .filter(|token| !token.trim().is_empty()) + }; + Ok(AwsCredentials { + access_key_id, + secret_access_key, + session_token, + region: self.region(), + }) + } - let mut config = Self::new(access_key_id, secret_access_key, region); - if let Ok(token) = std::env::var("AWS_SESSION_TOKEN") - && !token.trim().is_empty() - { - config = config.with_session_token(token); + /// The credentials to sign one request with. + async fn credentials(&self) -> Result { + if let Some(provider) = &self.credential_provider { + let mut credentials = provider.resolve().await?; + if self.region.is_some() || std::env::var(REGION_ENV_VAR).is_ok() { + credentials.region = self.region(); + } else if credentials.region.is_empty() { + credentials.region = DEFAULT_REGION.to_string(); + } + return Ok(credentials); } - Ok(config) + self.static_credentials() } - /// Build the [`AwsCredentials`] used for signing. - fn credentials(&self) -> AwsCredentials { - AwsCredentials { - access_key_id: self.access_key_id.clone(), - secret_access_key: self.secret_access_key.clone(), - session_token: self.session_token.clone(), - region: self.region.clone(), + /// Fail now, with the typed setting error, when signing could not find its + /// inputs: the transport reports a credential failure as a transport + /// error, which would hide which setting is missing. A credential provider + /// is left to the signer. + fn preflight(&self) -> Result<(), AiMuxError> { + if self.credential_provider.is_none() { + self.static_credentials()?; } + Ok(()) } } -impl Clone for AwsPollyConfig { - fn clone(&self) -> Self { - Self { - access_key_id: self.access_key_id.clone(), - secret_access_key: self.secret_access_key.clone(), - region: self.region.clone(), - session_token: self.session_token.clone(), - base_url: self.base_url.clone(), +/// The provider headers: only the caller's, after checking that signing has +/// its credentials (the signature itself is made by the transport). +fn signing_headers(auth: Auth, user: Option) -> HeadersFn { + Resolvable::from_async_fn(move || { + let auth = auth.clone(); + let user = user.clone(); + async move { + auth.preflight()?; + Ok(user.unwrap_or_default()) } - } + }) } -impl std::fmt::Debug for AwsPollyConfig { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AwsPollyConfig") - .field("access_key_id", &"") - .field("secret_access_key", &"") - .field("region", &self.region) - .field( - "session_token", - &self.session_token.as_ref().map(|_| ""), - ) - .field("base_url", &self.base_url) - .finish() - } +/// Create an Amazon Polly provider. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` +/// URL with a host. That is the only way this fails: the region and the +/// credentials are loaded per request, not here. +pub fn create_aws_polly( + settings: AwsPollyProviderSettings, +) -> Result { + let base_url = settings + .base_url + .as_deref() + .map(validate_base_url) + .transpose()?; + let auth = Auth { + region: settings.region, + access_key_id: settings.access_key_id, + secret_access_key: settings.secret_access_key, + session_token: settings.session_token, + credential_provider: settings.credential_provider, + }; + let signing_auth = auth.clone(); + let headers_auth = auth.clone(); + let fetch: FetchFunction = Arc::new(SigV4Fetch::new( + settings.fetch.unwrap_or_else(default_fetch), + Resolvable::from_async_fn(move || { + let auth = signing_auth.clone(); + async move { auth.credentials().await } + }), + SERVICE_NAME, + )); + Ok(AwsPollyProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + auth, + base_url, + // SigV4 signs in the transport; the provider headers carry no credential. + headers: signing_headers(headers_auth, settings.headers), + fetch, + }) } -// ── Provider ───────────────────────────────────────────────────────────────── +/// The default provider: `create_aws_polly` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// missing credentials surface from the first request instead. +pub fn aws_polly() -> &'static AwsPollyProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_aws_polly(AwsPollyProviderSettings::default()) + .expect("default Amazon Polly settings are always valid") + }) +} -/// Amazon Polly provider — creates [`AwsPollySpeechModel`] instances. +/// An Amazon Polly provider. Speech only; it holds no HTTP client. pub struct AwsPollyProvider { - config: AwsPollyConfig, + name: String, + auth: Auth, + base_url: Option, + headers: HeadersFn, + fetch: FetchFunction, } impl AwsPollyProvider { - #[must_use] - pub fn new(config: AwsPollyConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + let auth = self.auth.clone(); + let base_url = self.base_url.clone(); + let headers = self.headers.clone(); + EndpointConfig::dynamic( + format!("{}.{method}", self.name), + Arc::new(move || { + let auth = auth.clone(); + let base_url = base_url.clone(); + let headers = headers.clone(); + Box::pin(async move { + Ok(Endpoint { + base_url: base_url.unwrap_or_else(|| { + format!("https://polly.{}.amazonaws.com", auth.region()) + }), + headers: headers.resolve().await?, + }) + }) + }), + Some(self.fetch.clone()), + ) } - /// Create a speech (TTS) model instance for the given model id - /// (e.g. `"aws_polly/neural"` or `"neural"`). + /// A speech (TTS) model (an engine id such as `"neural"`); `provider()` is + /// `"{name}.speech"`. #[must_use] pub fn speech(&self, model_id: &str) -> AwsPollySpeechModel { - AwsPollySpeechModel::new(model_id.to_string(), self.config.clone()) + AwsPollySpeechModel::from_config(model_id.to_string(), self.model_config("speech")) } } -impl Provider for AwsPollyProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } -} +crate::impl_single_modality_provider!(AwsPollyProvider, speech_model, |p, id| p.speech(id)); // ── Speech model ───────────────────────────────────────────────────────────── /// An Amazon Polly speech (TTS) model. pub struct AwsPollySpeechModel { model_id: String, - config: AwsPollyConfig, + config: EndpointConfig, } impl AwsPollySpeechModel { - #[must_use] - pub fn new(model_id: String, config: AwsPollyConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn endpoint(&self) -> String { - format!("{}/v1/speech", self.config.base_url) - } } #[async_trait] impl SpeechModel for AwsPollySpeechModel { fn provider(&self) -> &str { - PROVIDER_NAME + &self.config.provider } fn model_id(&self) -> &str { @@ -255,39 +372,13 @@ impl SpeechModel for AwsPollySpeechModel { let (body, warnings) = build_request(options, &self.model_id); let body_str = serde_json::to_string(&Value::Object(body.clone())) .map_err(|e| AiMuxError::JsonParse(e.to_string()))?; - let url = self.endpoint(); - - // SigV4 sign the request. User-supplied extra headers are included in - // the canonical headers; `Content-Type` is added separately (mirrors the - // Bedrock provider). The signed body bytes are sent verbatim via - // `HttpBody::Bytes` so the `x-amz-content-sha256` hash still matches — - // re-serializing the JSON would invalidate the signature. - let creds = self.config.credentials(); - let mut extra_headers: Vec<(String, String)> = Vec::new(); - if let Some(ref headers) = options.headers { - for (k, v) in headers { - extra_headers.push((k.clone(), v.clone())); - } - } - let signed = sign_request( - &creds, - SERVICE_NAME, - "POST", - &url, - &body_str, - &extra_headers, - ); + // The request is signed by the transport (`SigV4Fetch`), over the exact + // bytes sent: user-supplied headers and `Content-Type` are part of the + // signature. + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_to_api( - HttpRequest { - url, - headers: signed.headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(exchange.url("/v1/speech"), options), HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), aimux_provider_utils::create_binary_response_handler(), aws_failed_response_handler(), @@ -446,7 +537,7 @@ struct PollySpeechProviderOptions { fn parse_polly_provider_options( options: Option<&SharedProviderOptions>, ) -> Option { - let opts = options.and_then(|opts| opts.get("aws_polly"))?; + let opts = options::aws_polly_options(options)?; Some(PollySpeechProviderOptions { engine: opts @@ -506,19 +597,6 @@ mod tests { assert!(warning.is_some()); } - #[test] - fn config_debug_redacts_credentials() { - let config = AwsPollyConfig::new("AKIAEXAMPLE", "super-secret-key", "us-west-2") - .with_session_token("session-token"); - let debug = format!("{config:?}"); - assert!(debug.contains("")); - assert!(!debug.contains("AKIAEXAMPLE")); - assert!(!debug.contains("super-secret-key")); - assert!(!debug.contains("session-token")); - // Non-secret fields are still visible. - assert!(debug.contains("us-west-2")); - } - #[test] fn extracts_aws_error_fields() { let data = serde_json::json!({ diff --git a/aimux-providers/src/aws_polly/options.rs b/aimux-providers/src/aws_polly/options.rs new file mode 100644 index 000000000..3eb1e501a --- /dev/null +++ b/aimux-providers/src/aws_polly/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Amazon Polly. +//! +//! Amazon Polly reads `providerOptions.aws_polly`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "aws_polly"; + +/// The Amazon Polly options in a providerOptions map. +pub(crate) fn aws_polly_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/azure/mod.rs b/aimux-providers/src/azure/mod.rs index 3b3e39637..9d1deb243 100644 --- a/aimux-providers/src/azure/mod.rs +++ b/aimux-providers/src/azure/mod.rs @@ -1,16 +1,519 @@ //! # Azure OpenAI provider //! -//! Implements the Azure OpenAI chat-completions API. Azure speaks the OpenAI -//! wire format but uses a deployment-based URL with an `api-version` query -//! parameter and authenticates via either an `api-key` header or an Azure AD -//! (Microsoft Entra ID) bearer token. +//! [`create_azure`] is the Rust form of the AI SDK's `createAzure`: it takes +//! [`AzureOpenAIProviderSettings`], rejects an `api_key` given together 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 +//! (`AZURE_RESOURCE_NAME`) are resolved on every request, so a missing one +//! fails that request with `AiMuxError::LoadApiKey` / `LoadSetting`. +//! [`azure()`] is the default instance. //! -//! See [`AzureProvider`] and [`AzureConfig`]. +//! Azure speaks the OpenAI wire format, so its models are the OpenAI package's +//! (`crate::openai`) configured for Azure: provider strings +//! `azure.chat` / `azure.responses` / `azure.embeddings` / `azure.image` / +//! `azure.transcription` / `azure.speech`, a URL that follows the AI SDK's rules +//! (`{base}/v1{path}?api-version=` or the deployment form +//! `{base}/deployments/{id}{path}?api-version=`), `api-key` or +//! `Authorization: Bearer` authentication, and the Responses model's `azure` +//! 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. + +pub(crate) mod options; + +pub use crate::shared::TransformRequestBody; -pub mod model; -pub mod responses; +use std::sync::{Arc, OnceLock}; -pub use model::{ - AzureAuth, AzureConfig, AzureModel, AzureProvider, DEFAULT_API_VERSION, TokenProvider, +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::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_core::speech_model::SpeechModel; +use aimux_core::transcription_model::TranscriptionModel; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, Resolvable, combine_headers, load_setting, + validate_base_url, without_trailing_slash, }; -pub use responses::AzureResponsesModel; + +use crate::openai::config::OpenAIModelConfig; +use crate::openai::responses::ResponsesProfile; +use crate::openai::{ + OpenAIEmbeddingModel, OpenAIImageModel, OpenAIModel, OpenAIResponsesModel, OpenAISpeechModel, + OpenAITranscriptionModel, +}; +use crate::shared::{AuthScheme, Credential, credential_headers, is_valid_hostname_part}; + +/// The chat-completions model of the Azure package (`provider.chat(id)`): the +/// OpenAI one, configured for Azure. +pub type AzureChatModel = OpenAIModel; +/// The Responses model of the Azure package: the OpenAI one, configured for +/// Azure. +pub type AzureResponsesModel = OpenAIResponsesModel; + +pub(crate) const API_KEY_ENV_VAR: &str = "AZURE_API_KEY"; +const RESOURCE_NAME_ENV_VAR: &str = "AZURE_RESOURCE_NAME"; +/// The `api-version` of the AI SDK's v1 URL form. +const DEFAULT_API_VERSION: &str = "v1"; +/// The `api-version` the deployments listing needs: the v1 API has no such +/// call, so discovery keeps the dated version unless the caller sets one. +const DEPLOYMENTS_API_VERSION: &str = "2024-10-21"; + +/// Settings of [`create_azure`] (the AI SDK's `AzureOpenAIProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and the `api_key` / `token_provider` conflict; +/// the credential, the resource name and `headers` are evaluated on every +/// request. +#[derive(Clone, Default)] +pub struct AzureOpenAIProviderSettings { + /// Name of the Azure OpenAI resource, the `{resource}` in + /// `https://{resource}.openai.azure.com/openai`. Either this or `base_url` + /// can be used; `None` (and no `base_url`) loads `AZURE_RESOURCE_NAME` + /// when a request is made and fails that request with + /// `AiMuxError::LoadSetting` if it is unset. It must be a single DNS label. + pub resource_name: Option, + /// A different URL prefix for API calls, e.g. a proxy. Wins over + /// `resource_name`. With an unversioned Azure OpenAI base URL the request + /// URL is `{base_url}/v1{path}`; one that already ends in `/openai/v1` is + /// used as is; a non-Azure gateway gets `{base_url}{path}`. A trailing + /// slash is removed. + pub 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. + 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 + /// every request. A caller-set `Authorization` header wins over it. + pub token_provider: Option>, + /// Extra headers on every request. A `None` value removes the header. + /// Per-call headers win over these. + pub headers: Option, + /// The transport: a mock, a signing decorator, a proxy-aware client. + /// `None` uses the process default, resolved per request. + pub fetch: Option, + /// The `api-version` query parameter. Default `"v1"`. + pub api_version: Option, + /// Use the legacy deployment URL form + /// `{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 { + /// Never prints credentials or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AzureOpenAIProviderSettings") + .field("resource_name", &self.resource_name) + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field("token_provider", &self.token_provider) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .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() + } +} + +/// What the AI SDK's `getAzureOpenAIBaseURLInfo` derives from a base URL. +#[derive(Debug, Clone, Copy)] +struct BaseUrlInfo { + is_azure_openai: bool, + is_foundry_project: bool, + is_versioned: bool, +} + +fn base_url_info(base_url: Option<&str>) -> Result { + let Some(base_url) = base_url else { + return Ok(BaseUrlInfo { + is_azure_openai: true, + is_foundry_project: false, + is_versioned: false, + }); + }; + let url = url::Url::parse(base_url) + .map_err(|e| AiMuxError::InvalidArgument(format!("invalid Azure base URL: {e}")))?; + let host = url.host_str().unwrap_or_default().to_ascii_lowercase(); + let is_azure_openai = host.ends_with(".openai.azure.com") + || host.ends_with(".services.ai.azure.com") + || host.ends_with(".cognitiveservices.azure.com"); + let path = url.path().trim_end_matches('/').to_string(); + Ok(BaseUrlInfo { + is_azure_openai, + is_foundry_project: host.ends_with(".services.ai.azure.com") + && path.starts_with("/api/projects/"), + is_versioned: is_azure_openai && path.to_ascii_lowercase().ends_with("/openai/v1"), + }) +} + +/// Create an Azure OpenAI provider. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when both `api_key` and +/// `token_provider` are given or `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() { + return Err(AiMuxError::InvalidArgument( + "Both apiKey and tokenProvider were provided. Please use only one authentication \ + method." + .to_string(), + )); + } + let base_url = settings + .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()])), + None => credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Azure OpenAI"), + AuthScheme::Header("api-key"), + Vec::new(), + settings.headers, + ), + }; + Ok(AzureOpenAIProvider { + resource_name: settings.resource_name, + base_url, + info, + api_version: settings.api_version, + use_deployment_based_urls: settings.use_deployment_based_urls, + 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, + }) +} + +/// 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. +pub fn azure() -> &'static AzureOpenAIProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_azure(AzureOpenAIProviderSettings::default()) + .expect("default Azure settings are always valid") + }) +} + +/// An Azure OpenAI provider (the AI SDK's `AzureOpenAIProvider`). Models are +/// addressed by deployment name. Cheap to clone the models out of; it holds no +/// HTTP client. +pub struct AzureOpenAIProvider { + resource_name: Option, + base_url: Option, + 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. +#[derive(Clone)] +struct UrlRules { + resource_name: Option, + base_url: Option, + info: BaseUrlInfo, + api_version: String, + use_deployment_based_urls: bool, +} + +impl UrlRules { + /// The URL prefix: the base URL, or the resource's `/openai` endpoint. + fn prefix(&self) -> Result { + if let Some(base_url) = &self.base_url { + return Ok(without_trailing_slash(base_url)); + } + let resource = load_setting( + self.resource_name.as_deref(), + RESOURCE_NAME_ENV_VAR, + "resourceName", + )?; + // The resource name becomes part of the request host, so only a DNS + // label is accepted (`user@internal:8080/#` would rewrite the host). + if !is_valid_hostname_part(&resource) { + return Err(AiMuxError::InvalidArgument( + "Invalid Azure resource name. Expected a single DNS label (letters, digits, and \ + hyphens). Use `base_url` for custom endpoints." + .to_string(), + )); + } + Ok(format!("https://{resource}.openai.azure.com/openai")) + } + + /// The AI SDK's `url({ path, modelId })`. + fn url(&self, path: &str, model_id: &str) -> Result { + let prefix = self.prefix()?; + let info = self.info; + let full = if self.use_deployment_based_urls { + format!("{prefix}/deployments/{model_id}{path}") + } else if !info.is_azure_openai || info.is_versioned { + // Custom gateways own Azure routing and versioning themselves, and + // complete Azure OpenAI v1 URLs own theirs. + format!("{prefix}{path}") + } else { + format!("{prefix}/v1{path}") + }; + let mut url = url::Url::parse(&full) + .map_err(|e| AiMuxError::InvalidArgument(format!("invalid Azure URL `{full}`: {e}")))?; + if self.use_deployment_based_urls + || (info.is_azure_openai && !info.is_versioned && !info.is_foundry_project) + { + set_query_param(&mut url, "api-version", &self.api_version); + } + Ok(url.to_string()) + } +} + +/// `URLSearchParams.set`: replace the parameter, or append it. +fn set_query_param(url: &mut url::Url, name: &str, value: &str) { + let mut pairs: Vec<(String, String)> = url + .query_pairs() + .map(|(k, v)| (k.into_owned(), v.into_owned())) + .collect(); + match pairs.iter().position(|(k, _)| k == name) { + Some(index) => { + pairs[index].1 = value.to_string(); + let mut seen = false; + pairs.retain(|(k, _)| { + if k == name { + let keep = !seen; + seen = true; + keep + } else { + true + } + }); + } + None => pairs.push((name.to_string(), value.to_string())), + } + url.query_pairs_mut().clear().extend_pairs(pairs); +} + +impl AzureOpenAIProvider { + fn url_rules(&self) -> UrlRules { + UrlRules { + resource_name: self.resource_name.clone(), + base_url: self.base_url.clone(), + info: self.info, + api_version: self + .api_version + .clone() + .unwrap_or_else(|| DEFAULT_API_VERSION.to_string()), + use_deployment_based_urls: self.use_deployment_based_urls, + } + } + + /// The model configuration of a deployment, reporting `provider` as its + /// identity. + fn model_config(&self, provider: &str, deployment: &str) -> OpenAIModelConfig { + let rules = self.url_rules(); + let deployment = deployment.to_string(); + OpenAIModelConfig { + provider: provider.to_string(), + url: Arc::new(move |path| rules.url(path, &deployment)), + headers: self.headers.clone(), + token_provider: self.token_provider.clone(), + fetch: self.fetch.clone(), + supported_urls: aimux_core::language_model::SupportedUrls::default(), + transform_request_body: self.transform_request_body.clone(), + responses: ResponsesProfile::default(), + chat_options: options::CHAT_OPTIONS, + } + } + + /// A chat-completions model for a deployment; `provider()` is + /// `"azure.chat"`. + #[must_use] + pub fn chat(&self, deployment: &str) -> AzureChatModel { + OpenAIModel::from_config( + deployment.to_string(), + self.model_config("azure.chat", deployment), + ) + } + + /// A Responses API model for a deployment; `provider()` is + /// `"azure.responses"`. Reads providerOptions from `azure` (then `openai`), + /// writes metadata under `azure`, and treats `assistant-` file data as + /// uploaded file ids. + #[must_use] + pub fn responses(&self, deployment: &str) -> AzureResponsesModel { + let mut config = self.model_config("azure.responses", deployment); + config.responses = ResponsesProfile { + namespace: options::RESPONSES, + file_id_prefixes: vec!["assistant-"], + }; + OpenAIResponsesModel::from_config(deployment.to_string(), config) + } + + /// An embedding model for a deployment; `provider()` is + /// `"azure.embeddings"`. + #[must_use] + pub fn embedding(&self, deployment: &str) -> OpenAIEmbeddingModel { + OpenAIEmbeddingModel::from_config( + deployment.to_string(), + self.model_config("azure.embeddings", deployment), + ) + } + + /// An image model for a deployment (e.g. DALL-E); `provider()` is + /// `"azure.image"`. + #[must_use] + pub fn image(&self, deployment: &str) -> OpenAIImageModel { + OpenAIImageModel::from_config( + deployment.to_string(), + self.model_config("azure.image", deployment), + ) + } + + /// A transcription model for a deployment; `provider()` is + /// `"azure.transcription"`. + #[must_use] + pub fn transcription(&self, deployment: &str) -> OpenAITranscriptionModel { + OpenAITranscriptionModel::from_config( + deployment.to_string(), + self.model_config("azure.transcription", deployment), + ) + } + + /// A speech model for a deployment; `provider()` is `"azure.speech"`. + #[must_use] + pub fn speech(&self, deployment: &str) -> OpenAISpeechModel { + OpenAISpeechModel::from_config( + deployment.to_string(), + self.model_config("azure.speech", deployment), + ) + } + + /// The provider as a function: the default language model for a + /// deployment. The AI SDK's callable provider returns the Responses model, + /// the same one as [`responses`](Self::responses) and + /// [`language_model`](Provider::language_model). + #[must_use] + pub fn call(&self, deployment: &str) -> Arc { + Arc::new(self.responses(deployment)) + } +} + +impl Provider for AzureOpenAIProvider { + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) + } + + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.image(model_id))) + } + + fn transcription_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.transcription(model_id)))) + } + + fn speech_model(&self, model_id: &str) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.speech(model_id)))) + } +} + +impl ProviderDiscovery for AzureOpenAIProvider { + /// `GET {prefix}/deployments?api-version=...`: one exchange, no retry. + /// Azure lists *deployments*, not models; each deployment's `id` is the + /// model id to pass to [`Provider::language_model`]. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let mut rules = self.url_rules(); + if self.api_version.is_none() { + rules.api_version = DEPLOYMENTS_API_VERSION.to_string(); + } + let config = self.model_config("azure.responses", ""); + let fetch = self.fetch.clone(); + Box::pin(async move { + // Azure response: { data: [{ id, model, modelName, ... }] } + #[derive(serde::Deserialize)] + struct Resp { + #[serde(default)] + data: Vec, + } + #[derive(serde::Deserialize)] + struct Entry { + id: String, + #[serde(default)] + model: Option, + #[serde(default, rename = "modelName")] + model_name: Option, + } + + let provider_headers = config.request_headers(None).await?; + let mut url = url::Url::parse(&format!("{}/deployments", rules.prefix()?)) + .map_err(|e| AiMuxError::InvalidArgument(format!("invalid Azure URL: {e}")))?; + set_query_param(&mut url, "api-version", &rules.api_version); + let url = url.to_string(); + let resp = aimux_provider_utils::get_from_api( + HttpRequest { + credentialed_origin: Some(url.clone()), + fetch, + url, + headers: provider_headers, + ..Default::default() + }, + aimux_provider_utils::create_json_response_handler(), + crate::openai::openai_failed_response_handler(), + ) + .await?; + let parsed: Resp = resp.value; + Ok(parsed + .data + .into_iter() + .map(|entry| RuntimeModel { + id: entry.id, + owned_by: entry.model_name.or(entry.model), + created: None, + }) + .collect()) + }) + } +} diff --git a/aimux-providers/src/azure/model.rs b/aimux-providers/src/azure/model.rs deleted file mode 100644 index b7376abe3..000000000 --- a/aimux-providers/src/azure/model.rs +++ /dev/null @@ -1,618 +0,0 @@ -//! Azure OpenAI language model — implements `LanguageModel` trait. -//! -//! Azure OpenAI speaks the OpenAI chat-completions wire format but differs in -//! two respects: -//! -//! 1. **URL construction** — Azure addresses a *deployment* rather than a -//! model, and pins the API version via a `?api-version=` query parameter. -//! The classic deployment-based form is -//! `https://{resource}.openai.azure.com/openai/deployments/{deployment}/chat/completions?api-version={api_version}`. -//! A newer `v1` form (`…/openai/v1/chat/completions?api-version=…`) is also -//! supported via [`AzureConfig::use_v1_urls`]. -//! -//! 2. **Authentication** — either an `api-key` header (API key) or a -//! `Bearer` token supplied by a [`TokenProvider`] (Azure AD / Microsoft -//! Entra ID). The two are mutually exclusive. -//! -//! The actual request/response conversion is shared with the OpenAI provider -//! via [`crate::openai::model::execute_generate`] / -//! [`crate::openai::model::execute_stream`]. - -use std::collections::HashMap; -use std::sync::Arc; - -use async_trait::async_trait; - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::options::CallOptions; -use aimux_core::provider::Provider; -use aimux_core::result::{GenerateResult, StreamResult}; - -use aimux_provider_utils::{ - RetryConfig, load_api_key, with_user_agent_suffix, without_trailing_slash, -}; - -use crate::azure::responses::AzureResponsesModel; -use crate::openai::model::{execute_generate, execute_stream}; - -/// Default Azure OpenAI API version used for the deployment-based endpoint. -pub const DEFAULT_API_VERSION: &str = "2024-10-21"; - -// ── Token provider (Azure AD / Microsoft Entra ID) ─────────────────────────── - -/// A source of Azure AD (Microsoft Entra ID) bearer tokens. -/// -/// Implementors return a fresh access token on each call; the token is sent as -/// `Authorization: Bearer `. This mirrors the TS `tokenProvider: -/// () => Promise` option. -#[async_trait] -pub trait TokenProvider: Send + Sync { - /// Return a bearer access token. - async fn get_token(&self) -> Result; -} - -/// Authentication strategy for Azure OpenAI. -#[derive(Clone)] -pub enum AzureAuth { - /// API-key auth — sent as the `api-key` header. - ApiKey(String), - /// Azure AD (Microsoft Entra ID) auth — sent as `Authorization: Bearer `. - TokenProvider(Arc), -} - -// ── Config ─────────────────────────────────────────────────────────────────── - -/// Configuration for the Azure OpenAI provider. -/// -/// Either `resource_name` or `base_url` must be set (the former builds the -/// standard `https://{resource}.openai.azure.com/openai` prefix; the latter -/// overrides it entirely, e.g. for a proxy/gateway). -#[derive(Clone)] -pub struct AzureConfig { - /// Azure OpenAI resource name (the `{resource}` in the hostname). - pub resource_name: Option, - /// Override the URL prefix (takes precedence over `resource_name`). - pub base_url: Option, - /// `api-version` query parameter. - pub api_version: String, - /// Use the deployment-based URL form - /// (`…/deployments/{deployment}/chat/completions`). Defaults to `true`. - pub use_deployment_based_urls: bool, - /// Authentication. Mutually exclusive at construction time: providing both - /// an API key and a token provider is rejected by [`AzureProvider::new`]. - pub auth: Option, - /// Provider-level extra headers merged into every request. - pub extra_headers: HashMap, - /// api_key 来源(RFC-0023):`None` = explicit (api-key 或 token provider); - /// `Some("env:VAR")` = 环境变量。不存明文。 - pub api_key_source: Option, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, -} - -impl AzureConfig { - /// Create a new config with default settings (deployment-based URLs, - /// `api-version = 2024-10-21`, no auth yet). - #[must_use] - pub fn new() -> Self { - Self { - resource_name: None, - base_url: None, - api_version: DEFAULT_API_VERSION.to_string(), - use_deployment_based_urls: true, - auth: None, - extra_headers: HashMap::new(), - api_key_source: None, - retry_config: RetryConfig::default(), - } - } - - /// Set the Azure resource name. - #[must_use] - pub fn with_resource_name(mut self, name: impl Into) -> Self { - self.resource_name = Some(name.into()); - self - } - - /// Override the base URL prefix (takes precedence over `resource_name`). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = Some(without_trailing_slash(&url.into())); - self - } - - /// Set the `api-version` query parameter. - #[must_use] - pub fn with_api_version(mut self, version: impl Into) -> Self { - self.api_version = version.into(); - self - } - - /// Authenticate with an API key (sent as the `api-key` header). - #[must_use] - pub fn with_api_key(mut self, api_key: impl Into) -> Self { - self.auth = Some(AzureAuth::ApiKey(api_key.into())); - self - } - - /// Authenticate with an Azure AD token provider (sent as `Bearer`). - #[must_use] - pub fn with_token_provider(mut self, provider: Arc) -> Self { - self.auth = Some(AzureAuth::TokenProvider(provider)); - self - } - - /// 标注 api_key 来源(RFC-0023 回放重建用)。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self - } - - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self - } - - /// Add a provider-level extra header. - #[must_use] - pub fn with_header(mut self, key: impl Into, value: impl Into) -> Self { - self.extra_headers.insert(key.into(), value.into()); - self - } - - /// Use the deployment-based URL form (the default). - #[must_use] - pub fn use_deployment_based_urls(mut self) -> Self { - self.use_deployment_based_urls = true; - self - } - - /// Use the newer `v1` URL form (`…/openai/v1/chat/completions?api-version=…`). - #[must_use] - pub fn use_v1_urls(mut self) -> Self { - self.use_deployment_based_urls = false; - self - } - - /// Create a config from environment variables. - /// - /// Reads `AZURE_API_KEY` (required) and `AZURE_RESOURCE_NAME` (optional — - /// may also be supplied via `with_resource_name`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `AZURE_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "AZURE_API_KEY", "Azure OpenAI")?; - let mut config = Self::new() - .with_api_key(api_key) - .with_api_key_source(Some("env:AZURE_API_KEY")); - if let Ok(resource) = std::env::var("AZURE_RESOURCE_NAME") - && !resource.is_empty() - { - config = config.with_resource_name(resource); - } - Ok(config) - } -} - -impl Default for AzureConfig { - fn default() -> Self { - Self::new() - } -} - -// ── Provider ───────────────────────────────────────────────────────────────── - -/// Azure OpenAI provider — creates [`AzureModel`] instances for a deployment. -pub struct AzureProvider { - config: AzureConfig, -} - -impl AzureProvider { - /// Create a provider from a config. - /// - /// Returns [`AiMuxError::InvalidArgument`] if both an API key and a token - /// provider were supplied, or if neither `resource_name` nor `base_url` is - /// set. - /// - /// # Errors - /// - /// Returns [`AiMuxError::InvalidArgument`] when both an API key and a token - /// provider are supplied, or when neither `resource_name` nor `base_url` is - /// set. - pub fn new(config: AzureConfig) -> Result { - // Reject both auth methods — mirrors the TS `createAzure` guard. - // (The builder API makes it hard to set both, but a hand-built config - // could still carry conflicting state, so we validate here.) - if config.resource_name.is_none() && config.base_url.is_none() { - return Err(AiMuxError::InvalidArgument( - "Azure OpenAI requires either `resource_name` or `base_url` to be set.".to_string(), - )); - } - Ok(Self { config }) - } - - /// Create a model instance for the given deployment name - /// (e.g. `"gpt-4o-deployment"`). - #[must_use] - pub fn deployment(&self, deployment: &str) -> AzureModel { - AzureModel::new(deployment.to_string(), self.config.clone()) - } - - /// Alias for [`Self::deployment`] (matches the OpenAI provider's `model`). - #[must_use] - pub fn model(&self, deployment: &str) -> AzureModel { - self.deployment(deployment) - } - - /// Create a Responses API model instance for the given deployment name. - /// - /// Uses the `/responses` endpoint (via the Azure OpenAI v1 or - /// deployment-based URL form) instead of `/chat/completions`. - /// - /// Mirrors the TS `createResponsesModel` / `provider.responses()` factory. - #[must_use] - pub fn responses_model(&self, deployment: &str) -> AzureResponsesModel { - AzureResponsesModel::new(deployment.to_string(), self.config.clone()) - } - - /// Alias for [`Self::responses_model`] (matches the TS `provider.responses`). - #[must_use] - pub fn responses(&self, deployment: &str) -> AzureResponsesModel { - self.responses_model(deployment) - } -} - -impl Provider for AzureProvider { - fn name(&self) -> &str { - "azure" - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.deployment(model_id))) - } - - /// List deployments via `GET {prefix}/deployments?api-version=...` - /// (Azure OpenAI, RFC-0027). Azure lists *deployments*, not models — each - /// deployment's `id` is the model_id used by `language_model`. - fn list_models( - &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let config = self.config.clone(); - Box::pin(async move { - let prefix = match &config.base_url { - Some(url) => aimux_provider_utils::without_trailing_slash(url), - None => format!( - "https://{}.openai.azure.com/openai", - config.resource_name.as_deref().ok_or_else(|| { - AiMuxError::InvalidArgument( - "azure list_models: resource_name or base_url required".into(), - ) - })? - ), - }; - let url = format!("{prefix}/deployments?api-version={}", config.api_version); - - // Auth: api-key header or Bearer token. - let mut headers = vec![("Content-Type".to_string(), "application/json".to_string())]; - match &config.auth { - Some(AzureAuth::ApiKey(key)) => { - headers.push(("api-key".to_string(), key.clone())); - } - Some(AzureAuth::TokenProvider(tp)) => { - let token = tp.get_token().await.map_err(|e| { - AiMuxError::InvalidArgument(format!("azure token provider: {e}")) - })?; - headers.push(("Authorization".to_string(), format!("Bearer {token}"))); - } - None => { - return Err(AiMuxError::InvalidArgument( - "azure list_models: no auth configured".to_string(), - )); - } - } - - use aimux_provider_utils::HttpRequest; - // Retry rationale: see `openai::model::execute_list_models`. - let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: headers.clone(), - abort_signal: None, - call_id: None, - recording_context: None, - ..Default::default() - }, - aimux_provider_utils::create_json_response_handler(), - crate::openai::openai_failed_response_handler(), - ) - }) - .await?; - - // Azure response: { data: [{ id, model, ... }] } - #[derive(serde::Deserialize)] - struct Resp { - #[serde(default)] - data: Vec, - } - #[derive(serde::Deserialize)] - struct Entry { - id: String, - #[serde(default)] - model: Option, - #[serde(default, rename = "modelName")] - model_name: Option, - } - let parsed: Resp = resp.value; - let runtime: Vec = parsed - .data - .into_iter() - .map(|e| aimux_core::model_catalogue::RuntimeModel { - id: e.id, - owned_by: e.model_name.or(e.model), - created: None, - }) - .collect(); - Ok(runtime) - }) - } -} - -// ── Model ──────────────────────────────────────────────────────────────────── - -/// An Azure OpenAI language model (addressed by deployment name). -pub struct AzureModel { - deployment: String, - config: AzureConfig, -} - -impl AzureModel { - #[must_use] - pub fn new(deployment: String, config: AzureConfig) -> Self { - Self { deployment, config } - } - - /// Build the chat-completions endpoint URL for this deployment. - /// - /// Deployment-based (default): - /// `{base}/openai/deployments/{deployment}/chat/completions?api-version={v}` - /// - /// v1 form (`use_v1_urls`): - /// `{base}/openai/v1/chat/completions?api-version={v}` - /// - /// where `base` is `https://{resource}.openai.azure.com/openai` (derived - /// from `resource_name`) unless overridden by `base_url`. When `base_url` - /// is set to a non-Azure gateway, `api-version` is omitted (the gateway - /// owns its own versioning), matching the TS behaviour. - fn endpoint(&self) -> String { - let prefix = match &self.config.base_url { - Some(url) => without_trailing_slash(url), - None => format!( - "https://{}.openai.azure.com/openai", - self.config - .resource_name - .as_deref() - .expect("resource_name or base_url required (validated at construction)") - ), - }; - - let path = if self.config.use_deployment_based_urls { - format!("/deployments/{}/chat/completions", self.deployment) - } else { - "/v1/chat/completions".to_string() - }; - - // Only append api-version for the Azure OpenAI endpoint or the - // deployment-based form. A custom (non-Azure) gateway baseURL with the - // v1 form owns its own versioning — mirror the TS `useAzureOpenAIEndpoint` - // gate. - let is_azure = prefix.contains(".openai.azure.com"); - if is_azure || self.config.use_deployment_based_urls { - format!("{}{}?api-version={}", prefix, path, self.config.api_version) - } else { - format!("{prefix}{path}") - } - } - - /// Build the request headers: auth + provider extra headers + per-request - /// headers, plus the `ai-sdk/azure/…` user-agent suffix. - async fn build_headers( - &self, - extra: Option<&HashMap>, - ) -> Result, AiMuxError> { - let mut headers = HashMap::new(); - - match &self.config.auth { - Some(AzureAuth::ApiKey(key)) => { - headers.insert("api-key".to_string(), key.clone()); - } - Some(AzureAuth::TokenProvider(tp)) => { - let token = tp.get_token().await?; - headers.insert("Authorization".to_string(), format!("Bearer {token}")); - } - None => { - return Err(AiMuxError::InvalidArgument( - "Azure OpenAI has no authentication configured. \ - Provide an `api_key` or a `token_provider`." - .to_string(), - )); - } - } - - // Provider-level extra headers. - for (k, v) in &self.config.extra_headers { - headers.insert(k.clone(), v.clone()); - } - // Per-request headers override provider headers. - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - - with_user_agent_suffix(&mut headers, "azure"); - Ok(headers) - } -} - -#[async_trait] -impl LanguageModel for AzureModel { - fn provider(&self) -> &str { - "azure" - } - - fn model_id(&self) -> &str { - &self.deployment - } - - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id().to_string(), - base_url: self.config.base_url.clone(), - // Azure 不记录明文 key/token;来源按构造方式标注(env/explicit)。 - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: Some(serde_json::json!({ - "resource_name": self.config.resource_name, - "api_version": self.config.api_version, - "use_deployment_based_urls": self.config.use_deployment_based_urls, - })), - } - } - - async fn do_generate(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()).await?; - execute_generate( - &self.endpoint(), - &headers, - &self.deployment, - options, - "azure", - &crate::openai::OpenAICompatProfile::full(), - ) - .await - } - - async fn do_stream(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()).await?; - execute_stream( - &self.endpoint(), - &headers, - &self.deployment, - options, - "azure", - &crate::openai::OpenAICompatProfile::full(), - ) - .await - } -} - -// ── Unit tests for URL construction ────────────────────────────────────────── -#[cfg(test)] -mod tests { - use super::*; - - fn model(resource: &str) -> AzureModel { - AzureModel::new( - "gpt-4o".to_string(), - AzureConfig::new() - .with_resource_name(resource) - .with_api_key("k"), - ) - } - - /// Deployment-based (default) URL with a resource name: - /// `https://{resource}.openai.azure.com/openai/deployments/{deployment}/chat/completions?api-version=2024-10-21`. - #[test] - fn endpoint_deployment_based_with_resource_name() { - let m = model("my-resource"); - assert_eq!( - m.endpoint(), - "https://my-resource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21" - ); - } - - /// The v1 form on a real Azure endpoint still appends `api-version`. - #[test] - fn endpoint_v1_with_resource_name_appends_api_version() { - let mut m = model("my-resource"); - m.config = m.config.use_v1_urls(); - assert_eq!( - m.endpoint(), - "https://my-resource.openai.azure.com/openai/v1/chat/completions?api-version=2024-10-21" - ); - } - - /// A custom `api_version` is reflected in the query parameter. - #[test] - fn endpoint_custom_api_version() { - let mut m = model("my-resource"); - m.config = m.config.with_api_version("2025-04-01-preview"); - assert_eq!( - m.endpoint(), - "https://my-resource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2025-04-01-preview" - ); - } - - /// A `base_url` override pointing at an Azure host still appends - /// `api-version`. - #[test] - fn endpoint_azure_base_url_appends_api_version() { - let m = AzureModel::new( - "gpt-4o".to_string(), - AzureConfig::new() - .with_base_url("https://other.openai.azure.com/openai") - .with_api_key("k") - .use_v1_urls(), - ); - assert_eq!( - m.endpoint(), - "https://other.openai.azure.com/openai/v1/chat/completions?api-version=2024-10-21" - ); - } - - /// RFC-0023:config_snapshot 反映 provider/model_id、base_url、api_key_source。 - #[test] - fn config_snapshot_matches_config() { - let m = AzureModel::new( - "gpt-4o".to_string(), - AzureConfig::new() - .with_base_url("https://gateway.example.com/openai") - .with_api_version("2025-04-01-preview") - .use_v1_urls() - .with_api_key("k"), - ); - let snap = m.config_snapshot(); - assert_eq!(snap.provider, "azure"); - assert_eq!(snap.model_id, "gpt-4o"); - assert_eq!( - snap.base_url.as_deref(), - Some("https://gateway.example.com/openai") - ); - assert_eq!(snap.api_key_source, "explicit"); - } -} diff --git a/aimux-providers/src/azure/options.rs b/aimux-providers/src/azure/options.rs new file mode 100644 index 000000000..7b713d90c --- /dev/null +++ b/aimux-providers/src/azure/options.rs @@ -0,0 +1,22 @@ +//! The providerOptions and providerMetadata namespace of Azure OpenAI. +//! +//! 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. + +use crate::openai::responses::ResponsesNamespace; + +/// The Azure providerOptions / providerMetadata key. +pub(crate) const NAMESPACE: &str = "azure"; + +/// What the Responses model reads and writes: `azure` first, then `openai`; +/// metadata under `azure`. +pub(crate) const RESPONSES: ResponsesNamespace = + ResponsesNamespace::new(&[NAMESPACE, "openai"], NAMESPACE); + +/// Chat uses the same option schema as the OpenAI package. +pub(crate) const CHAT_OPTIONS: crate::openai::options::ChatOptionsParser = + crate::openai::options::parse_chat_options; diff --git a/aimux-providers/src/azure/responses.rs b/aimux-providers/src/azure/responses.rs deleted file mode 100644 index 75b65e234..000000000 --- a/aimux-providers/src/azure/responses.rs +++ /dev/null @@ -1,384 +0,0 @@ -//! Azure OpenAI Responses API language model. -//! -//! Implements the [`LanguageModel`] trait against the Azure OpenAI -//! `/responses` endpoint. Azure speaks the same Responses API wire format as -//! OpenAI (requests carry an `input` array, streaming events are typed as -//! `response.output_text.delta`, etc.) but differs in: -//! -//! 1. **URL construction** — Azure addresses a *deployment* and pins the API -//! version via a `?api-version=` query parameter. The v1 form is -//! `{base}/v1/responses?api-version={v}`; the deployment-based form is -//! `{base}/deployments/{deployment}/responses?api-version={v}`. -//! 2. **Authentication** — either an `api-key` header (API key) or a `Bearer` -//! token supplied by a `TokenProvider` (Azure AD / Microsoft Entra ID). -//! The two are mutually exclusive. -//! 3. **File ID prefix** — Azure file IDs use the `assistant-` prefix. When a -//! file content part's data starts with `assistant-`, it is passed through -//! as a `file_id` instead of being base64-encoded. -//! -//! The request-body building logic is shared with the OpenAI Responses -//! provider via [`crate::openai::responses::convert`]; the non-streaming -//! output parsing and the streaming SSE event reducer are shared via -//! [`crate::openai::responses::responses_convert`] (RFC-0012 §3.5). -//! -//! Mirrors the TS `createResponsesModel` factory in -//! `azure-openai-provider.ts`. - -use std::collections::HashMap; - -use async_trait::async_trait; -use futures::StreamExt; -use serde_json::{Value, json}; - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateResult, StreamResult}; - -use aimux_provider_utils::{HttpRequest, with_user_agent_suffix, without_trailing_slash}; - -use crate::azure::{AzureAuth, AzureConfig}; -use crate::openai::responses::convert::build_responses_request_body; -use crate::openai::responses::responses_convert::{ - build_header_list, build_responses_event_stream, build_responses_generate_result, -}; - -/// The file-ID prefix Azure uses for uploaded files. -const AZURE_FILE_ID_PREFIX: &str = "assistant-"; - -/// An Azure OpenAI Responses API language model (addressed by deployment name). -/// -/// Created via [`AzureProvider::responses_model`](super::AzureProvider::responses_model). -/// -/// This mirrors the TS `OpenAIResponsesLanguageModel` configured with -/// `provider: 'azure.responses'` and `fileIdPrefixes: ['assistant-']`. -pub struct AzureResponsesModel { - deployment: String, - config: AzureConfig, -} - -impl AzureResponsesModel { - #[must_use] - pub fn new(deployment: String, config: AzureConfig) -> Self { - Self { deployment, config } - } - - /// Build the Responses endpoint URL for this deployment. - /// - /// v1 form (default): `{base}/v1/responses?api-version={v}` - /// - /// Deployment-based: `{base}/deployments/{deployment}/responses?api-version={v}` - /// - /// Custom gateway (non-Azure baseURL, v1 form): `{base}/responses` — the - /// gateway owns its own versioning, so `api-version` is omitted. - fn endpoint(&self) -> String { - let prefix = match &self.config.base_url { - Some(url) => without_trailing_slash(url), - None => format!( - "https://{}.openai.azure.com/openai", - self.config - .resource_name - .as_deref() - .expect("resource_name or base_url required (validated at construction)") - ), - }; - - let path = if self.config.use_deployment_based_urls { - format!("/deployments/{}/responses", self.deployment) - } else { - "/v1/responses".to_string() - }; - - // Only append api-version for the Azure OpenAI endpoint or the - // deployment-based form. A custom (non-Azure) gateway baseURL with the - // v1 form owns its own versioning. - let is_azure = prefix.contains(".openai.azure.com"); - if is_azure || self.config.use_deployment_based_urls { - format!("{}{}?api-version={}", prefix, path, self.config.api_version) - } else { - format!("{prefix}{path}") - } - } - - /// Build the request headers: auth + provider extra headers + per-request - /// headers, plus the `ai-sdk/azure/…` user-agent suffix. - async fn build_headers( - &self, - extra: Option<&HashMap>, - ) -> Result, AiMuxError> { - let mut headers = HashMap::new(); - - match &self.config.auth { - Some(AzureAuth::ApiKey(key)) => { - headers.insert("api-key".to_string(), key.clone()); - } - Some(AzureAuth::TokenProvider(tp)) => { - let token = tp.get_token().await?; - headers.insert("Authorization".to_string(), format!("Bearer {token}")); - } - None => { - return Err(AiMuxError::InvalidArgument( - "Azure OpenAI has no authentication configured. \ - Provide an `api_key` or a `token_provider`." - .to_string(), - )); - } - } - - // Provider-level extra headers. - for (k, v) in &self.config.extra_headers { - headers.insert(k.clone(), v.clone()); - } - // Per-request headers override provider headers. - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - - with_user_agent_suffix(&mut headers, "azure"); - Ok(headers) - } - - /// Post-process the request body to apply the Azure `assistant-` file-ID - /// prefix: when a file content part's base64 data starts with `assistant-`, - /// emit a `file_id` field instead of `image_url`/`file_data`. - /// - /// Mirrors the TS `fileIdPrefixes: ['assistant-']` option. - fn apply_file_id_prefixes(body: &mut Value) { - if let Some(input) = body.get_mut("input").and_then(|v| v.as_array_mut()) { - for msg in input.iter_mut() { - if let Some(content) = msg.get_mut("content").and_then(|v| v.as_array_mut()) { - for part in content.iter_mut() { - apply_prefix_to_part(part); - } - } - } - } - } -} - -/// Check if a content part's `image_url` or `file_data` contains an -/// `assistant-` prefixed data URL and, if so, replace it with a `file_id`. -/// -/// For `input_file` parts whose media type is an image, the type is also -/// changed to `input_image` — mirroring the TS behaviour where image files -/// are sent as `input_image` with a `file_id`. -fn apply_prefix_to_part(part: &mut Value) { - // input_image: { type: "input_image", image_url: "data:;base64," } - if part.get("type").and_then(|v| v.as_str()) == Some("input_image") { - let file_id = part - .get("image_url") - .and_then(|v| v.as_str()) - .and_then(extract_base64_data) - .filter(|data| data.starts_with(AZURE_FILE_ID_PREFIX)) - .map(std::string::ToString::to_string); - if let Some(file_id) = file_id - && let Some(obj) = part.as_object_mut() - { - obj.remove("image_url"); - obj.insert("file_id".to_string(), json!(file_id)); - } - } - - // input_file: { type: "input_file", file_data: "data:;base64," } - if part.get("type").and_then(|v| v.as_str()) == Some("input_file") { - let file_data = part - .get("file_data") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - if let Some(file_data) = file_data { - let media_type = extract_media_type(&file_data).unwrap_or(""); - let data = extract_base64_data(&file_data); - if let Some(data) = data - && data.starts_with(AZURE_FILE_ID_PREFIX) - && let Some(obj) = part.as_object_mut() - { - obj.remove("file_data"); - obj.remove("filename"); - obj.insert("file_id".to_string(), json!(data)); - // If the media type is an image, switch to input_image - // (mirrors TS behaviour). - if media_type.starts_with("image/") { - obj.insert("type".to_string(), json!("input_image")); - } - } - } - } -} - -/// Extract the MIME type from a `data:;base64,` URL. -fn extract_media_type(data_url: &str) -> Option<&str> { - let prefix = data_url.strip_prefix("data:")?; - let end = prefix.find(";base64,")?; - Some(&prefix[..end]) -} - -/// Extract the base64 payload from a `data:;base64,` URL. -/// Returns `None` if the URL doesn't match the expected format. -fn extract_base64_data(data_url: &str) -> Option<&str> { - let marker = ";base64,"; - data_url.split_once(marker).map(|(_, rest)| rest) -} - -/// The provider-metadata key: always `"azure"` for the Azure responses model. -fn provider_key() -> &'static str { - "azure" -} - -#[async_trait] -impl LanguageModel for AzureResponsesModel { - fn provider(&self) -> &str { - "azure.responses" - } - - fn model_id(&self) -> &str { - &self.deployment - } - - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id().to_string(), - base_url: self.config.base_url.clone(), - // Azure 不记录明文 key/token;来源按构造方式标注(env/explicit)。 - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: Some(serde_json::json!({ - "resource_name": self.config.resource_name, - "api_version": self.config.api_version, - "use_deployment_based_urls": self.config.use_deployment_based_urls, - })), - } - } - - async fn do_generate(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()).await?; - let request_result = build_responses_request_body(&self.deployment, options, false)?; - let mut body = request_result.body; - - // Apply Azure assistant- file ID prefix passthrough. - Self::apply_file_id_prefixes(&mut body); - - let provider_key = provider_key().to_string(); - let endpoint = self.endpoint(); - let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.clone(), build_header_list(&headers), options), - body.clone(), - aimux_provider_utils::create_json_response_handler::(), - crate::openai::openai_failed_response_handler(), - ) - .await?; - - let response_headers = resp.response_headers; - let raw_body = resp - .raw_value - .as_ref() - .map(ToString::to_string) - .unwrap_or_default(); - let data = resp.value; - - build_responses_generate_result( - &data, - &raw_body, - request_result.warnings, - provider_key, - endpoint, - body, - response_headers, - ) - } - - async fn do_stream(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()).await?; - let request_result = build_responses_request_body(&self.deployment, options, false)?; - let mut body = request_result.body; - let warnings = request_result.warnings; - let provider_key = provider_key().to_string(); - - // Apply Azure assistant- file ID prefix passthrough. - Self::apply_file_id_prefixes(&mut body); - - // The `store` request option (None by default). Used to decide when - // reasoning summary parts are concluded. - let store_flag = options - .provider_options - .as_ref() - .and_then(|m| m.get("openai")) - .and_then(|o| o.get("store")) - .and_then(serde_json::Value::as_bool) - == Some(true); - - let endpoint = self.endpoint(); - let mut stream_body = body.clone(); - stream_body["stream"] = Value::Bool(true); - let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.clone(), build_header_list(&headers), options), - stream_body, - aimux_provider_utils::create_event_source_response_handler::(), - crate::openai::openai_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, - }; - let stream = build_responses_event_stream( - first_event, - sse_stream, - provider_key, - warnings, - store_flag, - endpoint, - body.clone(), - response_headers.clone(), - )?; - - Ok(StreamResult { - stream, - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) - } -} - -// ── Unit tests ─────────────────────────────────────────────────────────────── -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn config_snapshot_matches_config() { - let m = AzureResponsesModel::new( - "gpt-4o".to_string(), - AzureConfig::new() - .with_base_url("https://gateway.example.com/openai") - .with_api_version("2025-04-01-preview") - .use_v1_urls() - .with_api_key("k"), - ); - let snap = m.config_snapshot(); - assert_eq!(snap.provider, "azure.responses"); - assert_eq!(snap.model_id, "gpt-4o"); - assert_eq!( - snap.base_url.as_deref(), - Some("https://gateway.example.com/openai") - ); - assert_eq!(snap.api_key_source, "explicit"); - } -} diff --git a/aimux-providers/src/bedrock/convert.rs b/aimux-providers/src/bedrock/convert.rs index 089392b60..5d24f8825 100644 --- a/aimux-providers/src/bedrock/convert.rs +++ b/aimux-providers/src/bedrock/convert.rs @@ -14,6 +14,7 @@ //! - Consecutive same-role messages are merged into a single message (matching //! the TS `groupIntoBlocks` behaviour). +use super::options; use aimux_core::language_model_message::{ AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, UserPart, @@ -168,11 +169,7 @@ pub fn convert_prompt_to_bedrock( text, provider_options, }) => { - let options = provider_options.as_ref().and_then(|options| { - options - .get("amazonBedrock") - .or_else(|| options.get("bedrock")) - }); + let options = options::read(provider_options.as_ref()); let Some(sig) = options .and_then(|options| options.get("signature")) .and_then(Value::as_str) @@ -457,34 +454,20 @@ fn push_file_block( } /// Extract a `{ cachePoint: {...} }` block from a part's `providerOptions` -/// (`bedrock.cachePoint` or `amazonBedrock.cachePoint`), if present. +/// (`amazonBedrock.cachePoint`), if present. fn cache_point(provider_options: &Option) -> Option { - let po = provider_options.as_ref()?; - for key in ["bedrock", "amazonBedrock"] { - if let Some(cp) = po.get(key).and_then(|v| v.get("cachePoint")) { - return Some(json!({ "cachePoint": cp.clone() })); - } - } - None + let cp = options::read(provider_options.as_ref())?.get("cachePoint")?; + Some(json!({ "cachePoint": cp.clone() })) } -/// Whether `citations.enabled` is set on a part's `bedrock`/`amazonBedrock` -/// provider options. +/// Whether `citations.enabled` is set on a part's `amazonBedrock` provider +/// options. fn citations_enabled(provider_options: &Option) -> bool { - let Some(po) = provider_options.as_ref() else { - return false; - }; - for key in ["bedrock", "amazonBedrock"] { - if let Some(enabled) = po - .get(key) - .and_then(|b| b.get("citations")) - .and_then(|c| c.get("enabled")) - .and_then(serde_json::Value::as_bool) - { - return enabled; - } - } - false + options::read(provider_options.as_ref()) + .and_then(|b| b.get("citations")) + .and_then(|c| c.get("enabled")) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) } /// Resolve a tool result `output` value into Bedrock's `toolResult.content` @@ -695,10 +678,7 @@ pub fn build_request_body( let (system, messages) = convert_prompt_to_bedrock(&options.prompt)?; // Extract Bedrock-specific provider options. - let bedrock_opts = options - .provider_options - .as_ref() - .and_then(|po| po.get("amazonBedrock").or_else(|| po.get("bedrock"))); + 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")); diff --git a/aimux-providers/src/bedrock/embedding.rs b/aimux-providers/src/bedrock/embedding.rs index 56db97f28..1694e5c70 100644 --- a/aimux-providers/src/bedrock/embedding.rs +++ b/aimux-providers/src/bedrock/embedding.rs @@ -9,8 +9,6 @@ //! request/response payloads. This implementation adapts based on the model ID, //! matching the TS reference. -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -20,11 +18,10 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::{HttpBody, HttpRequest}; +use aimux_provider_utils::HttpBody; -use super::BedrockAuth; -use super::model::BedrockConfig; -use super::sigv4::sign_request; +use super::options; +use crate::shared::EndpointConfig; /// An Amazon Bedrock embedding model (e.g. `"amazon.titan-embed-text-v2:0"`). /// @@ -32,52 +29,13 @@ use super::sigv4::sign_request; /// `Client` internally (RFC-0009 §4.1). pub struct BedrockEmbeddingModel { model_id: String, - config: BedrockConfig, + config: EndpointConfig, } impl BedrockEmbeddingModel { - #[must_use] - pub fn new(model_id: String, config: BedrockConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn endpoint(&self) -> String { - // 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. - format!("{}/model/{}/invoke", self.config.base_url, self.model_id) - } - - fn build_headers( - &self, - body: &str, - url: &str, - extra: Option<&HashMap>, - ) -> Result, AiMuxError> { - let mut extra_headers: Vec<(String, String)> = Vec::new(); - if let Some(extra) = extra { - for (k, v) in extra { - extra_headers.push((k.clone(), v.clone())); - } - } - - match &self.config.auth { - BedrockAuth::BearerToken(token) => { - let mut headers = vec![("Authorization".to_string(), format!("Bearer {token}"))]; - headers.extend(extra_headers); - Ok(headers) - } - BedrockAuth::SigV4(creds) => { - let signed = sign_request(creds, "bedrock", "POST", url, body, &extra_headers); - let mut headers: Vec<(String, String)> = Vec::new(); - for (k, v) in &signed.headers { - headers.push((k.clone(), v.clone())); - } - Ok(headers) - } - } - } } /// Returns `true` if the model ID is a Cohere embedding model. @@ -95,17 +53,13 @@ fn is_nova_embedding_model(model_id: &str) -> bool { #[async_trait] impl EmbeddingModel for BedrockEmbeddingModel { fn provider(&self) -> &str { - "amazon-bedrock" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - fn max_embeddings_per_call(&self) -> Option { if is_cohere_embedding_model(&self.model_id) { Some(96) @@ -195,20 +149,17 @@ impl EmbeddingModel for BedrockEmbeddingModel { Value::Object(body) }; + 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 = self.endpoint(); - let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; + // 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 resp = aimux_provider_utils::post_to_api( - HttpRequest { - url, - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(url, options), HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), aimux_provider_utils::create_json_response_handler(), super::bedrock_failed_response_handler(), @@ -335,14 +286,11 @@ struct BedrockEmbeddingProviderOptions { /// Parse Bedrock embedding provider options. /// -/// Tries the `"amazonBedrock"` key first, then falls back to the legacy -/// `"bedrock"` key for backward compatibility. +/// Reads the `amazonBedrock` key only. fn parse_bedrock_provider_options( options: Option<&SharedProviderOptions>, ) -> BedrockEmbeddingProviderOptions { - let provider_opts = options - .and_then(|o| o.get("amazonBedrock")) - .or_else(|| options.and_then(|o| o.get("bedrock"))); + let provider_opts = options::read(options); BedrockEmbeddingProviderOptions { dimensions: provider_opts diff --git a/aimux-providers/src/bedrock/image.rs b/aimux-providers/src/bedrock/image.rs index cbebdedba..0e7e32c0a 100644 --- a/aimux-providers/src/bedrock/image.rs +++ b/aimux-providers/src/bedrock/image.rs @@ -5,8 +5,6 @@ //! //! Endpoint: `POST {base_url}/model/{model_id}/invoke` -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -17,10 +15,10 @@ use aimux_core::image_model::{ }; use aimux_core::shared::Warning; -use aimux_provider_utils::{HttpBody, HttpRequest}; +use aimux_provider_utils::HttpBody; -use super::sigv4::sign_request; -use super::{BedrockAuth, BedrockConfig}; +use super::options; +use crate::shared::EndpointConfig; /// Returns the max images per call for the given model ID. fn get_max_images_per_call(model_id: &str) -> u32 { @@ -37,48 +35,14 @@ fn get_max_images_per_call(model_id: &str) -> u32 { /// `Client` internally (RFC-0009 §4.1). pub struct BedrockImageModel { model_id: String, - config: BedrockConfig, + config: EndpointConfig, } impl BedrockImageModel { - #[must_use] - pub fn new(model_id: String, config: BedrockConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - fn endpoint(&self) -> String { - format!("{}/model/{}/invoke", self.config.base_url, self.model_id) - } - - fn build_headers( - &self, - body: &str, - url: &str, - extra: Option<&HashMap>, - ) -> Result, AiMuxError> { - let mut extra_headers: Vec<(String, String)> = Vec::new(); - if let Some(extra) = extra { - for (k, v) in extra { - extra_headers.push((k.clone(), v.clone())); - } - } - match &self.config.auth { - BedrockAuth::BearerToken(token) => { - let mut headers = vec![("Authorization".into(), format!("Bearer {token}"))]; - headers.extend(extra_headers); - Ok(headers) - } - BedrockAuth::SigV4(creds) => { - let signed = sign_request(creds, "bedrock", "POST", url, body, &extra_headers); - let mut headers: Vec<(String, String)> = Vec::new(); - for (k, v) in &signed.headers { - headers.push((k.clone(), v.clone())); - } - Ok(headers) - } - } - } - fn get_base64_data(file: &ImageFile) -> Result { match file { ImageFile::Url { .. } => Err(AiMuxError::InvalidArgument( @@ -98,14 +62,11 @@ impl BedrockImageModel { #[async_trait] impl ImageModel for BedrockImageModel { fn provider(&self) -> &str { - "amazon-bedrock" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } fn max_images_per_call(&self) -> Option { Some(get_max_images_per_call(&self.model_id)) } @@ -118,11 +79,8 @@ impl ImageModel for BedrockImageModel { .unwrap_or((None, None)); let has_files = options.files.as_ref().is_some_and(|f| !f.is_empty()); - // Parse provider options (amazonBedrock or bedrock) - let ab_opts = options - .provider_options - .get("amazonBedrock") - .or_else(|| options.provider_options.get("bedrock")); + // Parse provider options (`amazonBedrock`) + let ab_opts = options::read(Some(&options.provider_options)); // Build image generation config let mut image_gen_config = Map::new(); @@ -277,21 +235,14 @@ impl ImageModel for BedrockImageModel { }); } - let body_str = serde_json::to_string(&Value::Object(args)) - .map_err(|e| AiMuxError::JsonParse(e.to_string()))?; - let url = self.endpoint(); - let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; + 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 resp = aimux_provider_utils::post_to_api( - HttpRequest { - url: url.clone(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(url.clone(), options), HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), aimux_provider_utils::create_json_response_handler::(), super::bedrock_failed_response_handler(), diff --git a/aimux-providers/src/bedrock/mod.rs b/aimux-providers/src/bedrock/mod.rs index df28bcfdd..63d1c69ef 100644 --- a/aimux-providers/src/bedrock/mod.rs +++ b/aimux-providers/src/bedrock/mod.rs @@ -1,36 +1,80 @@ //! Amazon Bedrock provider. //! -//! Implements the `LanguageModel` trait against the Bedrock Converse API -//! (`bedrock-runtime.{region}.amazonaws.com/model/{model-id}/converse`). The -//! Converse API provides a unified interface across all Bedrock-backed models -//! (Anthropic Claude, Meta Llama, Mistral, Amazon Nova, etc.). +//! [`create_amazon_bedrock`] is the Rust form of the AI SDK's +//! `createAmazonBedrock`: it takes [`AmazonBedrockProviderSettings`], validates +//! an explicit base URL and returns an [`AmazonBedrockProvider`]. Nothing is +//! read from the environment and no region, key or credential is needed until +//! a request is made. [`amazon_bedrock()`] is the default instance. //! -//! Authentication is either via: -//! - **AWS SigV4** signing (using `access_key_id` + `secret_access_key` + -//! `region`), or -//! - **Bearer token** (using `AWS_BEARER_TOKEN_BEDROCK` / `apiKey`), which -//! takes precedence when provided. +//! The models speak the Bedrock Converse API +//! (`bedrock-runtime.{region}.amazonaws.com/model/{model-id}/converse`), which +//! gives one interface to every Bedrock-hosted model family (Anthropic Claude, +//! Meta Llama, Mistral, Amazon Nova, ...), plus `invoke` for embeddings and +//! images and the Agent Runtime `rerank` endpoint. Claude's native Messages +//! protocol on AWS is a different package ([`crate::anthropic_aws`]) and the +//! OpenAI-compatible Mantle endpoint is a preset; neither is served from here. +//! +//! Authentication is decided for each request, the way `createAmazonBedrock` +//! decides it once at creation: +//! +//! - **API key** when one is available (the `api_key` setting, else +//! `AWS_BEARER_TOKEN_BEDROCK`): `Authorization: Bearer `, nothing is +//! signed. +//! - **AWS SigV4** otherwise: a [`SigV4Fetch`] transport decorator signs the +//! final method, URL, headers and body bytes. The credentials come from +//! `credential_provider`, else `access_key_id` + `secret_access_key` (+ +//! `session_token`), else `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY` and +//! `AWS_SESSION_TOKEN`. A missing region or credential fails the request +//! with `AiMuxError::LoadSetting`. pub mod convert; pub mod embedding; pub mod event_stream; pub mod image; mod model; +pub(crate) mod options; pub mod reranking; -pub mod sigv4; mod types; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::{RetryConfig, without_trailing_slash}; +use std::sync::{Arc, OnceLock}; + +use async_trait::async_trait; +use futures::future::BoxFuture; use serde_json::Value; +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_core::reranking_model::RerankingModel; +use aimux_provider_utils::{ + AwsCredentials, Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, + HeadersFn, Resolvable, SigV4Fetch, combine_headers, default_fetch, load_optional_setting, + load_setting, validate_base_url, without_trailing_slash, +}; + +use crate::shared::{Endpoint, EndpointConfig, TransformRequestBody, is_valid_hostname_part}; + pub use embedding::BedrockEmbeddingModel; pub use image::BedrockImageModel; -pub use model::{BedrockConfig, BedrockModel}; +pub use model::BedrockModel; pub use reranking::BedrockRerankingModel; -pub use sigv4::AwsCredentials; + +/// The region of a request, resolved when the request is made. +pub(crate) type RegionFn = Arc Result + Send + Sync>; + +const PROVIDER: &str = "amazon-bedrock"; +const SIGV4_SERVICE: &str = "bedrock"; + +const REGION_ENV_VAR: &str = "AWS_REGION"; +pub(crate) const BEARER_TOKEN_ENV_VAR: &str = "AWS_BEARER_TOKEN_BEDROCK"; +const ACCESS_KEY_ID_ENV_VAR: &str = "AWS_ACCESS_KEY_ID"; +const SECRET_ACCESS_KEY_ENV_VAR: &str = "AWS_SECRET_ACCESS_KEY"; +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 bedrock_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -53,284 +97,455 @@ pub(crate) fn bedrock_failed_response_handler() -> aimux_provider_utils::Respons }) } -/// Authentication method for the Bedrock provider. -#[derive(Debug, Clone)] -pub enum BedrockAuth { - /// Bearer token authentication (takes precedence when provided). - BearerToken(String), - /// AWS SigV4 signing with static credentials. - SigV4(AwsCredentials), +/// Settings of [`create_amazon_bedrock`] (the AI SDK's +/// `AmazonBedrockProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except an explicit `base_url`; everything else is evaluated on every +/// request. An explicit value is used as given and never falls back to the +/// environment. +#[derive(Clone, Default)] +pub struct AmazonBedrockProviderSettings { + /// The AWS region. `None` loads `AWS_REGION` when a request is made (else + /// the region of a `credential_provider` that is a plain value) and fails + /// that request with `AiMuxError::LoadSetting` if none is found. + pub region: Option, + /// API key for Bearer authentication; when one is available it is used + /// instead of SigV4. `None` loads `AWS_BEARER_TOKEN_BEDROCK` when a + /// request is made. An empty or whitespace-only key counts as no key. A + /// [`Resolvable::Future`] is awaited once, an [`Resolvable::AsyncFn`] on + /// every request. + pub api_key: Option>, + /// The AWS access key id of SigV4 signing. `None` loads `AWS_ACCESS_KEY_ID`. + pub access_key_id: Option, + /// The AWS secret access key of SigV4 signing. `None` loads + /// `AWS_SECRET_ACCESS_KEY`. + pub secret_access_key: Option, + /// The AWS session token of temporary credentials. When `access_key_id` + /// and `secret_access_key` are both given only this field is used; when + /// either comes from the environment the token also falls back to + /// `AWS_SESSION_TOKEN`. + 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. + pub base_url: Option, + /// Extra headers on every request. A `None` value removes the header. + /// Per-call headers win over these. + pub headers: Option, + /// The transport SigV4 signing (when used) wraps. `None` uses the process + /// default, resolved per request. + pub fetch: Option, + /// Dynamic AWS credentials, like the AWS SDK's credential providers. When + /// set, its credentials are used instead of `access_key_id`, + /// `secret_access_key` and `session_token`; it is resolved on every + /// request, so an [`Resolvable::AsyncFn`] can hand out rotating STS + /// credentials. Its `region` is used only when neither `region` nor + /// `AWS_REGION` names one. + pub credential_provider: Option>, + /// Rewrites every JSON request body once, after it is serialized and + /// before it is sent (and before it is signed). + pub transform_request_body: Option, } -/// Configuration for the Amazon Bedrock provider. -#[derive(Debug, Clone)] -pub struct BedrockProviderConfig { - pub base_url: String, - pub auth: BedrockAuth, - /// AWS region, used for constructing model ARNs and agent-runtime URLs. - pub region: String, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, - /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 - pub api_key_source: Option, +impl std::fmt::Debug for AmazonBedrockProviderSettings { + /// Never prints keys, credentials or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AmazonBedrockProviderSettings") + .field("region", &self.region) + .field("api_key", &self.api_key) + .field("access_key_id", &self.access_key_id.is_some()) + .field("secret_access_key", &self.secret_access_key.is_some()) + .field("session_token", &self.session_token.is_some()) + .field("base_url", &self.base_url) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("fetch", &self.fetch.is_some()) + .field("credential_provider", &self.credential_provider.is_some()) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ) + .finish() + } } -impl BedrockProviderConfig { - /// Create a config using AWS SigV4 credentials. - pub fn new( - access_key_id: impl Into, - secret_access_key: impl Into, - region: impl Into, - ) -> Self { - let region = region.into(); - let base_url = format!("https://bedrock-runtime.{region}.amazonaws.com"); - Self { - base_url, - auth: BedrockAuth::SigV4(AwsCredentials { - access_key_id: access_key_id.into(), - secret_access_key: secret_access_key.into(), - session_token: None, - region: region.clone(), - }), - region, - retry_config: RetryConfig::default(), - api_key_source: None, - } - } +/// Where the region, the API key and the AWS credentials of a request come +/// from. +struct Auth { + region: Option, + api_key: Option>, + access_key_id: Option, + secret_access_key: Option, + session_token: Option, + credential_provider: Option>, +} - /// Create a config using a Bearer token. - pub fn with_bearer_token(token: impl Into, region: impl Into) -> Self { - let region = region.into(); - let base_url = format!("https://bedrock-runtime.{region}.amazonaws.com"); - Self { - base_url, - auth: BedrockAuth::BearerToken(token.into()), - region, - retry_config: RetryConfig::default(), - api_key_source: None, - } +impl Auth { + /// The Bearer API key, when there is a non-blank one (trimmed). + async fn bearer(&self) -> Result, AiMuxError> { + let raw = match &self.api_key { + Some(key) => Some(key.resolve().await?), + None => load_optional_setting(None, BEARER_TOKEN_ENV_VAR), + }; + Ok(raw + .map(|key| key.trim().to_string()) + .filter(|key| !key.is_empty())) } - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self + /// The region: the setting, `AWS_REGION`, or the region of a plain-value + /// credential provider. + fn region(&self) -> Result { + let region = + if let Some(region) = load_optional_setting(self.region.as_deref(), REGION_ENV_VAR) { + region + } else if let Some(Resolvable::Value(credentials)) = &self.credential_provider + && !credentials.region.is_empty() + { + credentials.region.clone() + } else { + load_setting(None, REGION_ENV_VAR, "region")? + }; + if !is_valid_hostname_part(®ion) { + return Err(AiMuxError::InvalidArgument( + "Invalid Amazon Bedrock region. Expected a single DNS label (letters, digits, and \ + hyphens). Use `base_url` for custom endpoints." + .to_string(), + )); + } + Ok(region) } - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self + /// The credentials of the settings and the environment (no provider). + fn static_credentials(&self) -> Result { + let region = self.region()?; + let access_key_id = load_setting( + self.access_key_id.as_deref(), + ACCESS_KEY_ID_ENV_VAR, + "access_key_id", + )?; + let secret_access_key = load_setting( + self.secret_access_key.as_deref(), + SECRET_ACCESS_KEY_ENV_VAR, + "secret_access_key", + )?; + let session_token = if self.access_key_id.is_some() && self.secret_access_key.is_some() { + self.session_token.clone() + } else { + load_optional_setting(self.session_token.as_deref(), SESSION_TOKEN_ENV_VAR) + }; + Ok(AwsCredentials { + access_key_id, + secret_access_key, + session_token, + region, + }) } - /// 标注凭证来源(RFC-0023 回放重建用)。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self + /// The credentials to sign one request with. + async fn credentials(&self) -> Result { + match &self.credential_provider { + Some(provider) => { + let mut credentials = provider.resolve().await?; + if let Ok(region) = self.region() { + credentials.region = region; + } else if credentials.region.is_empty() { + return Err(AiMuxError::LoadSetting { + env_var: REGION_ENV_VAR.to_string(), + name: "region".to_string(), + }); + } + Ok(credentials) + } + None => self.static_credentials(), + } } - /// Add a session token for temporary STS credentials. - #[must_use] - pub fn with_session_token(mut self, token: impl Into) -> Self { - if let BedrockAuth::SigV4(ref mut creds) = self.auth { - creds.session_token = Some(token.into()); + /// Fail now, with the typed setting error, when signing could not find its + /// inputs: the transport reports a credential failure as a transport + /// error, which would hide which setting is missing. A credential provider + /// is the caller's own and is only called when the request is signed. + fn preflight(&self) -> Result<(), AiMuxError> { + if self.credential_provider.is_some() { + return Ok(()); } - self + self.static_credentials().map(|_| ()) } - /// Create from environment variables. - /// - /// Checks for `AWS_BEARER_TOKEN_BEDROCK` first (Bearer auth), then falls - /// back to `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` + `AWS_REGION`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when neither `AWS_BEARER_TOKEN_BEDROCK` - /// nor the `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` pair is present. - pub fn from_env() -> Result { - // Check for bearer token first. - if let Ok(token) = std::env::var("AWS_BEARER_TOKEN_BEDROCK") - && !token.trim().is_empty() + /// The base URL of `service`: the explicit one, the service's endpoint + /// environment variable, or the regional host. + fn base_url( + explicit: Option<&str>, + service: &str, + endpoint_env_var: &str, + region: &dyn Fn() -> Result, + ) -> Result { + if let Some(url) = explicit { + return Ok(url.to_string()); + } + if let Some(url) = load_optional_setting(None, endpoint_env_var) + .or_else(|| load_optional_setting(None, "AWS_ENDPOINT_URL")) { - let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string()); - return Ok(Self::with_bearer_token(token, region) - .with_api_key_source(Some("env:AWS_BEARER_TOKEN_BEDROCK"))); + return Ok(without_trailing_slash(&url)); } + Ok(format!("https://{service}.{}.amazonaws.com", region()?)) + } +} - let access_key_id = std::env::var("AWS_ACCESS_KEY_ID").map_err(|_| { - AiMuxError::InvalidArgument( - "AWS_ACCESS_KEY_ID environment variable is required for Bedrock SigV4 auth" - .to_string(), - ) - })?; - let secret_access_key = std::env::var("AWS_SECRET_ACCESS_KEY").map_err(|_| { - AiMuxError::InvalidArgument( - "AWS_SECRET_ACCESS_KEY environment variable is required for Bedrock SigV4 auth" - .to_string(), - ) - })?; - let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string()); +/// Signs a request with SigV4 unless it already carries an `Authorization` +/// header (the API-key path): one decision, made by the provider headers. +struct BedrockAuthFetch { + signing: SigV4Fetch, + inner: FetchFunction, +} - let mut config = Self::new(access_key_id, secret_access_key, region) - .with_api_key_source(Some("env:AWS_ACCESS_KEY_ID")); - if let Ok(token) = std::env::var("AWS_SESSION_TOKEN") { - config = config.with_session_token(token); +#[async_trait] +impl Fetch for BedrockAuthFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + if request.headers.contains_key(reqwest::header::AUTHORIZATION) { + self.inner.fetch(request).await + } else { + self.signing.fetch(request).await } - Ok(config) - } - - /// The Bedrock Agent Runtime base URL for this region (used for reranking). - #[must_use] - pub fn agent_runtime_url(&self) -> String { - format!( - "https://bedrock-agent-runtime.{}.amazonaws.com", - self.region - ) } } -/// Amazon Bedrock provider — creates [`BedrockModel`] instances. +/// Create an Amazon Bedrock provider. +/// +/// # Errors /// -/// 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 BedrockProvider { - config: BedrockProviderConfig, +/// Returns `AiMuxError::InvalidArgument` when an explicit `base_url` is not an +/// `http(s)` URL with a host. That is the only way this fails: the region, the +/// key and the credentials are loaded per request, not here. +pub fn create_amazon_bedrock( + settings: AmazonBedrockProviderSettings, +) -> Result { + let base_url = settings + .base_url + .as_deref() + .map(validate_base_url) + .transpose()?; + let auth = Arc::new(Auth { + region: settings.region, + api_key: settings.api_key, + access_key_id: settings.access_key_id, + secret_access_key: settings.secret_access_key, + session_token: settings.session_token, + credential_provider: settings.credential_provider, + }); + + let user_headers = settings.headers; + let header_auth = auth.clone(); + let headers: HeadersFn = Resolvable::from_async_fn(move || { + let auth = header_auth.clone(); + let user = user_headers.clone().unwrap_or_default(); + async move { + match auth.bearer().await? { + Some(key) => { + let mut bearer = HeaderMapOpt::new(); + bearer.insert("Authorization".to_string(), Some(format!("Bearer {key}"))); + Ok(combine_headers(&[&user, &bearer])) + } + None => { + auth.preflight()?; + Ok(user) + } + } + } + }); + + let credentials_auth = auth.clone(); + let inner = settings.fetch.unwrap_or_else(default_fetch); + let fetch: FetchFunction = Arc::new(BedrockAuthFetch { + signing: SigV4Fetch::new( + inner.clone(), + Resolvable::from_async_fn(move || { + let auth = credentials_auth.clone(); + async move { auth.credentials().await } + }), + SIGV4_SERVICE, + ), + inner, + }); + + Ok(AmazonBedrockProvider { + auth, + base_url: base_url.map(|url| without_trailing_slash(&url)), + headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( + headers, + "amazon-bedrock", + "5.0.100", + ), + fetch, + transform_request_body: settings.transform_request_body, + }) } -impl BedrockProvider { - #[must_use] - pub fn new(config: BedrockProviderConfig) -> Self { - Self { config } +/// The default provider: `create_amazon_bedrock` with default settings, +/// created on first use. Creating it reads nothing from the environment and +/// cannot fail; a missing region, key or credential surfaces from the first +/// request instead. +pub fn amazon_bedrock() -> &'static AmazonBedrockProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_amazon_bedrock(AmazonBedrockProviderSettings::default()) + .expect("default Bedrock settings are always valid") + }) +} + +/// An Amazon Bedrock provider (the AI SDK's `AmazonBedrockProvider`). Cheap to +/// clone the models out of; it holds no HTTP client. +pub struct AmazonBedrockProvider { + auth: Arc, + base_url: Option, + headers: HeadersFn, + fetch: FetchFunction, + transform_request_body: Option, +} + +impl AmazonBedrockProvider { + /// The configuration of a model that calls `service` (`bedrock-runtime` + /// or `bedrock-agent-runtime`). + fn model_config( + &self, + service: &'static str, + endpoint_env_var: &'static str, + ) -> EndpointConfig { + let auth = self.auth.clone(); + let base_url = self.base_url.clone(); + let headers = self.headers.clone(); + EndpointConfig { + provider: PROVIDER.to_string(), + endpoint: Arc::new(move || { + let auth = auth.clone(); + let base_url = base_url.clone(); + let headers = headers.clone(); + Box::pin(async move { + let region_auth = auth.clone(); + let base_url = Auth::base_url( + base_url.as_deref(), + service, + endpoint_env_var, + &move || region_auth.region(), + )?; + Ok(Endpoint { + base_url, + headers: headers.resolve().await?, + }) + }) + }), + fetch: Some(self.fetch.clone()), + supported_urls: Arc::new(|_| SupportedUrls::default()), + transform_request_body: self.transform_request_body.clone(), + } + } + + fn runtime_config(&self) -> EndpointConfig { + self.model_config("bedrock-runtime", RUNTIME_ENDPOINT_ENV_VAR) } - /// Create a model instance for the given Bedrock model id - /// (e.g. `"anthropic.claude-3-5-sonnet-20240620-v1:0"`). + /// A Converse language model for a Bedrock model id (e.g. + /// `"anthropic.claude-3-5-sonnet-20240620-v1:0"`); `provider()` is + /// `"amazon-bedrock"`. #[must_use] - pub fn model(&self, model_id: &str) -> BedrockModel { - BedrockModel::new( - model_id.to_string(), - BedrockConfig { - base_url: self.config.base_url.clone(), - auth: self.config.auth.clone(), - retry_config: self.config.retry_config, - api_key_source: self.config.api_key_source.clone(), - }, - ) + pub fn chat(&self, model_id: &str) -> BedrockModel { + BedrockModel::from_config(model_id.to_string(), self.runtime_config()) } - /// Create an embedding model instance for the given Bedrock model id - /// (e.g. `"amazon.titan-embed-text-v2:0"`). + /// An embedding model (e.g. `"amazon.titan-embed-text-v2:0"`); + /// `provider()` is `"amazon-bedrock"`. #[must_use] - pub fn embedding_model(&self, model_id: &str) -> BedrockEmbeddingModel { - BedrockEmbeddingModel::new( - model_id.to_string(), - BedrockConfig { - base_url: self.config.base_url.clone(), - auth: self.config.auth.clone(), - retry_config: self.config.retry_config, - api_key_source: self.config.api_key_source.clone(), - }, - ) + pub fn embedding(&self, model_id: &str) -> BedrockEmbeddingModel { + BedrockEmbeddingModel::from_config(model_id.to_string(), self.runtime_config()) } - /// Create an image generation model instance for the given Bedrock model id - /// (e.g. `"amazon.titan-image-generator-v1"` or `"amazon.nova-canvas-v1:0"`). + /// An image generation model (e.g. `"amazon.titan-image-generator-v1"` or + /// `"amazon.nova-canvas-v1:0"`); `provider()` is `"amazon-bedrock"`. #[must_use] pub fn image(&self, model_id: &str) -> BedrockImageModel { - BedrockImageModel::new( - model_id.to_string(), - BedrockConfig { - base_url: self.config.base_url.clone(), - auth: self.config.auth.clone(), - retry_config: self.config.retry_config, - api_key_source: self.config.api_key_source.clone(), - }, - ) + BedrockImageModel::from_config(model_id.to_string(), self.runtime_config()) } - /// Create a reranking model instance for the given Bedrock model id - /// (e.g. `"cohere.rerank-v3-5:0"`). Uses the Bedrock Agent Runtime - /// `/rerank` endpoint. + /// A reranking model (e.g. `"cohere.rerank-v3-5:0"`) on the Agent Runtime + /// `rerank` endpoint; `provider()` is `"amazon-bedrock"`. #[must_use] - pub fn reranking_model(&self, model_id: &str) -> BedrockRerankingModel { - // Derive the agent-runtime URL from the base URL by replacing - // `bedrock-runtime` with `bedrock-agent-runtime`. When the base URL - // is overridden (e.g. for tests), the replacement is a no-op. - let reranking_base_url = self - .config - .base_url - .replace("bedrock-runtime", "bedrock-agent-runtime"); - BedrockRerankingModel::new( + pub fn reranking(&self, model_id: &str) -> BedrockRerankingModel { + let auth = self.auth.clone(); + BedrockRerankingModel::from_config( model_id.to_string(), - reranking_base_url, - self.config.region.clone(), - self.config.auth.clone(), + self.model_config("bedrock-agent-runtime", AGENT_RUNTIME_ENDPOINT_ENV_VAR), + Arc::new(move || auth.region()), ) - .with_retry_config(self.config.retry_config) + } + + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; the same model as [`chat`](Self::chat). + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.chat(model_id)) } } -impl Provider for BedrockProvider { - fn name(&self) -> &str { - "amazon-bedrock" +impl Provider for AmazonBedrockProvider { + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) + } + + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.image(model_id))) } - /// List foundation models via AWS Bedrock `ListFoundationModels` API - /// (SigV4-signed GET, RFC-0027). - fn list_models( + fn reranking_model( &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let region = self.config.region.clone(); - let auth = self.config.auth.clone(); - let retry_config = self.config.retry_config; + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.reranking(model_id)))) + } +} + +impl ProviderDiscovery for AmazonBedrockProvider { + /// `GET /foundation-models` (the Bedrock `ListFoundationModels` API): one + /// exchange, no retry, authenticated like any other request. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let auth = self.auth.clone(); + let base_url = self.base_url.clone(); + let config = self.runtime_config(); Box::pin(async move { - let url = format!("https://bedrock.{region}.api.amazonaws.com/foundation-models"); - // Sign the request (empty body for GET). - let signed = match &auth { - BedrockAuth::SigV4(creds) => { - sigv4::sign_request(creds, "bedrock", "GET", &url, "", &[]) - } - BedrockAuth::BearerToken(_) => { - return Err(AiMuxError::UnsupportedFunctionality( - "list_models via BearerToken not supported for Bedrock; use SigV4 credentials".into(), - )); - } + // `ListFoundationModels` is on the control plane, whose host is + // `bedrock.{region}.amazonaws.com` (the runtime is + // `bedrock-runtime.{region}.amazonaws.com`); an explicit base URL + // replaces it. + let url = match &base_url { + Some(base) => format!("{base}/foundation-models"), + None => format!( + "https://bedrock.{}.amazonaws.com/foundation-models", + auth.region()? + ), }; - let mut headers = signed.headers; + let exchange = config.exchange(None).await?; + let mut headers = exchange.headers(); headers.push(("Accept".to_string(), "application/json".to_string())); - - use aimux_provider_utils::HttpRequest; - // Retry rationale: see `openai::model::execute_list_models`. - let resp = aimux_core::retry::prepare_retries(None, retry_config, None) - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: headers.clone(), - abort_signal: None, - call_id: None, - recording_context: None, - ..Default::default() - }, - aimux_provider_utils::create_json_response_handler(), - bedrock_failed_response_handler(), - ) - }) - .await?; + let resp = aimux_provider_utils::get_from_api( + exchange.with_transport(aimux_provider_utils::HttpRequest { + url, + headers, + ..Default::default() + }), + aimux_provider_utils::create_json_response_handler(), + bedrock_failed_response_handler(), + ) + .await?; // AWS response: { modelSummaries: [{ modelId, modelName, ... }] } #[derive(serde::Deserialize)] @@ -346,16 +561,15 @@ impl Provider for BedrockProvider { name: Option, } let parsed: Resp = resp.value; - let runtime: Vec = parsed + Ok(parsed .summaries .into_iter() - .map(|e| aimux_core::model_catalogue::RuntimeModel { - id: e.id, - owned_by: e.name.or(Some("amazon".to_string())), + .map(|entry| RuntimeModel { + id: entry.id, + owned_by: entry.name.or(Some("amazon".to_string())), created: None, }) - .collect(); - Ok(runtime) + .collect()) }) } } diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs index eea84b214..6d3546ccc 100644 --- a/aimux-providers/src/bedrock/model.rs +++ b/aimux-providers/src/bedrock/model.rs @@ -16,7 +16,6 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; -use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, @@ -24,12 +23,12 @@ use aimux_core::types::{ use serde_json::json; -use aimux_provider_utils::{HttpBody, HttpRequest, RetryConfig}; +use aimux_provider_utils::HttpBody; -use super::BedrockAuth; use super::convert::{build_request_body, convert_usage, map_finish_reason}; -use super::sigv4::sign_request; +use super::options; use super::types::{BedrockContentBlock, BedrockConverseResponse}; +use crate::shared::EndpointConfig; fn bedrock_event_stream_response_handler() -> aimux_provider_utils::ResponseHandler>> { @@ -88,28 +87,16 @@ fn bedrock_event_stream_response_handler() /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct BedrockModel { model_id: String, - config: BedrockConfig, -} - -/// Configuration for a Bedrock model instance (cloned from the provider). -#[derive(Debug, Clone)] -pub struct BedrockConfig { - pub base_url: String, - pub auth: BedrockAuth, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, - /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 - /// SigV4 记 access-key 来源;BearerToken 记 bearer-token 来源。不存明文。 - pub api_key_source: Option, + config: EndpointConfig, } impl BedrockModel { - #[must_use] - pub fn new(model_id: String, config: BedrockConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - fn endpoint(&self, stream: bool) -> String { + /// `/model/{model-id}/converse` or `/converse-stream`. + fn path(&self, stream: bool) -> String { let suffix = if stream { "converse-stream" } else { @@ -119,90 +106,28 @@ impl BedrockModel { // `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.config.base_url, self.model_id, suffix - ) - } - - fn build_headers( - &self, - body: &str, - url: &str, - extra: Option<&HashMap>, - ) -> Result, AiMuxError> { - let mut extra_headers: Vec<(String, String)> = Vec::new(); - if let Some(extra) = extra { - for (k, v) in extra { - extra_headers.push((k.clone(), v.clone())); - } - } - - match &self.config.auth { - BedrockAuth::BearerToken(token) => { - let mut headers = vec![("Authorization".to_string(), format!("Bearer {token}"))]; - headers.extend(extra_headers); - Ok(headers) - } - BedrockAuth::SigV4(creds) => { - let signed = sign_request(creds, "bedrock", "POST", url, body, &extra_headers); - let mut headers: Vec<(String, String)> = Vec::new(); - for (k, v) in &signed.headers { - headers.push((k.clone(), v.clone())); - } - Ok(headers) - } - } + format!("/model/{}/{}", self.model_id, suffix) } } #[async_trait] impl LanguageModel for BedrockModel { fn provider(&self) -> &str { - "amazon-bedrock" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - // M2b: record identity + credential source + region (encoded in - // base_url) + retry. Never serialize the SigV4 credentials or bearer - // token plaintext — only the auth *kind* and source marker. - let auth_kind = match &self.config.auth { - BedrockAuth::BearerToken(_) => "bearer_token", - BedrockAuth::SigV4(_) => "sigv4", - }; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.base_url.clone()), - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: Some(serde_json::json!({ - "auth_kind": auth_kind, - "max_retries": self.config.retry_config.max_retries, - })), - } - } - async fn do_generate(&self, options: &CallOptions) -> Result { let body = build_request_body(&self.model_id, options)?; + 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 = self.endpoint(false); - let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; + let url = exchange.url(&self.path(false)); let resp = aimux_provider_utils::post_to_api( - HttpRequest::new(url, headers, options), + exchange.request(url, options), HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), aimux_provider_utils::create_json_response_handler(), super::bedrock_failed_response_handler(), @@ -259,11 +184,12 @@ impl LanguageModel for BedrockModel { async fn do_stream(&self, options: &CallOptions) -> Result { let body = build_request_body(&self.model_id, options)?; + 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 = self.endpoint(true); - let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; + let url = exchange.url(&self.path(true)); let resp = aimux_provider_utils::post_to_api( - HttpRequest::new(url, headers, options), + exchange.request(url, options), HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), bedrock_event_stream_response_handler(), super::bedrock_failed_response_handler(), @@ -570,8 +496,8 @@ impl LanguageModel for BedrockModel { } // #26: merge the stop sentinel into the metadata payload, then build - // the Finish provider_metadata under both `amazonBedrock` and - // `bedrock` keys (mirrors the TS `doStream` flush handler). + // the Finish provider_metadata under `amazonBedrock` (mirrors the + // TS `doStream` flush handler). if let Some(seq) = stop_sequence { finish_meta.insert( "stopSequence".to_string(), @@ -581,10 +507,8 @@ impl LanguageModel for BedrockModel { let provider_metadata = if finish_meta.is_empty() { None } else { - Some(aimux_core::shared::NamespaceMap::from([ - ("amazonBedrock".to_string(), finish_meta.clone()), - ("bedrock".to_string(), finish_meta), - ])) + let payload = serde_json::Value::Object(finish_meta); + Some(options::metadata(payload)) }; yield Ok(StreamPart::Finish { @@ -608,18 +532,9 @@ impl LanguageModel for BedrockModel { } /// Wrap the accumulated reasoning signature as provider_metadata in the same -/// dual-key shape the non-streaming path emits (`amazonBedrock` + `bedrock`), -/// so consumers reading either key see it. +/// shape the non-streaming path emits (under `amazonBedrock`). fn reasoning_signature_meta(sig: Option) -> Option { - sig.map(|s| { - let mut metadata = provider_namespace("amazonBedrock", json!({ "signature": &s })) - .expect("provider metadata must be an object"); - metadata.extend( - provider_namespace("bedrock", json!({ "signature": s })) - .expect("provider metadata must be an object"), - ); - metadata - }) + sig.map(|s| options::metadata(json!({ "signature": s }))) } /// Extract `GenerateContent` items from a non-streaming content block. @@ -628,10 +543,10 @@ fn reasoning_signature_meta(sig: Option) -> Option { /// `text`, `toolUse`, or `reasoningContent`. Empty `text` blocks are preserved /// (matching the TS SDK) so that empty text between reasoning blocks survives. /// `reasoningContent.reasoningText` yields a `Reasoning` item whose -/// `provider_metadata` carries the `signature` under both `amazonBedrock` and -/// `bedrock` keys (or `None` when no signature is present). +/// `provider_metadata` carries the `signature` under `amazonBedrock` (or +/// `None` when no signature is present). /// `reasoningContent.redactedReasoning` yields a `Reasoning` item with empty -/// text and `redactedData` under both metadata keys. +/// text and `redactedData` under the same key. fn extract_content(block: &BedrockContentBlock, content: &mut Vec) { if let Some(text) = &block.text { content.push(GenerateContent::Text { @@ -656,15 +571,10 @@ fn extract_content(block: &BedrockContentBlock, content: &mut Vec) -> Option<&JsonObject> { + let options = provider_options?; + options + .get(AMAZON_BEDROCK) + .or_else(|| options.get("bedrock")) +} + +/// 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") +} diff --git a/aimux-providers/src/bedrock/reranking.rs b/aimux-providers/src/bedrock/reranking.rs index a68b446e8..3570a9f53 100644 --- a/aimux-providers/src/bedrock/reranking.rs +++ b/aimux-providers/src/bedrock/reranking.rs @@ -5,8 +5,6 @@ //! //! Uses the Bedrock Agent Runtime `/rerank` endpoint. -use std::collections::HashMap; - use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; @@ -18,10 +16,11 @@ use aimux_core::reranking_model::{ }; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::{HttpBody, HttpRequest, RetryConfig}; +use aimux_provider_utils::HttpBody; -use super::BedrockAuth; -use super::sigv4::sign_request; +use super::RegionFn; +use super::options; +use crate::shared::EndpointConfig; /// Bedrock provider-specific reranking options. #[derive(Debug, Clone, Default)] @@ -35,8 +34,7 @@ fn parse_bedrock_reranking_options( ) -> BedrockRerankingOptions { let mut opts = BedrockRerankingOptions::default(); if let Some(po) = provider_options { - // Prefer "amazonBedrock"; fall back to "bedrock". - let bedrock = po.get("amazonBedrock").or_else(|| po.get("bedrock")); + let bedrock = options::read(Some(po)); if let Some(bedrock) = bedrock { if let Some(token) = bedrock.get("nextToken").and_then(|v| v.as_str()) { opts.next_token = Some(token.to_string()); @@ -68,60 +66,16 @@ struct BedrockRerankingResult { /// `Client` internally (RFC-0009 §4.1). pub struct BedrockRerankingModel { model_id: String, - base_url: String, - region: String, - auth: BedrockAuth, - retry_config: RetryConfig, + config: EndpointConfig, + region: RegionFn, } impl BedrockRerankingModel { - #[must_use] - pub fn new(model_id: String, base_url: String, region: String, auth: BedrockAuth) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig, region: RegionFn) -> Self { Self { model_id, - base_url, + config, region, - auth, - retry_config: RetryConfig::default(), - } - } - - pub(crate) fn with_retry_config(mut self, retry_config: RetryConfig) -> Self { - self.retry_config = retry_config; - self - } - - fn endpoint(&self) -> String { - format!("{}/rerank", self.base_url) - } - - fn build_headers( - &self, - body: &str, - url: &str, - extra: Option<&HashMap>, - ) -> Result, AiMuxError> { - let mut extra_headers: Vec<(String, String)> = Vec::new(); - if let Some(extra) = extra { - for (k, v) in extra { - extra_headers.push((k.clone(), v.clone())); - } - } - - match &self.auth { - BedrockAuth::BearerToken(token) => { - let mut headers = vec![("Authorization".to_string(), format!("Bearer {token}"))]; - headers.extend(extra_headers); - Ok(headers) - } - BedrockAuth::SigV4(creds) => { - let signed = sign_request(creds, "bedrock", "POST", url, body, &extra_headers); - let mut headers: Vec<(String, String)> = Vec::new(); - for (k, v) in &signed.headers { - headers.push((k.clone(), v.clone())); - } - Ok(headers) - } } } } @@ -129,26 +83,23 @@ impl BedrockRerankingModel { #[async_trait] impl RerankingModel for BedrockRerankingModel { fn provider(&self) -> &str { - "amazon-bedrock" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.retry_config - } - async fn do_rerank( &self, options: &RerankingCallOptions, ) -> Result { let bedrock_options = parse_bedrock_reranking_options(options.provider_options.as_ref()); + let region = (self.region)()?; let model_arn = format!( - "arn:aws:bedrock:{}::foundation-model/{}", - self.region, self.model_id + "arn:aws:bedrock:{region}::foundation-model/{}", + self.model_id ); // Build sources array. @@ -206,20 +157,13 @@ impl RerankingModel for BedrockRerankingModel { ["additionalModelRequestFields"] = fields.clone(); } + 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 = self.endpoint(); - let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; + let url = exchange.url("/rerank"); let resp = aimux_provider_utils::post_to_api( - HttpRequest { - url, - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(url, options), HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), aimux_provider_utils::create_json_response_handler::(), super::bedrock_failed_response_handler(), diff --git a/aimux-providers/src/bedrock/sigv4.rs b/aimux-providers/src/bedrock/sigv4.rs deleted file mode 100644 index 681f91f3b..000000000 --- a/aimux-providers/src/bedrock/sigv4.rs +++ /dev/null @@ -1,450 +0,0 @@ -//! Simplified AWS Signature Version 4 (SigV4) signer. -//! -//! Implements the core SigV4 signing algorithm for authenticating requests to -//! AWS services (Amazon Bedrock, Claude Platform on AWS). The signer produces -//! the `Authorization` header and supporting `X-Amz-*` headers that AWS -//! expects. -//! -//! Reference: -//! -//! This is a self-contained implementation using only `sha2` + `hmac` (no AWS -//! SDK dependency). It supports: -//! - Static credentials (`access_key_id` + `secret_access_key` + optional -//! `session_token`). -//! - Any AWS service name and region. -//! - POST requests with a JSON body. - -use hmac::{Hmac, Mac}; -use sha2::{Digest, Sha256}; - -/// A set of AWS credentials for SigV4 signing. -/// -/// `Debug` is hand-written so the secret key and session token never reach -/// logs or error output; the access key id and region stay visible because -/// they identify the credentials without granting access. -#[derive(Clone)] -pub struct AwsCredentials { - pub access_key_id: String, - pub secret_access_key: String, - /// Optional temporary session token (STS). - pub session_token: Option, - pub region: String, -} - -impl std::fmt::Debug for AwsCredentials { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AwsCredentials") - .field("access_key_id", &self.access_key_id) - .field("secret_access_key", &"[REDACTED]") - .field( - "session_token", - &self.session_token.as_ref().map(|_| "[REDACTED]"), - ) - .field("region", &self.region) - .finish() - } -} - -/// The result of signing a request. -#[derive(Debug, Clone)] -pub struct SignedRequest { - /// All headers to send with the request, including the `Authorization` - /// header and any `X-Amz-*` headers. - pub headers: Vec<(String, String)>, -} - -type HmacSha256 = Hmac; - -/// Sign an HTTP request using AWS SigV4. -/// -/// # Arguments -/// - `credentials` - The AWS credentials to sign with. -/// - `service` - The AWS service name (e.g. `"bedrock"`, `"aws-external-anthropic"`). -/// - `method` - HTTP method (e.g. `"POST"`). -/// - `url` - The full request URL. -/// - `body` - The request body as a string (e.g. JSON). -/// - `extra_headers` - Additional headers to include in the request (these are -/// also included in the canonical headers for signing). -#[must_use] -pub fn sign_request( - credentials: &AwsCredentials, - service: &str, - method: &str, - url: &str, - body: &str, - extra_headers: &[(String, String)], -) -> SignedRequest { - sign_request_at( - credentials, - service, - method, - url, - body, - extra_headers, - chrono::Utc::now(), - ) -} - -/// [`sign_request`] with an explicit signing time. -fn sign_request_at( - credentials: &AwsCredentials, - service: &str, - method: &str, - url: &str, - body: &str, - extra_headers: &[(String, String)], - now: chrono::DateTime, -) -> SignedRequest { - let parsed = - url::Url::parse(url).unwrap_or_else(|_| url::Url::parse("http://localhost").unwrap()); - // The signed host must match the wire host: keep the port whenever the - // URL carries a non-default one (Url::port() is Some only for non-default - // ports). Signing a bare host diverges from the Host header reqwest would - // otherwise so gateways and proxies on non-standard ports reject the - // signature (observed as 502 when an env proxy forwarded a signed - // loopback request whose host lacked the port). - let host = match parsed.port() { - Some(port) => format!("{}:{}", parsed.host_str().unwrap_or(""), port), - None => parsed.host_str().unwrap_or("").to_string(), - }; - let path = canonical_uri(parsed.path()); - let query = canonical_query(parsed.query().unwrap_or("")); - - // Timestamp:yyyyMMddTHHmmssZ and date:yyyyMMdd - let amz_date = now.format("%Y%m%dT%H%M%SZ").to_string(); - let date_stamp = now.format("%Y%m%d").to_string(); - - // Body hash - let body_hash = hex::encode(Sha256::digest(body.as_bytes())); - - // Build canonical headers (sorted, lowercased, trimmed). - // Always include host, x-amz-date, and x-amz-content-sha256. - let mut canonical_headers: Vec<(String, String)> = vec![ - ("host".to_string(), host.to_string()), - ("x-amz-content-sha256".to_string(), body_hash.clone()), - ("x-amz-date".to_string(), amz_date.clone()), - ]; - - if let Some(ref token) = credentials.session_token { - canonical_headers.push(("x-amz-security-token".to_string(), token.clone())); - } - - for (k, v) in extra_headers { - let lower = k.to_lowercase(); - // Don't override the required headers. - if lower != "host" - && lower != "x-amz-date" - && lower != "x-amz-content-sha256" - && lower != "x-amz-security-token" - { - canonical_headers.push((lower, canonical_header_value(v))); - } - } - - canonical_headers.sort_by(|a, b| a.0.cmp(&b.0)); - - let canonical_headers_str: String = canonical_headers - .iter() - .map(|(k, v)| format!("{k}:{v}\n")) - .collect(); - let signed_headers: String = canonical_headers - .iter() - .map(|(k, _)| k.as_str()) - .collect::>() - .join(";"); - - // Canonical request - let canonical_request = format!( - "{}\n{}\n{}\n{}\n{}\n{}", - method.to_uppercase(), - path, - query, - canonical_headers_str, - signed_headers, - body_hash - ); - - let canonical_request_hash = hex::encode(Sha256::digest(canonical_request.as_bytes())); - - // Credential scope - let credential_scope = format!( - "{}/{}/{}/aws4_request", - date_stamp, credentials.region, service - ); - - // String to sign - let string_to_sign = - format!("AWS4-HMAC-SHA256\n{amz_date}\n{credential_scope}\n{canonical_request_hash}"); - - // Signing key: derived through chained HMAC, each step keyed by the - // previous one and starting from "AWS4" + secret. - let k_date = hmac_sha256( - format!("AWS4{}", credentials.secret_access_key).as_bytes(), - date_stamp.as_bytes(), - ); - let k_region = hmac_sha256(&k_date, credentials.region.as_bytes()); - let k_service = hmac_sha256(&k_region, service.as_bytes()); - let k_signing = hmac_sha256(&k_service, b"aws4_request"); - - let signature = hex::encode(hmac_sha256(&k_signing, string_to_sign.as_bytes())); - - let authorization = format!( - "AWS4-HMAC-SHA256 Credential={}/{}, SignedHeaders={}, Signature={}", - credentials.access_key_id, credential_scope, signed_headers, signature - ); - - // Build the final header list. - let mut headers: Vec<(String, String)> = canonical_headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - headers.push(("Authorization".to_string(), authorization)); - - SignedRequest { headers } -} - -/// SigV4 `UriEncode`: every byte except the unreserved set -/// (`A-Z a-z 0-9 - _ . ~`) becomes `%XX` with uppercase hex. -fn uri_encode(input: &[u8]) -> String { - let mut out = String::with_capacity(input.len()); - for &byte in input { - if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b'~') { - out.push(byte as char); - } else { - out.push_str(&format!("%{byte:02X}")); - } - } - out -} - -/// Canonical URI for every service but S3: each segment of the path as sent -/// is URI-encoded again, so `:` signs as `%3A` and an already-escaped `%3A` -/// signs as `%253A`. -fn canonical_uri(path: &str) -> String { - if path.is_empty() { - return "/".to_string(); - } - path.split('/') - .map(|segment| uri_encode(segment.as_bytes())) - .collect::>() - .join("/") -} - -/// Canonical query string: each name and value percent-decoded, then -/// URI-encoded, and the pairs sorted by name and then value. -fn canonical_query(query: &str) -> String { - let mut pairs: Vec<(String, String)> = query - .split('&') - .filter(|pair| !pair.is_empty()) - .map(|pair| { - let (name, value) = pair.split_once('=').unwrap_or((pair, "")); - ( - uri_encode(&percent_decode(name)), - uri_encode(&percent_decode(value)), - ) - }) - .collect(); - pairs.sort(); - pairs - .iter() - .map(|(name, value)| format!("{name}={value}")) - .collect::>() - .join("&") -} - -fn percent_decode(input: &str) -> Vec { - let bytes = input.as_bytes(); - let mut out = Vec::with_capacity(bytes.len()); - let mut i = 0; - while i < bytes.len() { - let escaped = (bytes[i] == b'%') - .then(|| bytes.get(i + 1..i + 3)) - .flatten() - .and_then(|hex| std::str::from_utf8(hex).ok()) - .and_then(|hex| u8::from_str_radix(hex, 16).ok()); - match escaped { - Some(byte) => { - out.push(byte); - i += 3; - } - None => { - out.push(bytes[i]); - i += 1; - } - } - } - out -} - -/// Canonical header value: trimmed, with runs of spaces collapsed to one. -fn canonical_header_value(value: &str) -> String { - value.split_whitespace().collect::>().join(" ") -} - -fn hmac_sha256(key: &[u8], data: &[u8]) -> Vec { - let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length"); - mac.update(data); - mac.finalize().into_bytes().to_vec() -} - -#[cfg(test)] -mod tests { - use super::*; - - fn creds() -> AwsCredentials { - AwsCredentials { - access_key_id: "AKIAIOSFODNN7EXAMPLE".to_string(), - secret_access_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(), - session_token: None, - region: "us-east-1".to_string(), - } - } - - #[test] - fn debug_redacts_secret_key_and_session_token() { - let mut credentials = creds(); - credentials.session_token = Some("FwoGZXIvYXdzEXAMPLETOKEN".to_string()); - let printed = format!("{credentials:?}"); - assert!(!printed.contains("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY")); - assert!(!printed.contains("FwoGZXIvYXdzEXAMPLETOKEN")); - assert!(printed.contains("[REDACTED]")); - assert!(printed.contains("us-east-1")); - } - - fn host_header(signed: &SignedRequest) -> String { - signed - .headers - .iter() - .find(|(k, _)| k.eq_ignore_ascii_case("host")) - .map(|(_, v)| v.clone()) - .expect("signed request carries a host header") - } - - #[test] - fn host_keeps_non_default_port() { - let signed = sign_request( - &creds(), - "bedrock", - "POST", - "http://127.0.0.1:54321/model/x/invoke", - "{}", - &[], - ); - assert_eq!(host_header(&signed), "127.0.0.1:54321"); - } - - #[test] - fn host_omits_default_port() { - let signed = sign_request( - &creds(), - "bedrock", - "POST", - "https://bedrock-runtime.us-east-1.amazonaws.com/model/x/invoke", - "{}", - &[], - ); - assert_eq!( - host_header(&signed), - "bedrock-runtime.us-east-1.amazonaws.com" - ); - } - - fn fixed_time() -> chrono::DateTime { - use chrono::TimeZone; - chrono::Utc.with_ymd_and_hms(2024, 1, 2, 3, 4, 5).unwrap() - } - - fn signature(signed: &SignedRequest) -> String { - signed - .headers - .iter() - .find(|(k, _)| k == "Authorization") - .and_then(|(_, v)| v.rsplit_once("Signature=")) - .expect("signed request carries a signature") - .1 - .to_string() - } - - /// Expected values come from botocore 1.43's `SigV4Auth` (AWS's own - /// Python signer) for the same credentials, headers, body and fixed - /// time, not from this implementation. - #[test] - fn signatures_match_botocore() { - let signed = sign_request_at( - &creds(), - "bedrock", - "POST", - "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-5-sonnet-20240620-v1:0/converse", - r#"{"messages":[]}"#, - &[("X-Custom".to_string(), " value ".to_string())], - fixed_time(), - ); - assert_eq!( - signature(&signed), - "538444fb4b17f3cdd777bbd5f1cd4e3e1f53d6ac97167da96c81ec8d14d6468a" - ); - - let session = AwsCredentials { - session_token: Some("SESSION".to_string()), - ..creds() - }; - let signed = sign_request_at( - &session, - "bedrock", - "POST", - "http://127.0.0.1:54321/model/x/invoke?a=1", - "{}", - &[], - fixed_time(), - ); - assert_eq!( - signature(&signed), - "dd6a2384cb9c8f0db8ea8ee383a61bf06dd85da679cec54d5c944b1b0cb8d9f6" - ); - - let signed = sign_request_at( - &creds(), - "bedrock", - "POST", - "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-3-5-sonnet-20240620-v1%3A0/converse-stream", - "{}", - &[ - ("content-type".to_string(), "application/json".to_string()), - ("X-Multi".to_string(), "a b c".to_string()), - ], - fixed_time(), - ); - assert_eq!( - signature(&signed), - "a6abc7f26ad04c0922b03e05c9184a20a8fcf4b68b7dfce358e9a6436c8a9489" - ); - - let signed = sign_request_at( - &creds(), - "bedrock", - "POST", - "https://polly.us-east-1.amazonaws.com/v1/speech?b=2&a=x%20y&a=1", - "{}", - &[], - fixed_time(), - ); - assert_eq!( - signature(&signed), - "70c6e5794c4f1456f33f2b41cebc483ee0d19c98b95660411f4b51ccac464006" - ); - } - - #[test] - fn canonical_uri_and_query_follow_sigv4() { - assert_eq!(canonical_uri(""), "/"); - assert_eq!( - canonical_uri("/model/a.b-v1:0/converse"), - "/model/a.b-v1%3A0/converse" - ); - assert_eq!(canonical_uri("/a+b/c%2Fd"), "/a%2Bb/c%252Fd"); - assert_eq!( - canonical_query("b=2&a=x%20y&a=1&e=&z"), - "a=1&a=x%20y&b=2&e=&z=" - ); - } -} diff --git a/aimux-providers/src/bedrock_mantle.rs b/aimux-providers/src/bedrock_mantle.rs deleted file mode 100644 index 75d847c57..000000000 --- a/aimux-providers/src/bedrock_mantle.rs +++ /dev/null @@ -1,115 +0,0 @@ -//! Bedrock Mantle provider — a thin OpenAI-compatible wrapper. -//! -//! Exposes an OpenAI-compatible Chat Completions API at -//! `https://bedrock-mantle.{region}.api.aws/v1` (region defaults to `us-east-1`). -//! Provider-specific details are the region-aware base URL and the -//! `BEDROCK_MANTLE_API_KEY` environment variable; everything else is delegated -//! to the shared [`OpenAIProvider`]. -//! -//! # Authentication -//! -//! This thin wrapper only supports **Bearer API key** auth via the -//! `BEDROCK_MANTLE_API_KEY` environment variable (the shared OpenAI layer sends -//! it as `Authorization: Bearer `). Bedrock Mantle also accepts AWS SigV4 -//! signing, but that cannot ride on the OpenAI shared layer's Bearer auth — -//! SigV4 callers should use the dedicated [`bedrock`](crate::bedrock) provider -//! instead, which signs requests with `service = "bedrock"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::load_api_key; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_REGION: &str = "us-east-1"; -const ENV_VAR: &str = "BEDROCK_MANTLE_API_KEY"; -const PROVIDER_NAME: &str = "bedrock_mantle"; - -/// Build the Bedrock Mantle base URL for a region: -/// `https://bedrock-mantle.{region}.api.aws/v1`. -fn base_url_for_region(region: &str) -> String { - format!("https://bedrock-mantle.{region}.api.aws/v1") -} - -/// Resolve the Bedrock Mantle region from the environment. -/// -/// Precedence: `BEDROCK_MANTLE_REGION` → `AWS_REGION` → `us-east-1`. -fn resolve_region() -> String { - std::env::var("BEDROCK_MANTLE_REGION") - .ok() - .filter(|r| !r.trim().is_empty()) - .or_else(|| { - std::env::var("AWS_REGION") - .ok() - .filter(|r| !r.trim().is_empty()) - }) - .unwrap_or_else(|| DEFAULT_REGION.to_string()) -} - -/// Configuration for the Bedrock Mantle provider (wraps [`OpenAIConfig`]). -pub struct BedrockMantleConfig(OpenAIConfig); - -impl BedrockMantleConfig { - /// Create from an API key, deriving the region-aware base URL from the - /// environment (`BEDROCK_MANTLE_REGION` / `AWS_REGION` / `us-east-1`). - pub fn new(api_key: impl Into) -> Self { - let region = resolve_region(); - Self( - OpenAIConfig::new(api_key) - .with_base_url(base_url_for_region(®ion)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()), - ) - } - - /// Create from the `BEDROCK_MANTLE_API_KEY` environment variable. - /// - /// The region is resolved from `BEDROCK_MANTLE_REGION` / `AWS_REGION` - /// (defaulting to `us-east-1`). - /// - /// # Errors - /// - /// Returns [`AiMuxError::InvalidArgument`] when `BEDROCK_MANTLE_API_KEY` is - /// not set. - pub fn from_env() -> Result { - let key = load_api_key(None, ENV_VAR, "Bedrock Mantle")?; - Ok(Self::new(key)) - } - - /// Override the base URL (useful for tests / self-hosted endpoints). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// Bedrock Mantle provider — creates [`OpenAIModel`] instances. -pub struct BedrockMantleProvider(OpenAIProvider); - -impl BedrockMantleProvider { - #[must_use] - pub fn new(config: BedrockMantleConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given model id - /// (e.g. `"bedrock_mantle/openai.gpt-oss-120b"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for BedrockMantleProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/black_forest_labs.rs b/aimux-providers/src/black_forest_labs.rs index bef61f9dc..9407ab55d 100644 --- a/aimux-providers/src/black_forest_labs.rs +++ b/aimux-providers/src/black_forest_labs.rs @@ -5,8 +5,15 @@ //! //! Uses an async submit + poll pattern: POST to submit, then GET poll_url //! until status is "Ready", then download the image. +//! +//! [`create_black_forest_labs`] takes [`BlackForestLabsProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`BlackForestLabsProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `BFL_API_KEY`. +//! [`black_forest_labs()`] is the default instance; it reads nothing and cannot fail. use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::Duration; use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -17,9 +24,14 @@ use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; -use aimux_core::retry; use aimux_core::shared::Warning; -use aimux_provider_utils::{HttpRequest, load_api_key, sleep_or_abort, without_trailing_slash}; +use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{ + Credential, EndpointConfig, POLL_INTERVAL_MILLIS_KEY, PollStep, is_poll_control_key, + poll_interval_ms, poll_until, provider_headers, retry_download, +}; /// AI SDK's `isTrustedUrl` (black-forest-labs-api.ts): credentials may go to /// the configured origin, or over HTTPS to `bfl.ai` and its subdomains — BFL @@ -62,95 +74,146 @@ fn bfl_failed_response_handler() -> aimux_provider_utils::ResponseHandler, + start_time: Option, + end_time: Option, + duration: Option, +} -/// Configuration for the Black Forest Labs provider. -#[derive(Debug, Clone)] -pub struct BlackForestLabsConfig { - pub api_key: String, - pub base_url: String, - pub headers: Option>, +pub(crate) mod options; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.bfl.ai"; +pub(crate) const API_KEY_ENV_VAR: &str = "BFL_API_KEY"; +const DEFAULT_NAME: &str = "blackForestLabs"; + +/// Settings of [`create_black_forest_labs`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct BlackForestLabsProviderSettings { + /// Base URL for the API calls. Default `https://api.bfl.ai`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `BFL_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.image"`). + /// Default `"blackForestLabs"`. The providerOptions key stays `blackForestLabs`. + 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, } -impl BlackForestLabsConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.bfl.ai".to_string(), - headers: None, - } - } - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } - /// Create from the `BFL_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "BFL_API_KEY", "Black Forest Labs")?; - Ok(Self::new(api_key)) +impl std::fmt::Debug for BlackForestLabsProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BlackForestLabsProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } +/// Create a Black Forest Labs provider. +/// +/// # 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. +pub fn create_black_forest_labs( + settings: BlackForestLabsProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(BlackForestLabsProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Black Forest Labs"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_black_forest_labs` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn black_forest_labs() -> &'static BlackForestLabsProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_black_forest_labs(BlackForestLabsProviderSettings::default()) + .expect("default Black Forest Labs settings are always valid") + }) +} + +/// A Black Forest Labs provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct BlackForestLabsProvider { - config: BlackForestLabsConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } + impl BlackForestLabsProvider { - #[must_use] - pub fn new(config: BlackForestLabsConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + + /// An image model (e.g. `"flux-pro-1.1"`); `provider()` is `"{name}.image"`. #[must_use] pub fn image(&self, model_id: &str) -> BlackForestLabsImageModel { - BlackForestLabsImageModel::new(model_id.to_string(), self.config.clone()) + BlackForestLabsImageModel::from_config(model_id.to_string(), self.model_config("image")) } } -/// A Black Forest Labs image generation model. +crate::impl_single_modality_provider!(BlackForestLabsProvider, image_model, |p, id| p.image(id)); + pub struct BlackForestLabsImageModel { model_id: String, - config: BlackForestLabsConfig, + config: EndpointConfig, } impl BlackForestLabsImageModel { - #[must_use] - pub fn new(model_id: String, config: BlackForestLabsConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut h = HashMap::new(); - h.insert( - "Authorization".into(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - h.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - h.insert(k.clone(), v.clone()); - } - } - h - } - - fn submit_endpoint(&self) -> String { - format!("{}/{}", self.config.base_url, self.model_id) - } - fn file_to_string(file: &ImageFile) -> Result { match file { ImageFile::Url { url } => Ok(url.clone()), @@ -181,7 +244,7 @@ fn gcd(a: u32, b: u32) -> u32 { #[async_trait] impl ImageModel for BlackForestLabsImageModel { fn provider(&self) -> &str { - "blackForestLabs" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id @@ -210,7 +273,7 @@ impl ImageModel for BlackForestLabsImageModel { warnings.push(Warning::Unsupported { feature: "size".into(), details: Some("Black Forest Labs ignores size when aspectRatio is provided. Use the width and height provider options to specify dimensions for models that support them".into()) }); } - let bfl_opts = options.provider_options.get("blackForestLabs"); + let bfl_opts = options::black_forest_labs_options(Some(&options.provider_options)); let (width_str, height_str) = options .size .map(|s| (s.width().to_string(), s.height().to_string())) @@ -300,10 +363,7 @@ impl ImageModel for BlackForestLabsImageModel { ("webhookUrl", "webhook_url"), ]; for (key, value) in bfl { - if matches!( - key.as_str(), - "width" | "height" | "pollIntervalMillis" | "pollTimeoutMillis" - ) { + if is_poll_control_key(key) || matches!(key.as_str(), "width" | "height") { continue; } let ak = map @@ -321,12 +381,12 @@ impl ImageModel for BlackForestLabsImageModel { body.insert("mask".into(), json!(m)); } - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; - // Submit + // Submit. This is the only request that creates a generation: nothing + // below sends it again. let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.submit_endpoint(), header_list.clone(), options), + exchange.request(exchange.url(&format!("/{}", self.model_id)), options), Value::Object(body), aimux_provider_utils::create_json_response_handler(), bfl_failed_response_handler(), @@ -346,22 +406,12 @@ impl ImageModel for BlackForestLabsImageModel { .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); - let retries = retry::prepare_retries( - options.max_retries, - self.retry_config(), - options.abort_signal.clone(), - ); - - // Poll for result - let poll_interval = bfl_opts - .and_then(|o| o.get("pollIntervalMillis")) - .and_then(serde_json::Value::as_u64) - .unwrap_or(DEFAULT_POLL_INTERVAL_MS); - let poll_timeout = bfl_opts - .and_then(|o| o.get("pollTimeoutMillis")) - .and_then(serde_json::Value::as_u64) - .unwrap_or(DEFAULT_POLL_TIMEOUT_MS); - let max_attempts = (poll_timeout / poll_interval.max(1)) as usize; + + let poll_interval = Duration::from_millis(poll_interval_ms( + bfl_opts, + POLL_INTERVAL_MILLIS_KEY, + POLL_INTERVAL_MS, + )); let mut poll_url_with_id = url::Url::parse(&poll_url) .map_err(|e| AiMuxError::InvalidResponseData(format!("invalid BFL poll URL: {e}")))?; @@ -370,127 +420,98 @@ impl ImageModel for BlackForestLabsImageModel { .query_pairs_mut() .append_pair("id", &request_id); } - - let mut image_url = None; - let mut result_seed = None; - let mut result_start_time = None; - let mut result_end_time = None; - let mut result_duration = None; + let poll_url = poll_url_with_id.to_string(); // AI SDK gates headers per URL via isTrustedUrl; response-supplied // targets outside the BFL allowlist get none. - let gated_headers = |url: &str| -> Vec<(String, String)> { - if bfl_trusted_url(url, &self.config.base_url) { - header_list.clone() - } else { - vec![] + let gated_request = |url: &str| -> HttpRequest { + let mut request = exchange.request(url.to_string(), options); + // Headers are gated per URL here, not by the credentialed origin. + request.credentialed_origin = None; + if !bfl_trusted_url(url, exchange.base_url()) { + request.headers = Vec::new(); } + request.validate_url = true; + request.trusted_origin = Some(exchange.base_url().to_string()); + request }; - for _ in 0..max_attempts { - // AI SDK polls polling_url with validateUrl: true and gates the - // headers itself via isTrustedUrl (base_url origin or HTTPS - // *.bfl.ai), so credentialed_origin is None here. - let poll_url = poll_url_with_id.to_string(); - let pr = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: poll_url.clone(), - headers: gated_headers(&poll_url), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: true, - trusted_origin: Some(self.config.base_url.clone()), - // Headers are gated per URL by gated_headers above. - credentialed_origin: None, - }, - aimux_provider_utils::create_json_response_handler::(), - bfl_failed_response_handler(), - ) - }) + // Poll for the result. AI SDK polls polling_url with validateUrl: true + // and gates the headers itself via isTrustedUrl (base_url origin or + // HTTPS *.bfl.ai). + let ready = poll_until( + &format!("blackForestLabs task {request_id}"), + options.abort_signal.as_ref(), + poll_interval, + MAX_POLL_ATTEMPTS, + || async { + let pr = aimux_provider_utils::get_from_api( + gated_request(&poll_url), + aimux_provider_utils::create_json_response_handler::(), + bfl_failed_response_handler(), + ) .await?; - let response_body = pr.raw_value.as_ref().map(ToString::to_string); - let pv = pr.value; - - let poll_status = pv - .get("status") - .and_then(|v| v.as_str()) - .or_else(|| pv.get("state").and_then(|v| v.as_str())) - .unwrap_or(""); - if poll_status == "Ready" { - if let Some(result) = pv.get("result") { - image_url = result - .get("sample") + let response_body = pr.raw_value.as_ref().map(ToString::to_string); + let pv = pr.value; + + let poll_status = pv + .get("status") + .and_then(|v| v.as_str()) + .or_else(|| pv.get("state").and_then(|v| v.as_str())) + .unwrap_or(""); + if poll_status == "Ready" { + let result = pv.get("result"); + let sample = result + .and_then(|r| r.get("sample")) .and_then(|v| v.as_str()) - .map(String::from); - result_seed = result.get("seed").cloned(); - result_start_time = result.get("start_time").cloned(); - result_end_time = result.get("end_time").cloned(); - result_duration = result.get("duration").cloned(); + .map(String::from) + .ok_or_else(|| { + AiMuxError::InvalidResponseData( + "BFL poll reported Ready without result.sample".to_string(), + ) + })?; + return Ok(PollStep::Ready(BflResult { + image_url: sample, + seed: result.and_then(|r| r.get("seed")).cloned(), + start_time: result.and_then(|r| r.get("start_time")).cloned(), + end_time: result.and_then(|r| r.get("end_time")).cloned(), + duration: result.and_then(|r| r.get("duration")).cloned(), + })); } - if image_url.is_none() { - return Err(AiMuxError::InvalidResponseData( - "BFL poll reported Ready without result.sample".to_string(), - )); + if poll_status == "Error" || poll_status == "Failed" { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some(poll_status.to_string()), + message: "Black Forest Labs generation failed.".into(), + response_body, + ..ApiCallError::new( + "Black Forest Labs generation failed.", + poll_url.clone(), + serde_json::json!({}), + ) + }))); } - break; - } - if poll_status == "Error" || poll_status == "Failed" { - return Err(AiMuxError::ApiCall(Box::new(ApiCallError { - status_code: Some(200), - provider_code: Some(poll_status.to_string()), - message: "Black Forest Labs generation failed.".into(), - response_body, - ..ApiCallError::new( - "Black Forest Labs generation failed.", - poll_url_with_id.to_string(), - serde_json::json!({}), - ) - }))); - } - sleep_or_abort( - std::time::Duration::from_millis(poll_interval), - options.abort_signal.as_ref(), - ) - .await?; - } - - let image_url = image_url.ok_or_else(|| { - AiMuxError::Timeout(format!( - "blackForestLabs task {request_id} polling timed out after {poll_timeout}ms" - )) - })?; + Ok(PollStep::Pending) + }, + ) + .await?; + let image_url = ready.image_url; + let result_seed = ready.seed; + let result_start_time = ready.start_time; + let result_end_time = ready.end_time; + let result_duration = ready.duration; // Download image; result.sample is a URL from the poll response body. // AI SDK sends its headers to trusted BFL hosts on the download too, // gated by the same allowlist. - let ir = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: image_url.clone(), - headers: gated_headers(&image_url), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: true, - trusted_origin: Some(self.config.base_url.clone()), - // Headers are gated per URL by gated_headers above. - credentialed_origin: None, - }, - aimux_provider_utils::create_binary_response_handler(), - bfl_failed_response_handler(), - ) - }) - .await?; + let ir = retry_download(options.abort_signal.as_ref(), poll_interval, || { + aimux_provider_utils::get_from_api( + gated_request(&image_url), + aimux_provider_utils::create_binary_response_handler(), + bfl_failed_response_handler(), + ) + }) + .await?; let image_bytes = ir.value.to_vec(); let download_headers: HashMap = HashMap::new(); @@ -520,7 +541,7 @@ impl ImageModel for BlackForestLabsImageModel { img_meta.insert("outputMegapixels".into(), o.clone()); } bfl_meta.insert("images".into(), json!([Value::Object(img_meta)])); - metadata.insert("blackForestLabs".into(), bfl_meta); + metadata.insert(options::NAMESPACE.into(), bfl_meta); Ok(ImageResult { images: ImageOutputs::Binary(vec![image_bytes]), diff --git a/aimux-providers/src/black_forest_labs/options.rs b/aimux-providers/src/black_forest_labs/options.rs new file mode 100644 index 000000000..656dbf673 --- /dev/null +++ b/aimux-providers/src/black_forest_labs/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Black Forest Labs. +//! +//! Black Forest Labs reads `providerOptions.blackForestLabs`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "blackForestLabs"; + +/// The Black Forest Labs options in a providerOptions map. +pub(crate) fn black_forest_labs_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/cartesia.rs b/aimux-providers/src/cartesia.rs index a6123a18e..00ed79c0d 100644 --- a/aimux-providers/src/cartesia.rs +++ b/aimux-providers/src/cartesia.rs @@ -17,8 +17,13 @@ //! //! `instructions` is not supported and emits a warning. `speed` must be //! between 0.6 and 1.5 (inclusive) or it is ignored with a warning. +//! +//! [`create_cartesia`] takes [`CartesiaProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`CartesiaProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `CARTESIA_API_KEY`. +//! [`cartesia()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -28,8 +33,10 @@ use aimux_core::shared::{SharedProviderOptions, Warning}; use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; +use aimux_provider_utils::HttpBody; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; -use aimux_provider_utils::{HttpBody, HttpRequest, load_api_key}; +use crate::shared::{Credential, EndpointConfig, provider_headers}; fn cartesia_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -62,86 +69,163 @@ const SAMPLE_RATES: &[u32] = &[8000, 16000, 22050, 24000, 44100, 48000]; // ── Config ─────────────────────────────────────────────────────────────────── -/// Configuration for the Cartesia provider. -#[derive(Debug, Clone)] -pub struct CartesiaConfig { - pub api_key: String, - /// Base URL for the Cartesia API (no trailing slash). Defaults to - /// `https://api.cartesia.ai`. - pub base_url: String, - /// The Cartesia API version (sent via the `Cartesia-Version` header). - pub version: String, - /// Extra headers merged into every request. - pub headers: Option>, -} - -impl CartesiaConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.cartesia.ai".to_string(), - version: CARTESIA_API_VERSION.to_string(), - headers: None, - } - } +pub(crate) mod options; - /// Override the base URL (for testing or proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = url.into().trim_end_matches('/').to_string(); - self - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.cartesia.ai"; +pub(crate) const API_KEY_ENV_VAR: &str = "CARTESIA_API_KEY"; +const DEFAULT_NAME: &str = "cartesia"; - /// Override the Cartesia API version. - #[must_use] - pub fn with_version(mut self, version: impl Into) -> Self { - self.version = version.into(); - self - } +/// Settings of [`create_cartesia`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct CartesiaProviderSettings { + /// Base URL for the API calls. Default `https://api.cartesia.ai`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `CARTESIA_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.speech"`, `"{name}.transcription"`). + /// Default `"cartesia"`. The providerOptions key stays `cartesia`. + 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, +} - /// Attach extra headers merged into every request. - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self +impl std::fmt::Debug for CartesiaProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CartesiaProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Create from environment variable `CARTESIA_API_KEY`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `CARTESIA_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "CARTESIA_API_KEY", "Cartesia")?; - Ok(Self::new(api_key)) - } +/// Create a Cartesia provider. +/// +/// # 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. +pub fn create_cartesia(settings: CartesiaProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(CartesiaProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Cartesia"), + vec![( + "Cartesia-Version".to_string(), + CARTESIA_API_VERSION.to_string(), + )], + settings.headers, + ), + fetch: settings.fetch, + }) } -// ── Provider ───────────────────────────────────────────────────────────────── +/// The default provider: `create_cartesia` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn cartesia() -> &'static CartesiaProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_cartesia(CartesiaProviderSettings::default()) + .expect("default Cartesia settings are always valid") + }) +} -/// Cartesia provider — creates `CartesiaSpeechModel` instances. +/// A Cartesia provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct CartesiaProvider { - config: CartesiaConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl CartesiaProvider { - #[must_use] - pub fn new(config: CartesiaConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a speech (TTS) model instance for the given model name (e.g. - /// `"sonic-3.5"`). + /// A speech (TTS) model (e.g. `"sonic-3.5"`); `provider()` is `"{name}.speech"`. #[must_use] pub fn speech(&self, model_id: &str) -> CartesiaSpeechModel { - CartesiaSpeechModel::new(model_id.to_string(), self.config.clone()) + CartesiaSpeechModel::from_config(model_id.to_string(), self.model_config("speech")) } - /// Create a transcription (STT) model instance for the given model name - /// (e.g. `"best"`). Uses the `/stt` endpoint. + /// A transcription (STT) model (e.g. `"ink-whisper"`); `provider()` is `"{name}.transcription"`. #[must_use] pub fn transcription(&self, model_id: &str) -> CartesiaTranscriptionModel { - CartesiaTranscriptionModel::new(model_id.to_string(), self.config.clone()) + CartesiaTranscriptionModel::from_config( + model_id.to_string(), + self.model_config("transcription"), + ) + } +} + +impl ::aimux_core::Provider for CartesiaProvider { + fn language_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "languageModel")) + } + + fn embedding_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } + + fn transcription_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(::std::sync::Arc::new(self.transcription(model_id)))) + } + + fn speech_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(::std::sync::Arc::new(self.speech(model_id)))) } } @@ -150,44 +234,19 @@ impl CartesiaProvider { /// A Cartesia speech (TTS) model. pub struct CartesiaSpeechModel { model_id: String, - config: CartesiaConfig, + config: EndpointConfig, } impl CartesiaSpeechModel { - #[must_use] - pub fn new(model_id: String, config: CartesiaConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - headers.insert("Cartesia-Version".to_string(), self.config.version.clone()); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/tts/bytes", self.config.base_url) - } } #[async_trait] impl SpeechModel for CartesiaSpeechModel { fn provider(&self) -> &str { - "cartesia.speech" + &self.config.provider } fn model_id(&self) -> &str { @@ -197,10 +256,10 @@ impl SpeechModel for CartesiaSpeechModel { async fn do_generate(&self, options: &SpeechCallOptions) -> Result { let (body, warnings) = build_request(options, &self.model_id)?; - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), headers.into_iter().collect(), options), + exchange.request(exchange.url("/tts/bytes"), options), Value::Object(body.clone()), aimux_provider_utils::create_binary_response_handler(), cartesia_failed_response_handler(), @@ -576,7 +635,7 @@ struct CartesiaSpeechProviderOptions { fn parse_cartesia_provider_options( options: Option<&SharedProviderOptions>, ) -> Option { - let opts = options.and_then(|opts| opts.get("cartesia"))?; + let opts = options::cartesia_options(options)?; Some(CartesiaSpeechProviderOptions { container: opts @@ -653,38 +712,13 @@ struct CartesiaTranscriptionResponse { /// Endpoint: `POST {base_url}/stt` (multipart form-data) pub struct CartesiaTranscriptionModel { model_id: String, - config: CartesiaConfig, + config: EndpointConfig, } impl CartesiaTranscriptionModel { - #[must_use] - pub fn new(model_id: String, config: CartesiaConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - headers.insert("Cartesia-Version".to_string(), self.config.version.clone()); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/stt", self.config.base_url) - } } fn audio_input_to_bytes_stt(audio: &AudioInput) -> Result, AiMuxError> { @@ -694,13 +728,14 @@ fn audio_input_to_bytes_stt(audio: &AudioInput) -> Result, AiMuxError> { base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid base64: {e}"))) } + AudioInput::Url(_) => Err(AudioInput::url_error()), } } #[async_trait] impl TranscriptionModel for CartesiaTranscriptionModel { fn provider(&self) -> &str { - "cartesia" + &self.config.provider } fn model_id(&self) -> &str { @@ -723,9 +758,7 @@ impl TranscriptionModel for CartesiaTranscriptionModel { // Parse provider options. let mut language: Option = None; let mut timestamp_granularities: Option> = None; - if let Some(ref po) = options.provider_options - && let Some(cartesia) = po.get("cartesia") - { + if let Some(cartesia) = options::cartesia_options(options.provider_options.as_ref()) { if let Some(l) = cartesia.get("language").and_then(|v| v.as_str()) { language = Some(l.to_string()); } @@ -769,10 +802,10 @@ impl TranscriptionModel for CartesiaTranscriptionModel { let (body_bytes, content_type) = form.finish(); - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_to_api( - HttpRequest::new(self.endpoint(), headers.into_iter().collect(), options), + exchange.request(exchange.url("/stt"), options), HttpBody::Bytes(body_bytes, content_type), aimux_provider_utils::create_json_response_handler::(), cartesia_failed_response_handler(), diff --git a/aimux-providers/src/cartesia/options.rs b/aimux-providers/src/cartesia/options.rs new file mode 100644 index 000000000..59fa1cb62 --- /dev/null +++ b/aimux-providers/src/cartesia/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Cartesia. +//! +//! Cartesia reads `providerOptions.cartesia`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "cartesia"; + +/// The Cartesia options in a providerOptions map. +pub(crate) fn cartesia_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/catalogue.rs b/aimux-providers/src/catalogue.rs index 4e18bea4c..4b5815e93 100644 --- a/aimux-providers/src/catalogue.rs +++ b/aimux-providers/src/catalogue.rs @@ -3,7 +3,7 @@ //! Provides [`get_model_specs`] — a thin function that fetches the community //! model catalogue (anya2a `dist/all.json`) and returns a [`Catalogue`] of //! [`ModelSpec`] entries. This is **deliberately decoupled** from -//! `Provider::list_models`: the host calls both independently and merges them +//! `ProviderDiscovery::list_models`: the host calls both independently and merges them //! as needed. //! //! `get_model_specs` does **no caching, no FS writes, no TTL** — it is a pure @@ -13,7 +13,6 @@ use std::collections::HashMap; use std::time::Duration; use aimux_core::AiMuxError; -use aimux_core::ApiCallError; use aimux_core::model_catalogue::{ CatalogueSource, Modality, ModelCapabilities, ModelCost, ModelLimits, ModelModalities, ModelSpec, ModelType, ReasoningMode, ReasoningSpec, ReasoningVisibility, @@ -407,53 +406,22 @@ fn build_cost(m: &Value) -> Option { }) } +/// Whole-exchange bound for the catalogue download (the aggregate file is a +/// few MiB; the helper default is tuned for API calls). +const FETCH_TIMEOUT: Duration = Duration::from_secs(60); + async fn fetch_text(url: &str) -> Result { - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(60)) - .build() - .map_err(|e| { - AiMuxError::ApiCall(Box::new(ApiCallError { - is_retryable: true, - ..ApiCallError::new( - format!("get_model_specs: build client: {e}"), - url, - serde_json::json!({}), - ) - })) - })?; - let resp = client.get(url).send().await.map_err(|e| { - AiMuxError::ApiCall(Box::new(ApiCallError { - is_retryable: true, - ..ApiCallError::new( - format!("get_model_specs: fetch {url}: {e}"), - url, - serde_json::json!({}), - ) - })) - })?; - if !resp.status().is_success() { - let status = resp.status().as_u16(); - return Err(AiMuxError::ApiCall(Box::new(ApiCallError { - status_code: Some(status), - response_body: resp.text().await.ok().filter(|b| !b.is_empty()), - is_retryable: status == 429 || status >= 500, - ..ApiCallError::new( - format!("get_model_specs: fetch {url}"), - url, - serde_json::json!({}), - ) - }))); - } - resp.text().await.map_err(|e| { - AiMuxError::ApiCall(Box::new(ApiCallError { - is_retryable: true, - ..ApiCallError::new( - format!("get_model_specs: read body: {e}"), - url, - serde_json::json!({}), - ) - })) - }) + let output = aimux_provider_utils::get_from_api( + aimux_provider_utils::HttpRequest { + url: url.to_string(), + response_timeout: Some(FETCH_TIMEOUT), + ..Default::default() + }, + aimux_provider_utils::create_binary_response_handler(), + aimux_provider_utils::create_status_code_error_response_handler(), + ) + .await?; + Ok(String::from_utf8_lossy(&output.value).into_owned()) } #[cfg(test)] diff --git a/aimux-providers/src/codex.rs b/aimux-providers/src/codex.rs index 91eae9309..290a630ff 100644 --- a/aimux-providers/src/codex.rs +++ b/aimux-providers/src/codex.rs @@ -1,14 +1,21 @@ //! Codex provider (RFC-0018) — OpenAI Codex models over the Responses API. //! +//! [`create_codex`] takes [`CodexProviderSettings`], validates the base URL, +//! fixes the provider name and returns a [`CodexProvider`]. Credentials are not +//! read there; they are resolved in the request headers of every call, so a +//! host that refreshes the ChatGPT access token supplies an +//! [`Resolvable::AsyncFn`] or rebuilds the provider. The AI SDK has no Codex +//! package; the factory follows the shape of the OpenAI one. +//! //! Codex models (`gpt-5.2-codex`, `gpt-5.1-codex`, `gpt-5.1-codex-mini`, //! `gpt-5-codex`, …) are available **only** through the Responses API //! (verified 2026-08-03, RFC-0018 §2 V1). //! -//! Two access modes ([`CodexMode`]): +//! Two access modes ([`CodexMode`], chosen in the settings): //! -//! - **ApiKey** (default): `https://api.openai.com/v1/responses` — the +//! - **`ApiKey`** (default): `https://api.openai.com/v1/responses` — the //! official, documented, usage-based channel. Recommended for CI/automation. -//! - **Subscription**: `https://chatgpt.com/backend-api/codex/responses` — +//! - **`ChatGptAccount`**: `https://chatgpt.com/backend-api/codex/responses` — //! the ChatGPT-subscription channel (single personal account, per OpenAI //! ToU). Undocumented endpoint — **best-effort, no SLA**. Mirrors the //! official Codex client: **always streams** (`stream: true`, `store: false`). @@ -17,261 +24,310 @@ //! integrator performs the device-code login, persists tokens, and orchestrates //! refresh. The library exposes [`codex_refresh`] as a stateless protocol //! helper and maps subscription-channel 401 responses to -//! [`AiMuxError::TokenExpired`] so the integrator can refresh and retry. +//! [`AiMuxError::TokenExpired`] so the integrator can refresh and rebuild the +//! provider. //! //! Scope boundaries (RFC-0018): no account pools, no quota rotation, no //! reselling — single account, single user. -use std::collections::HashMap; +use std::sync::{Arc, OnceLock}; use async_trait::async_trait; use futures::StreamExt; +use futures::future::BoxFuture; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; +use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; use aimux_core::language_model::LanguageModel; +use aimux_core::model_catalogue::RuntimeModel; use aimux_core::options::CallOptions; -use aimux_core::provider::Provider; +use aimux_core::provider::{Provider, ProviderDiscovery}; use aimux_core::result::{GenerateResult, StreamResult}; use aimux_core::types::Warning; -use aimux_provider_utils::{HttpRequest, RetryConfig}; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, Resolvable, validate_base_url, +}; +use crate::openai::config::OpenAIModelConfig; use crate::openai::responses::responses_convert::{ - build_header_list, build_responses_event_stream, build_responses_generate_result, + build_responses_event_stream, build_responses_generate_result, +}; +use crate::openai::responses::{ + OpenAIResponsesModel, ResponsesNamespace, build_responses_request_body, }; -use crate::openai::responses::{OpenAIResponsesModel, build_responses_request_body}; -use crate::openai::{OpenAICompatProfile, OpenAIConfig}; +pub use crate::shared::TransformRequestBody; +use crate::shared::{AuthScheme, Credential, credential_headers}; /// Default API-key base URL (official OpenAI Responses endpoint). pub const CODEX_API_BASE_URL: &str = "https://api.openai.com/v1"; -/// Default subscription base URL (undocumented ChatGPT backend endpoint). -pub const CODEX_SUBSCRIPTION_BASE_URL: &str = "https://chatgpt.com/backend-api"; -/// Environment variable read by [`CodexConfig::from_env`]. +/// Default subscription base URL (undocumented ChatGPT backend endpoint). The +/// Responses path is appended, giving `https://chatgpt.com/backend-api/codex/responses` +/// (the endpoint the `chatgpt` cassettes record). +pub const CODEX_SUBSCRIPTION_BASE_URL: &str = "https://chatgpt.com/backend-api/codex"; +/// Environment variable the API-key mode reads when no key is given. pub const CODEX_API_KEY_ENV_VAR: &str = "CODEX_API_KEY"; /// OAuth token endpoint used by [`codex_refresh`] (RFC-0018 §2 V3). pub const CODEX_OAUTH_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; -/// Access mode for the Codex provider. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +const DEFAULT_NAME: &str = "codex"; +const DEFAULT_ORIGINATOR: &str = "aimux"; + +/// Access mode of the Codex provider, with the credential of that channel. +#[derive(Debug, Clone)] pub enum CodexMode { - /// Official API-key channel (`api.openai.com`), usage-based. - ApiKey, + /// Official API-key channel (`api.openai.com`), usage-based. `None` loads + /// `CODEX_API_KEY` when a request is made and fails that request with + /// `AiMuxError::LoadApiKey` if it is unset; an explicit key is used as + /// given, `""` included. + ApiKey(Option>), /// ChatGPT-subscription channel (`chatgpt.com/backend-api`), single /// personal account; always streams. - Subscription, -} - -/// Configuration for the Codex provider. -#[derive(Debug, Clone)] -pub struct CodexConfig { - openai: OpenAIConfig, - mode: CodexMode, - /// Subscription mode: `ChatGPT-Account-Id` header value (optional). - chatgpt_account_id: Option, - /// Subscription mode: `Originator` header value (client identifier; - /// reference clients use e.g. `"opencode"` / `"codex_cli_rs"`). - originator: String, -} - -impl CodexConfig { - /// API-key mode with the official base URL. - pub fn new(api_key: impl Into) -> Self { - Self { - openai: OpenAIConfig::new(api_key) - .with_base_url(CODEX_API_BASE_URL) - .with_provider("codex") - .with_profile(OpenAICompatProfile::full()), - mode: CodexMode::ApiKey, - chatgpt_account_id: None, - originator: "aimux".to_string(), - } - } - - /// Subscription mode with the ChatGPT backend base URL. /// - /// `account_token` is the access token minted by the integrator's OAuth - /// device-code flow — the library never performs login, persists tokens, - /// or refreshes automatically (RFC-0018 §3.2). - pub fn subscription(account_token: impl Into) -> Self { - Self { - openai: OpenAIConfig::new(account_token) - .with_base_url(CODEX_SUBSCRIPTION_BASE_URL) - .with_provider("codex") - .with_profile(OpenAICompatProfile::full()), - mode: CodexMode::Subscription, - chatgpt_account_id: None, - originator: "aimux".to_string(), - } - } - - /// API-key mode reading `CODEX_API_KEY` from the environment. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `CODEX_API_KEY` is not set. - pub fn from_env() -> Result { - let key = aimux_provider_utils::load_api_key(None, CODEX_API_KEY_ENV_VAR, "Codex")?; - Ok(Self::new(key).with_api_key_source(Some("env:CODEX_API_KEY"))) - } + /// `token` is the access token minted by the integrator's OAuth + /// device-code flow: the library never performs login, persists tokens, or + /// refreshes automatically (RFC-0018 §3.2). A [`Resolvable::AsyncFn`] is + /// evaluated on every request. + ChatGptAccount { + /// The account access token, sent as `Authorization: Bearer`. + token: Resolvable, + /// The `ChatGPT-Account-Id` header value (optional). + account_id: Option, + }, +} - /// 标注 api_key 来源(RFC-0023 回放重建用)。透传到内部 `OpenAIConfig`。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.openai = self.openai.with_api_key_source(source); - self +impl Default for CodexMode { + /// The API-key channel, key from `CODEX_API_KEY`. + fn default() -> Self { + Self::ApiKey(None) } +} - /// Override the base URL (tests / self-hosted endpoints). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.openai = self.openai.with_base_url(url); - self - } +/// Settings of [`create_codex`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; the credential and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct CodexProviderSettings { + /// The access mode and its credential. Default: API key from + /// `CODEX_API_KEY`. + pub mode: CodexMode, + /// Base URL for the API calls. Default [`CODEX_API_BASE_URL`] for + /// `ApiKey` and [`CODEX_SUBSCRIPTION_BASE_URL`] for `ChatGptAccount`; a + /// trailing slash is removed. + pub base_url: Option, + /// Extra headers on every request. A `None` value removes the header, + /// including `Authorization`. They win over the channel headers + /// (`Originator`, `ChatGPT-Account-Id`); per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of `provider()` (`"{name}.responses"`). + /// Default `"codex"`. + pub name: Option, + /// `ChatGptAccount` mode: the `Originator` header value (client + /// identifier; reference clients use e.g. `"opencode"` / `"codex_cli_rs"`). + /// Default `"aimux"`. + pub originator: Option, + /// The transport: a mock, a signing decorator, a proxy-aware client. + /// `None` uses the process default, resolved per request. + pub fetch: Option, + /// Rewrites every JSON request body once, after it is serialized (and + /// after the `store: false` rule of the subscription channel) and before it + /// is sent. + pub transform_request_body: Option, +} - /// Subscription mode: set the `ChatGPT-Account-Id` header. - #[must_use] - pub fn with_chatgpt_account_id(mut self, id: impl Into) -> Self { - self.chatgpt_account_id = Some(id.into()); - self +impl std::fmt::Debug for CodexProviderSettings { + /// Never prints the credential or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CodexProviderSettings") + .field("mode", &self.mode) + .field("base_url", &self.base_url) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("originator", &self.originator) + .field("fetch", &self.fetch.is_some()) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ) + .finish() } +} - /// Subscription mode: override the `Originator` header (default `"aimux"`). - #[must_use] - pub fn with_originator(mut self, originator: impl Into) -> Self { - self.originator = originator.into(); - self - } +/// Which channel a provider talks to; the model rules differ. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Channel { + ApiKey, + ChatGptAccount, +} - /// Override the retry settings used by Core model operations. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.openai = self.openai.with_retry_config(config); - self - } +/// Create a Codex provider. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` +/// URL with a host. That is the only way this fails: credentials are resolved +/// per request, not here. +pub fn create_codex(settings: CodexProviderSettings) -> Result { + let (channel, default_base_url, credential, fixed) = match settings.mode { + CodexMode::ApiKey(key) => ( + Channel::ApiKey, + CODEX_API_BASE_URL, + Credential::explicit_or_env(key, CODEX_API_KEY_ENV_VAR, "Codex"), + Vec::new(), + ), + CodexMode::ChatGptAccount { token, account_id } => { + let mut fixed = vec![( + "Originator".to_string(), + settings + .originator + .unwrap_or_else(|| DEFAULT_ORIGINATOR.to_string()), + )]; + if let Some(id) = account_id { + fixed.push(("ChatGPT-Account-Id".to_string(), id)); + } + ( + Channel::ChatGptAccount, + CODEX_SUBSCRIPTION_BASE_URL, + Credential::Explicit(token), + fixed, + ) + } + }; + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => default_base_url.to_string(), + }; + Ok(CodexProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + channel, + base_url, + headers: credential_headers(credential, AuthScheme::Bearer, fixed, settings.headers), + fetch: settings.fetch, + transform_request_body: settings.transform_request_body, + }) +} - /// Extra headers merged into every request (user-supplied values win over - /// the subscription-mode defaults). - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.openai = self.openai.with_headers(headers); - self - } +/// The default provider: `create_codex` with default settings (API key from +/// `CODEX_API_KEY`), created on first use. Creating it reads nothing from the +/// environment and cannot fail; a missing key surfaces from the first request +/// instead. +pub fn codex() -> &'static CodexProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_codex(CodexProviderSettings::default()) + .expect("default Codex settings are always valid") + }) } -/// Codex provider — creates [`CodexModel`] instances. +/// A Codex provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct CodexProvider { - config: CodexConfig, + name: String, + channel: Channel, + base_url: String, + headers: HeadersFn, + fetch: Option, + transform_request_body: Option, } impl CodexProvider { - #[must_use] - pub fn new(config: CodexConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> OpenAIModelConfig { + OpenAIModelConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + self.transform_request_body.clone(), + ) } - /// Create a model instance for the given Codex model id - /// (e.g. `"gpt-5.2-codex"`). + /// A Codex model over the Responses API (e.g. `"gpt-5.2-codex"`); + /// `provider()` is `"{name}.responses"`. #[must_use] - pub fn model(&self, model_id: &str) -> CodexModel { + pub fn responses(&self, model_id: &str) -> CodexModel { CodexModel { model_id: model_id.to_string(), - config: self.config.clone(), + channel: self.channel, + config: self.model_config("responses"), } } + + /// The provider as a function: the default language model for an id. The + /// same model as [`responses`](Self::responses) and + /// [`language_model`](Provider::language_model). + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.responses(model_id)) + } } impl Provider for CodexProvider { - fn name(&self) -> &str { - "codex" + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - /// List models via `GET {base_url}/models` (OpenAI-compatible), enriched - /// with the community catalogue portrait when available (RFC-0027). - fn list_models( - &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let config = self.config.openai.clone(); - Box::pin(async move { - let headers = crate::openai::model::build_auth_headers(&config); - let runtime = crate::openai::model::execute_list_models( - &config.base_url, - &headers, - config.retry_config, - ) - .await?; - Ok(runtime) - }) + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } +} + +impl ProviderDiscovery for CodexProvider { + /// `GET {base_url}/models` (OpenAI-compatible): one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config("models"); + Box::pin(async move { crate::openai::model::list_models_once(&config).await }) } } /// A Codex language model over the Responses API. -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct CodexModel { model_id: String, - config: CodexConfig, + channel: Channel, + config: OpenAIModelConfig, } impl CodexModel { + /// The shared Responses model, for the API-key channel. fn inner(&self) -> OpenAIResponsesModel { - OpenAIResponsesModel::new(self.model_id.clone(), self.config.openai.clone()) - } - - fn endpoint(&self) -> String { - format!("{}/responses", self.config.openai.base_url) - } - - /// Subscription mode: fold the channel headers (`Originator`, - /// `ChatGPT-Account-Id`) into the per-call options. User-supplied header - /// values win (`entry().or_insert`). - fn subscription_options(&self, options: &CallOptions) -> CallOptions { - let mut opts = options.clone(); - let mut headers = opts.headers.take().unwrap_or_default(); - headers - .entry("Originator".to_string()) - .or_insert_with(|| self.config.originator.clone()); - if let Some(id) = &self.config.chatgpt_account_id { - headers - .entry("ChatGPT-Account-Id".to_string()) - .or_insert_with(|| id.clone()); - } - opts.headers = Some(headers); - opts + OpenAIResponsesModel::from_config(self.model_id.clone(), self.config.clone()) } /// Build the request body; the subscription channel always streams and - /// never stores (mirrors the official Codex client). - fn body( + /// never stores (mirrors the official Codex client). That rule is part of + /// this package's request construction, not a body override. + fn subscription_body( &self, options: &CallOptions, - stream: bool, ) -> Result<(Value, Vec), AiMuxError> { - let mut result = build_responses_request_body(&self.model_id, options, stream)?; - if self.config.mode == CodexMode::Subscription { - result.body["store"] = Value::Bool(false); - } - Ok((result.body, result.warnings)) + let mut result = build_responses_request_body(&self.model_id, options, true)?; + result.body["store"] = Value::Bool(false); + Ok((self.config.transform_body(result.body), result.warnings)) } /// Map a subscription-channel authentication failure to /// [`AiMuxError::TokenExpired`] so the integrator can refresh the access - /// token and retry (RFC-0018 §3.2). The only credential on the - /// subscription channel is the account token, so any 401 from the + /// token and rebuild the provider (RFC-0018 §3.2). The only credential on + /// the subscription channel is the account token, so any 401 from the /// endpoint (read from the `status_code` field) means the token is /// invalid/expired. - fn map_subscription_401(&self, error: AiMuxError) -> AiMuxError { + fn map_subscription_401(error: AiMuxError) -> AiMuxError { match error { AiMuxError::ApiCall(d) if d.status_code == Some(401) => { AiMuxError::TokenExpired(d.message) @@ -288,20 +344,21 @@ impl CodexModel { &self, options: &CallOptions, ) -> Result { - let opts = self.subscription_options(options); - let headers = self.inner().build_headers(opts.headers.as_ref()); - let (body, warnings) = self.body(&opts, true)?; + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; + let (body, warnings) = self.subscription_body(options)?; - let endpoint = self.endpoint(); + let endpoint = self.config.url("/responses")?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.clone(), build_header_list(&headers), &opts), + self.config.http_request(endpoint.clone(), headers, options), body.clone(), aimux_provider_utils::create_event_source_response_handler::(), crate::openai::openai_failed_response_handler(), ) .await - .map_err(|e| self.map_subscription_401(e))?; - + .map_err(Self::map_subscription_401)?; let response_headers = resp.response_headers; let mut sse = resp.value; // (parsed response object, raw event payload) of the terminal event. @@ -370,7 +427,7 @@ impl CodexModel { &data, &raw, warnings, - "openai".to_string(), + ResponsesNamespace::OPENAI.write_key().to_string(), endpoint, body, response_headers, @@ -385,19 +442,21 @@ impl CodexModel { /// Subscription `do_stream`: always streams with `store: false`. async fn subscription_stream(&self, options: &CallOptions) -> Result { - let opts = self.subscription_options(options); - let headers = self.inner().build_headers(opts.headers.as_ref()); - let (body, warnings) = self.body(&opts, true)?; + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; + let (body, warnings) = self.subscription_body(options)?; - let endpoint = self.endpoint(); + let endpoint = self.config.url("/responses")?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.clone(), build_header_list(&headers), &opts), + self.config.http_request(endpoint.clone(), headers, options), body.clone(), aimux_provider_utils::create_event_source_response_handler::(), crate::openai::openai_failed_response_handler(), ) .await - .map_err(|e| self.map_subscription_401(e))?; + .map_err(Self::map_subscription_401)?; let response_headers = resp.response_headers; let mut sse_stream = resp.value; @@ -408,7 +467,7 @@ impl CodexModel { let stream = build_responses_event_stream( first_event, sse_stream, - "openai".to_string(), + ResponsesNamespace::OPENAI.write_key().to_string(), warnings, false, endpoint, @@ -429,45 +488,24 @@ impl CodexModel { #[async_trait] impl LanguageModel for CodexModel { fn provider(&self) -> &str { - "codex" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.openai.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.openai.base_url.clone()), - api_key_source: self - .config - .openai - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: None, - } - } - async fn do_generate(&self, options: &CallOptions) -> Result { - match self.config.mode { - CodexMode::ApiKey => self.inner().do_generate(options).await, - CodexMode::Subscription => self.subscription_generate(options).await, + match self.channel { + Channel::ApiKey => self.inner().do_generate(options).await, + Channel::ChatGptAccount => self.subscription_generate(options).await, } } async fn do_stream(&self, options: &CallOptions) -> Result { - match self.config.mode { - CodexMode::ApiKey => self.inner().do_stream(options).await, - CodexMode::Subscription => self.subscription_stream(options).await, + match self.channel { + Channel::ApiKey => self.inner().do_stream(options).await, + Channel::ChatGptAccount => self.subscription_stream(options).await, } } } @@ -478,7 +516,7 @@ impl LanguageModel for CodexModel { /// output. #[derive(Clone, Serialize, Deserialize)] pub struct CodexTokens { - /// Access token to pass as `CodexConfig::subscription(..)`. + /// Access token to pass as the `ChatGptAccount` token of [`CodexProviderSettings`]. pub access_token: String, /// Rotated refresh token — **single use**; the caller must persist the /// new value (the library never persists credentials). @@ -575,21 +613,7 @@ pub async fn codex_refresh_at( #[cfg(test)] mod tests { - use super::*; - - #[test] - fn config_snapshot_matches_model_identity() { - let config = CodexConfig::new("sk-test"); - let model = CodexModel { - model_id: "gpt-5.2-codex".to_string(), - config, - }; - let snap = model.config_snapshot(); - assert_eq!(snap.provider, model.provider()); - assert_eq!(snap.model_id, model.model_id()); - assert!(snap.base_url.is_some()); - assert_eq!(snap.api_key_source, "explicit"); - } + use super::CodexTokens; #[test] fn tokens_debug_redacts_secrets() { diff --git a/aimux-providers/src/cohere/convert.rs b/aimux-providers/src/cohere/convert.rs index 2d7868298..e6256e444 100644 --- a/aimux-providers/src/cohere/convert.rs +++ b/aimux-providers/src/cohere/convert.rs @@ -12,6 +12,7 @@ use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; use aimux_core::tool::Tool; 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}; // ── Prepared tools ────────────────────────────────────────────────────────── @@ -128,7 +129,7 @@ pub struct ConvertedPrompt { /// - Tool messages are flat `{role:"tool", content, tool_call_id}`. /// /// # Errors -/// Returns an error for unsupported text file data. +/// Returns an error for image text data or an unresolved inline image media type. pub fn convert_prompt_to_cohere( prompt: &LanguageModelPrompt, ) -> Result { @@ -153,7 +154,7 @@ pub fn convert_prompt_to_cohere( } UserPart::File(file) => { use base64::Engine; - if file.media_type.split('/').next() == Some("image") { + if get_top_level_media_type(&file.media_type) == "image" { let url = match &file.data { FileData::Data { data } => { let b64 = match data { @@ -163,9 +164,11 @@ pub fn convert_prompt_to_cohere( } FileBytes::Base64(data) => data.clone(), }; - let media_type = - crate::google::convert::tool_file_media_type(file)?; - format!("data:{media_type};base64,{b64}") + format!( + "data:{};base64,{}", + resolve_full_media_type(file)?, + b64 + ) } FileData::Url { url, .. } => url.clone(), FileData::Reference { .. } => { @@ -175,7 +178,7 @@ pub fn convert_prompt_to_cohere( } FileData::Text { .. } => { return Err(AiMuxError::UnsupportedFunctionality( - "image file parts with text data".to_string(), + "image file parts with text data".into(), )); } }; @@ -324,7 +327,7 @@ pub struct RequestBodyResult { /// and the `thinking` config resolved from `reasoning` / provider options. /// /// # Errors -/// Returns an error for unsupported text file data. +/// Returns an error for image text data or an unresolved inline image media type. pub fn build_request_body( model_id: &str, options: &CallOptions, @@ -437,15 +440,16 @@ pub fn resolve_cohere_thinking( provider_options: &Option, ) -> Result, AiMuxError> { // Provider options take precedence. - if let Some(po) = provider_options - && let Some(cohere) = po.get("cohere") + 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, "cohere", |key, value| { + let thinking = crate::openai::convert::parse_option_fields( + thinking, + crate::cohere::options::NAMESPACE, + |key, value| { let valid = match key { "type" => value .as_str() @@ -454,7 +458,8 @@ pub fn resolve_cohere_thinking( _ => return None, }; Some(if valid { Ok(value.clone()) } else { Err(()) }) - })?; + }, + )?; let t_type = thinking .get("type") .and_then(|v| v.as_str()) diff --git a/aimux-providers/src/cohere/embedding.rs b/aimux-providers/src/cohere/embedding.rs index 1c21a849f..9ca536956 100644 --- a/aimux-providers/src/cohere/embedding.rs +++ b/aimux-providers/src/cohere/embedding.rs @@ -5,8 +5,6 @@ //! //! Endpoint: `POST {base_url}/embed` -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -16,55 +14,30 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::HttpRequest; - -use super::CohereConfig; +use crate::shared::EndpointConfig; /// A Cohere embedding model (e.g. `"embed-english-v3.0"`). pub struct CohereEmbeddingModel { model_id: String, - config: CohereConfig, + config: EndpointConfig, } impl CohereEmbeddingModel { - #[must_use] - pub fn new(model_id: String, config: CohereConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/embed", self.config.base_url) - } } #[async_trait] impl EmbeddingModel for CohereEmbeddingModel { fn provider(&self) -> &str { - "cohere" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - fn max_embeddings_per_call(&self) -> Option { Some(96) } @@ -99,17 +72,11 @@ impl EmbeddingModel for CohereEmbeddingModel { body.insert("output_dimension".to_string(), json!(output_dimension)); } - let headers = self.build_headers(options.headers.as_ref()); - - let mut header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - header_list.push(("Content-Type".to_string(), "application/json".to_string())); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), header_list, options), - Value::Object(body), + exchange.request(exchange.url("/embed"), options), + exchange.transform_body(Value::Object(body)), aimux_provider_utils::create_json_response_handler(), super::cohere_failed_response_handler(), ) @@ -172,7 +139,7 @@ struct CohereEmbeddingProviderOptions { fn parse_cohere_provider_options( options: Option<&SharedProviderOptions>, ) -> CohereEmbeddingProviderOptions { - let provider_opts = options.and_then(|opts| opts.get("cohere")); + let provider_opts = super::options::cohere_options(options); CohereEmbeddingProviderOptions { input_type: provider_opts .and_then(|o| o.get("inputType")) diff --git a/aimux-providers/src/cohere/mod.rs b/aimux-providers/src/cohere/mod.rs index b4cc75315..72c44351b 100644 --- a/aimux-providers/src/cohere/mod.rs +++ b/aimux-providers/src/cohere/mod.rs @@ -1,5 +1,11 @@ //! 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 +//! 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. +//! //! Implements Cohere's v2 chat API with its own message format (not //! OpenAI-compatible). Supports text generation, streaming, tool calls, //! and reasoning (thinking). @@ -7,17 +13,29 @@ pub mod convert; pub mod embedding; mod model; +pub(crate) mod options; pub mod reranking; mod types; pub use embedding::CohereEmbeddingModel; +pub use model::CohereModel; pub use reranking::CohereRerankingModel; +use std::sync::{Arc, OnceLock}; + +use futures::future::BoxFuture; +use serde_json::Value; + +use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::{RetryConfig, load_api_key, without_trailing_slash}; -use serde_json::Value; +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 crate::shared::{Credential, EndpointConfig, provider_headers}; pub(crate) fn cohere_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -33,141 +51,177 @@ pub(crate) fn cohere_failed_response_handler() -> aimux_provider_utils::Response }) } -/// Configuration for the Cohere provider. -#[derive(Debug, Clone)] -pub struct CohereConfig { - pub api_key: String, - pub base_url: String, - /// api_key 来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 - pub api_key_source: Option, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, -} +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.cohere.com/v2"; +pub(crate) const API_KEY_ENV_VAR: &str = "COHERE_API_KEY"; +const DEFAULT_NAME: &str = "cohere"; -impl CohereConfig { - /// Create from an API key (uses default Cohere base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.cohere.com/v2".to_string(), - api_key_source: None, - retry_config: RetryConfig::default(), - } - } - - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Settings of [`create_cohere`] (the AI SDK's `CohereProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct CohereProviderSettings { + /// Base URL for the API calls. Default `https://api.cohere.com/v2`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `COHERE_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. + pub api_key: Option, + /// Extra headers on every request. A `None` value removes the header, + /// including `Authorization`. Per-call headers win over these. + pub headers: Option, + /// The transport: a mock, a signing decorator, a proxy-aware client. + /// `None` uses the process default, resolved per request. + pub fetch: Option, + /// Generates unique IDs for chat output. + pub generate_id: Option String + Send + Sync>>, +} - /// 标注 api_key 来源(RFC-0023 回放重建用)。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self +impl std::fmt::Debug for CohereProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CohereProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key.is_some()) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("fetch", &self.fetch.is_some()) + .field("generate_id", &self.generate_id.is_some()) + .finish() } +} - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self - } +/// Create a Cohere provider. +/// +/// # 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. +pub fn create_cohere(settings: CohereProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(CohereProvider { + 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, + "Cohere", + ), + Vec::new(), + settings.headers, + ), + options::NAMESPACE, + "4.0.52", + ), + fetch: settings.fetch, + generate_id: settings.generate_id, + }) +} - /// Create from environment variable `COHERE_API_KEY`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `COHERE_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "COHERE_API_KEY", "Cohere")?; - Ok(Self::new(api_key).with_api_key_source(Some("env:COHERE_API_KEY"))) - } +/// The default provider: `create_cohere` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn cohere() -> &'static CohereProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_cohere(CohereProviderSettings::default()) + .expect("default Cohere settings are always valid") + }) } -/// Cohere provider — creates `CohereModel` instances. +/// A Cohere provider (the AI SDK's `CohereProvider`). Cheap to clone the +/// models out of; it holds no HTTP client. pub struct CohereProvider { - config: CohereConfig, + name: String, + 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), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) + } + + /// A chat model; `provider()` is `"{name}.chat"`. #[must_use] - pub fn new(config: CohereConfig) -> Self { - Self { config } + pub fn chat(&self, model_id: &str) -> CohereModel { + CohereModel::from_config(model_id.to_string(), self.model_config("chat")) + .with_generate_id(self.generate_id.clone()) } - /// Create a model instance for the given model name (e.g. `"command-r-plus"`). + /// An embedding model (e.g. `"embed-english-v3.0"`); `provider()` is + /// `"{name}.textEmbedding"`. #[must_use] - pub fn model(&self, model_id: &str) -> model::CohereModel { - model::CohereModel::new(model_id.to_string(), self.config.clone()) + pub fn embedding(&self, model_id: &str) -> CohereEmbeddingModel { + CohereEmbeddingModel::from_config(model_id.to_string(), self.model_config("textEmbedding")) } - /// Create a reranking model instance for the given model name (e.g. - /// `"rerank-english-v3.0"`). + /// A reranking model (e.g. `"rerank-english-v3.0"`); `provider()` is + /// `"{name}.reranking"`. #[must_use] - pub fn reranking_model(&self, model_id: &str) -> reranking::CohereRerankingModel { - reranking::CohereRerankingModel::new(model_id.to_string(), self.config.clone()) + pub fn reranking(&self, model_id: &str) -> CohereRerankingModel { + CohereRerankingModel::from_config(model_id.to_string(), self.model_config("reranking")) } - /// Create an embedding model instance for the given model name (e.g. - /// `"embed-english-v3.0"`). + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; the same model as [`chat`](Self::chat) and + /// [`language_model`](Provider::language_model). #[must_use] - pub fn embedding_model(&self, model_id: &str) -> embedding::CohereEmbeddingModel { - embedding::CohereEmbeddingModel::new(model_id.to_string(), self.config.clone()) + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.chat(model_id)) } } impl Provider for CohereProvider { - fn name(&self) -> &str { - "cohere" + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) + } + + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) } - /// List models via `GET {base_url}/models` (Cohere v2, RFC-0027). - fn list_models( + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } + + fn reranking_model( &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let config = self.config.clone(); + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.reranking(model_id)))) + } +} + +impl ProviderDiscovery for CohereProvider { + /// `GET {base_url}/models` (Cohere v2): one exchange, no retry. Only the + /// chat-capable models are listed. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config("models"); Box::pin(async move { - let base = config.base_url.trim_end_matches('/'); - let url = format!("{base}/models"); - let headers = vec![ - ( - "Authorization".to_string(), - format!("Bearer {}", config.api_key), - ), - ("Content-Type".to_string(), "application/json".to_string()), - ]; - use aimux_provider_utils::HttpRequest; - // Retry rationale: see `openai::model::execute_list_models`. - let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: headers.clone(), - abort_signal: None, - call_id: None, - recording_context: None, - ..Default::default() - }, - aimux_provider_utils::create_json_response_handler(), - cohere_failed_response_handler(), - ) - }) - .await?; // Cohere v2: { models: [{ name, endpoints, ... }] } #[derive(serde::Deserialize)] struct Resp { @@ -180,23 +234,34 @@ impl Provider for CohereProvider { #[serde(default)] endpoints: Option>, } + + let exchange = config.exchange(None).await?; + let resp = aimux_provider_utils::get_from_api( + exchange.with_transport(aimux_provider_utils::HttpRequest { + url: exchange.url("/models"), + headers: exchange.headers(), + ..Default::default() + }), + aimux_provider_utils::create_json_response_handler(), + cohere_failed_response_handler(), + ) + .await?; let parsed: Resp = resp.value; - let runtime: Vec = parsed + Ok(parsed .models .into_iter() - // Only include chat-capable models (endpoints contain "chat"). - .filter(|e| { - e.endpoints + .filter(|entry| { + entry + .endpoints .as_ref() - .is_none_or(|eps| eps.iter().any(|e| e == "chat")) + .is_none_or(|endpoints| endpoints.iter().any(|e| e == "chat")) }) - .map(|e| aimux_core::model_catalogue::RuntimeModel { - id: e.name, - owned_by: Some("cohere".to_string()), + .map(|entry| RuntimeModel { + id: entry.name, + owned_by: Some(options::NAMESPACE.to_string()), created: None, }) - .collect(); - Ok(runtime) + .collect()) }) } } diff --git a/aimux-providers/src/cohere/model.rs b/aimux-providers/src/cohere/model.rs index 055585dc6..9c9a3685e 100644 --- a/aimux-providers/src/cohere/model.rs +++ b/aimux-providers/src/cohere/model.rs @@ -5,53 +5,46 @@ //! (`event: type\ndata: json`). use aimux_core::tool::RawToolCall; -use std::collections::HashMap; - use async_trait::async_trait; use futures::StreamExt; -use serde_json::{Value, json}; +use serde_json::Value; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; -use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::HttpRequest; +use crate::shared::EndpointConfig; -use super::CohereConfig; use super::convert::{build_request_body, parse_finish_reason}; use super::types::{ChatResponse, StreamEvent, TokenPair}; /// A Cohere language model. pub struct CohereModel { model_id: String, - config: CohereConfig, + config: EndpointConfig, + generate_id: std::sync::Arc String + Send + Sync>, } impl CohereModel { - pub fn new(model_id: String, config: CohereConfig) -> Self { - Self { model_id, config } - } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } + 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; } - headers + self } - fn endpoint(&self) -> String { - format!("{}/chat", self.config.base_url) + 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), + } } } @@ -94,46 +87,19 @@ struct PendingToolCall { #[async_trait] impl LanguageModel for CohereModel { fn provider(&self) -> &str { - "cohere" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.base_url.clone()), - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: None, - } - } - async fn do_generate(&self, options: &CallOptions) -> Result { let body_result = build_request_body(&self.model_id, options, false)?; - let body = body_result.body.clone(); - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let body = exchange.transform_body(body_result.body.clone()); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new( - self.endpoint(), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), + exchange.request(exchange.url("/chat"), options), body.clone(), aimux_provider_utils::create_json_response_handler(), super::cohere_failed_response_handler(), @@ -181,7 +147,7 @@ impl LanguageModel for CohereModel { // item. The per-citation metadata (start/end/text/sources/citationType) // is preserved in `provider_metadata`. if let Some(citations) = &data.message.citations { - for (i, citation) in citations.iter().enumerate() { + for citation in citations { let title = citation .get("sources") .and_then(|s| s.as_array()) @@ -212,14 +178,13 @@ impl LanguageModel for CohereModel { cohere_meta.insert("citationType".to_string(), Value::String(t.to_string())); } content.push(GenerateContent::Source(Source::Document { - id: format!("citation-{i}"), + id: (self.generate_id)(), media_type: "text/plain".to_string(), title, filename: None, - provider_metadata: Some(aimux_core::shared::NamespaceMap::from([( - "cohere".to_string(), + provider_metadata: Some(super::options::cohere_metadata(Value::Object( cohere_meta, - )])), + ))), })); } } @@ -268,18 +233,11 @@ impl LanguageModel for CohereModel { async fn do_stream(&self, options: &CallOptions) -> Result { let body_result = build_request_body(&self.model_id, options, true)?; - let body = body_result.body.clone(); - let headers = self.build_headers(options.headers.as_ref()); - let endpoint = self.endpoint(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let body = exchange.transform_body(body_result.body.clone()); + let endpoint = exchange.url("/chat"); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new( - endpoint.clone(), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), + exchange.request(endpoint.clone(), options), body.clone(), aimux_provider_utils::create_event_source_response_handler::(), super::cohere_failed_response_handler(), @@ -574,7 +532,7 @@ impl LanguageModel for CohereModel { } else { final_usage }, - provider_metadata: Some(provider_namespace("cohere", json!({})).expect("provider metadata must be an object")), + provider_metadata: Some(super::options::cohere_metadata(serde_json::json!({}))), }); }; @@ -617,21 +575,3 @@ fn cohere_stream_error( ) })) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn config_snapshot_reports_config() { - let model = CohereModel::new("command-r-plus".to_string(), CohereConfig::new("test-key")); - let snapshot = model.config_snapshot(); - assert_eq!(snapshot.provider, "cohere"); - assert_eq!(snapshot.model_id, "command-r-plus"); - assert_eq!( - snapshot.base_url.as_deref(), - Some("https://api.cohere.com/v2") - ); - assert_eq!(snapshot.api_key_source, "explicit"); - } -} diff --git a/aimux-providers/src/cohere/options.rs b/aimux-providers/src/cohere/options.rs new file mode 100644 index 000000000..d8d4741f3 --- /dev/null +++ b/aimux-providers/src/cohere/options.rs @@ -0,0 +1,24 @@ +//! The providerOptions and providerMetadata namespace of Cohere. +//! +//! `@ai-sdk/cohere` reads `providerOptions.cohere` and writes response +//! metadata under `cohere`, 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; + +/// The providerOptions / providerMetadata key. +pub(crate) const NAMESPACE: &str = "cohere"; + +/// The Cohere options in a providerOptions map. +pub(crate) fn cohere_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} + +/// `{ "cohere": payload }`. +pub(crate) fn cohere_metadata(payload: Value) -> ProviderMetadata { + provider_namespace(NAMESPACE, payload).expect("provider metadata must be an object") +} diff --git a/aimux-providers/src/cohere/reranking.rs b/aimux-providers/src/cohere/reranking.rs index e8b58a9a4..bd35ae9d2 100644 --- a/aimux-providers/src/cohere/reranking.rs +++ b/aimux-providers/src/cohere/reranking.rs @@ -3,8 +3,6 @@ //! Aligned with Vercel AI SDK `CohereRerankingModel` //! (`reference/ai/packages/cohere/src/reranking/cohere-reranking-model.ts`). -use std::collections::HashMap; - use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; @@ -17,9 +15,7 @@ use aimux_core::reranking_model::{ use aimux_core::shared::SharedProviderOptions; use aimux_core::types::Warning; -use aimux_provider_utils::HttpRequest; - -use super::CohereConfig; +use crate::shared::EndpointConfig; /// Cohere provider-specific reranking options. #[derive(Debug, Clone, Default)] @@ -32,9 +28,7 @@ fn parse_cohere_reranking_options( provider_options: Option<&SharedProviderOptions>, ) -> CohereRerankingOptions { let mut opts = CohereRerankingOptions::default(); - if let Some(po) = provider_options - && let Some(cohere) = po.get("cohere") - { + if let Some(cohere) = super::options::cohere_options(provider_options) { if let Some(v) = cohere .get("maxTokensPerDoc") .and_then(serde_json::Value::as_u64) @@ -65,48 +59,25 @@ struct CohereRerankingResult { /// A Cohere reranking model. pub struct CohereRerankingModel { model_id: String, - config: CohereConfig, + config: EndpointConfig, } impl CohereRerankingModel { - #[must_use] - pub fn new(model_id: String, config: CohereConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/rerank", self.config.base_url) - } } #[async_trait] impl RerankingModel for CohereRerankingModel { fn provider(&self) -> &str { - "cohere" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - async fn do_rerank( &self, options: &RerankingCallOptions, @@ -144,15 +115,11 @@ impl RerankingModel for CohereRerankingModel { body["priority"] = json!(priority); } - let headers = self.build_headers(options.headers.as_ref()); - let mut header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - header_list.push(("Content-Type".to_string(), "application/json".to_string())); + 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( - HttpRequest::new(self.endpoint(), header_list, options), + exchange.request(exchange.url("/rerank"), options), body.clone(), aimux_provider_utils::create_json_response_handler::(), super::cohere_failed_response_handler(), diff --git a/aimux-providers/src/cybertron.rs b/aimux-providers/src/cybertron.rs deleted file mode 100644 index c875f2356..000000000 --- a/aimux-providers/src/cybertron.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! Cybertron (Rust) provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `CYBERTRON_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "CYBERTRON_BASE_URL"; -const PROVIDER_NAME: &str = "cybertron"; -const PLACEHOLDER_API_KEY: &str = "cybertron"; - -pub struct CybertronConfig(OpenAIConfig); - -impl CybertronConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `CYBERTRON_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct CybertronProvider(OpenAIProvider); - -impl CybertronProvider { - #[must_use] - pub fn new(config: CybertronConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for CybertronProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/dataforseo.rs b/aimux-providers/src/dataforseo.rs index b0016adb1..5cb670193 100644 --- a/aimux-providers/src/dataforseo.rs +++ b/aimux-providers/src/dataforseo.rs @@ -7,8 +7,13 @@ //! `DATAFORSEO_PASSWORD` credentials. The request body is a JSON array of //! task objects and the response is deeply nested. DataForSEO is a //! search-only provider that does not support language models. +//! +//! [`create_dataforseo`] takes [`DataforseoProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`DataforseoProvider`]. The credential is not read +//! there: it is loaded for every request, from the settings or from `DATAFORSEO_LOGIN` and `DATAFORSEO_PASSWORD`. +//! [`dataforseo()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use base64::Engine; @@ -16,16 +21,14 @@ use serde::Deserialize; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_core::shared::SharedHeaders; - -use aimux_provider_utils::{HttpRequest, without_trailing_slash}; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, combine_headers, validate_base_url, +}; -/// Provider canonical name. -const PROVIDER_NAME: &str = "dataforseo"; +use crate::shared::{Credential, EndpointConfig}; /// Fixed model id for the DataForSEO search model. const MODEL_ID: &str = "dataforseo-search"; @@ -56,95 +59,150 @@ fn dataforseo_failed_response_handler() -> aimux_provider_utils::ResponseHandler // ── Config ─────────────────────────────────────────────────────────────────── -/// Configuration for the DataForSEO provider. -/// -/// Holds HTTP Basic credentials (`login` + `password`). The `Debug` -/// implementation redacts both so credentials never appear in logs or error -/// messages. -#[derive(Clone)] -pub struct DataforseoConfig { - login: String, - password: String, - base_url: String, -} - -impl DataforseoConfig { - /// Create from explicit login + password (uses the default DataForSEO base URL). - pub fn new(login: impl Into, password: impl Into) -> Self { - Self { - login: login.into(), - password: password.into(), - base_url: "https://api.dataforseo.com".to_string(), - } - } - - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.dataforseo.com"; +const LOGIN_ENV_VAR: &str = "DATAFORSEO_LOGIN"; +const PASSWORD_ENV_VAR: &str = "DATAFORSEO_PASSWORD"; +const DEFAULT_NAME: &str = "dataforseo"; - /// Create from the `DATAFORSEO_LOGIN` + `DATAFORSEO_PASSWORD` environment variables. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `DATAFORSEO_LOGIN` or - /// `DATAFORSEO_PASSWORD` is not set. - pub fn from_env() -> Result { - let login = std::env::var("DATAFORSEO_LOGIN").map_err(|_| { - AiMuxError::InvalidArgument( - "No DataForSEO login found. Please set the `DATAFORSEO_LOGIN` environment variable." - .to_string(), - ) - })?; - let password = std::env::var("DATAFORSEO_PASSWORD").map_err(|_| { - AiMuxError::InvalidArgument( - "No DataForSEO password found. Please set the `DATAFORSEO_PASSWORD` environment variable." - .to_string(), - ) - })?; - Ok(Self::new(login, password)) - } +/// Settings of [`create_dataforseo`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; the credentials and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct DataforseoProviderSettings { + /// Base URL for the API calls. Default `https://api.dataforseo.com`; a + /// trailing slash is removed. + pub base_url: Option, + /// The HTTP Basic login. `None` loads `DATAFORSEO_LOGIN` when a request is + /// made and fails that request with `AiMuxError::LoadApiKey` if it is + /// unset. An explicit value is used as given, `""` included. + pub login: Option>, + /// The HTTP Basic password. `None` loads `DATAFORSEO_PASSWORD` when a + /// request is made and fails that request with `AiMuxError::LoadApiKey` + /// if it is unset. An explicit value is used as given, `""` included. + pub password: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including `Authorization`. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` string + /// (`"{name}.search"`). Default `"dataforseo"`. + 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, } -impl std::fmt::Debug for DataforseoConfig { +impl std::fmt::Debug for DataforseoProviderSettings { + /// Never prints the credentials or header values. fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DataforseoConfig") - .field("login", &"") - .field("password", &"") + f.debug_struct("DataforseoProviderSettings") .field("base_url", &self.base_url) + .field("login", &self.login) + .field("password", &self.password) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) .finish() } } -// ── Provider ───────────────────────────────────────────────────────────────── +/// The provider headers: `Authorization: Basic base64(login:password)` with +/// both parts resolved for every request, then the caller's headers. +fn basic_auth_headers( + login: Credential, + password: Credential, + user: Option, +) -> HeadersFn { + Resolvable::from_async_fn(move || { + let login = login.clone(); + let password = password.clone(); + let user = user.clone(); + async move { + let login = login.secret().await?.unwrap_or_default(); + let password = password.secret().await?.unwrap_or_default(); + let encoded = base64::engine::general_purpose::STANDARD + .encode(format!("{login}:{password}").as_bytes()); + let mut layer = HeaderMapOpt::new(); + layer.insert( + "Authorization".to_string(), + Some(format!("Basic {encoded}")), + ); + Ok(match &user { + Some(user) => combine_headers(&[&layer, user]), + None => layer, + }) + } + }) +} -/// DataForSEO provider — creates [`DataforseoSearchModel`] instances. +/// Create a DataForSEO provider. /// -/// DataForSEO is a search-only provider; it does not support language models. +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` +/// URL with a host. That is the only way this fails: the credentials are +/// loaded per request, not here. +pub fn create_dataforseo( + settings: DataforseoProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(DataforseoProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: basic_auth_headers( + Credential::explicit_or_env(settings.login, LOGIN_ENV_VAR, "DataForSEO login"), + Credential::explicit_or_env(settings.password, PASSWORD_ENV_VAR, "DataForSEO password"), + settings.headers, + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_dataforseo` with default settings, created +/// on first use. Creating it reads nothing from the environment and cannot +/// fail; a missing login or password surfaces from the first request instead. +pub fn dataforseo() -> &'static DataforseoProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_dataforseo(DataforseoProviderSettings::default()) + .expect("default DataForSEO settings are always valid") + }) +} + +/// A DataForSEO provider. Search only; it holds no HTTP client. pub struct DataforseoProvider { - config: DataforseoConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl DataforseoProvider { - #[must_use] - pub fn new(config: DataforseoConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a search model instance. + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> DataforseoSearchModel { - DataforseoSearchModel::new(self.config.clone()) + DataforseoSearchModel::from_config(self.model_config("search")) } } -impl Provider for DataforseoProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } -} +crate::impl_single_modality_provider!(DataforseoProvider, search_model, |p, _id| p.search_model()); // ── Request builder ────────────────────────────────────────────────────────── @@ -218,45 +276,19 @@ fn map_result(r: DataforseoOrganic) -> SearchResultItem { /// A DataForSEO search model. pub struct DataforseoSearchModel { - config: DataforseoConfig, + config: EndpointConfig, } impl DataforseoSearchModel { - #[must_use] - pub fn new(config: DataforseoConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - /// Build the HTTP Basic `Authorization` header value. - fn basic_auth(&self) -> String { - let credentials = format!("{}:{}", self.config.login, self.config.password); - let encoded = base64::engine::general_purpose::STANDARD.encode(credentials.as_bytes()); - format!("Basic {encoded}") - } - - fn build_headers(&self, extra: Option<&SharedHeaders>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("Authorization".to_string(), self.basic_auth()); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!( - "{}/v3/serp/google/organic/live/advanced", - self.config.base_url - ) - } } #[async_trait] impl SearchModel for DataforseoSearchModel { fn provider(&self) -> &str { - PROVIDER_NAME + &self.config.provider } fn model_id(&self) -> &str { @@ -266,21 +298,13 @@ impl SearchModel for DataforseoSearchModel { async fn do_search(&self, options: &SearchCallOptions) -> Result { let depth = resolve_depth(options.max_results); let body = build_request_body(&options.query, depth); - let headers: Vec<(String, String)> = self - .build_headers(options.headers.as_ref()) - .into_iter() - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request( + exchange.url("/v3/serp/google/organic/live/advanced"), + options, + ), body, aimux_provider_utils::create_json_response_handler::(), dataforseo_failed_response_handler(), diff --git a/aimux-providers/src/decision_scoring.rs b/aimux-providers/src/decision_scoring.rs index a77bb78cb..fd6410878 100644 --- a/aimux-providers/src/decision_scoring.rs +++ b/aimux-providers/src/decision_scoring.rs @@ -83,7 +83,8 @@ pub(crate) async fn score( } encoded.push(entry); } - let base = model.config.base_url.trim_end_matches('/'); + let base = model.config.base_url()?; + let base = base.as_str(); if model.provider() == "sglang" { let body = json!({"model":model.model_id,"query":[],"items":encoded.iter().map(|e| &e.prompt_token_ids).collect::>(),"label_token_ids":encoded.iter().map(|e| &e.label_token_ids).collect::>(),"apply_softmax":true,"temperature":temperature,"return_token_logprobs":true}); return post( diff --git a/aimux-providers/src/decision_support.rs b/aimux-providers/src/decision_support.rs index c027c648f..a1136a87e 100644 --- a/aimux-providers/src/decision_support.rs +++ b/aimux-providers/src/decision_support.rs @@ -8,7 +8,59 @@ use aimux_provider_utils::HttpRequest; use base64::Engine; use serde_json::Value; -use crate::openai::OpenAIConfig; +use crate::openai::config::OpenAIModelConfig; +use crate::openai_compatible::config::CompatModelConfig; + +/// The connection a decision adapter borrows from the provider that made it: +/// native OpenAI, or an OpenAI-compatible preset (a local runtime, +/// Cloudflare). Headers are resolved per request, like every other model of +/// that provider. +#[derive(Clone)] +pub(crate) enum DecisionConfig { + OpenAI(OpenAIModelConfig), + Compatible(CompatModelConfig), +} + +impl DecisionConfig { + /// Resolve the full endpoint through the native model's URL callback. + pub(crate) fn url(&self, path: &str) -> Result { + match self { + Self::OpenAI(config) => config.url(path), + Self::Compatible(config) => Ok(format!( + "{}{}", + config.base_url.resolve()?.trim_end_matches('/'), + path + )), + } + } + + /// The provider's base URL; adapters append their own endpoint paths. + pub(crate) fn base_url(&self) -> Result { + let base = match self { + Self::OpenAI(config) => config.url("")?, + Self::Compatible(config) => config.base_url.resolve()?, + }; + Ok(base.trim_end_matches('/').to_string()) + } + + async fn request( + &self, + url: &str, + options: &DecisionCallOptions, + ) -> Result { + Ok(match self { + Self::OpenAI(config) => { + let headers = config.request_headers(options.headers.as_ref()).await?; + config.http_request(url.to_string(), headers, options) + } + Self::Compatible(config) => { + let headers = config.request_headers(options.headers.as_ref()).await?; + let origin = config.base_url.resolve()?; + config.with_transport(HttpRequest::new(url, headers, options), origin) + } + }) + } +} pub(crate) struct InlineImage { pub data_url: String, @@ -45,18 +97,14 @@ pub(crate) fn inline_image(image: &DecisionImage) -> Result( - config: &OpenAIConfig, + config: &DecisionConfig, options: &DecisionCallOptions, url: &str, body: Value, convert: impl FnOnce(Value, HashMap) -> Result, ) -> Result { - let mut headers = crate::openai::model::build_auth_headers(config); - if let Some(extra) = &options.headers { - headers.extend(extra.clone()); - } let response = aimux_provider_utils::post_json_to_api( - HttpRequest::new(url, headers.into_iter().collect(), options), + config.request(url, options).await?, body.clone(), aimux_provider_utils::create_json_response_handler::(), aimux_provider_utils::create_json_error_response_handler(|data| { @@ -129,3 +177,31 @@ pub(crate) fn options( Some(options) => Ok(options.clone()), } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use aimux_provider_utils::Resolvable; + + use super::*; + + #[test] + fn native_decision_url_passes_endpoint_to_url_callback() { + let mut config = OpenAIModelConfig::fixed( + "openai".into(), + "https://example.com/v1".into(), + Resolvable::Value(HashMap::new()), + None, + None, + ); + config.url = Arc::new(|path| { + assert_eq!(path, "/decisions"); + Ok("https://example.com/decisions?version=1".into()) + }); + assert_eq!( + DecisionConfig::OpenAI(config).url("/decisions").unwrap(), + "https://example.com/decisions?version=1" + ); + } +} diff --git a/aimux-providers/src/deepgram.rs b/aimux-providers/src/deepgram.rs index 1fcfa51c9..d6d1104b4 100644 --- a/aimux-providers/src/deepgram.rs +++ b/aimux-providers/src/deepgram.rs @@ -9,8 +9,13 @@ //! `Content-Type` header set to the audio media type) and query parameters for //! model configuration. It returns a JSON body with `results.channels[0]` //! containing the transcript, words, and detected language. +//! +//! [`create_deepgram`] takes [`DeepgramProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`DeepgramProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `DEEPGRAM_API_KEY`. +//! [`deepgram()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; @@ -22,7 +27,10 @@ use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::{HttpBody, HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::HttpBody; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; /// Deepgram errors: `{"err_code": "...", "err_msg": "...", "request_id": ...}` /// on most endpoints; some return `{"category": "...", "message": "...", @@ -59,64 +67,124 @@ fn deepgram_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, +pub(crate) mod options; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.deepgram.com"; +pub(crate) const API_KEY_ENV_VAR: &str = "DEEPGRAM_API_KEY"; +const DEFAULT_NAME: &str = "deepgram"; + +/// Settings of [`create_deepgram`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct DeepgramProviderSettings { + /// Base URL for the API calls. Default `https://api.deepgram.com`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `DEEPGRAM_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.transcription"`). + /// Default `"deepgram"`. The providerOptions key stays `deepgram`. + 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, } -impl DeepgramConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.deepgram.com".to_string(), - headers: None, - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self +impl std::fmt::Debug for DeepgramProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DeepgramProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } +/// Create a Deepgram provider. +/// +/// # 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. +pub fn create_deepgram(settings: DeepgramProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(DeepgramProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Deepgram"), + AuthScheme::Scheme("Token"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `DEEPGRAM_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "DEEPGRAM_API_KEY", "Deepgram")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_deepgram` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn deepgram() -> &'static DeepgramProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_deepgram(DeepgramProviderSettings::default()) + .expect("default Deepgram settings are always valid") + }) } -/// Deepgram provider — creates `DeepgramTranscriptionModel` instances. +/// A Deepgram provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct DeepgramProvider { - config: DeepgramConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl DeepgramProvider { - #[must_use] - pub fn new(config: DeepgramConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// A transcription (STT) model (e.g. `"nova-3"`); `provider()` is `"{name}.transcription"`. #[must_use] pub fn transcription(&self, model_id: &str) -> DeepgramTranscriptionModel { - DeepgramTranscriptionModel::new(model_id.to_string(), self.config.clone()) + DeepgramTranscriptionModel::from_config( + model_id.to_string(), + self.model_config("transcription"), + ) } } +crate::impl_single_modality_provider!(DeepgramProvider, transcription_model, |p, id| p + .transcription(id)); + // ── Provider options ──────────────────────────────────────────────────────── #[derive(Debug, Clone, Default)] @@ -138,9 +206,7 @@ struct DeepgramOptions { fn parse_deepgram_options(provider_options: Option<&SharedProviderOptions>) -> DeepgramOptions { let mut opts = DeepgramOptions::default(); - if let Some(po) = provider_options - && let Some(dg) = po.get("deepgram") - { + if let Some(dg) = options::deepgram_options(provider_options) { opts.detect_entities = dg .get("detectEntities") .and_then(serde_json::Value::as_bool); @@ -217,44 +283,25 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid base64: {e}"))) } + AudioInput::Url(_) => Err(AudioInput::url_error()), } } pub struct DeepgramTranscriptionModel { model_id: String, - config: DeepgramConfig, + config: EndpointConfig, } impl DeepgramTranscriptionModel { - #[must_use] - pub fn new(model_id: String, config: DeepgramConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Token {}", self.config.api_key), - ); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } } #[async_trait] impl TranscriptionModel for DeepgramTranscriptionModel { fn provider(&self) -> &str { - "deepgram" + &self.config.provider } fn model_id(&self) -> &str { @@ -318,22 +365,13 @@ impl TranscriptionModel for DeepgramTranscriptionModel { .collect::>() .join("&"); - let url = format!("{}/v1/listen?{query_string}", self.config.base_url); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let url = exchange.url(&format!("/v1/listen?{query_string}")); let audio_bytes = audio_input_to_bytes(&options.audio)?; - let headers = self.build_headers(options.headers.as_ref()); - let resp = aimux_provider_utils::post_to_api( - HttpRequest { - url, - headers: headers.into_iter().collect(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(url, options), HttpBody::Bytes(audio_bytes, options.media_type.clone()), aimux_provider_utils::create_json_response_handler::(), deepgram_failed_response_handler(), diff --git a/aimux-providers/src/deepgram/options.rs b/aimux-providers/src/deepgram/options.rs new file mode 100644 index 000000000..ba37c886b --- /dev/null +++ b/aimux-providers/src/deepgram/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Deepgram. +//! +//! Deepgram reads `providerOptions.deepgram`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "deepgram"; + +/// The Deepgram options in a providerOptions map. +pub(crate) fn deepgram_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/deepseek/convert.rs b/aimux-providers/src/deepseek/convert.rs new file mode 100644 index 000000000..2ba10b45d --- /dev/null +++ b/aimux-providers/src/deepseek/convert.rs @@ -0,0 +1,386 @@ +//! Prompt to DeepSeek chat messages (`convert-to-deepseek-chat-messages.ts`). + +use base64::Engine; +use serde_json::{Map, Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, +}; +use aimux_core::options::ResponseFormat; +use aimux_core::shared::{FileBytes, FileData}; +use aimux_core::types::Warning; +use aimux_provider_utils::resolve_full_media_type; + +use super::is_v4_model::is_deepseek_v4_model; +use super::options::{parse_file_part_options, parse_message_options}; + +const SUPPORTED_IMAGE_MEDIA_TYPES: [&str; 5] = [ + "image/gif", + "image/jpeg", + "image/jpg", + "image/png", + "image/webp", +]; + +fn resolve_deepseek_image_media_type(part: &FilePart) -> Result { + let resolved = resolve_full_media_type(part)?; + if !SUPPORTED_IMAGE_MEDIA_TYPES.contains(&resolved.as_str()) { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "DeepSeek image media type {resolved}: DeepSeek supports JPEG, PNG, GIF, and WebP image inputs." + ))); + } + Ok(resolved) +} + +/// The content part of an image file part. +fn convert_image_part( + part: &FilePart, + provider_options_name: &str, + allow_file_data: bool, +) -> 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 { + options.file_data = None; + } + let image_url = |url: &str| { + let mut image_url = json!({ "url": url }); + if let Some(detail) = &options.image_detail { + image_url["detail"] = json!(detail); + } + json!({ "type": "image_url", "image_url": image_url }) + }; + match &part.data { + FileData::Reference { reference } => { + let file_id = reference.get("deepseek").ok_or_else(|| { + AiMuxError::InvalidArgument( + "No provider reference found for provider 'deepseek'.".to_string(), + ) + })?; + Ok(json!({ "type": "file", "file_id": file_id })) + } + FileData::Url { url, .. } => { + resolve_deepseek_image_media_type(part)?; + if url.len() > 8192 { + return Err(AiMuxError::InvalidPrompt( + "DeepSeek image URLs must not exceed 8192 characters.".to_string(), + )); + } + if options.file_data == Some(true) { + return Err(AiMuxError::InvalidPrompt( + "DeepSeek `fileData` image parts require inline data, not a URL.".to_string(), + )); + } + Ok(image_url(url)) + } + FileData::Data { data } => { + let media_type = resolve_deepseek_image_media_type(part)?; + 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)); + } + if options.image_detail.is_some() { + return Err(AiMuxError::InvalidPrompt( + "DeepSeek `imageDetail` cannot be combined with `fileData`.".to_string(), + )); + } + let mut file = json!({ "type": "file", "file_data": data_url }); + if let Some(filename) = &part.filename { + file["filename"] = json!(filename); + } + Ok(file) + } + FileData::Text { .. } => Err(AiMuxError::UnsupportedFunctionality( + "text file parts".to_string(), + )), + } +} + +/// The messages of a request and the warnings raised while converting. +pub(crate) struct ConvertedMessages { + pub messages: Vec, + 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 +/// +/// `InvalidPrompt` / `UnsupportedFunctionality` / `InvalidArgument` for a +/// prompt DeepSeek cannot take (an assistant prefix that is not the last +/// message or needs a beta base URL, an unsupported image). +pub(crate) fn convert_to_deepseek_chat_messages( + prompt: &LanguageModelPrompt, + response_format: Option<&ResponseFormat>, + model_id: &str, + provider_options_name: &str, + supports_assistant_prefix_completion: bool, +) -> Result { + let is_deepseek_v4 = is_deepseek_v4_model(model_id); + let mut messages = Vec::new(); + let mut warnings = Vec::new(); + + // Inject a system message if the response format is JSON. The DeepSeek + // package never supports structured outputs, so a schema is injected too. + if let Some(ResponseFormat::Json { schema, .. }) = response_format { + match schema { + None => messages.push(json!({ "role": "system", "content": "Return JSON." })), + Some(schema) => { + messages.push(json!({ + "role": "system", + "content": format!( + "Return JSON that conforms to the following schema: {schema}" + ), + })); + warnings.push(Warning::Compatibility { + feature: "responseFormat JSON schema".to_string(), + details: Some( + "JSON response schema is injected into the system message.".to_string(), + ), + }); + } + } + } + let last_user_message_index = prompt + .iter() + .rposition(|message| matches!(message, LanguageModelMessage::User { .. })); + + for (index, message) in prompt.iter().enumerate() { + let provider_options = match message { + LanguageModelMessage::System { + provider_options, .. + } + | LanguageModelMessage::User { + provider_options, .. + } + | LanguageModelMessage::Assistant { + provider_options, .. + } + | LanguageModelMessage::Tool { + provider_options, .. + } => provider_options, + }; + let options = parse_message_options(provider_options.as_ref(), provider_options_name)?; + if options.prefix == Some(true) + && !matches!(message, LanguageModelMessage::Assistant { .. }) + { + return Err(AiMuxError::InvalidPrompt( + "DeepSeek assistant prefix completion requires `prefix: true` on an assistant message." + .to_string(), + )); + } + let name = |wire: &mut Map| { + if let Some(name) = &options.name { + wire.insert("name".into(), json!(name)); + } + }; + + match message { + LanguageModelMessage::System { content, .. } => { + let mut wire = Map::new(); + wire.insert("role".into(), json!("system")); + wire.insert("content".into(), json!(content)); + name(&mut wire); + messages.push(Value::Object(wire)); + } + + 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 mut content = Vec::new(); + for part in parts { + match part { + UserPart::Text(TextPart { text, .. }) => { + content.push(json!({ "type": "text", "text": 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(_) => warnings.push(Warning::Unsupported { + feature: "user message part type: file".to_string(), + details: None, + }), + } + } + let mut wire = Map::new(); + wire.insert("role".into(), json!("user")); + wire.insert( + "content".into(), + if has_image_part { + Value::Array(content) + } else { + json!(join_text(parts)) + }, + ); + name(&mut wire); + messages.push(Value::Object(wire)); + } + + LanguageModelMessage::Assistant { content, .. } => { + if options.prefix == Some(true) { + if index != prompt.len() - 1 { + return Err(AiMuxError::InvalidPrompt( + "DeepSeek assistant prefix completion requires the prefixed assistant message to be the final message." + .to_string(), + )); + } + if !supports_assistant_prefix_completion { + return Err(AiMuxError::UnsupportedFunctionality( + "DeepSeek assistant prefix completion requires a beta base URL ending in `/beta`." + .to_string(), + )); + } + } + + let mut text = String::new(); + let mut reasoning: Option = None; + let mut tool_calls = Vec::new(); + for part in content { + match part { + AssistantPart::Text(TextPart { text: t, .. }) => text.push_str(t), + AssistantPart::Reasoning(ReasoningPart { text: t, .. }) => { + // R1 must not receive prior reasoning; V4 requires it. + if last_user_message_index.is_some_and(|last| index <= last) + && !is_deepseek_v4 + { + continue; + } + reasoning.get_or_insert_with(String::new).push_str(t); + } + AssistantPart::ToolCall(ToolCallPart { + tool_call_id, + tool_name, + input, + .. + }) => tool_calls.push(json!({ + "id": tool_call_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": input.to_string(), + }, + })), + _ => {} + } + } + + let mut wire = Map::new(); + wire.insert("role".into(), json!("assistant")); + wire.insert("content".into(), json!(text)); + name(&mut wire); + if options.prefix == Some(true) { + wire.insert("prefix".into(), json!(true)); + } + // V4 demands the field on every assistant turn: back-fill an + // empty string when the message had no reasoning part at all. + if let Some(reasoning) = reasoning.or_else(|| is_deepseek_v4.then(String::new)) { + wire.insert("reasoning_content".into(), json!(reasoning)); + } + if !tool_calls.is_empty() { + wire.insert("tool_calls".into(), Value::Array(tool_calls)); + } + messages.push(Value::Object(wire)); + } + + LanguageModelMessage::Tool { content, .. } => { + if options.name.is_some() { + warnings.push(Warning::Unsupported { + feature: "message name on tool messages".to_string(), + details: None, + }); + } + for part in content { + let ToolPart::ToolResult(ToolResultPart { + tool_call_id, + output, + .. + }) = part + else { + continue; + }; + messages.push(json!({ + "role": "tool", + "tool_call_id": tool_call_id, + "content": convert_tool_result(output, provider_options_name, &mut warnings)?, + })); + } + } + } + } + + Ok(ConvertedMessages { messages, warnings }) +} + +fn convert_tool_result( + output: &ToolResultOutput, + provider_options_name: &str, + warnings: &mut Vec, +) -> 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 { .. }) + }; + if !value + .iter() + .any(|part| matches!(part, ToolResultContent::File(file) if is_image(file))) + { + return Ok(crate::openai::convert::tool_result_to_content(output)); + } + let mut parts = Vec::new(); + for part in value { + match part { + ToolResultContent::Text(text) => { + parts.push(json!({ "type": "text", "text": text.text })) + } + ToolResultContent::File(file) if is_image(file) => { + parts.push(convert_image_part(file, provider_options_name, false)?); + } + _ => 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)) +} diff --git a/aimux-providers/src/deepseek/finish_reason.rs b/aimux-providers/src/deepseek/finish_reason.rs new file mode 100644 index 000000000..ede5e2935 --- /dev/null +++ b/aimux-providers/src/deepseek/finish_reason.rs @@ -0,0 +1,15 @@ +//! DeepSeek finish reasons (`map-deepseek-finish-reason.ts`). + +use aimux_core::types::FinishReasonUnified; + +/// The unified finish reason of a DeepSeek `finish_reason`. +pub(crate) fn map_deepseek_finish_reason(finish_reason: Option<&str>) -> FinishReasonUnified { + match finish_reason { + Some("stop") => FinishReasonUnified::Stop, + Some("length") => FinishReasonUnified::Length, + Some("content_filter") => FinishReasonUnified::ContentFilter, + Some("tool_calls") => FinishReasonUnified::ToolCalls, + Some("insufficient_system_resource") => FinishReasonUnified::Error, + _ => FinishReasonUnified::Other, + } +} diff --git a/aimux-providers/src/deepseek/is_v4_model.rs b/aimux-providers/src/deepseek/is_v4_model.rs new file mode 100644 index 000000000..f7dbf8a26 --- /dev/null +++ b/aimux-providers/src/deepseek/is_v4_model.rs @@ -0,0 +1,10 @@ +//! Whether a model id is a DeepSeek V4 model (`is-deepseek-v4-model.ts`). + +/// Whether a model id is a DeepSeek V4-generation model (thinking mode on by +/// default, `reasoning_content` required on every assistant turn). Only the +/// legacy `deepseek-chat` / `deepseek-reasoner` ids predate V4. +pub(crate) fn is_deepseek_v4_model(model_id: &str) -> bool { + model_id.contains("deepseek-v4") + || model_id.starts_with("deepseek-flash") + || model_id.starts_with("deepseek-pro") +} diff --git a/aimux-providers/src/deepseek/mod.rs b/aimux-providers/src/deepseek/mod.rs new file mode 100644 index 000000000..0c9605bea --- /dev/null +++ b/aimux-providers/src/deepseek/mod.rs @@ -0,0 +1,190 @@ +//! DeepSeek provider. +//! +//! [`create_deepseek`] is the Rust form of the AI SDK's `createDeepSeek`: a +//! provider of [`DeepSeekChatLanguageModel`], the package's own chat model +//! (`@ai-sdk/deepseek` does not build on the OpenAI-compatible package). Its +//! models report `"deepseek.chat"` and read providerOptions and report provider +//! metadata under `deepseek`. +//! +//! As in the other packages the API key is not read when the provider is +//! created: `DEEPSEEK_API_KEY` is loaded on every request unless `api_key` is +//! given. [`deepseek()`] is the default instance. + +mod convert; +mod finish_reason; +mod is_v4_model; +mod model; +mod options; +mod prepare_tools; +mod types; +mod usage; + +pub use model::DeepSeekChatLanguageModel; + +use std::sync::{Arc, OnceLock}; + +use futures::future::BoxFuture; + +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; +use model::DeepSeekChatConfig; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.deepseek.com"; +pub(crate) const API_KEY_ENV_VAR: &str = "DEEPSEEK_API_KEY"; + +/// The URL patterns the chat model fetches itself (`supportedUrls` of the AI +/// SDK's `DeepSeekChatLanguageModel`): `http(s)` images. +fn chat_supported_urls() -> SupportedUrls { + let http = regex::Regex::new(r"^https?://.*$").expect("static pattern"); + SupportedUrls(std::iter::once(("image/*".to_string(), vec![http])).collect()) +} + +/// Settings of [`create_deepseek`] (the AI SDK's `DeepSeekProviderSettings`). +#[derive(Clone, Default)] +pub struct DeepSeekProviderSettings { + /// Base URL for the API calls. Default `https://api.deepseek.com`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `DEEPSEEK_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. + pub api_key: Option, + /// Extra headers on every request; a `None` value removes the header. + pub headers: Option, + /// The transport. `None` uses the process default, resolved per request. + pub fetch: Option, +} + +impl std::fmt::Debug for DeepSeekProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DeepSeekProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key.is_some()) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("fetch", &self.fetch.is_some()) + .finish() + } +} + +/// Create a DeepSeek provider. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` +/// URL with a host. The key is not read here. +pub fn create_deepseek(settings: DeepSeekProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(DeepSeekProvider { + name: "deepseek".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, + "DeepSeek", + ), + Vec::new(), + settings.headers, + ), + "deepseek", + "3.0.56", + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_deepseek` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail. +pub fn deepseek() -> &'static DeepSeekProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_deepseek(DeepSeekProviderSettings::default()) + .expect("default DeepSeek settings are always valid") + }) +} + +/// A DeepSeek provider (the AI SDK's `DeepSeekProvider`). Language models +/// only: embedding and image models are `NoSuchModel`. +pub struct DeepSeekProvider { + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, +} + +impl DeepSeekProvider { + fn endpoint_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) + } + + /// A chat model; `provider()` is `"{name}.chat"` (`"deepseek.chat"`). + #[must_use] + pub fn chat(&self, model_id: &str) -> DeepSeekChatLanguageModel { + DeepSeekChatLanguageModel::from_config( + model_id.to_string(), + DeepSeekChatConfig { + endpoint: self + .endpoint_config("chat") + .with_supported_urls(Arc::new(|_| chat_supported_urls())), + supports_assistant_prefix_completion: self.base_url.ends_with("/beta"), + supports_strict_tool_calls: self.base_url.ends_with("/beta"), + }, + ) + } + + /// The provider as a function: the default language model, the chat model. + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.chat(model_id)) + } +} + +impl Provider for DeepSeekProvider { + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) + } + + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } +} + +impl ProviderDiscovery for DeepSeekProvider { + /// `GET {base_url}/models`: one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.endpoint_config("models"); + Box::pin(async move { + crate::shared::list_data_models(&config, types::deepseek_failed_response_handler()) + .await + }) + } +} diff --git a/aimux-providers/src/deepseek/model.rs b/aimux-providers/src/deepseek/model.rs new file mode 100644 index 000000000..cc64f2ea0 --- /dev/null +++ b/aimux-providers/src/deepseek/model.rs @@ -0,0 +1,774 @@ +//! The DeepSeek chat language model (`deepseek-chat-language-model.ts`). + +use aimux_core::tool::RawToolCall; +use std::collections::{BTreeMap, HashMap}; + +use async_trait::async_trait; +use futures::StreamExt; +use serde_json::{Map, Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::options::{CallOptions, ResponseFormat}; +use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; +use aimux_core::shared::{RequestInfo, ResponseInfo, StreamResponseInfo, provider_namespace}; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::{ + FinishReason, FinishReasonUnified, ReasoningEffort, ResponseMetadata, Warning, +}; +use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker, generate_id}; + +use super::convert::convert_to_deepseek_chat_messages; +use super::finish_reason::map_deepseek_finish_reason; +use super::is_v4_model::is_deepseek_v4_model; +use super::options::{ProviderReasoningEffort, ThinkingType, parse_chat_options}; +use super::prepare_tools::prepare_tools; +use super::types::{ChatChunk, ChatResponse}; +use super::usage::convert_deepseek_usage; +use crate::shared::EndpointConfig; + +/// What a DeepSeek chat model needs to talk to the API (`DeepSeekChatConfig`). +pub(crate) struct DeepSeekChatConfig { + /// Identity (`"{name}.chat"`), endpoint, headers, transport. + pub(crate) endpoint: EndpointConfig, + /// Whether the base URL is the beta one, which accepts assistant prefix + /// completion. + pub(crate) supports_assistant_prefix_completion: bool, + /// Whether the base URL is the beta one, which accepts strict tool calls. + pub(crate) supports_strict_tool_calls: bool, +} + +/// A DeepSeek chat model. +pub struct DeepSeekChatLanguageModel { + model_id: String, + config: DeepSeekChatConfig, +} + +/// A request body and the warnings raised while building it. +struct RequestBodyResult { + /// The JSON body, before the provider's `transform_request_body`. + body: Value, + /// Warnings raised while building it. + warnings: Vec, +} + +fn timestamp(created: Option) -> Option { + created + .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) + .map(|dt| dt.to_rfc3339()) +} + +/// The HTTP status a stream error code stands for: a number or a three-digit +/// string in `400..=599`. +fn http_status_code(code: Option<&Value>) -> Option { + let status = match code? { + Value::String(code) if code.len() == 3 && code.bytes().all(|b| b.is_ascii_digit()) => { + code.parse::().ok()? + } + Value::Number(code) => code.as_u64()?, + _ => return None, + }; + u16::try_from(status) + .ok() + .filter(|status| (400..=599).contains(status)) +} + +/// The HTTP status a named error kind stands for. +fn status_of_error_kind(kind: &str) -> Option { + Some(match kind { + "rate_limit_exceeded" | "rate_limit_error" => 429, + "server_error" | "api_error" | "internal_server_error" => 500, + "overloaded_error" | "service_unavailable" => 503, + "timeout" | "timeout_error" => 504, + "authentication_error" | "invalid_api_key" => 401, + "permission_error" => 403, + "not_found_error" | "model_not_found" => 404, + "bad_request" | "context_length_exceeded" | "invalid_request_error" => 400, + _ => return None, + }) +} + +/// The error of an error event inside a stream (`createDeepSeekStreamError`). +fn deepseek_stream_error( + event: &Value, + url: &str, + request_body_values: Value, + response_headers: HashMap, +) -> AiMuxError { + let error = &event["error"]; + let code = error.get("code").filter(|code| !code.is_null()); + let kind = error.get("type").and_then(Value::as_str); + let names = [code.and_then(Value::as_str), kind]; + let quota_exhausted = names.contains(&Some("insufficient_quota")); + let status_code = if quota_exhausted { + Some(429) + } else { + http_status_code(code) + .or_else(|| names.into_iter().flatten().find_map(status_of_error_kind)) + }; + let provider_code = code + .or_else(|| error.get("type")) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + let message = error + .get("message") + .and_then(Value::as_str) + .unwrap_or("DeepSeek stream failed before any output was generated") + .to_owned(); + let mut error = aimux_provider_utils::stream_error_api_call( + message, + provider_code, + status_code, + event, + url, + request_body_values, + response_headers, + ); + if quota_exhausted && let AiMuxError::ApiCall(api_error) = &mut error { + api_error.is_retryable = false; + } + error +} + +impl DeepSeekChatLanguageModel { + pub(crate) fn from_config(model_id: String, config: DeepSeekChatConfig) -> Self { + Self { model_id, config } + } + + /// The providerOptions namespace: the provider name. + fn provider_options_name(&self) -> &str { + self.config + .endpoint + .provider + .split('.') + .next() + .unwrap_or_default() + .trim() + } + + /// The JSON body a call would send, before the provider's + /// `transform_request_body`, with the warnings raised while building it + /// (`getArgs`, plus the stream fields). + fn request_body( + &self, + options: &CallOptions, + stream: bool, + ) -> Result { + let provider_options_name = self.provider_options_name(); + let deepseek_options = + parse_chat_options(options.provider_options.as_ref(), provider_options_name)?; + + let converted = convert_to_deepseek_chat_messages( + &options.prompt, + options.response_format.as_ref(), + &self.model_id, + provider_options_name, + self.config.supports_assistant_prefix_completion, + )?; + let mut warnings = converted.warnings; + + if options.top_k.is_some() { + warnings.push(Warning::Unsupported { + feature: "topK".to_string(), + details: None, + }); + } + if options.seed.is_some() { + warnings.push(Warning::Unsupported { + feature: "seed".to_string(), + details: None, + }); + } + for (value, setting) in [ + (options.frequency_penalty, "frequencyPenalty"), + (options.presence_penalty, "presencePenalty"), + ] { + if value.is_some() { + warnings.push(Warning::Deprecated { + setting: setting.to_string(), + message: format!( + "{setting} is deprecated by DeepSeek and has been omitted. Remove {setting} from the request." + ), + }); + } + } + + let prepared = prepare_tools( + options.tools.as_ref(), + options.tool_choice.as_ref(), + self.config.supports_strict_tool_calls, + )?; + + let thinking_type = deepseek_options + .thinking + .as_ref() + .and_then(|thinking| thinking.kind); + if thinking_type == Some(ThinkingType::Adaptive) { + warnings.push(Warning::Compatibility { + feature: "thinking.type".to_string(), + details: Some( + "thinking.type \"adaptive\" is not a canonical DeepSeek value. mapped to \"enabled\"." + .to_string(), + ), + }); + } + + let reasoning = options.reasoning.filter(|effort| effort.is_custom()); + let thinking = match (thinking_type, reasoning) { + (Some(ThinkingType::Disabled), _) => Some("disabled"), + (Some(_), _) => Some("enabled"), + (None, Some(ReasoningEffort::None)) => Some("disabled"), + (None, Some(_)) => Some("enabled"), + (None, None) => None, + }; + + let is_thinking_enabled = thinking != Some("disabled") + && (thinking.is_some() + || self.model_id == "deepseek-reasoner" + || is_deepseek_v4_model(&self.model_id)); + + for (value, feature) in [ + (options.temperature, "temperature"), + (options.top_p, "topP"), + ] { + if is_thinking_enabled && value.is_some() { + warnings.push(Warning::Unsupported { + feature: feature.to_string(), + details: Some(format!( + "{feature} has no effect when DeepSeek thinking is enabled. Set providerOptions.deepseek.thinking.type to 'disabled' to use {feature}." + )), + }); + } + } + + let reasoning_effort = match (deepseek_options.reasoning_effort, reasoning) { + (Some(effort), _) => { + let mapped = match effort { + ProviderReasoningEffort::Medium => "high", + ProviderReasoningEffort::Xhigh => "max", + ProviderReasoningEffort::Low => "low", + ProviderReasoningEffort::High => "high", + ProviderReasoningEffort::Max => "max", + }; + let given = match effort { + ProviderReasoningEffort::Low => "low", + ProviderReasoningEffort::Medium => "medium", + ProviderReasoningEffort::High => "high", + ProviderReasoningEffort::Xhigh => "xhigh", + ProviderReasoningEffort::Max => "max", + }; + if mapped != given { + warnings.push(Warning::Compatibility { + feature: "reasoningEffort".to_string(), + details: Some(format!( + "reasoningEffort \"{given}\" is not a canonical DeepSeek value. mapped to \"{mapped}\"." + )), + }); + } + Some(mapped) + } + (None, Some(ReasoningEffort::None) | None) => None, + (None, Some(reasoning)) => { + let mapped = match reasoning { + ReasoningEffort::Minimal | ReasoningEffort::Low => "low", + ReasoningEffort::Medium | ReasoningEffort::High => "high", + _ => "max", + }; + if mapped != reasoning.to_string() { + warnings.push(Warning::Compatibility { + feature: "reasoning".to_string(), + details: Some(format!( + "reasoning \"{reasoning}\" is not directly supported by this model. mapped to effort \"{mapped}\"." + )), + }); + } + Some(mapped) + } + }; + + let mut body = Map::new(); + body.insert("model".into(), json!(self.model_id)); + if deepseek_options.logprobs == Some(true) || deepseek_options.top_logprobs.is_some() { + body.insert("logprobs".into(), json!(true)); + } + if let Some(top_logprobs) = deepseek_options.top_logprobs { + body.insert("top_logprobs".into(), json!(top_logprobs)); + } + if let Some(max_tokens) = options.max_output_tokens { + body.insert("max_tokens".into(), json!(max_tokens)); + } + if !is_thinking_enabled { + for (key, value) in [ + ("temperature", options.temperature), + ("top_p", options.top_p), + ] { + if let Some(value) = value { + body.insert(key.into(), json!(value)); + } + } + } + if matches!(options.response_format, Some(ResponseFormat::Json { .. })) { + body.insert("response_format".into(), json!({ "type": "json_object" })); + } + if let Some(stop) = &options.stop_sequences { + body.insert("stop".into(), json!(stop)); + } + body.insert("messages".into(), Value::Array(converted.messages)); + if let Some(tools) = prepared.tools { + body.insert("tools".into(), Value::Array(tools)); + } + if let Some(tool_choice) = prepared.tool_choice { + body.insert("tool_choice".into(), tool_choice); + } + if let Some(thinking) = thinking { + body.insert("thinking".into(), json!({ "type": thinking })); + } + if let Some(user_id) = deepseek_options.user_id { + body.insert("user_id".into(), json!(user_id)); + } + if let Some(reasoning_effort) = reasoning_effort + && thinking != Some("disabled") + { + body.insert("reasoning_effort".into(), json!(reasoning_effort)); + } + if stream { + body.insert("stream".into(), json!(true)); + body.insert("stream_options".into(), json!({ "include_usage": true })); + } + + warnings.extend(prepared.tool_warnings); + Ok(RequestBodyResult { + body: Value::Object(body), + warnings, + }) + } +} + +#[async_trait] +impl LanguageModel for DeepSeekChatLanguageModel { + /// `"{name}.chat"`. + fn provider(&self) -> &str { + &self.config.endpoint.provider + } + + fn model_id(&self) -> &str { + &self.model_id + } + + async fn supported_urls(&self) -> SupportedUrls { + (self.config.endpoint.supported_urls)(&self.model_id) + } + + async fn do_generate(&self, options: &CallOptions) -> Result { + let exchange = self + .config + .endpoint + .exchange(options.headers.as_ref()) + .await?; + let built = self.request_body(options, false)?; + let body = exchange.transform_body(built.body); + let resp = aimux_provider_utils::post_json_to_api( + exchange.request(exchange.url("/chat/completions"), options), + body.clone(), + aimux_provider_utils::create_json_response_handler::(), + super::types::deepseek_failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers; + let raw = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; + + let choice = data.choices.into_iter().next().ok_or_else(|| { + AiMuxError::InvalidResponseData("Response did not contain any choices.".to_string()) + })?; + + let mut metadata = Map::new(); + let usage_value = raw.get("usage").filter(|usage| !usage.is_null()); + for (key, name) in [ + ("prompt_cache_hit_tokens", "promptCacheHitTokens"), + ("prompt_cache_miss_tokens", "promptCacheMissTokens"), + ] { + if let Some(value) = usage_value.and_then(|usage| usage.get(key)) { + metadata.insert(name.into(), value.clone()); + } + } + if let Some(object) = &data.object { + metadata.insert("responseObject".into(), json!(object)); + } + if let Some(index) = choice.index { + metadata.insert("choiceIndex".into(), json!(index)); + } + if let Some(role) = &choice.message.role { + metadata.insert("messageRole".into(), json!(role)); + } + if let Some(tool_calls) = &choice.message.tool_calls { + let types: Vec<&String> = tool_calls + .iter() + .filter_map(|call| call.r#type.as_ref()) + .collect(); + metadata.insert("toolCallTypes".into(), json!(types)); + } + if let Some(logprobs) = choice.logprobs.filter(|logprobs| !logprobs.is_null()) { + metadata.insert("logprobs".into(), logprobs); + } + if let Some(system_fingerprint) = &data.system_fingerprint { + metadata.insert("systemFingerprint".into(), json!(system_fingerprint)); + } + + let mut content = Vec::new(); + + // reasoning content (before text): + if let Some(text) = choice + .message + .reasoning_content + .filter(|text| !text.is_empty()) + { + content.push(GenerateContent::Reasoning(ReasoningOutput { + text, + provider_metadata: None, + })); + } + + // tool calls: + for call in choice.message.tool_calls.into_iter().flatten() { + content.push(GenerateContent::ToolCall(RawToolCall { + tool_call_id: call + .id + .filter(|id| !id.is_empty()) + .unwrap_or_else(generate_id), + tool_name: call.function.name, + input: call.function.arguments, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + } + + // text content: + if let Some(text) = choice.message.content.filter(|text| !text.is_empty()) { + content.push(GenerateContent::Text { + text, + provider_metadata: None, + }); + } + + Ok(GenerateResult { + content, + finish_reason: FinishReason { + unified: map_deepseek_finish_reason(choice.finish_reason.as_deref()), + raw: choice.finish_reason, + }, + 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) }), + response: Some(ResponseInfo { + id: data.id, + timestamp: timestamp(data.created), + model_id: data.model, + headers: Some(response_headers), + body: Some(raw), + }), + }) + } + + async fn do_stream(&self, options: &CallOptions) -> Result { + let exchange = self + .config + .endpoint + .exchange(options.headers.as_ref()) + .await?; + let built = self.request_body(options, true)?; + let warnings = built.warnings; + let body = exchange.transform_body(built.body); + let endpoint = exchange.url("/chat/completions"); + let resp = aimux_provider_utils::post_json_to_api( + exchange.request(endpoint.clone(), options), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + super::types::deepseek_failed_response_handler(), + ) + .await?; + let response_headers = resp.response_headers; + let mut sse_stream = resp.value; + + let provider_options_name = self.provider_options_name().to_string(); + let emit_raw_chunks = options.include_raw_chunks == Some(true); + let stream_error_body = body.clone(); + let stream_response_headers = response_headers.clone(); + + let stream = async_stream::stream! { + yield Ok(StreamPart::StreamStart { warnings }); + + let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(generate_id); + let mut tool_parts = Vec::new(); + let mut finish_reason = FinishReason { unified: FinishReasonUnified::Other, raw: None }; + let mut usage: Option = None; + let mut system_fingerprint: Option = None; + let mut is_first_chunk = true; + let mut is_active_reasoning = false; + let mut is_active_text = false; + let mut response_object: Option = None; + let mut choice_index: Option = None; + let mut message_role: Option = None; + let mut tool_call_types: BTreeMap = BTreeMap::new(); + let mut content_logprobs: Vec = Vec::new(); + let mut reasoning_logprobs: Vec = Vec::new(); + + while let Some(event) = sse_stream.next().await { + let parsed = match event { + Ok(parsed) => parsed, + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } + continue; + } + }; + + // Emit the raw chunk if requested (before anything else). + if emit_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); + } + + // handle error chunks: + if parsed.get("error").is_some() { + finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; + yield Ok(StreamPart::Error { + error: deepseek_stream_error( + &parsed, + &endpoint, + stream_error_body.clone(), + stream_response_headers.clone(), + ), + }); + continue; + } + + // handle failed chunk validation: + let chunk: ChatChunk = match serde_json::from_value(parsed) { + Ok(chunk) => chunk, + Err(error) => { + finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; + yield Ok(StreamPart::Error { error: error.into() }); + continue; + } + }; + + if is_first_chunk { + is_first_chunk = false; + yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: timestamp(chunk.created), + model_id: chunk.model.clone(), + })); + } + + if let Some(chunk_usage) = chunk.usage.filter(|usage| !usage.is_null()) { + usage = Some(chunk_usage); + } + + if chunk.object.is_some() { + response_object = chunk.object; + } + + // The fingerprint is repeated on stream chunks; keep the + // latest non-null value in case it changes during the response. + if chunk.system_fingerprint.is_some() { + system_fingerprint = chunk.system_fingerprint; + } + + let Some(choice) = chunk.choices.into_iter().next() else { + continue; + }; + + if choice.index.is_some() { + choice_index = choice.index; + } + + if let Some(reason) = choice.finish_reason { + finish_reason = FinishReason { + unified: map_deepseek_finish_reason(Some(&reason)), + raw: Some(reason), + }; + } + + if let Some(logprobs) = &choice.logprobs { + for (key, sink) in [ + ("content", &mut content_logprobs), + ("reasoning_content", &mut reasoning_logprobs), + ] { + if let Some(entries) = logprobs.get(key).and_then(Value::as_array) { + sink.extend(entries.iter().cloned()); + } + } + } + + let Some(delta) = choice.delta else { + continue; + }; + + if delta.role.is_some() { + message_role = delta.role; + } + + // enqueue reasoning before text deltas: + if let Some(reasoning) = delta + .reasoning_content + .filter(|reasoning| !reasoning.is_empty()) + { + if !is_active_reasoning { + yield Ok(StreamPart::ReasoningStart { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + is_active_reasoning = true; + } + yield Ok(StreamPart::ReasoningDelta { + id: "reasoning-0".to_string(), + delta: reasoning, + provider_metadata: None, + }); + } + + if let Some(text) = delta.content.filter(|text| !text.is_empty()) { + if !is_active_text { + yield Ok(StreamPart::TextStart { + id: "txt-0".to_string(), + provider_metadata: None, + }); + is_active_text = true; + } + + // end reasoning when text starts: + if is_active_reasoning { + yield Ok(StreamPart::ReasoningEnd { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + is_active_reasoning = false; + } + + yield Ok(StreamPart::TextDelta { + id: "txt-0".to_string(), + delta: text, + provider_metadata: None, + }); + } + + if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { + // end reasoning when tool calls start: + if is_active_reasoning { + yield Ok(StreamPart::ReasoningEnd { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + is_active_reasoning = false; + } + + for tool_call in &deltas { + if let Some(kind) = &tool_call.r#type { + tool_call_types.insert(tool_call.index, kind.clone()); + } + let delta = StreamingToolCallDelta { + index: Some(tool_call.index), + id: tool_call.id.as_deref(), + r#type: tool_call.r#type.as_deref(), + name: tool_call.function.name.as_deref(), + arguments: tool_call.function.arguments.as_deref(), + provider_metadata: None, + }; + if let Err(error) = tool_calls.process(delta, &mut tool_parts) { + yield Err(error.into()); + return; + } + for part in tool_parts.drain(..) { + yield Ok(part); + } + } + } + } + + if is_active_reasoning { + yield Ok(StreamPart::ReasoningEnd { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + } + if is_active_text { + yield Ok(StreamPart::TextEnd { + id: "txt-0".to_string(), + provider_metadata: None, + }); + } + tool_calls.finish(&mut tool_parts); + for part in tool_parts.drain(..) { + yield Ok(part); + } + + let mut metadata = Map::new(); + for (key, name) in [ + ("prompt_cache_hit_tokens", "promptCacheHitTokens"), + ("prompt_cache_miss_tokens", "promptCacheMissTokens"), + ] { + if let Some(value) = usage + .as_ref() + .and_then(|usage| usage.get(key)) + .filter(|value| !value.is_null()) + { + metadata.insert(name.into(), value.clone()); + } + } + if let Some(object) = response_object { + metadata.insert("responseObject".into(), json!(object)); + } + if let Some(index) = choice_index { + metadata.insert("choiceIndex".into(), json!(index)); + } + if let Some(role) = message_role { + metadata.insert("messageRole".into(), json!(role)); + } + if !tool_call_types.is_empty() { + metadata.insert( + "toolCallTypes".into(), + json!(tool_call_types.into_values().collect::>()), + ); + } + let mut logprobs = Map::new(); + if !content_logprobs.is_empty() { + logprobs.insert("content".into(), Value::Array(content_logprobs)); + } + if !reasoning_logprobs.is_empty() { + logprobs.insert("reasoning_content".into(), Value::Array(reasoning_logprobs)); + } + if !logprobs.is_empty() { + metadata.insert("logprobs".into(), Value::Object(logprobs)); + } + if let Some(system_fingerprint) = system_fingerprint { + metadata.insert("systemFingerprint".into(), json!(system_fingerprint)); + } + yield Ok(StreamPart::Finish { + finish_reason, + usage: convert_deepseek_usage(usage.as_ref()), + provider_metadata: Some(provider_namespace(&provider_options_name, Value::Object(metadata)).expect("provider metadata must be an object")), + }); + }; + + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} diff --git a/aimux-providers/src/deepseek/options.rs b/aimux-providers/src/deepseek/options.rs new file mode 100644 index 000000000..62d808722 --- /dev/null +++ b/aimux-providers/src/deepseek/options.rs @@ -0,0 +1,208 @@ +//! The providerOptions of the DeepSeek chat model +//! (`deepseek-chat-language-model-options.ts`, `deepseek-file-part-options.ts`). +//! +//! They are read from the namespace named by the provider (`deepseek`). + +use serde::{Deserialize, Deserializer}; +use serde_json::Value; + +use aimux_core::error::AiMuxError; +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// `thinking` of the call options. `adaptive` is accepted for backwards +/// compatibility and mapped to `enabled`. +#[derive(Debug, Default, Deserialize)] +pub(crate) struct DeepSeekThinkingOptions { + #[serde(rename = "type", default, deserialize_with = "deserialize_optional")] + pub kind: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum ThinkingType { + Adaptive, + Enabled, + Disabled, +} + +/// `reasoningEffort` of the call options. `medium` and `xhigh` are accepted for +/// backwards compatibility and mapped to canonical DeepSeek values. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum ProviderReasoningEffort { + Low, + Medium, + High, + Xhigh, + Max, +} + +/// `providerOptions.` of a call (`deepseekLanguageModelChatOptions`). +/// Fields the schema does not know are dropped. +#[derive(Debug, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct DeepSeekChatOptions { + /// Whether to return log probabilities for generated tokens. + #[serde(default, deserialize_with = "deserialize_optional")] + pub logprobs: Option, + /// Number of most likely tokens to return at each token position (0 to + /// 20); setting it enables `logprobs`. + #[serde(default, deserialize_with = "deserialize_optional")] + pub top_logprobs: Option, + /// An opaque identifier for the end user: ASCII letters, numbers, + /// underscores and hyphens, at most 512 characters. + #[serde(default, deserialize_with = "deserialize_optional")] + pub user_id: Option, + /// The thinking configuration. + #[serde(default, deserialize_with = "deserialize_optional")] + pub thinking: Option, + /// The thinking strength. + #[serde(default, deserialize_with = "deserialize_optional")] + pub reasoning_effort: Option, + #[serde( + rename = "strictJsonSchema", + default, + deserialize_with = "deserialize_optional" + )] + pub _strict_json_schema: Option, +} + +/// The `providerOptions` of a message (`name`, and for an assistant message +/// `prefix`). +#[derive(Debug, Default, Deserialize)] +pub(crate) struct DeepSeekMessageOptions { + /// The name of the participant represented by the message. + pub name: Option, + /// Whether the assistant message is a prefix DeepSeek should continue. + pub prefix: Option, +} + +/// The `providerOptions` of a file part. +#[derive(Debug, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct DeepSeekFilePartOptions { + /// How DeepSeek processes an image sent as an `image_url` part + /// (`low`, `high`, `original`, `auto`). + pub image_detail: Option, + /// Send inline image data as a `file` part instead of an `image_url`. + pub file_data: Option, +} + +// An optional schema field may be absent, but an explicit null is invalid. +fn deserialize_optional<'de, D, T>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, + T: Deserialize<'de>, +{ + T::deserialize(deserializer).map(Some) +} + +/// The options of `namespace` in a providerOptions value; no namespace is the +/// defaults. +fn parse_namespace Deserialize<'de>>( + namespace: Option<&JsonObject>, + name: &str, +) -> Result { + match namespace { + None => Ok(T::default()), + Some(options) => serde_json::from_value(Value::Object(options.clone())).map_err(|error| { + AiMuxError::InvalidArgument(format!("invalid provider options for \"{name}\": {error}")) + }), + } +} + +fn invalid(name: &str, message: &str) -> AiMuxError { + AiMuxError::InvalidArgument(format!( + "invalid provider options for \"{name}\": {message}" + )) +} + +/// The chat options of a call. +/// +/// # Errors +/// +/// `InvalidArgument` when an option has the wrong type or is out of range. +pub(crate) fn parse_chat_options( + provider_options: Option<&SharedProviderOptions>, + name: &str, +) -> Result { + crate::openai::convert::parse_chat_provider_options(provider_options, name, |namespace| { + if namespace + .get("thinking") + .is_some_and(|thinking| !thinking.is_object()) + { + return Err(invalid(name, "thinking must be an object")); + } + let mut normalized = namespace.clone(); + if let Some(value) = normalized.get_mut("topLogprobs") + && let Some(number) = value.as_f64() + && number.fract() == 0.0 + && (0.0..=20.0).contains(&number) + { + *value = Value::from(number as u8); + } + let options: DeepSeekChatOptions = parse_namespace(Some(&normalized), name)?; + if options.top_logprobs.is_some_and(|n| n > 20) { + return Err(invalid(name, "topLogprobs must be at most 20")); + } + if let Some(user_id) = &options.user_id { + if user_id.is_empty() + || !user_id + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-') + { + return Err(invalid(name, "userId must match /^[a-zA-Z0-9_-]+$/")); + } + if user_id.len() > 512 { + return Err(invalid(name, "userId must be at most 512 characters long")); + } + } + Ok(options) + }) + .map(Option::unwrap_or_default) +} + +/// The options of a message. +/// +/// # Errors +/// +/// `InvalidArgument` when `name` is not a string or `prefix` is not `true`. +pub(crate) fn parse_message_options( + provider_options: Option<&SharedProviderOptions>, + name: &str, +) -> Result { + let options: DeepSeekMessageOptions = + parse_namespace(provider_options.and_then(|all| all.get(name)), name)?; + if options.prefix == Some(false) { + return Err(invalid(name, "prefix must be true")); + } + Ok(options) +} + +/// The options of a file part. +/// +/// # Errors +/// +/// `InvalidArgument` when `imageDetail` is not a known value or `fileData` is +/// not `true`. +pub(crate) fn parse_file_part_options( + provider_options: Option<&SharedProviderOptions>, + name: &str, +) -> Result { + let options: DeepSeekFilePartOptions = + parse_namespace(provider_options.and_then(|all| all.get(name)), name)?; + if options.file_data == Some(false) { + return Err(invalid(name, "fileData must be true")); + } + if options + .image_detail + .as_deref() + .is_some_and(|detail| !["low", "high", "original", "auto"].contains(&detail)) + { + return Err(invalid( + name, + "imageDetail must be low, high, original or auto", + )); + } + Ok(options) +} diff --git a/aimux-providers/src/deepseek/prepare_tools.rs b/aimux-providers/src/deepseek/prepare_tools.rs new file mode 100644 index 000000000..7af9b0969 --- /dev/null +++ b/aimux-providers/src/deepseek/prepare_tools.rs @@ -0,0 +1,96 @@ +//! DeepSeek tool preparation (`deepseek-prepare-tools.ts`). + +use serde_json::{Map, Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::options::LanguageModelV4ToolChoice; +use aimux_core::tool::Tool; +use aimux_core::types::Warning; + +/// `tools` / `tool_choice` ready for the body, with the warnings raised while +/// preparing them. +pub(crate) struct PreparedTools { + pub tools: Option>, + pub tool_choice: Option, + pub tool_warnings: Vec, +} + +/// Function tools go to the body; a provider-defined tool is unsupported and +/// warned about. With no tools there is no tool choice. +/// +/// # Errors +/// +/// `UnsupportedFunctionality` when a tool is `strict` and the base URL is not +/// the beta one, or when strict and non-strict function tools are mixed. +pub(crate) fn prepare_tools( + tools: Option<&Vec>, + tool_choice: Option<&LanguageModelV4ToolChoice>, + supports_strict_tool_calls: bool, +) -> Result { + let mut tool_warnings = Vec::new(); + let Some(tools) = tools.filter(|tools| !tools.is_empty()) else { + return Ok(PreparedTools { + tools: None, + tool_choice: None, + tool_warnings, + }); + }; + + let function_tools = || { + tools.iter().filter_map(|tool| match tool { + Tool::Function(tool) => Some(tool), + Tool::Provider(_) => None, + }) + }; + let has_strict_tool = function_tools().any(|tool| tool.strict == Some(true)); + if has_strict_tool && !supports_strict_tool_calls { + return Err(AiMuxError::UnsupportedFunctionality( + "DeepSeek strict tool calls require a beta base URL ending in `/beta`.".to_string(), + )); + } + if has_strict_tool + && supports_strict_tool_calls + && function_tools().any(|tool| tool.strict != Some(true)) + { + return Err(AiMuxError::UnsupportedFunctionality( + "DeepSeek strict mode requires every function tool in the request to set `strict: true`." + .to_string(), + )); + } + + let mut deepseek_tools = Vec::new(); + for tool in tools { + match tool { + Tool::Provider(tool) => tool_warnings.push(Warning::Unsupported { + feature: format!("provider-defined tool {}", tool.id), + details: None, + }), + Tool::Function(tool) => { + let mut function = Map::new(); + function.insert("name".into(), json!(tool.name)); + if let Some(description) = &tool.description { + function.insert("description".into(), json!(description)); + } + function.insert("parameters".into(), tool.input_schema.clone()); + if let Some(strict) = tool.strict { + function.insert("strict".into(), json!(strict)); + } + deepseek_tools.push(json!({ "type": "function", "function": function })); + } + } + } + + let tool_choice = tool_choice.map(|tool_choice| match tool_choice { + LanguageModelV4ToolChoice::Auto => json!("auto"), + LanguageModelV4ToolChoice::None => json!("none"), + LanguageModelV4ToolChoice::Required => json!("required"), + LanguageModelV4ToolChoice::Tool { tool_name } => { + json!({ "type": "function", "function": { "name": tool_name } }) + } + }); + Ok(PreparedTools { + tools: Some(deepseek_tools), + tool_choice, + tool_warnings, + }) +} diff --git a/aimux-providers/src/deepseek/types.rs b/aimux-providers/src/deepseek/types.rs new file mode 100644 index 000000000..534cb7445 --- /dev/null +++ b/aimux-providers/src/deepseek/types.rs @@ -0,0 +1,154 @@ +//! DeepSeek chat API response types (`deepseek-chat-api-types.ts`). +//! +//! Limited versions of the schemas, focused on what the model needs. The +//! `usage` object stays a `Value`: it is returned as `usage.raw`. + +use serde::{Deserialize, Deserializer}; +use serde_json::Value; + +/// `deepSeekErrorSchema` turned into the failed-response handler, as the +/// chat model builds it (deepseek@3.0.56 `deepseek-chat-api-types.ts:117-124`, +/// `deepseek-chat-language-model.ts:213-217`): the message is +/// `error.message`; `error.code` (a string or a number), else `error.type`, +/// is the provider code. +pub(crate) fn deepseek_failed_response_handler() +-> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error"); + aimux_provider_utils::ProviderErrorParts { + message: error + .and_then(|error| error.get("message")) + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .and_then(|error| error.get("code").or_else(|| error.get("type"))) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }), + } + }) +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ChatResponse { + pub id: Option, + pub created: Option, + pub model: Option, + #[serde(default, deserialize_with = "deserialize_response_object")] + pub object: Option, + pub system_fingerprint: Option, + pub choices: Vec, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ResponseChoice { + pub index: Option, + pub message: ResponseMessage, + pub logprobs: Option, + pub finish_reason: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ResponseMessage { + #[serde(default, deserialize_with = "deserialize_role")] + pub role: Option, + pub content: Option, + pub reasoning_content: Option, + pub tool_calls: Option>, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ResponseToolCall { + pub id: Option, + #[serde(default, deserialize_with = "deserialize_tool_type")] + pub r#type: Option, + pub function: ResponseToolCallFunction, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ResponseToolCallFunction { + pub name: String, + pub arguments: String, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ChatChunk { + pub id: Option, + pub created: Option, + pub model: Option, + #[serde(default, deserialize_with = "deserialize_chunk_object")] + pub object: Option, + pub system_fingerprint: Option, + pub choices: Vec, + pub usage: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ChunkChoice { + pub index: Option, + pub delta: Option, + pub logprobs: Option, + pub finish_reason: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ChunkDelta { + #[serde(default, deserialize_with = "deserialize_role")] + pub role: Option, + pub content: Option, + pub reasoning_content: Option, + pub tool_calls: Option>, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ChunkToolCall { + pub index: usize, + pub id: Option, + #[serde(default, deserialize_with = "deserialize_tool_type")] + pub r#type: Option, + pub function: ChunkToolCallFunction, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ChunkToolCallFunction { + pub name: Option, + pub arguments: Option, +} + +fn deserialize_literal<'de, D>(deserializer: D, literal: &str) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let value = Option::::deserialize(deserializer)?; + if value.as_deref().is_some_and(|value| value != literal) { + return Err(serde::de::Error::custom(format!("expected {literal}"))); + } + Ok(value) +} + +fn deserialize_response_object<'de, D: Deserializer<'de>>( + deserializer: D, +) -> Result, D::Error> { + deserialize_literal(deserializer, "chat.completion") +} + +fn deserialize_chunk_object<'de, D: Deserializer<'de>>( + deserializer: D, +) -> Result, D::Error> { + deserialize_literal(deserializer, "chat.completion.chunk") +} + +fn deserialize_role<'de, D: Deserializer<'de>>( + deserializer: D, +) -> Result, D::Error> { + deserialize_literal(deserializer, "assistant") +} + +fn deserialize_tool_type<'de, D: Deserializer<'de>>( + deserializer: D, +) -> Result, D::Error> { + deserialize_literal(deserializer, "function") +} diff --git a/aimux-providers/src/deepseek/usage.rs b/aimux-providers/src/deepseek/usage.rs new file mode 100644 index 000000000..6c9ca7ae5 --- /dev/null +++ b/aimux-providers/src/deepseek/usage.rs @@ -0,0 +1,39 @@ +//! DeepSeek usage (`convert-to-deepseek-usage.ts`). + +use serde_json::Value; + +use aimux_core::types::{InputTokenUsage, OutputTokenUsage, Usage}; + +fn count(usage: &Value, pointer: &str) -> u32 { + usage + .pointer(pointer) + .and_then(Value::as_u64) + .map_or(0, |n| u32::try_from(n).unwrap_or(u32::MAX)) +} + +/// The usage of a DeepSeek `usage` object; no object is null usage. The object +/// is kept as `raw`. +pub(crate) fn convert_deepseek_usage(usage: Option<&Value>) -> Usage { + let Some(usage) = usage.filter(|usage| !usage.is_null()) else { + return Usage::default(); + }; + let prompt_tokens = count(usage, "/prompt_tokens"); + let completion_tokens = count(usage, "/completion_tokens"); + let cache_read_tokens = count(usage, "/prompt_cache_hit_tokens"); + let reasoning_tokens = count(usage, "/completion_tokens_details/reasoning_tokens"); + + Usage { + input_tokens: InputTokenUsage { + total: Some(prompt_tokens), + no_cache: Some(prompt_tokens.saturating_sub(cache_read_tokens)), + cache_read: Some(cache_read_tokens), + cache_write: None, + }, + output_tokens: OutputTokenUsage { + total: Some(completion_tokens), + text: Some(completion_tokens.saturating_sub(reasoning_tokens)), + reasoning: Some(reasoning_tokens), + }, + raw: usage.as_object().cloned(), + } +} diff --git a/aimux-providers/src/default_providers.rs b/aimux-providers/src/default_providers.rs new file mode 100644 index 000000000..8969b87cc --- /dev/null +++ b/aimux-providers/src/default_providers.rs @@ -0,0 +1,698 @@ +//! The built-in providers. +//! +//! The AI SDK ships no vendor list: a caller creates providers with the +//! package factories and hands a map of them to `createProviderRegistry`. +//! aimux knows its vendors, so it does that assembly once, the way a +//! models.dev consumer does: a name selects a factory (a vendor package, or +//! the OpenAI-compatible factory with a row of `provider_registry.json`) and +//! the same few settings go to whichever factory it is. +//! +//! - [`create_provider`] is that one step for one name. +//! - [`default_providers`] is every name with default settings, ready for +//! [`aimux_core::create_provider_registry`]. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use aimux_core::error::AiMuxError; +use aimux_core::provider::Provider; + +use crate::preset::{self, PresetSettings}; + +/// The vendor packages, by the name of their default instance. +const PACKAGES: [&str; 45] = [ + "amazon_bedrock", + "anthropic", + "anthropic_aws", + "assemblyai", + "aws_polly", + "azure", + "black_forest_labs", + "cartesia", + "codex", + "cohere", + "dataforseo", + "deepgram", + "deepseek", + "elevenlabs", + "exa_ai", + "fal", + "firecrawl", + "gladia", + "google", + "google_pse", + "google_vertex", + "groq", + "huggingface", + "hume", + "jina_ai", + "klingai", + "linkup", + "lmnt", + "luma", + "mistral", + "openai", + "parallel_ai", + "prodia", + "recraft", + "replicate", + "revai", + "runwayml", + "searxng", + "serper", + "stability", + "tavily", + "tinyfish", + "voyage", + "xai", + "you_com", +]; + +/// The names [`create_provider`] accepts: the vendor packages and every +/// registry row, sorted. +#[must_use] +pub fn provider_names() -> Vec<&'static str> { + let mut names: Vec<&'static str> = PACKAGES.to_vec(); + names.extend(preset::names().filter(|name| !PACKAGES.contains(name))); + names.sort_unstable(); + names +} + +/// Create the built-in provider `name` with `settings` (key, base URL, extra +/// headers, transport) applied to its factory. A vendor package wins over a +/// registry row of the same name. +/// +/// # Errors +/// +/// [`AiMuxError::NoSuchProvider`] for an unknown name (there is no fallback), +/// `InvalidArgument` for a setting the selected factory does not have, and +/// whatever the factory itself reports. +pub fn create_provider( + name: &str, + settings: PresetSettings, +) -> Result, AiMuxError> { + if !PACKAGES.contains(&name) { + // Checked here, not left to `preset::create`: its error lists only the + // registry rows, and a misspelled package name must list the packages. + if preset::lookup(name).is_none() { + return Err(AiMuxError::NoSuchProvider { + provider_id: name.to_string(), + model_id: name.to_string(), + model_type: String::new(), + available_providers: provider_names().into_iter().map(str::to_owned).collect(), + }); + } + return Ok(Arc::new(preset::create(name, settings)?)); + } + let PresetSettings { + api_key, + base_url, + headers, + fetch, + params, + transform_request_body, + } = settings; + let unsupported = |field: &str| { + Err(AiMuxError::InvalidArgument(format!( + "provider '{name}' has no `{field}` setting" + ))) + }; + if !params.is_empty() { + return unsupported("params"); + } + // The packages whose settings also support a body transform. + macro_rules! package { + ($module:ident, $create:ident, $settings:ident) => { + Ok(Arc::new(crate::$module::$create( + crate::$module::$settings { + api_key, + base_url, + headers, + fetch, + transform_request_body, + ..Default::default() + }, + )?)) + }; + } + macro_rules! package_without_transform { + ($module:ident, $create:ident, $settings:ident $(, $api_key:ident)?) => {{ + if transform_request_body.is_some() { + return unsupported("transform_request_body"); + } + Ok(Arc::new(crate::$module::$create( + crate::$module::$settings { + $($api_key,)? + base_url, + headers, + fetch, + ..Default::default() + }, + )?)) + }}; + } + macro_rules! package_with_string_key { + ($module:ident, $create:ident, $settings:ident $(, $defaults:expr)?) => {{ + if transform_request_body.is_some() { + return unsupported("transform_request_body"); + } + let api_key = match api_key { + Some(aimux_provider_utils::Resolvable::Value(key)) => Some(key), + None => None, + Some(_) => return unsupported("resolvable api_key"), + }; + Ok(Arc::new(crate::$module::$create( + crate::$module::$settings { + api_key, + base_url, + headers, + fetch, + $(..$defaults)? + }, + )?)) + }}; + } + match name { + "amazon_bedrock" => package!( + bedrock, + create_amazon_bedrock, + AmazonBedrockProviderSettings + ), + "anthropic" => { + package_with_string_key!( + anthropic, + create_anthropic, + AnthropicProviderSettings, + Default::default() + ) + } + "azure" => package!(azure, create_azure, AzureOpenAIProviderSettings), + "cohere" => package_with_string_key!( + cohere, + create_cohere, + CohereProviderSettings, + Default::default() + ), + "deepseek" => package_with_string_key!(deepseek, create_deepseek, DeepSeekProviderSettings), + "elevenlabs" => package!(elevenlabs, create_elevenlabs, ElevenLabsProviderSettings), + "google" => package_with_string_key!( + google, + create_google, + GoogleProviderSettings, + Default::default() + ), + "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() + ), + "voyage" => package!(voyage, create_voyage, VoyageProviderSettings), + "xai" => package_with_string_key!(xai, create_xai, XAIProviderSettings), + "assemblyai" => package_without_transform!( + assemblyai, + create_assemblyai, + AssemblyAIProviderSettings, + api_key + ), + "black_forest_labs" => package_without_transform!( + black_forest_labs, + create_black_forest_labs, + BlackForestLabsProviderSettings, + api_key + ), + "cartesia" => { + package_without_transform!(cartesia, create_cartesia, CartesiaProviderSettings, api_key) + } + "deepgram" => { + package_without_transform!(deepgram, create_deepgram, DeepgramProviderSettings, api_key) + } + "exa_ai" => { + package_without_transform!(exa_ai, create_exa_ai, ExaAiProviderSettings, api_key) + } + "fal" => package_without_transform!(fal, create_fal, FalProviderSettings, api_key), + "firecrawl" => package_without_transform!( + firecrawl, + create_firecrawl, + FirecrawlProviderSettings, + api_key + ), + "gladia" => { + package_without_transform!(gladia, create_gladia, GladiaProviderSettings, api_key) + } + "google_pse" => package_without_transform!( + google_pse, + create_google_pse, + GooglePseProviderSettings, + api_key + ), + "hume" => package_without_transform!(hume, create_hume, HumeProviderSettings, api_key), + "jina_ai" => { + package_without_transform!(jina_ai, create_jina_ai, JinaAiProviderSettings, api_key) + } + "klingai" => { + package_without_transform!(klingai, create_klingai, KlingAIProviderSettings, api_key) + } + "linkup" => { + package_without_transform!(linkup, create_linkup, LinkupProviderSettings, api_key) + } + "lmnt" => package_without_transform!(lmnt, create_lmnt, LMNTProviderSettings, api_key), + "luma" => package_without_transform!(luma, create_luma, LumaProviderSettings, api_key), + "parallel_ai" => package_without_transform!( + parallel_ai, + create_parallel_ai, + ParallelAiProviderSettings, + api_key + ), + "prodia" => { + package_without_transform!(prodia, create_prodia, ProdiaProviderSettings, api_key) + } + "recraft" => { + package_without_transform!(recraft, create_recraft, RecraftProviderSettings, api_key) + } + "replicate" => package_without_transform!( + replicate, + create_replicate, + ReplicateProviderSettings, + api_key + ), + "revai" => package_without_transform!(revai, create_revai, RevaiProviderSettings, api_key), + "runwayml" => { + package_without_transform!(runwayml, create_runwayml, RunwaymlProviderSettings, api_key) + } + "searxng" => { + package_without_transform!(searxng, create_searxng, SearxngProviderSettings, api_key) + } + "serper" => { + package_without_transform!(serper, create_serper, SerperProviderSettings, api_key) + } + "stability" => package_without_transform!( + stability, + create_stability, + StabilityProviderSettings, + api_key + ), + "tavily" => { + package_without_transform!(tavily, create_tavily, TavilyProviderSettings, api_key) + } + "tinyfish" => { + package_without_transform!(tinyfish, create_tinyfish, TinyfishProviderSettings, api_key) + } + "you_com" => { + package_without_transform!(you_com, create_you_com, YouComProviderSettings, api_key) + } + "aws_polly" => { + if api_key.is_some() { + return unsupported("api_key"); + } + package_without_transform!(aws_polly, create_aws_polly, AwsPollyProviderSettings) + } + "dataforseo" => { + if api_key.is_some() { + return unsupported("api_key"); + } + package_without_transform!(dataforseo, create_dataforseo, DataforseoProviderSettings) + } + "google_vertex" => Ok(Arc::new(crate::vertex::create_google_vertex( + crate::vertex::VertexProviderSettings { + api_key, + base_url, + headers: headers.map(Into::into), + fetch, + transform_request_body, + ..Default::default() + }, + )?)), + "anthropic_aws" | "codex" => { + if api_key.is_some() { + return unsupported("api_key"); + } + if name == "codex" { + Ok(Arc::new(crate::codex::create_codex( + crate::codex::CodexProviderSettings { + base_url, + headers, + fetch, + transform_request_body, + ..Default::default() + }, + )?)) + } else { + Ok(Arc::new(crate::anthropic_aws::create_anthropic_aws( + crate::anthropic_aws::AnthropicAwsProviderSettings { + base_url, + headers, + fetch, + transform_request_body, + ..Default::default() + }, + )?)) + } + } + _ => unreachable!("every name in PACKAGES has an arm"), + } +} + +/// The environment variable `name`'s API key is read from when +/// [`create_provider`] gets no `api_key`: the vendor package's key variable, +/// or the registry row's `env_var`. `None` for an unknown name, for a +/// registry row without a variable, and for the packages whose `api_key` +/// setting has no variable (`searxng`) or that take no `api_key` at all +/// (`anthropic_aws`, `aws_polly`, `codex`, `dataforseo`). +/// +/// Hosts that let a caller pick a key by variable name (the aimux-web +/// console's `env:VAR`) use it to allow only the provider's own variable. +#[must_use] +pub fn provider_api_key_env(name: &str) -> Option<&'static str> { + use crate::{ + anthropic, assemblyai, azure, bedrock, black_forest_labs, cartesia, cohere, deepgram, + deepseek, elevenlabs, exa_ai, fal, firecrawl, gladia, google, google_pse, groq, + huggingface, hume, jina_ai, klingai, linkup, lmnt, luma, mistral, openai, parallel_ai, + prodia, recraft, replicate, revai, runwayml, serper, stability, tavily, tinyfish, vertex, + voyage, xai, you_com, + }; + Some(match name { + "amazon_bedrock" => bedrock::BEARER_TOKEN_ENV_VAR, + "anthropic" => anthropic::API_KEY_ENV_VAR, + "assemblyai" => assemblyai::API_KEY_ENV_VAR, + "azure" => azure::API_KEY_ENV_VAR, + "black_forest_labs" => black_forest_labs::API_KEY_ENV_VAR, + "cartesia" => cartesia::API_KEY_ENV_VAR, + "cohere" => cohere::API_KEY_ENV_VAR, + "deepgram" => deepgram::API_KEY_ENV_VAR, + "deepseek" => deepseek::API_KEY_ENV_VAR, + "elevenlabs" => elevenlabs::API_KEY_ENV_VAR, + "exa_ai" => exa_ai::API_KEY_ENV_VAR, + "fal" => fal::API_KEY_ENV_VAR, + "firecrawl" => firecrawl::API_KEY_ENV_VAR, + "gladia" => gladia::API_KEY_ENV_VAR, + "google" => google::API_KEY_ENV_VAR, + "google_pse" => google_pse::API_KEY_ENV_VAR, + "google_vertex" => vertex::API_KEY_ENV_VAR, + "groq" => groq::API_KEY_ENV_VAR, + "huggingface" => huggingface::API_KEY_ENV_VAR, + "hume" => hume::API_KEY_ENV_VAR, + "jina_ai" => jina_ai::API_KEY_ENV_VAR, + "klingai" => klingai::API_KEY_ENV_VAR, + "linkup" => linkup::API_KEY_ENV_VAR, + "lmnt" => lmnt::API_KEY_ENV_VAR, + "luma" => luma::API_KEY_ENV_VAR, + "mistral" => mistral::API_KEY_ENV_VAR, + "openai" => openai::API_KEY_ENV_VAR, + "parallel_ai" => parallel_ai::API_KEY_ENV_VAR, + "prodia" => prodia::API_KEY_ENV_VAR, + "recraft" => recraft::API_KEY_ENV_VAR, + "replicate" => replicate::API_KEY_ENV_VAR, + "revai" => revai::API_KEY_ENV_VAR, + "runwayml" => runwayml::API_KEY_ENV_VAR, + "serper" => serper::API_KEY_ENV_VAR, + "stability" => stability::API_KEY_ENV_VAR, + "tavily" => tavily::API_KEY_ENV_VAR, + "tinyfish" => tinyfish::API_KEY_ENV_VAR, + "voyage" => voyage::API_KEY_ENV_VAR, + "xai" => xai::API_KEY_ENV_VAR, + "you_com" => you_com::API_KEY_ENV_VAR, + _ if PACKAGES.contains(&name) => return None, + _ => preset::lookup(name) + .map(|entry| entry.descriptor.env_var.as_str()) + .filter(|var| !var.is_empty())?, + }) +} + +/// The endpoint `name`'s factory calls when [`create_provider`] gets no +/// `base_url` and no base-URL variable is set: the vendor package's default +/// base URL, or the registry row's `base_url`. `None` for an unknown name and +/// where the endpoint is not one fixed URL: built from settings (a region, +/// resource, project or a registry row's template parameters) or read only +/// from the environment. +/// +/// Hosts that let a caller choose the base URL (the aimux-web console) use it +/// to tell the provider's own endpoint from one the caller picked. +#[must_use] +pub fn provider_default_base_url(name: &str) -> Option<&'static str> { + use crate::{ + anthropic, assemblyai, black_forest_labs, cartesia, cohere, dataforseo, deepgram, deepseek, + elevenlabs, exa_ai, fal, firecrawl, gladia, google, google_pse, groq, huggingface, hume, + jina_ai, klingai, linkup, lmnt, luma, mistral, openai, parallel_ai, prodia, recraft, + replicate, revai, runwayml, serper, stability, tavily, tinyfish, voyage, xai, you_com, + }; + Some(match name { + "anthropic" => anthropic::DEFAULT_BASE_URL, + "assemblyai" => assemblyai::DEFAULT_BASE_URL, + "black_forest_labs" => black_forest_labs::DEFAULT_BASE_URL, + "cartesia" => cartesia::DEFAULT_BASE_URL, + "cohere" => cohere::DEFAULT_BASE_URL, + "dataforseo" => dataforseo::DEFAULT_BASE_URL, + "deepgram" => deepgram::DEFAULT_BASE_URL, + "deepseek" => deepseek::DEFAULT_BASE_URL, + "elevenlabs" => elevenlabs::DEFAULT_BASE_URL, + "exa_ai" => exa_ai::DEFAULT_BASE_URL, + "fal" => fal::DEFAULT_BASE_URL, + "firecrawl" => firecrawl::DEFAULT_BASE_URL, + "gladia" => gladia::DEFAULT_BASE_URL, + "google" => google::DEFAULT_BASE_URL, + "google_pse" => google_pse::DEFAULT_BASE_URL, + "groq" => groq::DEFAULT_BASE_URL, + "huggingface" => huggingface::DEFAULT_BASE_URL, + "hume" => hume::DEFAULT_BASE_URL, + "jina_ai" => jina_ai::DEFAULT_BASE_URL, + "klingai" => klingai::DEFAULT_BASE_URL, + "linkup" => linkup::DEFAULT_BASE_URL, + "lmnt" => lmnt::DEFAULT_BASE_URL, + "luma" => luma::DEFAULT_BASE_URL, + "mistral" => mistral::DEFAULT_BASE_URL, + "openai" => openai::DEFAULT_BASE_URL, + "parallel_ai" => parallel_ai::DEFAULT_BASE_URL, + "prodia" => prodia::DEFAULT_BASE_URL, + "recraft" => recraft::DEFAULT_BASE_URL, + "replicate" => replicate::DEFAULT_BASE_URL, + "revai" => revai::DEFAULT_BASE_URL, + "runwayml" => runwayml::DEFAULT_BASE_URL, + "serper" => serper::DEFAULT_BASE_URL, + "stability" => stability::DEFAULT_BASE_URL, + "tavily" => tavily::DEFAULT_BASE_URL, + "tinyfish" => tinyfish::DEFAULT_BASE_URL, + "voyage" => voyage::DEFAULT_BASE_URL, + "xai" => xai::DEFAULT_BASE_URL, + "you_com" => you_com::DEFAULT_BASE_URL, + // amazon_bedrock, anthropic_aws, aws_polly, azure, codex, + // google_vertex and searxng. + _ if PACKAGES.contains(&name) => return None, + _ => preset::lookup(name) + .filter(|entry| entry.descriptor.params.is_empty()) + .map(|entry| entry.descriptor.base_url.as_str())?, + }) +} + +/// Every built-in provider with default settings, keyed by name: pass it to +/// [`aimux_core::create_provider_registry`], after adding or replacing entries +/// as needed. +/// +/// No entry can fail to be created. Keys are read per request, so a missing +/// key fails the call that needs it. The `openai` and `anthropic` packages +/// read `OPENAI_BASE_URL` / `ANTHROPIC_BASE_URL` once, here, as the AI SDK +/// does when a provider is created; an invalid value fails each call of that +/// provider with `InvalidArgument` naming the variable, not this function. +/// Registry rows read their base-URL variables per request. +#[must_use] +pub fn default_providers() -> BTreeMap> { + provider_names() + .into_iter() + .map(|name| { + let provider = create_provider(name, PresetSettings::default()) + .expect("default settings have nothing a factory can reject"); + (name.to_string(), provider) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Sets an environment variable for the test's duration and restores + /// the previous value on drop. + struct EnvVar(&'static str, Option); + + impl EnvVar { + fn set(name: &'static str, value: &str) -> Self { + let previous = std::env::var_os(name); + // SAFETY: the tests touching these variables are `#[serial]`. + unsafe { std::env::set_var(name, value) }; + Self(name, previous) + } + } + + impl Drop for EnvVar { + fn drop(&mut self) { + // SAFETY: as in `set`. + unsafe { + match &self.1 { + Some(value) => std::env::set_var(self.0, value), + None => std::env::remove_var(self.0), + } + } + } + } + + /// An invalid `OPENAI_BASE_URL` / `ANTHROPIC_BASE_URL` does not fail + /// creating the default providers. Each call of that provider reports it + /// instead, as the AI SDK's factories read the variable without failing + /// (openai@4.0.83 `openai-provider.ts:207-215`, anthropic@4.0.71 + /// `anthropic-provider.ts:127-133`) and the first request rejects the URL. + #[tokio::test] + #[serial_test::serial] + async fn an_invalid_env_base_url_fails_each_call_not_creation() { + let _openai = EnvVar::set("OPENAI_BASE_URL", "not a url"); + let _anthropic = EnvVar::set("ANTHROPIC_BASE_URL", "not a url"); + let _openai_key = EnvVar::set("OPENAI_API_KEY", "test-key"); + let _anthropic_key = EnvVar::set("ANTHROPIC_API_KEY", "test-key"); + let providers = default_providers(); + for (name, var) in [ + ("openai", "OPENAI_BASE_URL"), + ("anthropic", "ANTHROPIC_BASE_URL"), + ] { + let model = providers[name].language_model("some-model").unwrap(); + let options = aimux_core::options::CallOptions::new(vec![ + aimux_core::language_model_message::LanguageModelMessage::user_text("hi"), + ]); + let error = model.do_generate(&options).await.expect_err(name); + assert!( + matches!(&error, AiMuxError::InvalidArgument(message) if message.starts_with(var)), + "{name}: {error:?}" + ); + } + let _ = create_provider("openai", PresetSettings::default()).unwrap(); + } + + /// A misspelled name lists every name `create_provider` accepts, the + /// vendor packages included, not only the registry rows. + #[test] + fn an_unknown_name_lists_every_built_in_provider() { + let Err(AiMuxError::NoSuchProvider { + provider_id, + available_providers, + .. + }) = create_provider("opeanai", PresetSettings::default()) + else { + panic!("expected NoSuchProvider"); + }; + assert_eq!(provider_id, "opeanai"); + assert_eq!(available_providers, provider_names()); + assert!(available_providers.iter().any(|name| name == "openai")); + } + + /// The key variable is the one each factory reads with no `api_key`: + /// the package's own constant, else the registry row's `env_var`. + #[test] + fn provider_api_key_env_names_the_variable_the_factory_reads() { + assert_eq!(provider_api_key_env("openai"), Some("OPENAI_API_KEY")); + assert_eq!(provider_api_key_env("groq"), Some("GROQ_API_KEY")); + assert_eq!( + provider_api_key_env("google"), + Some("GOOGLE_GENERATIVE_AI_API_KEY") + ); + assert_eq!( + provider_api_key_env("amazon_bedrock"), + Some("AWS_BEARER_TOKEN_BEDROCK") + ); + assert_eq!(provider_api_key_env("fal"), Some("FAL_KEY")); + // Registry rows. + assert_eq!( + provider_api_key_env("openrouter"), + Some("OPENROUTER_API_KEY") + ); + assert_eq!(provider_api_key_env("ollama"), Some("OLLAMA_API_KEY")); + // Every package that takes a key names its variable; the rest none. + let keyless = [ + "anthropic_aws", + "aws_polly", + "codex", + "dataforseo", + "searxng", + ]; + for name in PACKAGES { + assert_eq!( + provider_api_key_env(name).is_some(), + !keyless.contains(&name), + "{name}" + ); + } + // A registry row without a variable, and an unknown name. + let without = preset::entries() + .find(|entry| entry.descriptor.env_var.is_empty()) + .map(|entry| entry.descriptor.name.as_str()); + if let Some(name) = without { + assert_eq!(provider_api_key_env(name), None, "{name}"); + } + assert_eq!(provider_api_key_env("no-such-provider"), None); + } + + /// The default endpoint is the package constant or the registry row's + /// URL; a URL built from settings has none. + #[test] + fn provider_default_base_url_is_the_factorys_fixed_endpoint() { + assert_eq!( + provider_default_base_url("openai"), + Some("https://api.openai.com/v1") + ); + assert_eq!( + provider_default_base_url("groq"), + Some("https://api.groq.com/openai/v1") + ); + assert_eq!( + provider_default_base_url("openrouter"), + Some("https://openrouter.ai/api/v1") + ); + for name in [ + "amazon_bedrock", + "anthropic_aws", + "aws_polly", + "azure", + "codex", + "google_vertex", + "searxng", + // A registry row with template parameters. + "bedrock_mantle", + "no-such-provider", + ] { + assert_eq!(provider_default_base_url(name), None, "{name}"); + } + } + + /// Providers and their models report the V4 specification version, as + /// the AI SDK's `ProviderV4` and model V4 interfaces do: a vendor package + /// and a registry row from `create_provider`, and the same providers + /// reached through a provider registry. + #[test] + fn providers_and_models_report_specification_version_v4() { + for (name, model_id) in [("openai", "gpt-4o"), ("fireworks", "some-model")] { + let provider = create_provider(name, PresetSettings::default()).unwrap(); + assert_eq!(provider.specification_version(), "v4", "{name}"); + let model = provider.language_model(model_id).unwrap(); + assert_eq!(model.specification_version(), "v4", "{name}"); + } + let openai = create_provider("openai", PresetSettings::default()).unwrap(); + let files = openai.files().expect("openai offers files"); + assert_eq!(files.specification_version(), "v4"); + + let registry = aimux_core::create_provider_registry( + default_providers(), + aimux_core::ProviderRegistryOptions::default(), + ); + assert_eq!( + registry.provider("openai").unwrap().specification_version(), + "v4" + ); + let model = registry.language_model("openai:gpt-4o").unwrap(); + assert_eq!(model.specification_version(), "v4"); + } +} diff --git a/aimux-providers/src/docker_model_runner.rs b/aimux-providers/src/docker_model_runner.rs deleted file mode 100644 index de1fc1529..000000000 --- a/aimux-providers/src/docker_model_runner.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! Docker Model Runner provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://model-runner.docker.internal/engines/llama.cpp/v1`. The `DOCKER_MODEL_RUNNER_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://model-runner.docker.internal/engines/llama.cpp/v1"; -const ENV_VAR: &str = "DOCKER_MODEL_RUNNER_BASE_URL"; -const PROVIDER_NAME: &str = "docker_model_runner"; -const PLACEHOLDER_API_KEY: &str = "docker_model_runner"; - -pub struct DockerModelRunnerConfig(OpenAIConfig); - -impl DockerModelRunnerConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `DOCKER_MODEL_RUNNER_BASE_URL` environment variable, falling - /// back to `http://model-runner.docker.internal/engines/llama.cpp/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct DockerModelRunnerProvider(OpenAIProvider); - -impl DockerModelRunnerProvider { - #[must_use] - pub fn new(config: DockerModelRunnerConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for DockerModelRunnerProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/elevenlabs.rs b/aimux-providers/src/elevenlabs.rs index 5664f933a..6fcaac431 100644 --- a/aimux-providers/src/elevenlabs.rs +++ b/aimux-providers/src/elevenlabs.rs @@ -13,7 +13,7 @@ //! returns raw binary audio. The `output_format` is passed as a query parameter. //! `instructions` is not supported and emits a warning. -use std::collections::HashMap; +use std::sync::{Arc, OnceLock}; use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -24,7 +24,19 @@ use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::{HttpBody, HttpRequest, load_api_key}; +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::LanguageModel; +use aimux_core::provider::Provider; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, HeadersFn, HttpBody, Resolvable, validate_base_url, +}; + +use crate::shared::{ + AuthScheme, Credential, EndpointConfig, TransformRequestBody, credential_headers, +}; + +pub(crate) mod options; /// ElevenLabs errors: `{"detail": {"status": "invalid_api_key", "message": ...}}` /// where `detail.status` is the machine code. FastAPI validation errors carry @@ -64,78 +76,181 @@ fn elevenlabs_failed_response_handler() -> aimux_provider_utils::ResponseHandler aimux_provider_utils::create_json_error_response_handler(elevenlabs_error_parts) } -// ── Config ─────────────────────────────────────────────────────────────────── - -/// Configuration for the ElevenLabs provider. -#[derive(Debug, Clone)] -pub struct ElevenLabsConfig { - pub api_key: String, - /// Base URL for the ElevenLabs API (no trailing slash). Defaults to - /// `https://api.elevenlabs.io`. - pub base_url: String, - /// Extra headers merged into every request. - pub headers: Option>, -} +// ── Settings ───────────────────────────────────────────────────────────────── -impl ElevenLabsConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.elevenlabs.io".to_string(), - headers: None, - } - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.elevenlabs.io"; +pub(crate) const API_KEY_ENV_VAR: &str = "ELEVENLABS_API_KEY"; +const DEFAULT_NAME: &str = "elevenlabs"; - /// Override the base URL (for testing or proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = url.into().trim_end_matches('/').to_string(); - self - } +/// Settings of [`create_elevenlabs`] (the AI SDK's `ElevenLabsProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request, including the WebSocket handshake of realtime +/// transcription. +#[derive(Clone, Default)] +pub struct ElevenLabsProviderSettings { + /// Base URL for the API calls. Default `https://api.elevenlabs.io`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key, sent as `xi-api-key`. `None` loads `ELEVENLABS_API_KEY` + /// when a request is made and fails that request with + /// `AiMuxError::LoadApiKey` if it is unset. An explicit value is used as + /// given, `""` included: it never falls back to the environment. A + /// [`Resolvable::Future`] is awaited once, an [`Resolvable::AsyncFn`] on + /// every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including `xi-api-key`. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings + /// (`"{name}.speech"`, `"{name}.transcription"`). Default `"elevenlabs"`. + /// The providerOptions key stays `elevenlabs`. + 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, + /// Rewrites every JSON request body once, after it is serialized and + /// before it is sent. + pub transform_request_body: Option, + /// Opens the WebSocket of realtime transcription (the AI SDK's + /// `webSocket`). `None` uses the built-in tungstenite connector. + #[cfg(feature = "realtime")] + pub web_socket: Option>, +} - /// Attach extra headers merged into every request. - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self +impl std::fmt::Debug for ElevenLabsProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut debug = f.debug_struct("ElevenLabsProviderSettings"); + debug + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ); + #[cfg(feature = "realtime")] + debug.field("web_socket", &self.web_socket.is_some()); + debug.finish() } +} - /// Create from environment variable `ELEVENLABS_API_KEY`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `ELEVENLABS_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "ELEVENLABS_API_KEY", "ElevenLabs")?; - Ok(Self::new(api_key)) - } +/// Create an ElevenLabs provider. +/// +/// # 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. +pub fn create_elevenlabs( + settings: ElevenLabsProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(ElevenLabsProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "ElevenLabs"), + AuthScheme::Header("xi-api-key"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + transform_request_body: settings.transform_request_body, + #[cfg(feature = "realtime")] + web_socket: settings.web_socket, + }) +} + +/// The default provider: `create_elevenlabs` with default settings, created +/// on first use. Creating it reads nothing from the environment and cannot +/// fail; a missing key surfaces from the first request instead. +pub fn elevenlabs() -> &'static ElevenLabsProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_elevenlabs(ElevenLabsProviderSettings::default()) + .expect("default ElevenLabs settings are always valid") + }) } // ── Provider ───────────────────────────────────────────────────────────────── -/// ElevenLabs provider — creates `ElevenLabsSpeechModel` instances. +/// ElevenLabs provider (the AI SDK's `ElevenLabsProvider`) — creates speech +/// and transcription models. Cheap to clone the models out of; it holds no +/// HTTP client. pub struct ElevenLabsProvider { - config: ElevenLabsConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, + transform_request_body: Option, + #[cfg(feature = "realtime")] + web_socket: Option>, } impl ElevenLabsProvider { - #[must_use] - pub fn new(config: ElevenLabsConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + self.transform_request_body.clone(), + ) } - /// Create a speech (TTS) model instance for the given model name (e.g. - /// `"eleven_multilingual_v2"`). + /// A speech (TTS) model (e.g. `"eleven_multilingual_v2"`); `provider()` + /// is `"{name}.speech"`. #[must_use] pub fn speech(&self, model_id: &str) -> ElevenLabsSpeechModel { - ElevenLabsSpeechModel::new(model_id.to_string(), self.config.clone()) + ElevenLabsSpeechModel::from_config(model_id.to_string(), self.model_config("speech")) } - /// Create a transcription (STT) model instance for the given model name - /// (e.g. `"scribe_v1"`). Uses the `/v1/speech-to-text` endpoint. + /// A transcription (STT) model (e.g. `"scribe_v1"`); `provider()` is + /// `"{name}.transcription"`. Uses the `/v1/speech-to-text` endpoint. #[must_use] pub fn transcription(&self, model_id: &str) -> ElevenLabsTranscriptionModel { - ElevenLabsTranscriptionModel::new(model_id.to_string(), self.config.clone()) + ElevenLabsTranscriptionModel::from_config( + model_id.to_string(), + self.model_config("transcription"), + #[cfg(feature = "realtime")] + self.web_socket.clone(), + ) + } +} + +impl Provider for ElevenLabsProvider { + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "languageModel")) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } + + fn transcription_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.transcription(model_id)))) + } + + fn speech_model(&self, model_id: &str) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.speech(model_id)))) } } @@ -150,40 +265,19 @@ const DEFAULT_OUTPUT_FORMAT: &str = "mp3_44100_128"; /// An ElevenLabs speech (TTS) model. pub struct ElevenLabsSpeechModel { model_id: String, - config: ElevenLabsConfig, + config: EndpointConfig, } impl ElevenLabsSpeechModel { - #[must_use] - pub fn new(model_id: String, config: ElevenLabsConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("xi-api-key".to_string(), self.config.api_key.clone()); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self, voice_id: &str) -> String { - format!("{}/v1/text-to-speech/{}", self.config.base_url, voice_id) - } } #[async_trait] impl SpeechModel for ElevenLabsSpeechModel { fn provider(&self) -> &str { - "elevenlabs.speech" + &self.config.provider } fn model_id(&self) -> &str { @@ -193,32 +287,22 @@ impl SpeechModel for ElevenLabsSpeechModel { async fn do_generate(&self, options: &SpeechCallOptions) -> Result { let (body, query_params, warnings, voice_id) = build_request(options, &self.model_id)?; - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let endpoint = exchange.url(&format!("/v1/text-to-speech/{voice_id}")); let url = if query_params.is_empty() { - self.endpoint(&voice_id) + endpoint } else { let qs: Vec = query_params .iter() .map(|(k, v)| format!("{k}={v}")) .collect(); - format!("{}?{}", self.endpoint(&voice_id), qs.join("&")) + format!("{endpoint}?{}", qs.join("&")) }; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url, - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, - Value::Object(body.clone()), + exchange.request(url, options), + exchange.transform_body(Value::Object(body.clone())), aimux_provider_utils::create_binary_response_handler(), elevenlabs_failed_response_handler(), ) @@ -457,7 +541,7 @@ struct ElevenLabsPronunciationLocator { fn parse_elevenlabs_provider_options( options: Option<&SharedProviderOptions>, ) -> Option { - let opts = options.and_then(|opts| opts.get("elevenlabs"))?; + let opts = options::elevenlabs_options(options)?; let voice_settings = opts .get("voiceSettings") @@ -580,33 +664,25 @@ struct ElevenLabsTranscriptionResponse { /// Endpoint: `POST {base_url}/v1/speech-to-text` (multipart form-data) pub struct ElevenLabsTranscriptionModel { model_id: String, - config: ElevenLabsConfig, + config: EndpointConfig, + #[cfg(feature = "realtime")] + web_socket: Option>, } impl ElevenLabsTranscriptionModel { - #[must_use] - pub fn new(model_id: String, config: ElevenLabsConfig) -> Self { - Self { model_id, config } - } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("xi-api-key".to_string(), self.config.api_key.clone()); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } + pub(crate) fn from_config( + model_id: String, + config: EndpointConfig, + #[cfg(feature = "realtime")] web_socket: Option< + Arc, + >, + ) -> Self { + Self { + model_id, + config, + #[cfg(feature = "realtime")] + web_socket, } - headers - } - - fn endpoint(&self) -> String { - format!("{}/v1/speech-to-text", self.config.base_url) } } @@ -617,6 +693,7 @@ fn audio_input_to_bytes_stt(audio: &AudioInput) -> Result, AiMuxError> { base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid base64: {e}"))) } + AudioInput::Url(_) => Err(AudioInput::url_error()), } } @@ -630,7 +707,7 @@ fn is_realtime_transcription_model_id(model_id: &str) -> bool { #[async_trait] impl TranscriptionModel for ElevenLabsTranscriptionModel { fn provider(&self) -> &str { - "elevenlabs" + &self.config.provider } fn model_id(&self) -> &str { @@ -661,9 +738,7 @@ impl TranscriptionModel for ElevenLabsTranscriptionModel { form.text("diarize", "true")?; // Parse provider options. - if let Some(ref po) = options.provider_options - && let Some(el) = po.get("elevenlabs") - { + if let Some(el) = options::elevenlabs_options(options.provider_options.as_ref()) { if let Some(v) = el.get("diarize").and_then(serde_json::Value::as_bool) { form.text("diarize", &v.to_string())?; } @@ -689,17 +764,10 @@ impl TranscriptionModel for ElevenLabsTranscriptionModel { let (body_bytes, content_type) = form.finish(); - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_to_api( - HttpRequest::new( - self.endpoint(), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), + exchange.request(exchange.url("/v1/speech-to-text"), options), HttpBody::Bytes(body_bytes, content_type), aimux_provider_utils::create_json_response_handler::(), elevenlabs_failed_response_handler(), @@ -784,9 +852,7 @@ impl TranscriptionModel for ElevenLabsTranscriptionModel { // languageCode + includeTimestamps. Anything else waits for a user. let mut language_code: Option = None; let mut include_timestamps = false; - if let Some(ref po) = options.provider_options - && let Some(el) = po.get("elevenlabs") - { + if let Some(el) = options::elevenlabs_options(options.provider_options.as_ref()) { if let Some(v) = el.get("languageCode").and_then(serde_json::Value::as_str) { language_code = Some(v.to_string()); } @@ -815,7 +881,8 @@ impl TranscriptionModel for ElevenLabsTranscriptionModel { } }; - let base = self.config.base_url.trim_end_matches('/'); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let base = exchange.base_url(); let (scheme, host) = if let Some(rest) = base.strip_prefix("https://") { ("wss", rest) } else if let Some(rest) = base.strip_prefix("http://") { @@ -843,10 +910,7 @@ impl TranscriptionModel for ElevenLabsTranscriptionModel { ) }; - let header_list: Vec<(String, String)> = self - .build_headers(options.headers.as_ref()) - .into_iter() - .collect(); + let header_list = exchange.headers(); // Connect BEFORE the stream so connect failures surface from // do_stream's Result (same contract as the OpenAI realtime path). @@ -856,6 +920,7 @@ impl TranscriptionModel for ElevenLabsTranscriptionModel { subprotocols: Vec::new(), abort_signal: options.abort_signal.clone(), timeout: options.timeout, + connector: self.web_socket.clone(), }; let mut ws = ws_connect(&req).await?; diff --git a/aimux-providers/src/elevenlabs/options.rs b/aimux-providers/src/elevenlabs/options.rs new file mode 100644 index 000000000..3c3fedfa4 --- /dev/null +++ b/aimux-providers/src/elevenlabs/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of ElevenLabs. +//! +//! `@ai-sdk/elevenlabs` reads `providerOptions.elevenlabs`, whatever the +//! provider is named. This module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "elevenlabs"; + +/// The ElevenLabs options in a providerOptions map. +pub(crate) fn elevenlabs_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/exa_ai.rs b/aimux-providers/src/exa_ai.rs index 61b8d4edf..8b59450a8 100644 --- a/aimux-providers/src/exa_ai.rs +++ b/aimux-providers/src/exa_ai.rs @@ -3,19 +3,25 @@ //! Implements the `SearchModel` trait against the Exa search API //! (`POST https://api.exa.ai/search`). Uses `x-api-key` header auth //! via `EXA_API_KEY`. +//! +//! [`create_exa_ai`] takes [`ExaAiProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`ExaAiProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `EXA_API_KEY`. +//! [`exa_ai()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; fn exa_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -36,61 +42,117 @@ fn exa_failed_response_handler() -> aimux_provider_utils::ResponseHandler, + /// The API key. `None` loads `EXA_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.search"`). + /// Default `"exa_ai"`. The providerOptions key stays `exa_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, } -impl ExaAiConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.exa.ai".to_string(), - } +impl std::fmt::Debug for ExaAiProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ExaAiProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Create a Exa AI provider. +/// +/// # 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. +pub fn create_exa_ai(settings: ExaAiProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(ExaAiProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Exa AI"), + AuthScheme::Header("x-api-key"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `EXA_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "EXA_API_KEY", "Exa AI")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_exa_ai` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn exa_ai() -> &'static ExaAiProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_exa_ai(ExaAiProviderSettings::default()) + .expect("default Exa AI settings are always valid") + }) } -/// Exa AI provider — search-only. +/// A Exa AI provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct ExaAiProvider { - config: ExaAiConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl ExaAiProvider { - #[must_use] - pub fn new(config: ExaAiConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> ExaAiSearchModel { - ExaAiSearchModel::new(self.config.clone()) + ExaAiSearchModel::from_config(self.model_config("search")) } } -impl Provider for ExaAiProvider { - fn name(&self) -> &str { - "exa_ai" - } -} +crate::impl_single_modality_provider!(ExaAiProvider, search_model, |p, _id| p.search_model()); fn build_request_body(options: &SearchCallOptions) -> Value { let mut body = json!({ @@ -138,39 +200,21 @@ fn map_results(entries: Vec) -> Vec { .collect() } -/// Exa AI search model — implements `SearchModel`. +/// A Exa AI search model. pub struct ExaAiSearchModel { - config: ExaAiConfig, + config: EndpointConfig, } impl ExaAiSearchModel { - #[must_use] - pub fn new(config: ExaAiConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn endpoint(&self) -> String { - format!("{}/search", self.config.base_url) - } - - fn build_headers(&self, extra: Option<&HashMap>) -> Vec<(String, String)> { - let mut headers = vec![ - ("Content-Type".to_string(), "application/json".to_string()), - ("x-api-key".to_string(), self.config.api_key.clone()), - ]; - if let Some(extra) = extra { - for (k, v) in extra { - headers.push((k.clone(), v.clone())); - } - } - headers - } } #[async_trait] impl SearchModel for ExaAiSearchModel { fn provider(&self) -> &str { - "exa_ai" + &self.config.provider } fn model_id(&self) -> &str { @@ -179,18 +223,10 @@ impl SearchModel for ExaAiSearchModel { async fn do_search(&self, options: &SearchCallOptions) -> Result { let body = build_request_body(options); - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(exchange.url("/search"), options), body, aimux_provider_utils::create_json_response_handler(), exa_failed_response_handler(), diff --git a/aimux-providers/src/fal.rs b/aimux-providers/src/fal.rs index f274c64c8..13286d010 100644 --- a/aimux-providers/src/fal.rs +++ b/aimux-providers/src/fal.rs @@ -5,21 +5,32 @@ //! //! Fal uses an async queue pattern: POST to submit, then GET to poll until //! the result is ready. The audio is sent as a base64 data URL. +//! +//! [`create_fal`] takes [`FalProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`FalProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `FAL_KEY`. +//! [`fal()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::Duration; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::retry; use aimux_core::shared::Warning; use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::{HttpRequest, load_api_key, sleep_or_abort, without_trailing_slash}; +use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{ + AuthScheme, Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, credential_headers, + is_poll_control_key, poll_interval_ms, poll_until, retry_download, +}; /// fal errors are FastAPI-style: `{"detail": "..."}` or /// `{"detail": [{"loc": [...], "msg": "...", "type": "..."}]}` where `type` @@ -60,73 +71,173 @@ fn fal_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, -} +/// Milliseconds between two polls of a queued transcription request +/// (`providerOptions.fal.pollIntervalMs` overrides it for one call). +const POLL_INTERVAL_MS: u64 = 100; +/// How many times a queued transcription request is polled before the call +/// gives up (ten minutes at the default interval). +const MAX_POLL_ATTEMPTS: u32 = 6_000; +/// Milliseconds between two attempts to download a generated image +/// (`providerOptions.fal.pollIntervalMs` overrides it for one call). +const DOWNLOAD_RETRY_INTERVAL_MS: u64 = 1_000; -impl FalConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://queue.fal.run".to_string(), - headers: None, - } - } +pub(crate) mod options; - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://queue.fal.run"; +pub(crate) const API_KEY_ENV_VAR: &str = "FAL_KEY"; +const DEFAULT_NAME: &str = "fal"; - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } +/// Settings of [`create_fal`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct FalProviderSettings { + /// Base URL for the API calls. Default `https://queue.fal.run`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `FAL_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.transcription"`, `"{name}.video"`, `"{name}.image"`). + /// Default `"fal"`. The providerOptions key stays `fal`. + 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, +} - /// Create from the `FAL_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "FAL_KEY", "Fal")?; - Ok(Self::new(api_key)) +impl std::fmt::Debug for FalProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("FalProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } +/// Create a Fal provider. +/// +/// # 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. +pub fn create_fal(settings: FalProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(FalProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Fal"), + AuthScheme::Scheme("Key"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_fal` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn fal() -> &'static FalProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_fal(FalProviderSettings::default()).expect("default Fal settings are always valid") + }) +} + +/// A Fal provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct FalProvider { - config: FalConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl FalProvider { - #[must_use] - pub fn new(config: FalConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// A transcription (STT) model (e.g. `"wizper"`); `provider()` is `"{name}.transcription"`. #[must_use] pub fn transcription(&self, model_id: &str) -> FalTranscriptionModel { - FalTranscriptionModel::new(model_id.to_string(), self.config.clone()) + FalTranscriptionModel::from_config(model_id.to_string(), self.model_config("transcription")) } - /// Create a video generation model instance for the given model name - /// (e.g. `"fal-ai/kling-video"`). + /// A video model (e.g. `"fal-ai/kling-video"`); `provider()` is `"{name}.video"`. #[must_use] pub fn video(&self, model_id: &str) -> FalVideoModel { - FalVideoModel::new(model_id.to_string(), self.config.clone()) + FalVideoModel::from_config(model_id.to_string(), self.model_config("video")) } - /// Create an image generation model instance for the given model name - /// (e.g. `"fal-ai/flux/schnell"`). + /// An image model (e.g. `"fal-ai/flux/schnell"`); `provider()` is `"{name}.image"`. #[must_use] pub fn image(&self, model_id: &str) -> FalImageModel { - FalImageModel::new(model_id.to_string(), self.config.clone()) + FalImageModel::from_config(model_id.to_string(), self.model_config("image")) + } +} + +impl ::aimux_core::Provider for FalProvider { + fn language_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "languageModel")) + } + + fn embedding_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Ok(::std::sync::Arc::new(self.image(model_id))) + } + + fn transcription_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(::std::sync::Arc::new(self.transcription(model_id)))) + } + + fn video_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(::std::sync::Arc::new(self.video(model_id)))) } } @@ -162,55 +273,33 @@ fn audio_input_to_base64(audio: &AudioInput) -> Result { &base64::engine::general_purpose::STANDARD, bytes, )), + AudioInput::Url(_) => Err(AudioInput::url_error()), } } pub struct FalTranscriptionModel { model_id: String, - config: FalConfig, + config: EndpointConfig, } impl FalTranscriptionModel { - #[must_use] - pub fn new(model_id: String, config: FalConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Key {}", self.config.api_key), - ); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers + fn submit_path(&self) -> String { + format!("/fal-ai/{}", self.model_id) } - fn submit_url(&self) -> String { - format!("{}/fal-ai/{}", self.config.base_url, self.model_id) - } - - fn poll_url(&self, request_id: &str) -> String { - format!( - "{}/fal-ai/{}/requests/{}", - self.config.base_url, self.model_id, request_id - ) + fn poll_path(&self, request_id: &str) -> String { + format!("/fal-ai/{}/requests/{request_id}", self.model_id) } } #[async_trait] impl TranscriptionModel for FalTranscriptionModel { fn provider(&self) -> &str { - "fal" + &self.config.provider } fn model_id(&self) -> &str { @@ -233,9 +322,8 @@ impl TranscriptionModel for FalTranscriptionModel { body.insert("audio_url".to_string(), json!(audio_url)); // Parse provider options. - if let Some(ref po) = options.provider_options - && let Some(fal) = po.get("fal") - { + let fal_options = options::fal_options(options.provider_options.as_ref()); + if let Some(fal) = fal_options { if let Some(v) = fal.get("language") { body.insert("language".to_string(), v.clone()); } @@ -256,26 +344,11 @@ impl TranscriptionModel for FalTranscriptionModel { } } - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; // Submit job. let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.submit_url(), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, + exchange.request(exchange.url(&self.submit_path()), options), Value::Object(body), aimux_provider_utils::create_json_response_handler(), fal_failed_response_handler(), @@ -283,71 +356,44 @@ impl TranscriptionModel for FalTranscriptionModel { .await?; let job: FalJobResponse = resp.value; - let retries = retry::prepare_retries( - options.max_retries, - self.retry_config(), - options.abort_signal.clone(), - ); - - // Poll for result. - let raw_body: Value; - let parsed: FalTranscriptionResponse; - let response_headers: HashMap; - loop { - // Fal returns 400/404 while a queued request is still registering. - // Normalize that state inside the retry attempt so a preceding 5xx - // cannot wrap the pending response in RetryError. - let resp = retries - .retry(|| { - let request = aimux_provider_utils::get_from_api( - HttpRequest { - url: self.poll_url(&job.request_id), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, - aimux_provider_utils::create_json_response_handler::< - FalTranscriptionResponse, - >(), - fal_failed_response_handler(), - ); - async move { - match request.await { - Ok(response) => Ok(Some(response)), - Err(AiMuxError::ApiCall(detail)) - if matches!(detail.status_code, Some(400 | 404)) => - { - Ok(None) - } - Err(error) => Err(error), - } - } - }) - .await?; - let Some(resp) = resp else { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - continue; - }; - response_headers = resp.response_headers; - raw_body = resp.raw_value.unwrap_or(Value::Null); - parsed = resp.value; - break; - } + // Poll for the result. Submitting is over: only this status request + // repeats. + let poll_url = exchange.url(&self.poll_path(&job.request_id)); + let resp = poll_until( + &format!("fal request {}", job.request_id), + options.abort_signal.as_ref(), + Duration::from_millis(poll_interval_ms( + fal_options, + POLL_INTERVAL_MS_KEY, + POLL_INTERVAL_MS, + )), + MAX_POLL_ATTEMPTS, + || async { + // Fal returns 400/404 while a queued request is still + // registering; that is pending, not a failure. + match aimux_provider_utils::get_from_api( + exchange.request(poll_url.clone(), options), + aimux_provider_utils::create_json_response_handler::( + ), + fal_failed_response_handler(), + ) + .await + { + Ok(response) => Ok(PollStep::Ready(response)), + Err(AiMuxError::ApiCall(detail)) + if matches!(detail.status_code, Some(400 | 404)) => + { + Ok(PollStep::Pending) + } + Err(error) => Err(error), + } + }, + ) + .await?; + let response_headers = resp.response_headers; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; let segments: Vec = parsed .chunks @@ -414,36 +460,16 @@ use aimux_core::image_model::{ /// An fal.ai image generation model. pub struct FalImageModel { model_id: String, - config: FalConfig, + config: EndpointConfig, } impl FalImageModel { - #[must_use] - pub fn new(model_id: String, config: FalConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Key {}", self.config.api_key), - ); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/{}", self.config.base_url, self.model_id) + fn endpoint_path(&self) -> String { + format!("/{}", self.model_id) } fn file_to_data_uri(file: &ImageFile) -> Result { @@ -480,7 +506,7 @@ impl FalImageModel { #[async_trait] impl ImageModel for FalImageModel { fn provider(&self) -> &str { - "fal" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id @@ -509,7 +535,7 @@ impl ImageModel for FalImageModel { body.insert("image_size".into(), image_size); body.insert("num_images".into(), json!(options.n)); - let fal_opts = options.provider_options.get("fal"); + let fal_opts = options::fal_options(Some(&options.provider_options)); if let Some(ref files) = options.files && !files.is_empty() @@ -552,7 +578,10 @@ impl ImageModel for FalImageModel { ("safetyTolerance", "safety_tolerance"), ]; for (key, value) in fal { - if key == "__deprecatedKeys" || key == "useMultipleImages" { + if key == "__deprecatedKeys" + || key == "useMultipleImages" + || is_poll_control_key(key) + { continue; } let ak = map @@ -564,25 +593,10 @@ impl ImageModel for FalImageModel { } } - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, + exchange.request(exchange.url(&self.endpoint_path()), options), Value::Object(body), aimux_provider_utils::create_json_response_handler(), fal_failed_response_handler(), @@ -591,12 +605,6 @@ impl ImageModel for FalImageModel { let rh = resp.response_headers; let rb: Value = resp.value; - let retries = retry::prepare_retries( - options.max_retries, - self.retry_config(), - options.abort_signal.clone(), - ); - let target_images: Vec = if let Some(i) = rb.get("images").and_then(|v| v.as_array()) { i.clone() @@ -611,27 +619,29 @@ impl ImageModel for FalImageModel { if let Some(url) = img.get("url").and_then(|v| v.as_str()) { // images[].url comes from the queue result response body, so // it goes through the SSRF download guard. - let ir = retries - .retry(|| { + let ir = retry_download( + options.abort_signal.as_ref(), + Duration::from_millis(poll_interval_ms( + fal_opts, + POLL_INTERVAL_MS_KEY, + DOWNLOAD_RETRY_INTERVAL_MS, + )), + || { aimux_provider_utils::get_from_api( HttpRequest { url: url.to_string(), - headers: vec![], - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, validate_url: true, - trusted_origin: Some(self.config.base_url.clone()), - credentialed_origin: Some(self.config.base_url.clone()), + trusted_origin: Some(exchange.base_url().to_string()), + credentialed_origin: Some(exchange.base_url().to_string()), + ..Default::default() }, aimux_provider_utils::create_binary_response_handler(), fal_failed_response_handler(), ) - }) - .await?; + }, + ) + .await?; downloaded.push(ir.value.to_vec()); } } @@ -678,7 +688,7 @@ impl ImageModel for FalImageModel { } } } - metadata.insert("fal".into(), fm); + metadata.insert(options::NAMESPACE.into(), fm); Ok(ImageResult { images: ImageOutputs::Binary(downloaded), @@ -712,12 +722,11 @@ use aimux_core::video_model::{ /// polling is driven by Core via `do_status`. pub struct FalVideoModel { model_id: String, - config: FalConfig, + config: EndpointConfig, } impl FalVideoModel { - #[must_use] - pub fn new(model_id: String, config: FalConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } @@ -729,41 +738,16 @@ impl FalVideoModel { .to_string() } - fn submit_url(&self) -> String { - format!( - "{}/fal-ai/{}", - self.config.base_url, - self.normalized_model_id() - ) + fn submit_path(&self) -> String { + format!("/fal-ai/{}", self.normalized_model_id()) } - fn poll_url(&self, request_id: &str) -> String { + fn poll_path(&self, request_id: &str) -> String { format!( - "{}/fal-ai/{}/requests/{}", - self.config.base_url, - self.normalized_model_id(), - request_id + "/fal-ai/{}/requests/{request_id}", + self.normalized_model_id() ) } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Key {}", self.config.api_key), - ); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } } fn video_file_to_data_uri(file: &VideoFile) -> Result { @@ -782,7 +766,7 @@ fn video_file_to_data_uri(file: &VideoFile) -> Result { #[async_trait] impl VideoModel for FalVideoModel { fn provider(&self) -> &str { - "fal" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id @@ -817,26 +801,11 @@ impl VideoModel for FalVideoModel { body.insert("seed".to_string(), json!(seed)); } - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; // Submit. let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.submit_url(), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, + exchange.request(exchange.url(&self.submit_path()), options), Value::Object(body), aimux_provider_utils::create_json_response_handler(), fal_failed_response_handler(), @@ -872,25 +841,10 @@ impl VideoModel for FalVideoModel { ) })?; - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::get_from_api( - HttpRequest { - url: self.poll_url(request_id), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, + exchange.request(exchange.url(&self.poll_path(request_id)), options), aimux_provider_utils::create_json_response_handler::(), fal_failed_response_handler(), ) diff --git a/aimux-providers/src/fal/options.rs b/aimux-providers/src/fal/options.rs new file mode 100644 index 000000000..1bf45646e --- /dev/null +++ b/aimux-providers/src/fal/options.rs @@ -0,0 +1,14 @@ +//! The providerOptions namespace of Fal. +//! +//! Fal reads `providerOptions.fal`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "fal"; + +/// The Fal options in a providerOptions map. +pub(crate) fn fal_options(provider_options: Option<&SharedProviderOptions>) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/firecrawl.rs b/aimux-providers/src/firecrawl.rs index f917f64f3..3380781eb 100644 --- a/aimux-providers/src/firecrawl.rs +++ b/aimux-providers/src/firecrawl.rs @@ -3,19 +3,25 @@ //! Implements the `SearchModel` trait against the Firecrawl search API //! (`POST https://api.firecrawl.dev/v2/search`). Bearer auth via //! `FIRECRAWL_API_KEY`. +//! +//! [`create_firecrawl`] takes [`FirecrawlProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`FirecrawlProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `FIRECRAWL_API_KEY`. +//! [`firecrawl()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; const MODEL_ID: &str = "firecrawl-search"; @@ -33,61 +39,118 @@ fn firecrawl_failed_response_handler() -> aimux_provider_utils::ResponseHandler< }) } -/// Configuration for the Firecrawl provider. -#[derive(Debug, Clone)] -pub struct FirecrawlConfig { - pub api_key: String, - pub base_url: String, +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.firecrawl.dev"; +pub(crate) const API_KEY_ENV_VAR: &str = "FIRECRAWL_API_KEY"; +const DEFAULT_NAME: &str = "firecrawl"; + +/// Settings of [`create_firecrawl`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct FirecrawlProviderSettings { + /// Base URL for the API calls. Default `https://api.firecrawl.dev`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `FIRECRAWL_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.search"`). + /// Default `"firecrawl"`. The providerOptions key stays `firecrawl`. + 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, } -impl FirecrawlConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.firecrawl.dev".to_string(), - } +impl std::fmt::Debug for FirecrawlProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("FirecrawlProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Create a Firecrawl provider. +/// +/// # 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. +pub fn create_firecrawl( + settings: FirecrawlProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(FirecrawlProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Firecrawl"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `FIRECRAWL_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "FIRECRAWL_API_KEY", "Firecrawl")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_firecrawl` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn firecrawl() -> &'static FirecrawlProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_firecrawl(FirecrawlProviderSettings::default()) + .expect("default Firecrawl settings are always valid") + }) } -/// Firecrawl provider — search-only. +/// A Firecrawl provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct FirecrawlProvider { - config: FirecrawlConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl FirecrawlProvider { - #[must_use] - pub fn new(config: FirecrawlConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> FirecrawlSearchModel { - FirecrawlSearchModel::new(self.config.clone()) + FirecrawlSearchModel::from_config(self.model_config("search")) } } -impl Provider for FirecrawlProvider { - fn name(&self) -> &str { - "firecrawl" - } -} +crate::impl_single_modality_provider!(FirecrawlProvider, search_model, |p, _id| p.search_model()); fn build_request_body(options: &SearchCallOptions) -> Value { let mut body = json!({ @@ -142,42 +205,21 @@ fn map_results(entries: Vec) -> Vec { .collect() } -/// Firecrawl search model — implements `SearchModel`. +/// A Firecrawl search model. pub struct FirecrawlSearchModel { - config: FirecrawlConfig, + config: EndpointConfig, } impl FirecrawlSearchModel { - #[must_use] - pub fn new(config: FirecrawlConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn endpoint(&self) -> String { - format!("{}/v2/search", self.config.base_url) - } - - fn build_headers(&self, extra: Option<&HashMap>) -> Vec<(String, String)> { - let mut headers = vec![ - ("Content-Type".to_string(), "application/json".to_string()), - ( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ), - ]; - if let Some(extra) = extra { - for (k, v) in extra { - headers.push((k.clone(), v.clone())); - } - } - headers - } } #[async_trait] impl SearchModel for FirecrawlSearchModel { fn provider(&self) -> &str { - "firecrawl" + &self.config.provider } fn model_id(&self) -> &str { @@ -186,18 +228,10 @@ impl SearchModel for FirecrawlSearchModel { async fn do_search(&self, options: &SearchCallOptions) -> Result { let body = build_request_body(options); - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(exchange.url("/v2/search"), options), body, aimux_provider_utils::create_json_response_handler(), firecrawl_failed_response_handler(), diff --git a/aimux-providers/src/gaudi.rs b/aimux-providers/src/gaudi.rs deleted file mode 100644 index 23ab9e791..000000000 --- a/aimux-providers/src/gaudi.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! Intel Gaudi provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `GAUDI_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "GAUDI_BASE_URL"; -const PROVIDER_NAME: &str = "gaudi"; -const PLACEHOLDER_API_KEY: &str = "gaudi"; - -pub struct GaudiConfig(OpenAIConfig); - -impl GaudiConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `GAUDI_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct GaudiProvider(OpenAIProvider); - -impl GaudiProvider { - #[must_use] - pub fn new(config: GaudiConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for GaudiProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/gladia.rs b/aimux-providers/src/gladia.rs index a8dcf356d..9aaa7cb3b 100644 --- a/aimux-providers/src/gladia.rs +++ b/aimux-providers/src/gladia.rs @@ -7,23 +7,31 @@ //! 1. POST `/v2/upload` with multipart form (audio file) → returns `audio_url` //! 2. POST `/v2/pre-recorded` with JSON body (audio_url + options) → returns `result_url` //! 3. GET `result_url` polling until status is `done` +//! +//! [`create_gladia`] takes [`GladiaProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`GladiaProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `GLADIA_API_KEY`. +//! [`gladia()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::Duration; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Map, Value, json}; use aimux_core::error::{AiMuxError, ApiCallError}; -use aimux_core::retry; use aimux_core::shared::{SharedProviderMetadata, Warning, provider_namespace}; use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::{ - HttpBody, HttpRequest, MultipartForm, load_api_key, media_type_to_extension, sleep_or_abort, - without_trailing_slash, +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{HttpBody, MultipartForm, media_type_to_extension}; + +use crate::shared::{ + Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, is_poll_control_key, + poll_interval_ms, poll_until, provider_headers, }; fn gladia_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -48,62 +56,130 @@ fn gladia_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, +/// Milliseconds between two polls of a transcription job +/// (`providerOptions.gladia.pollIntervalMs` overrides it for one call). +const POLL_INTERVAL_MS: u64 = 100; +/// How many times a transcription job is polled before the call gives up (ten +/// minutes at the default interval). +const MAX_POLL_ATTEMPTS: u32 = 6_000; + +pub(crate) mod options; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.gladia.io"; +pub(crate) const API_KEY_ENV_VAR: &str = "GLADIA_API_KEY"; +const DEFAULT_NAME: &str = "gladia"; + +/// Settings of [`create_gladia`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct GladiaProviderSettings { + /// Base URL for the API calls. Default `https://api.gladia.io`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `GLADIA_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.transcription"`). + /// Default `"gladia"`. The providerOptions key stays `gladia`. + 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, } -impl GladiaConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.gladia.io".to_string(), - headers: None, - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self +impl std::fmt::Debug for GladiaProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GladiaProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } +/// Create a Gladia provider. +/// +/// # 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. +pub fn create_gladia(settings: GladiaProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(GladiaProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Gladia"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `GLADIA_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "GLADIA_API_KEY", "Gladia")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_gladia` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn gladia() -> &'static GladiaProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_gladia(GladiaProviderSettings::default()) + .expect("default Gladia settings are always valid") + }) } +/// A Gladia provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct GladiaProvider { - config: GladiaConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl GladiaProvider { - #[must_use] - pub fn new(config: GladiaConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// A transcription (STT) model; `provider()` is `"{name}.transcription"`. #[must_use] pub fn transcription(&self, model_id: &str) -> GladiaTranscriptionModel { - GladiaTranscriptionModel::new(model_id.to_string(), self.config.clone()) + GladiaTranscriptionModel::from_config( + model_id.to_string(), + self.model_config("transcription"), + ) } } +crate::impl_single_modality_provider!(GladiaProvider, transcription_model, |p, id| p + .transcription(id)); + // ── Response schema ───────────────────────────────────────────────────────── #[derive(Debug, Deserialize)] @@ -157,52 +233,25 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid base64: {e}"))) } + AudioInput::Url(_) => Err(AudioInput::url_error()), } } pub struct GladiaTranscriptionModel { model_id: String, - config: GladiaConfig, + config: EndpointConfig, } impl GladiaTranscriptionModel { - #[must_use] - pub fn new(model_id: String, config: GladiaConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn upload_url(&self) -> String { - format!("{}/v2/upload", self.config.base_url) - } - - fn initiate_url(&self) -> String { - format!("{}/v2/pre-recorded", self.config.base_url) - } } #[async_trait] impl TranscriptionModel for GladiaTranscriptionModel { fn provider(&self) -> &str { - "gladia" + &self.config.provider } fn model_id(&self) -> &str { @@ -219,175 +268,136 @@ impl TranscriptionModel for GladiaTranscriptionModel { let file_extension = media_type_to_extension(&options.media_type); let filename = format!("audio.{file_extension}"); - let headers = self.build_headers(options.headers.as_ref()); - - // Per-stage retry, not one retry around `do_generate`: Core's outer - // retry would re-upload the audio to retry a later stage. An - // exhausted inner retry returns `AiMuxError::Retry`, which the outer - // retry passes through, so `do_generate` is never replayed (§6.2). - let retries = retry::prepare_retries( - options.max_retries, - self.retry_config(), - options.abort_signal.clone(), - ); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let gladia_options = options::gladia_options(options.provider_options.as_ref()); - // Step 1: Upload audio. + // Step 1: Upload audio. One attempt: an upload is never replayed + // implicitly. let mut form = MultipartForm::new(); form.file("audio", &filename, &options.media_type, &audio_bytes)?; let (body_bytes, content_type) = form.finish(); - let resp = retries - .retry(|| { - aimux_provider_utils::post_to_api( - HttpRequest::new( - self.upload_url(), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), - HttpBody::Bytes(body_bytes.clone(), content_type.clone()), - aimux_provider_utils::create_json_response_handler(), - gladia_failed_response_handler(), - ) - }) - .await?; + let resp = aimux_provider_utils::post_to_api( + exchange.request(exchange.url("/v2/upload"), options), + HttpBody::Bytes(body_bytes, content_type), + aimux_provider_utils::create_json_response_handler(), + gladia_failed_response_handler(), + ) + .await?; let upload: GladiaUploadResponse = resp.value; - // Step 2: Initiate transcription. + // Step 2: Initiate transcription. This creates the job and is sent + // exactly once; only the status poll below repeats. let mut body = Map::new(); body.insert("audio_url".to_string(), json!(upload.audio_url)); - // Parse provider options. - if let Some(ref po) = options.provider_options - && let Some(gladia) = po.get("gladia") - { - // Forward all provider options as-is (the API uses snake_case). - for (k, v) in gladia { - body.insert(k.clone(), v.clone()); + // Forward all provider options as-is (the API uses snake_case), except + // the poll interval, which is read here. + if let Some(obj) = gladia_options { + for (k, v) in obj { + if !is_poll_control_key(k) { + body.insert(k.clone(), v.clone()); + } } } - let body = Value::Object(body); - let resp = retries - .retry(|| { - aimux_provider_utils::post_json_to_api( - HttpRequest::new( - self.initiate_url(), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), - body.clone(), - aimux_provider_utils::create_json_response_handler(), - gladia_failed_response_handler(), - ) - }) - .await?; + let resp = aimux_provider_utils::post_json_to_api( + exchange.request(exchange.url("/v2/pre-recorded"), options), + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + gladia_failed_response_handler(), + ) + .await?; let init: GladiaInitResponse = resp.value; - // Step 3: Poll for result. - let mut raw_body: Value; - let mut response_headers: HashMap; - loop { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - - // AI SDK polls result_url with validateUrl: true and - // credentialedOrigin = the API origin: the target is validated - // and headers survive only while it stays on base_url's origin. - let resp = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: init.result_url.clone(), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: true, - trusted_origin: Some(self.config.base_url.clone()), - credentialed_origin: Some(self.config.base_url.clone()), - }, - aimux_provider_utils::create_json_response_handler::( - ), - gladia_failed_response_handler(), - ) - }) + // Step 3: Poll for the result. + // + // AI SDK polls result_url with validateUrl: true and credentialedOrigin + // = the API origin: the target is validated and headers survive only + // while it stays on base_url's origin. + let resp = poll_until( + "gladia transcription", + options.abort_signal.as_ref(), + Duration::from_millis(poll_interval_ms( + gladia_options, + POLL_INTERVAL_MS_KEY, + POLL_INTERVAL_MS, + )), + MAX_POLL_ATTEMPTS, + || async { + let mut request = exchange.request(init.result_url.clone(), options); + request.validate_url = true; + request.trusted_origin = Some(exchange.base_url().to_string()); + let resp = aimux_provider_utils::get_from_api( + request, + aimux_provider_utils::create_json_response_handler::(), + gladia_failed_response_handler(), + ) .await?; - - response_headers = resp.response_headers; - raw_body = resp.raw_value.unwrap_or(Value::Null); - let parsed = resp.value; - - if parsed.status == "done" { - let result = parsed.result.ok_or_else(|| { - AiMuxError::InvalidResponseData( - "Gladia transcription completed without a result".to_string(), - ) - })?; - - let segments: Vec = result - .transcription - .utterances - .iter() - .map(|u| TranscriptionSegment { - text: u.text.clone(), - start_second: u.start, - end_second: u.end, - }) - .collect(); - - let language = result.transcription.languages.first().cloned(); - let duration_in_seconds = Some(result.metadata.audio_duration); - - let timestamp = chrono::Utc::now().to_rfc3339(); - - let provider_metadata: Option = - Some(provider_namespace("gladia", raw_body.clone())?); - return Ok(TranscriptionResult { - text: result.transcription.full_transcript, - segments, - language, - duration_in_seconds, - warnings, - request: Some(TranscriptionRequest { body: None }), - response: TranscriptionResponse { - timestamp: Some(timestamp), - model_id: Some(self.model_id.clone()), - headers: Some(response_headers), - body: Some(raw_body), - }, - provider_metadata, - }); - } - - if parsed.status == "error" { - return Err(AiMuxError::ApiCall(Box::new(ApiCallError { - status_code: Some(200), - provider_code: Some(parsed.status.clone()), - response_body: Some(raw_body.to_string()), - ..ApiCallError::new( - "Transcription job failed", - init.result_url.clone(), - serde_json::json!({}), - ) - }))); - } - } + match resp.value.status.clone().as_str() { + "done" => Ok(PollStep::Ready(resp)), + "error" => Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some(resp.value.status.clone()), + response_body: Some( + resp.raw_value.clone().unwrap_or(Value::Null).to_string(), + ), + ..ApiCallError::new( + "Transcription job failed", + init.result_url.clone(), + serde_json::json!({}), + ) + }))), + _ => Ok(PollStep::Pending), + } + }, + ) + .await?; + + let response_headers = resp.response_headers; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; + + let result = parsed.result.ok_or_else(|| { + AiMuxError::InvalidResponseData( + "Gladia transcription completed without a result".to_string(), + ) + })?; + + let segments: Vec = result + .transcription + .utterances + .iter() + .map(|u| TranscriptionSegment { + text: u.text.clone(), + start_second: u.start, + end_second: u.end, + }) + .collect(); + + let language = result.transcription.languages.first().cloned(); + let duration_in_seconds = Some(result.metadata.audio_duration); + + let timestamp = chrono::Utc::now().to_rfc3339(); + + let provider_metadata: Option = + Some(provider_namespace(options::NAMESPACE, raw_body.clone())?); + Ok(TranscriptionResult { + text: result.transcription.full_transcript, + segments, + language, + duration_in_seconds, + warnings, + request: Some(TranscriptionRequest { body: None }), + response: TranscriptionResponse { + timestamp: Some(timestamp), + model_id: Some(self.model_id.clone()), + headers: Some(response_headers), + body: Some(raw_body), + }, + provider_metadata, + }) } } diff --git a/aimux-providers/src/gladia/options.rs b/aimux-providers/src/gladia/options.rs new file mode 100644 index 000000000..15a227f5e --- /dev/null +++ b/aimux-providers/src/gladia/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Gladia. +//! +//! Gladia reads `providerOptions.gladia`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "gladia"; + +/// The Gladia options in a providerOptions map. +pub(crate) fn gladia_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/google/convert.rs b/aimux-providers/src/google/convert.rs index 3e42fba14..9f5453a93 100644 --- a/aimux-providers/src/google/convert.rs +++ b/aimux-providers/src/google/convert.rs @@ -18,6 +18,7 @@ //! fields we actually read back are `text`, function tool parts, and native //! provider-executed tool parts. +use super::options::{GOOGLE, Namespace}; use aimux_core::error::AiMuxError; use aimux_core::language_model_message::{ AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, @@ -25,7 +26,7 @@ use aimux_core::language_model_message::{ }; use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat}; use aimux_core::result::{GenerateContent, Source}; -use aimux_core::shared::{FileBytes, FileData, JsonObject, SharedProviderOptions}; +use aimux_core::shared::{FileBytes, FileData}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified, Warning}; use base64::Engine; @@ -61,33 +62,6 @@ pub struct GooglePrompt { pub contents: Vec, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ProviderMetadataNamespace { - Google, - Vertex, -} - -/// Resolve per-part provider options in the same precedence order as the AI -/// SDK. Vertex prefers `googleVertex`, then legacy `vertex`; `google` is only a -/// cross-provider fallback. The public Google provider uses the inverse -/// fallback so transcripts survive gateway/provider failover. -fn read_provider_options( - provider_options: Option<&SharedProviderOptions>, - namespace: ProviderMetadataNamespace, -) -> Option<&JsonObject> { - let provider_options = provider_options?; - match namespace { - ProviderMetadataNamespace::Google => provider_options - .get("google") - .or_else(|| provider_options.get("googleVertex")) - .or_else(|| provider_options.get("vertex")), - ProviderMetadataNamespace::Vertex => provider_options - .get("googleVertex") - .or_else(|| provider_options.get("vertex")) - .or_else(|| provider_options.get("google")), - } -} - // ── convertToGoogleMessages ────────────────────────────────────────────────── /// Convert a provider-facing prompt into Google's `{ systemInstruction, contents }`. @@ -107,12 +81,12 @@ fn read_provider_options( /// `thoughtSignature` sibling of the `functionCall` part. #[must_use] pub fn convert_to_google_messages(prompt: &LanguageModelPrompt) -> GooglePrompt { - convert_to_google_messages_for_namespace(prompt, ProviderMetadataNamespace::Google, true) + convert_to_google_messages_for_namespace(prompt, Namespace::Google, true) } fn convert_to_google_messages_for_namespace( prompt: &LanguageModelPrompt, - namespace: ProviderMetadataNamespace, + namespace: Namespace, supports_function_response_parts: bool, ) -> GooglePrompt { let mut system_parts: Vec = Vec::new(); @@ -156,8 +130,7 @@ fn convert_to_google_messages_for_namespace( let mut ordinary = Vec::new(); for part in content { if let ToolPart::ToolResult(result) = part - && let Some(options) = - read_provider_options(result.provider_options.as_ref(), namespace) + && let Some(options) = namespace.read(result.provider_options.as_ref()) && let (Some(server_id), Some(server_type)) = ( options.get("serverToolCallId"), options.get("serverToolType"), @@ -201,7 +174,7 @@ fn convert_to_google_messages_for_namespace( } /// Convert user-role content parts into Google parts. -fn convert_user_parts(content: &[UserPart], namespace: ProviderMetadataNamespace) -> Vec { +fn convert_user_parts(content: &[UserPart], namespace: Namespace) -> Vec { let mut parts = Vec::new(); for part in content { match part { @@ -212,9 +185,7 @@ fn convert_user_parts(content: &[UserPart], namespace: ProviderMetadataNamespace data, media_type, .. }) => { let inline_media_type = if matches!(data, FileData::Text { .. }) - && !media_type - .split_once('/') - .is_some_and(|(_, subtype)| !subtype.is_empty() && subtype != "*") + && !aimux_provider_utils::is_full_media_type(media_type) { "text/plain" } else { @@ -242,8 +213,8 @@ fn convert_user_parts(content: &[UserPart], namespace: ProviderMetadataNamespace continue; } FileData::Reference { reference } => { - if namespace == ProviderMetadataNamespace::Google - && let Some(reference) = reference.get("google") + if namespace == Namespace::Google + && let Some(reference) = reference.get(GOOGLE) { parts.push(json!({ "fileData": { "mimeType": media_type, "fileUri": reference } @@ -268,10 +239,7 @@ fn convert_user_parts(content: &[UserPart], namespace: ProviderMetadataNamespace /// /// - `Text` → `{ text }` (skipped when empty, matching the TS SDK). /// - `ToolCall` → `{ functionCall: { id?, name, args } }`. -fn convert_assistant_parts( - content: &[AssistantPart], - namespace: ProviderMetadataNamespace, -) -> Vec { +fn convert_assistant_parts(content: &[AssistantPart], namespace: Namespace) -> Vec { let mut parts = Vec::new(); for part in content { match part { @@ -283,7 +251,8 @@ fn convert_assistant_parts( let mut p = json!({ "text": text }); // Echo thoughtSignature from provider_options if present // (upstream convert-to-google-messages.ts:355-377). - if let Some(sig) = read_provider_options(provider_options.as_ref(), namespace) + if let Some(sig) = namespace + .read(provider_options.as_ref()) .and_then(|g| g.get("thoughtSignature")) .and_then(|v| v.as_str()) { @@ -298,7 +267,8 @@ fn convert_assistant_parts( }) => { if !text.is_empty() { let mut p = json!({ "text": text, "thought": true }); - if let Some(sig) = read_provider_options(provider_options.as_ref(), namespace) + if let Some(sig) = namespace + .read(provider_options.as_ref()) .and_then(|g| g.get("thoughtSignature")) .and_then(Value::as_str) { @@ -314,7 +284,7 @@ fn convert_assistant_parts( provider_options, .. }) => { - let google_options = read_provider_options(provider_options.as_ref(), namespace); + let google_options = namespace.read(provider_options.as_ref()); let server_tool_call_id = google_options .and_then(|options| options.get("serverToolCallId")) .and_then(|value| value.as_str()); @@ -347,7 +317,7 @@ fn convert_assistant_parts( } } else { let mut function_call = Map::new(); - if !tool_call_id.is_empty() && namespace != ProviderMetadataNamespace::Vertex { + if !tool_call_id.is_empty() && namespace != Namespace::Vertex { function_call.insert("id".to_string(), json!(tool_call_id)); } function_call.insert("name".to_string(), json!(tool_name)); @@ -378,7 +348,7 @@ fn convert_assistant_parts( // 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) = read_provider_options(provider_options.as_ref(), namespace) { + if let Some(opts) = namespace.read(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) { @@ -412,9 +382,9 @@ fn convert_assistant_parts( FileBytes::Base64(data) => data.clone(), }; let mut value = json!({ "inlineData": { "mimeType": part.media_type, "data": data }, "thought": true }); - if let Some(signature) = - read_provider_options(part.provider_options.as_ref(), namespace) - .and_then(|options| options.get("thoughtSignature")) + if let Some(signature) = namespace + .read(part.provider_options.as_ref()) + .and_then(|options| options.get("thoughtSignature")) { value["thoughtSignature"] = signature.clone(); } @@ -433,10 +403,10 @@ fn convert_assistant_parts( json!({ "inlineData": { "mimeType": file.media_type, "data": data } }) } FileData::Text { text } => { - json!({ "inlineData": { "mimeType": if file.media_type.split_once('/').is_some_and(|(_, subtype)| !subtype.is_empty() && subtype != "*") { file.media_type.as_str() } else { "text/plain" }, "data": base64::engine::general_purpose::STANDARD.encode(text.as_bytes()) } }) + 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") { + if let Some(uri) = reference.get(GOOGLE) { json!({ "fileData": { "mimeType": file.media_type, "fileUri": uri } }) } else { continue; @@ -444,9 +414,7 @@ fn convert_assistant_parts( } FileData::Url { .. } => continue, }; - if let Some(options) = - read_provider_options(file.provider_options.as_ref(), namespace) - { + if let Some(options) = namespace.read(file.provider_options.as_ref()) { if options.get("thought").and_then(Value::as_bool) == Some(true) { value["thought"] = json!(true); } @@ -471,17 +439,7 @@ fn convert_assistant_parts( pub(crate) fn tool_file_media_type( file: &FilePart, ) -> Result { - if let FileData::Data { data } = &file.data { - let bytes = match data { - FileBytes::Binary(bytes) => bytes.clone(), - FileBytes::Base64(data) => base64::engine::general_purpose::STANDARD - .decode(data) - .unwrap_or_default(), - }; - crate::anthropic::convert::resolve_full_media_type(&file.media_type, &bytes) - } else { - crate::anthropic::convert::resolve_full_media_type(&file.media_type, &[]) - } + aimux_provider_utils::resolve_full_media_type(file) } pub(crate) fn validate_tool_result_files( @@ -519,7 +477,7 @@ fn contains_schema_reference(value: &Value) -> bool { fn convert_tool_parts( content: &[ToolPart], - namespace: ProviderMetadataNamespace, + namespace: Namespace, supports_function_response_parts: bool, ) -> Vec { let mut parts = Vec::new(); @@ -535,7 +493,7 @@ fn convert_tool_parts( }; let response = |content: Value| { let mut value = json!({ "functionResponse": { "name": tool_name, "response": { "name": tool_name, "content": content } } }); - if namespace != ProviderMetadataNamespace::Vertex && !tool_call_id.is_empty() { + if namespace != Namespace::Vertex && !tool_call_id.is_empty() { value["functionResponse"]["id"] = json!(tool_call_id); } value @@ -574,7 +532,7 @@ fn convert_tool_parts( } } FileData::Url { url, original_url } - if namespace == ProviderMetadataNamespace::Vertex + if namespace == Namespace::Vertex && url.starts_with("gs:") && original_url .as_deref() @@ -683,16 +641,7 @@ pub fn prepare_tools( if t.strict == Some(true) { has_strict = true; } - let mut decl = json!({ - "name": t.name, - "parameters": convert_json_schema_to_openapi_schema(&t.input_schema, true), - }); - if let Some(desc) = &t.description { - decl["description"] = json!(desc); - } else { - decl["description"] = json!(""); - } - decl + build_function_declaration(t) }) .collect(); @@ -1044,9 +993,10 @@ fn push_provider_tool( /// Build a single `functionDeclarations` entry from a `FunctionTool`. /// -/// `parameters` is omitted when the converted schema is null (empty object -/// schemas at the root), matching the TS `convertJSONSchemaToOpenAPISchema` -/// returning `undefined`. +/// 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. fn build_function_declaration(ft: &FunctionTool) -> Value { let mut decl = Map::new(); decl.insert("name".to_string(), json!(ft.name)); @@ -1054,9 +1004,10 @@ fn build_function_declaration(ft: &FunctionTool) -> Value { "description".to_string(), json!(ft.description.as_deref().unwrap_or("")), ); - let params = convert_json_schema_to_openapi_schema(&ft.input_schema, true); - if !params.is_null() { - decl.insert("parameters".to_string(), params); + // 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()); } Value::Object(decl) } @@ -1274,10 +1225,11 @@ pub(crate) fn image_generation_option(key: &str, value: &Value) -> Value { /// Build the Gemini `generateContent` request body from `CallOptions`. /// -/// Mirrors `getArgs` in `google-language-model.ts`. Provider-specific options -/// (`thinkingConfig`, `safetySettings`, `cachedContent`, `labels`, -/// `serviceTier`, …) are not yet surfaced through `CallOptions` in the Rust -/// port and are therefore omitted. +/// Mirrors `getArgs` in `google-language-model.ts`: the sampling settings, the +/// response format, and the provider options the SDK maps (`thinkingConfig`, +/// `responseModalities`, `audioTimestamp`, `mediaResolution`, `imageConfig` +/// into `generationConfig`; `safetySettings`, `cachedContent`, `labels`, +/// `serviceTier` and `retrievalConfig` next to it). /// /// This is the request-body-only entry point; warnings about unsupported tools /// are discarded. Use [`build_request_body_with_warnings`] to surface them. @@ -1297,12 +1249,8 @@ pub(crate) fn build_vertex_request_body( model_id: &str, options: &CallOptions, ) -> Result { - build_request_body_with_warnings_for_namespace( - model_id, - options, - ProviderMetadataNamespace::Vertex, - ) - .map(|(body, _)| body) + build_request_body_with_warnings_for_namespace(model_id, options, Namespace::Vertex) + .map(|(body, _)| body) } /// Build the Gemini `generateContent` request body **and** collect the tool @@ -1315,24 +1263,17 @@ pub fn build_request_body_with_warnings( model_id: &str, options: &CallOptions, ) -> Result<(Value, Vec), AiMuxError> { - build_request_body_with_warnings_for_namespace( - model_id, - options, - ProviderMetadataNamespace::Google, - ) + build_request_body_with_warnings_for_namespace(model_id, options, Namespace::Google) } fn build_request_body_with_warnings_for_namespace( model_id: &str, options: &CallOptions, - namespace: ProviderMetadataNamespace, + namespace: Namespace, ) -> Result<(Value, Vec), AiMuxError> { - let names: &[&str] = match namespace { - ProviderMetadataNamespace::Google => &["google"], - ProviderMetadataNamespace::Vertex => &["googleVertex", "vertex", "google"], - }; let provider_options = options.provider_options.as_ref().and_then(|options| { - names + namespace + .read_keys() .iter() .find_map(|name| options.get(*name).map(|value| (*name, value))) }); @@ -1399,23 +1340,56 @@ fn build_request_body_with_warnings_for_namespace( } } + // 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", + ] { + if let Some(value) = option(name) { + generation_config.insert(name.to_string(), value.clone()); + } + } + let mut body = Map::new(); body.insert("contents".to_string(), Value::Array(contents)); if let Some(sys) = system_instruction { body.insert("systemInstruction".to_string(), sys); } - if !generation_config.is_empty() { - body.insert( - "generationConfig".to_string(), - Value::Object(generation_config), - ); + // Always present, even when empty (`generationConfig: {}`), like the SDK. + body.insert( + "generationConfig".to_string(), + Value::Object(generation_config), + ); + for name in ["safetySettings", "cachedContent", "labels", "serviceTier"] { + if let Some(value) = option(name) { + body.insert(name.to_string(), value.clone()); + } } 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)); } - if let Some(tc) = prepared.tool_config { + let mut tool_config = prepared.tool_config; + 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()); + tool_config = Some(Value::Object(config)); + } + if let Some(tc) = tool_config { body.insert("toolConfig".to_string(), tc); } diff --git a/aimux-providers/src/google/embedding.rs b/aimux-providers/src/google/embedding.rs index 4eb061c43..733842289 100644 --- a/aimux-providers/src/google/embedding.rs +++ b/aimux-providers/src/google/embedding.rs @@ -7,8 +7,6 @@ //! - Single value: `POST {base_url}/models/{model}:embedContent` //! - Multiple values: `POST {base_url}/models/{model}:batchEmbedContents` -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -18,9 +16,8 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::HttpRequest; - -use super::GoogleConfig; +use super::options::google_options; +use crate::shared::EndpointConfig; /// A Google Gemini embedding model (e.g. `"gemini-embedding-001"`). /// @@ -28,51 +25,25 @@ use super::GoogleConfig; /// `Client` internally (RFC-0009 §4.1). pub struct GoogleEmbeddingModel { model_id: String, - config: GoogleConfig, + config: EndpointConfig, } impl GoogleEmbeddingModel { - #[must_use] - pub fn new(model_id: String, config: GoogleConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("x-goog-api-key".to_string(), self.config.api_key.clone()); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } -} - -/// Build the header list for a JSON POST: auth/extra headers + `Content-Type`. -fn build_header_list(headers: &HashMap) -> Vec<(String, String)> { - let mut list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - list.push(("Content-Type".to_string(), "application/json".to_string())); - list } #[async_trait] impl EmbeddingModel for GoogleEmbeddingModel { fn provider(&self) -> &str { - "google.generative-ai" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - fn max_embeddings_per_call(&self) -> Option { Some(100) } @@ -87,8 +58,7 @@ impl EmbeddingModel for GoogleEmbeddingModel { ) -> Result { let google_options = parse_google_provider_options(options.provider_options.as_ref()); - let headers = self.build_headers(options.headers.as_ref()); - let header_list = build_header_list(&headers); + let exchange = self.config.exchange(options.headers.as_ref()).await?; // For single embeddings, use the single endpoint. if options.values.len() == 1 { @@ -115,22 +85,11 @@ impl EmbeddingModel for GoogleEmbeddingModel { body.insert("taskType".to_string(), json!(task_type)); } - let url = format!( - "{}/models/{}:embedContent", - self.config.base_url, self.model_id - ); + let url = exchange.url(&format!("/models/{}:embedContent", self.model_id)); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url, - headers: header_list, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, - Value::Object(body), + exchange.request(url, options), + exchange.transform_body(Value::Object(body)), aimux_provider_utils::create_json_response_handler(), super::google_failed_response_handler(), ) @@ -198,22 +157,11 @@ impl EmbeddingModel for GoogleEmbeddingModel { let mut body = Map::new(); body.insert("requests".to_string(), Value::Array(requests)); - let url = format!( - "{}/models/{}:batchEmbedContents", - self.config.base_url, self.model_id - ); + let url = exchange.url(&format!("/models/{}:batchEmbedContents", self.model_id)); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url, - headers: header_list, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, - Value::Object(body), + exchange.request(url, options), + exchange.transform_body(Value::Object(body)), aimux_provider_utils::create_json_response_handler(), super::google_failed_response_handler(), ) @@ -266,7 +214,7 @@ struct GoogleEmbeddingProviderOptions { fn parse_google_provider_options( options: Option<&SharedProviderOptions>, ) -> GoogleEmbeddingProviderOptions { - let provider_opts = options.and_then(|opts| opts.get("google")); + let provider_opts = google_options(options); GoogleEmbeddingProviderOptions { output_dimensionality: provider_opts .and_then(|o| o.get("outputDimensionality")) diff --git a/aimux-providers/src/google/files.rs b/aimux-providers/src/google/files.rs index 644a59eea..48aa7b365 100644 --- a/aimux-providers/src/google/files.rs +++ b/aimux-providers/src/google/files.rs @@ -13,7 +13,6 @@ //! 3. If the file state is `PROCESSING`, poll `{base_url}/{file.name}` until //! the state becomes `ACTIVE` or `FAILED`. -use std::collections::HashMap; use std::time::{Duration, Instant}; use async_trait::async_trait; @@ -27,7 +26,8 @@ use aimux_core::types::Warning; use aimux_provider_utils::{HttpBody, HttpRequest, sleep_or_abort}; -use super::GoogleConfig; +use super::options::GOOGLE; +use crate::shared::EndpointConfig; /// Google-specific error structure: `{ "error": { "message": "..." } }`. /// Google provider-specific file upload options. @@ -44,7 +44,7 @@ fn parse_google_files_options( ) -> GoogleFilesUploadOptions { let mut opts = GoogleFilesUploadOptions::default(); if let Some(po) = provider_options - && let Some(google) = po.get("google") + && let Some(google) = po.get(GOOGLE) { if let Some(display_name) = google.get("displayName").and_then(|v| v.as_str()) { opts.display_name = Some(display_name.to_string()); @@ -112,35 +112,19 @@ struct UploadResponse { /// Aligned with TS `GoogleFiles`. Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers /// uses the process-wide shared `Client` internally (RFC-0009 §4.1). pub struct GoogleFiles { - config: GoogleConfig, + config: EndpointConfig, } impl GoogleFiles { - #[must_use] - pub fn new(config: GoogleConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn build_headers(&self) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("x-goog-api-key".to_string(), self.config.api_key.clone()); - headers - } - - /// The base origin is `base_url` with `/v1beta` stripped. - fn base_origin(&self) -> String { - self.config.base_url.replace("/v1beta", "") - } - - fn init_endpoint(&self) -> String { - format!("{}/upload/v1beta/files", self.base_origin()) - } } #[async_trait] impl Files for GoogleFiles { fn provider(&self) -> &str { - "google.generative-ai" + &self.config.provider } async fn upload_file( @@ -148,7 +132,13 @@ impl Files for GoogleFiles { options: &UploadFileCallOptions, ) -> Result { let google_options = parse_google_files_options(options.provider_options.as_ref()); - let resolved_headers = self.build_headers(); + 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", "") + ); let mut warnings = Vec::new(); if options.filename.is_some() { @@ -170,10 +160,7 @@ impl Files for GoogleFiles { json!({ "file": {} }) }; - let mut init_headers: Vec<(String, String)> = resolved_headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + let mut init_headers = exchange.headers(); init_headers.push(( "X-Goog-Upload-Protocol".to_string(), "resumable".to_string(), @@ -189,55 +176,38 @@ impl Files for GoogleFiles { )); init_headers.push(("Content-Type".to_string(), "application/json".to_string())); - // Nothing above `upload_file` retries it (§9.4), so the retry lives - // here — per stage, not around the whole upload: a failure in the - // upload or poll stage must not replay the init exchange that minted - // `upload_url`, nor resend the file body. - let retries = aimux_core::retry::prepare_retries( - None, - self.config.retry_config, - options.abort_signal.clone(), - ); - let init_resp = retries - .retry(|| async { - aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.init_endpoint(), - headers: init_headers.clone(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, - init_body_value.clone(), - aimux_provider_utils::ResponseHandler::new(|input| async move { - let headers = aimux_provider_utils::extract_response_headers::extract_response_headers( - input.response.headers(), - ); - Ok(aimux_provider_utils::ResponseHandlerOutput { - value: (), - raw_value: None, - response_headers: headers, - }) - }), - super::google_failed_response_handler(), - ) - .await - // Per attempt, so a `RetryError`'s saved `errors` carry it too. - .map_err(|e| match e { - AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError { - message: format!("Failed to initiate resumable upload: {}", d.message), - ..*d - })), - e => e, + // Nothing retries a file upload: the init, upload and poll exchanges + // each run once, and a failure is reported as it happened (a replayed + // upload would resend the file body). + let init_resp = aimux_provider_utils::post_json_to_api( + exchange.with_transport(HttpRequest { + url: init_endpoint, + headers: init_headers, + abort_signal: options.abort_signal.clone(), + ..Default::default() + }), + exchange.transform_body(init_body_value), + aimux_provider_utils::ResponseHandler::new(|input| async move { + let headers = + aimux_provider_utils::extract_response_headers::extract_response_headers( + input.response.headers(), + ); + Ok(aimux_provider_utils::ResponseHandlerOutput { + value: (), + raw_value: None, + response_headers: headers, }) - }) - .await?; + }), + super::google_failed_response_handler(), + ) + .await + .map_err(|e| match e { + AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError { + message: format!("Failed to initiate resumable upload: {}", d.message), + ..*d + })), + e => e, + })?; let upload_url = init_resp .response_headers @@ -262,47 +232,41 @@ impl Files for GoogleFiles { // The upload URL comes from the init response's x-goog-upload-url // header and receives the user's file bytes; validate it. (AI SDK // fetches this URL unvalidated — kept stricter here deliberately.) - let upload_resp = retries - .retry(|| async { - aimux_provider_utils::post_to_api( - HttpRequest { - url: upload_url.clone(), - headers: upload_headers.clone(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: true, - trusted_origin: Some(self.config.base_url.clone()), - credentialed_origin: Some(self.config.base_url.clone()), - }, - HttpBody::Bytes(file_bytes.clone(), media_type.clone()), - aimux_provider_utils::create_json_response_handler::(), - super::google_failed_response_handler(), - ) - .await - .map_err(|e| match e { - AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError { - message: format!("Failed to upload file data: {}", d.message), - ..*d - })), - e => e, - }) - }) - .await?; + let upload_resp = aimux_provider_utils::post_to_api( + HttpRequest { + url: upload_url, + headers: upload_headers, + abort_signal: options.abort_signal.clone(), + validate_url: true, + 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(), + ) + .await + .map_err(|e| match e { + AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError { + message: format!("Failed to upload file data: {}", d.message), + ..*d + })), + e => e, + })?; let mut file = upload_resp.value.file; + let mut raw_file = upload_resp + .raw_value + .as_ref() + .and_then(|value| value.get("file")) + .cloned() + .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 start_time = Instant::now(); - let poll_header_list: Vec<(String, String)> = resolved_headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); // Seed evidence from the upload response so a file that is already // FAILED (never polled) still carries the observed status + raw body. @@ -323,31 +287,22 @@ impl Files for GoogleFiles { ) .await?; - let poll_url = format!("{}/{}", self.config.base_url, file.name); - - let poll_resp = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: poll_url.clone(), - headers: poll_header_list.clone(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, - aimux_provider_utils::create_json_response_handler::(), - super::google_failed_response_handler(), - ) - }) - .await?; + let poll_url = exchange.url(&format!("/{}", file.name)); + + let poll_resp = aimux_provider_utils::get_from_api( + exchange.with_transport(HttpRequest { + url: poll_url.clone(), + headers: exchange.headers(), + abort_signal: options.abort_signal.clone(), + ..Default::default() + }), + aimux_provider_utils::create_json_response_handler::(), + super::google_failed_response_handler(), + ) + .await?; file = poll_resp.value; + raw_file = poll_resp.raw_value.clone().unwrap_or(Value::Null); last_poll_body = poll_resp.raw_value.map(|value| value.to_string()); last_poll_status = Some(200); last_poll_url = poll_url; @@ -371,9 +326,13 @@ impl Files for GoogleFiles { // Build provider metadata. let mut metadata = serde_json::Map::new(); metadata.insert("name".to_string(), json!(file.name)); - metadata.insert("displayName".to_string(), json!(file.display_name)); + if file.display_name.is_some() || raw_file.get("displayName").is_some() { + metadata.insert("displayName".to_string(), json!(file.display_name)); + } metadata.insert("mimeType".to_string(), json!(file.mime_type)); - metadata.insert("sizeBytes".to_string(), json!(file.size_bytes)); + if file.size_bytes.is_some() || raw_file.get("sizeBytes").is_some() { + metadata.insert("sizeBytes".to_string(), json!(file.size_bytes)); + } metadata.insert("state".to_string(), json!(file.state)); metadata.insert("uri".to_string(), json!(file.uri)); if let Some(ref create_time) = file.create_time { @@ -390,7 +349,7 @@ impl Files for GoogleFiles { } let mut provider_ref = aimux_core::shared::SharedProviderReference::new(); - provider_ref.insert("google".to_string(), file.uri.clone()); + provider_ref.insert(GOOGLE.to_string(), file.uri.clone()); let result_media_type = if file.mime_type.is_empty() { Some(options.media_type.clone()) @@ -403,7 +362,7 @@ impl Files for GoogleFiles { media_type: result_media_type, filename: None, provider_metadata: Some(aimux_core::shared::NamespaceMap::from([( - "google".to_string(), + GOOGLE.to_string(), metadata, )])), warnings, diff --git a/aimux-providers/src/google/image.rs b/aimux-providers/src/google/image.rs index b2ad74471..053148954 100644 --- a/aimux-providers/src/google/image.rs +++ b/aimux-providers/src/google/image.rs @@ -7,8 +7,6 @@ //! - **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}; @@ -19,9 +17,8 @@ use aimux_core::image_model::{ }; use aimux_core::shared::{SharedProviderMetadata, SharedProviderOptions, Warning}; -use aimux_provider_utils::HttpRequest; - -use super::GoogleConfig; +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. @@ -43,12 +40,15 @@ pub struct GoogleImageSettings { pub struct GoogleImageModel { model_id: String, settings: GoogleImageSettings, - config: GoogleConfig, + config: EndpointConfig, } impl GoogleImageModel { - #[must_use] - pub fn new(model_id: String, settings: GoogleImageSettings, config: GoogleConfig) -> Self { + pub(crate) fn from_config( + model_id: String, + settings: GoogleImageSettings, + config: EndpointConfig, + ) -> Self { Self { model_id, settings, @@ -56,39 +56,24 @@ impl GoogleImageModel { } } - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("x-goog-api-key".to_string(), self.config.api_key.clone()); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn predict_endpoint(&self) -> String { - format!("{}/models/{}:predict", self.config.base_url, self.model_id) + fn predict_path(&self) -> String { + format!("/models/{}:predict", self.model_id) } - fn generate_content_endpoint(&self) -> String { + 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!("{}/{}:generateContent", self.config.base_url, model_path) + format!("/{model_path}:generateContent") } fn max_images(&self) -> u32 { if let Some(max) = self.settings.max_images_per_call { return max; } - if is_gemini_model(&self.model_id) { - 10 - } else { - 4 - } + 1 } // ── Imagen path ───────────────────────────────────────────────────────── @@ -165,12 +150,10 @@ impl GoogleImageModel { "parameters": parameters, }); - let headers = self.build_headers(options.headers.as_ref()); - let header_list = build_header_list(&headers); - + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.predict_endpoint(), header_list, options), - body, + exchange.request(exchange.url(&self.predict_path()), options), + exchange.transform_body(body), aimux_provider_utils::create_json_response_handler(), super::google_failed_response_handler(), ) @@ -284,7 +267,7 @@ impl GoogleImageModel { } // Only declared generation-config options survive the upstream language schema. - if let Some(google) = options.provider_options.get("google") { + if let Some(google) = read_google_options(Some(&options.provider_options)) { for key in [ "audioTimestamp", "thinkingConfig", @@ -312,12 +295,10 @@ impl GoogleImageModel { body.insert("tools".to_string(), json!([{ "googleSearch": gs }])); } - let headers = self.build_headers(options.headers.as_ref()); - let header_list = build_header_list(&headers); - + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.generate_content_endpoint(), header_list, options), - Value::Object(body), + exchange.request(exchange.url(&self.generate_content_path()), options), + exchange.transform_body(Value::Object(body)), aimux_provider_utils::create_json_response_handler(), super::google_failed_response_handler(), ) @@ -346,17 +327,13 @@ impl GoogleImageModel { #[async_trait] impl ImageModel for GoogleImageModel { fn provider(&self) -> &str { - "google.generative-ai" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - fn max_images_per_call(&self) -> Option { Some(self.max_images()) } @@ -372,16 +349,6 @@ impl ImageModel for GoogleImageModel { // ── Helpers ───────────────────────────────────────────────────────────────── -/// Build the header list for a JSON POST: auth/extra headers + `Content-Type`. -fn build_header_list(headers: &HashMap) -> Vec<(String, String)> { - let mut list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - list.push(("Content-Type".to_string(), "application/json".to_string())); - list -} - /// Parsed Google image provider options. struct GoogleImageOptions { person_generation: Option, @@ -390,7 +357,7 @@ struct GoogleImageOptions { /// Parse Google image provider options from the `"google"` key. fn parse_google_image_options(provider_options: &SharedProviderOptions) -> GoogleImageOptions { - let google = provider_options.get("google"); + let google = read_google_options(Some(provider_options)); GoogleImageOptions { person_generation: google .and_then(|g| g.get("personGeneration")) @@ -430,7 +397,7 @@ fn extract_imagen_metadata(response: &Value) -> SharedProviderMetadata { 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.insert(GOOGLE.to_string(), google_meta); metadata } @@ -496,7 +463,7 @@ fn extract_gemini_result( if let Some(gm) = grounding_metadata { google_meta.insert("groundingMetadata".to_string(), gm); } - metadata.insert("google".to_string(), google_meta); + 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 0a3fc5765..67af738cd 100644 --- a/aimux-providers/src/google/mod.rs +++ b/aimux-providers/src/google/mod.rs @@ -1,7 +1,16 @@ //! Google Gemini provider. //! -//! Implements the `LanguageModel` trait against Google's Generative Language -//! API (`generativelanguage.googleapis.com/v1beta`). +//! [`create_google`] is the Rust form of the AI SDK's `createGoogle`: it takes +//! [`GoogleProviderSettings`], validates the base URL, fixes the provider name +//! and returns a [`GoogleProvider`]. The API key is not read there; it is +//! loaded in the request headers of every call, from the setting or from +//! `GOOGLE_GENERATIVE_AI_API_KEY`. [`google()`] is the default instance. +//! +//! The models speak Google's Generative Language API +//! (`generativelanguage.googleapis.com/v1beta`). providerOptions are read from +//! `providerOptions.google` and response metadata is written under `google` +//! (see `options`). Vertex AI serves the same request format and shares the +//! conversion in [`convert`]. //! //! Gemini's API shape is fundamentally different from OpenAI/Anthropic: //! - The model id is part of the URL path, not the request body. @@ -17,20 +26,34 @@ pub mod embedding; pub mod files; pub mod image; mod model; +pub(crate) mod options; pub mod types; pub mod utils; pub mod video; pub use embedding::GoogleEmbeddingModel; +pub use files::GoogleFiles; pub use image::{GoogleImageModel, GoogleImageSettings}; +pub use model::GoogleModel; pub use video::GoogleVideoModel; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::{RetryConfig, load_api_key, without_trailing_slash}; +use std::sync::{Arc, OnceLock}; + +use futures::future::BoxFuture; use serde_json::Value; +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::error::AiMuxError; +use aimux_core::files_model::Files; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_core::video_model::VideoModel; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{AuthScheme, Credential, Endpoint, EndpointConfig, credential_headers}; + pub(crate) fn google_stream_error( error: &types::GoogleError, url: &str, @@ -77,205 +100,335 @@ pub(crate) fn google_failed_response_handler() -> aimux_provider_utils::Response }) } -/// Configuration for the Google Gemini provider. -#[derive(Debug, Clone)] -pub struct GoogleConfig { - pub api_key: String, - pub base_url: String, - /// api_key 来源(RFC-0023 `ProviderRecord.api_key_source`):`None` = 显式 - /// (config_snapshot 记为 "explicit");`Some("env:VAR")` = 来自环境变量。 - /// 不存明文之外的信息,仅用于回放重建。 - pub api_key_source: Option, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, +pub(crate) const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta"; +pub(crate) const API_KEY_ENV_VAR: &str = "GOOGLE_GENERATIVE_AI_API_KEY"; +const DEFAULT_NAME: &str = "google.generative-ai"; + +/// The media types the API fetches from an external `https` URL itself +/// (`supportedExternalUrlMediaTypes` in `createGoogle`). +const EXTERNAL_URL_MEDIA_TYPES: [&str; 22] = [ + "text/html", + "text/css", + "text/plain", + "text/xml", + "text/csv", + "text/rtf", + "text/javascript", + "application/json", + "application/pdf", + "image/bmp", + "image/jpeg", + "image/png", + "image/webp", + "video/mp4", + "video/mpeg", + "video/quicktime", + "video/avi", + "video/x-flv", + "video/mpg", + "video/webm", + "video/wmv", + "video/3gpp", +]; + +/// `supportsExternalFileUrls`: `/(^|\/)gemini-/` matches and +/// `/(^|\/)gemini-2\.0/` does not. +fn supports_external_file_urls(model_id: &str) -> bool { + let starts_a_segment = + |prefix: &str| model_id.starts_with(prefix) || model_id.contains(&format!("/{prefix}")); + starts_a_segment("gemini-") && !starts_a_segment("gemini-2.0") } -impl GoogleConfig { - /// Create from an API key (uses the default Gemini API base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://generativelanguage.googleapis.com/v1beta".to_string(), - api_key_source: None, - retry_config: RetryConfig::default(), +/// The URL patterns the API fetches itself (`getSupportedUrls` in +/// `createGoogle`): Files API URLs and YouTube links for every media type, +/// plus external `https` URLs for the models that accept them. +fn supported_urls(base_url: &str, model_id: Option<&str>) -> SupportedUrls { + let include_external = model_id.is_none_or(supports_external_file_urls); + let pattern = |source: &str| regex::Regex::new(source).expect("static pattern"); + let mut urls = std::collections::HashMap::new(); + urls.insert( + "*".to_string(), + vec![ + pattern(r"^https://generativelanguage\.googleapis\.com/v1beta/files/.*$"), + pattern(&format!("^{}/files/.*$", regex::escape(base_url))), + pattern(r"^https://(?:www\.)?youtube\.com/watch\?v=[\w-]+(?:&[\w=&.-]*)?$"), + pattern(r"^https://youtu\.be/[\w-]+(?:\?[\w=&.-]*)?$"), + ], + ); + if include_external { + let https = pattern(r"^https://.*$"); + for media_type in EXTERNAL_URL_MEDIA_TYPES { + urls.insert(media_type.to_string(), vec![https.clone()]); } } + SupportedUrls(urls) +} - /// Use a custom base URL (e.g. for Vertex AI or a proxy). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Settings of [`create_google`] (the AI SDK's `GoogleProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct GoogleProviderSettings { + /// Base URL for the API calls. Default + /// `https://generativelanguage.googleapis.com/v1beta`; a trailing slash is + /// removed. + pub base_url: Option, + /// The API key, sent as `x-goog-api-key`. `None` loads + /// `GOOGLE_GENERATIVE_AI_API_KEY` when a request is made and fails that + /// request with `AiMuxError::LoadApiKey` if it is unset. An explicit value + /// is used as given, `""` included: it never falls back to the + /// environment. + pub api_key: Option, + /// Extra headers on every request. A `None` value removes the header, + /// including `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"`. + 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, +} - /// 标注 api_key 来源(RFC-0023 回放重建用)。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self +impl std::fmt::Debug for GoogleProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GoogleProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key.is_some()) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self - } +/// Create a Google Generative AI provider. +/// +/// # 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. +pub fn create_google(settings: GoogleProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(GoogleProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( + credential_headers( + Credential::explicit_or_env( + settings.api_key.map(Resolvable::Value), + API_KEY_ENV_VAR, + "Google Generative AI", + ), + AuthScheme::Header("x-goog-api-key"), + Vec::new(), + settings.headers, + ), + options::GOOGLE, + "4.0.85", + ), + fetch: settings.fetch, + }) +} - /// Create from the `GOOGLE_GENERATIVE_AI_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when - /// `GOOGLE_GENERATIVE_AI_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "GOOGLE_GENERATIVE_AI_API_KEY", "Google Generative AI")?; - Ok(Self::new(api_key).with_api_key_source(Some("env:GOOGLE_GENERATIVE_AI_API_KEY"))) - } +/// The default provider: `create_google` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn google() -> &'static GoogleProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_google(GoogleProviderSettings::default()) + .expect("default Google settings are always valid") + }) } -/// Google Gemini provider — creates `GoogleModel` instances. -/// -/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the -/// process-wide shared `Client` internally (RFC-0009 §4.1). +/// A Google Generative AI provider (the AI SDK's `GoogleProvider`). Cheap to +/// clone the models out of; it holds no HTTP client. pub struct GoogleProvider { - config: GoogleConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl GoogleProvider { - #[must_use] - pub fn new(config: GoogleConfig) -> Self { - Self { config } + /// The model configuration reporting `provider` as its identity. + fn model_config(&self, provider: String) -> EndpointConfig { + let base_url = self.base_url.clone(); + let headers = self.headers.clone(); + let urls_base = self.base_url.clone(); + EndpointConfig { + provider, + endpoint: Arc::new(move || { + let base_url = base_url.clone(); + let headers = headers.clone(); + Box::pin(async move { + Ok(Endpoint { + base_url, + headers: headers.resolve().await?, + }) + }) + }), + fetch: self.fetch.clone(), + supported_urls: Arc::new(move |model_id| supported_urls(&urls_base, Some(model_id))), + transform_request_body: None, + } } - /// Create a model instance for the given model name (e.g. `"gemini-2.0-flash"`). + /// A language model; `provider()` is the provider name + /// (`"google.generative-ai"` by default). #[must_use] - pub fn model(&self, model_id: &str) -> model::GoogleModel { - model::GoogleModel::new(model_id.to_string(), self.config.clone()) + pub fn chat(&self, model_id: &str) -> GoogleModel { + GoogleModel::from_config(model_id.to_string(), self.model_config(self.name.clone())) } - /// Create a Files interface for uploading files to Google. + /// A text embedding model (e.g. `"gemini-embedding-001"`); `provider()` is + /// the provider name. #[must_use] - pub fn files(&self) -> files::GoogleFiles { - files::GoogleFiles::new(self.config.clone()) + pub fn embedding(&self, model_id: &str) -> GoogleEmbeddingModel { + GoogleEmbeddingModel::from_config( + model_id.to_string(), + self.model_config(self.name.clone()), + ) } - /// Create an embedding model instance for the given model name (e.g. - /// `"gemini-embedding-001"`). + /// An image generation model (e.g. `"imagen-3.0-generate-002"` or + /// `"gemini-2.5-flash-image"`); `provider()` is the provider name. #[must_use] - pub fn embedding_model(&self, model_id: &str) -> embedding::GoogleEmbeddingModel { - embedding::GoogleEmbeddingModel::new(model_id.to_string(), self.config.clone()) + pub fn image(&self, model_id: &str) -> GoogleImageModel { + self.image_with_settings(model_id, GoogleImageSettings::default()) } - /// Create an image generation model instance for the given model name - /// (e.g. `"imagen-3.0-generate-002"` or `"gemini-2.5-flash-image"`). + /// [`image`](Self::image) with model settings. #[must_use] - pub fn image(&self, model_id: &str) -> image::GoogleImageModel { - image::GoogleImageModel::new( + pub fn image_with_settings( + &self, + model_id: &str, + settings: GoogleImageSettings, + ) -> GoogleImageModel { + GoogleImageModel::from_config( model_id.to_string(), - image::GoogleImageSettings::default(), - self.config.clone(), + settings, + self.model_config(self.name.clone()), ) } - /// Create a video generation model instance for the given model name - /// (e.g. `"veo-3.0-generate-001"`). + /// A video generation model (e.g. `"veo-3.0-generate-001"`); + /// `provider()` is the provider name. #[must_use] - pub fn video(&self, model_id: &str) -> video::GoogleVideoModel { - video::GoogleVideoModel::new(model_id.to_string(), self.config.clone()) + pub fn video(&self, model_id: &str) -> GoogleVideoModel { + GoogleVideoModel::from_config(model_id.to_string(), self.model_config(self.name.clone())) } - /// Create an image generation model instance with custom settings. + /// The files interface; `provider()` is the provider name. #[must_use] - pub fn image_with_settings( - &self, - model_id: &str, - settings: image::GoogleImageSettings, - ) -> image::GoogleImageModel { - image::GoogleImageModel::new(model_id.to_string(), settings, self.config.clone()) + pub fn files(&self) -> GoogleFiles { + GoogleFiles::from_config(self.model_config(self.name.clone())) + } + + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; the same model as [`chat`](Self::chat) and + /// [`language_model`](Provider::language_model). + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.chat(model_id)) } } impl Provider for GoogleProvider { - fn name(&self) -> &str { - "google" + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - /// List models via `GET {base_url}/models` (Gemini native), enriched with - /// the community catalogue portrait when available (RFC-0027). - fn list_models( - &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let config = self.config.clone(); - Box::pin(async move { - let base = config.base_url.trim_end_matches('/'); - let url = format!("{base}/models"); - let mut headers = std::collections::HashMap::new(); - headers.insert("x-goog-api-key".to_string(), config.api_key.clone()); - let mut header_list: Vec<(String, String)> = headers.into_iter().collect(); - header_list.push(("Content-Type".to_string(), "application/json".to_string())); - - use aimux_provider_utils::HttpRequest; - // Retry rationale: see `openai::model::execute_list_models`. - let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: header_list.clone(), - abort_signal: None, - call_id: None, - recording_context: None, - ..Default::default() - }, - aimux_provider_utils::create_json_response_handler(), - google_failed_response_handler(), - ) - }) - .await?; - - // Gemini response: { models: [{ name: "models/gemini-...", displayName, ... }] } - #[derive(serde::Deserialize)] - struct Resp { - #[serde(default)] - models: Vec, - } - #[derive(serde::Deserialize)] - struct Entry { - name: String, - #[serde(default)] - display_name: Option, - } - let parsed: Resp = resp.value; - let runtime: Vec = parsed - .models - .into_iter() - .map(|e| { - // Strip "models/" prefix → bare model id. - let id = e - .name - .strip_prefix("models/") - .unwrap_or(&e.name) - .to_string(); - aimux_core::model_catalogue::RuntimeModel { - id, - owned_by: e.display_name, - created: None, - } - }) - .collect(); - Ok(runtime) - }) + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.image(model_id))) + } + + fn video_model(&self, model_id: &str) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.video(model_id)))) + } + + fn files(&self) -> Option> { + Some(Arc::new(self.files())) + } +} + +impl ProviderDiscovery for GoogleProvider { + /// `GET {base_url}/models` (Gemini native): one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config(format!("{}.models", self.name)); + Box::pin(async move { list_models_once(&config).await }) } } + +/// One `GET {base_url}/models` exchange against a Gemini-shaped API: no retry, +/// no recording. Discovery is not a Core operation and the AI SDK has no +/// equivalent, so a failure is reported to the caller as it happened. Shared +/// by Vertex. +/// +/// # Errors +/// +/// Returns the header-resolution error (a missing key is `LoadApiKey`), +/// `ApiCall` for HTTP/transport failures and `JsonParse` when the body does +/// not deserialize into the models list. +pub(crate) async fn list_models_once( + config: &EndpointConfig, +) -> Result, AiMuxError> { + // Gemini response: { models: [{ name: "models/gemini-...", displayName, ... }] } + #[derive(serde::Deserialize)] + struct Resp { + #[serde(default)] + models: Vec, + } + #[derive(serde::Deserialize)] + struct Entry { + name: String, + #[serde(default)] + display_name: Option, + } + + let exchange = config.exchange(None).await?; + let resp = aimux_provider_utils::get_from_api( + exchange.with_transport(aimux_provider_utils::HttpRequest { + url: exchange.url("/models"), + headers: exchange.headers(), + ..Default::default() + }), + aimux_provider_utils::create_json_response_handler(), + google_failed_response_handler(), + ) + .await?; + let parsed: Resp = resp.value; + Ok(parsed + .models + .into_iter() + .map(|entry| RuntimeModel { + // Strip "models/" prefix -> bare model id. + id: entry + .name + .strip_prefix("models/") + .unwrap_or(&entry.name) + .to_string(), + owned_by: entry.display_name, + created: None, + }) + .collect()) +} diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index 8d0ebe62f..f0bf07c9f 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -2,8 +2,6 @@ use aimux_core::tool::RawToolCall; use aimux_core::tool::ToolResult; -use std::collections::HashMap; - use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; @@ -14,20 +12,21 @@ use aimux_core::options::CallOptions; use aimux_core::result::{ GenerateContent, GenerateResult, GeneratedFile, ReasoningOutput, Source, StreamResult, }; -use aimux_core::shared::{FileBytes, GeneratedFileData, provider_namespace}; +use aimux_core::shared::{FileBytes, GeneratedFileData}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, }; -use aimux_provider_utils::HttpRequest; +use aimux_core::language_model::SupportedUrls; -use super::GoogleConfig; use super::convert::{ build_request_body_with_warnings, code_execution_tool_name, convert_usage, extract_sources, parse_finish_reason, }; +use super::options::google_metadata; use super::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; +use crate::shared::EndpointConfig; /// A Google Gemini language model. /// @@ -35,91 +34,37 @@ use super::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct GoogleModel { model_id: String, - config: GoogleConfig, + config: EndpointConfig, } impl GoogleModel { - pub fn new(model_id: String, config: GoogleConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - // The TS SDK sends the API key via `x-goog-api-key`. The query-param - // form (`?key=…`) is also supported but the header form is preferred. - headers.insert("x-goog-api-key".to_string(), self.config.api_key.clone()); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - /// `…/models/{model}:generateContent` (or `{model}:generateContent` when - /// the model id already contains a `/`, e.g. a fine-tuned path). - fn generate_endpoint(&self) -> String { - let model_path = if self.model_id.contains('/') { + /// `models/{model}`, or the id itself when it already contains a `/` + /// (e.g. a fine-tuned path). + fn model_path(&self) -> String { + if self.model_id.contains('/') { self.model_id.clone() } else { format!("models/{}", self.model_id) - }; - format!("{}/{}:generateContent", self.config.base_url, model_path) - } - - fn stream_endpoint(&self) -> String { - let model_path = if self.model_id.contains('/') { - self.model_id.clone() - } else { - format!("models/{}", self.model_id) - }; - format!( - "{}/{}:streamGenerateContent?alt=sse", - self.config.base_url, model_path - ) + } } } -/// Build the header list for a JSON POST: auth/extra headers + `Content-Type`. -/// -/// Returns a `Vec<(String, String)>` for `HttpRequest` — no reqwest types. -fn build_header_list(headers: &HashMap) -> Vec<(String, String)> { - let mut list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - list.push(("Content-Type".to_string(), "application/json".to_string())); - list -} - #[async_trait] impl LanguageModel for GoogleModel { fn provider(&self) -> &str { - "google.generative-ai" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.base_url.clone()), - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: None, - } + async fn supported_urls(&self) -> SupportedUrls { + (self.config.supported_urls)(&self.model_id) } async fn do_generate(&self, options: &CallOptions) -> Result { @@ -158,13 +103,11 @@ impl LanguageModel for GoogleModel { crate::google::convert::validate_tool_result_files(&options.prompt)?; 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 headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let body = exchange.transform_body(body); + let url = exchange.url(&format!("/{}:generateContent", self.model_path())); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new( - self.generate_endpoint(), - build_header_list(&headers), - options, - ), + exchange.request(url, options), body.clone(), aimux_provider_utils::create_json_response_handler(), super::google_failed_response_handler(), @@ -200,20 +143,14 @@ impl LanguageModel for GoogleModel { // Provider metadata: wrap the raw Google metadata under a `google` // key (matching the TS `wrapProviderMetadata`). - let provider_metadata = Some( - provider_namespace( - "google", - 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, - }), - ) - .expect("provider metadata must be an object"), - ); + let provider_metadata = Some(google_metadata(serde_json::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, + }))); Ok(GenerateResult { content, @@ -224,7 +161,7 @@ impl LanguageModel for GoogleModel { response: Some(aimux_core::shared::ResponseInfo { id: data.response_id, timestamp: None, - model_id: None, + model_id: data.model_version, headers: Some(response_headers), body: response_body, }), @@ -268,10 +205,14 @@ impl LanguageModel for GoogleModel { crate::google::convert::validate_tool_result_files(&options.prompt)?; 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 headers = self.build_headers(options.headers.as_ref()); - let endpoint = self.stream_endpoint(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let body = exchange.transform_body(body); + let endpoint = exchange.url(&format!( + "/{}:streamGenerateContent?alt=sse", + self.model_path() + )); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.clone(), build_header_list(&headers), options), + exchange.request(endpoint.clone(), options), body.clone(), aimux_provider_utils::create_event_source_response_handler::(), super::google_failed_response_handler(), @@ -401,7 +342,7 @@ impl LanguageModel for GoogleModel { let thought_sig_meta: Option = part .get("thoughtSignature") .and_then(|v| v.as_str()) - .map(|s| provider_namespace("google", json!({ "thoughtSignature": s })).expect("provider metadata must be an object")); + .map(|s| google_metadata(json!({ "thoughtSignature": s }))); // text part if let Some(text) = part.get("text").and_then(|v| v.as_str()) { @@ -695,14 +636,14 @@ impl LanguageModel for GoogleModel { yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None}); } - let provider_metadata = Some(provider_namespace("google", json!({ + let provider_metadata = Some(google_metadata(serde_json::json!({ "promptFeedback": last_prompt_feedback, "groundingMetadata": last_grounding_metadata, "urlContextMetadata": last_url_context_metadata, "safetyRatings": last_safety_ratings, "usageMetadata": last_usage_metadata_value, "finishMessage": last_finish_message, - })).expect("provider metadata must be an object")); + }))); yield Ok(StreamPart::Finish { finish_reason: if stream_errored { @@ -745,7 +686,7 @@ fn server_tool_metadata( if let Some(signature) = thought_signature { payload["thoughtSignature"] = json!(signature); } - provider_namespace("google", payload).expect("provider metadata must be an object") + google_metadata(payload) } /// Extract `GenerateContent` items from a non-streaming candidate. @@ -779,10 +720,7 @@ fn extract_content_from_candidate( let thought_sig_meta: Option = part .get("thoughtSignature") .and_then(|v| v.as_str()) - .map(|s| { - provider_namespace("google", json!({ "thoughtSignature": s })) - .expect("provider metadata must be an object") - }); + .map(|s| google_metadata(json!({ "thoughtSignature": s }))); // Branch order matches upstream (google-language-model.ts:420-534): // executableCode → codeExecutionResult → text → functionCall @@ -999,26 +937,3 @@ fn set_provider_metadata(item: &mut GenerateContent, meta: ProviderMetadata) { } } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::google::GoogleConfig; - - #[test] - fn config_snapshot_records_provider_identity() { - let config = GoogleConfig::new("sk-test"); - let model = GoogleModel::new("gemini-2.0-flash".to_string(), config); - - let snap = model.config_snapshot(); - assert_eq!(snap.provider, "google.generative-ai"); - assert_eq!(snap.model_id, "gemini-2.0-flash"); - assert_eq!( - snap.base_url.as_deref(), - Some("https://generativelanguage.googleapis.com/v1beta") - ); - assert_eq!(snap.api_key_source, "explicit"); - assert_eq!(snap.profile, None); - assert_eq!(snap.provider_options, None); - } -} diff --git a/aimux-providers/src/google/options.rs b/aimux-providers/src/google/options.rs new file mode 100644 index 000000000..03580dfd3 --- /dev/null +++ b/aimux-providers/src/google/options.rs @@ -0,0 +1,100 @@ +//! The providerOptions and providerMetadata namespaces of the Gemini family. +//! +//! `@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. + +use aimux_core::shared::{JsonObject, SharedProviderOptions, provider_namespace}; +use aimux_core::types::ProviderMetadata; +use serde_json::Value; + +/// The Gemini API namespace. +pub(crate) const GOOGLE: &str = "google"; + +/// The Vertex AI namespace. +pub(crate) const GOOGLE_VERTEX: &str = "googleVertex"; +/// Keys the Vertex video model writes its metadata under: the current +/// namespace plus the two legacy keys upstream keeps for compatibility. +pub(crate) const VERTEX_VIDEO_METADATA_KEYS: [&str; 3] = [GOOGLE_VERTEX, "google-vertex", "vertex"]; + +/// Which provider's namespaces a model reads and writes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Namespace { + /// The public Gemini API: reads and writes `google`. + Google, + /// Vertex AI: reads `googleVertex`, then `google`; writes `googleVertex`. + Vertex, +} + +impl Namespace { + /// The keys to read, in order of precedence. + pub(crate) fn read_keys(self) -> &'static [&'static str] { + match self { + Self::Google => &[GOOGLE], + Self::Vertex => &[GOOGLE_VERTEX, GOOGLE], + } + } + + /// The keys response metadata is written under. + pub(crate) fn write_keys(self) -> &'static [&'static str] { + match self { + Self::Google => &[GOOGLE], + Self::Vertex => &[GOOGLE_VERTEX], + } + } + + /// The first present options object among the read keys. + pub(crate) fn read( + self, + provider_options: Option<&SharedProviderOptions>, + ) -> Option<&JsonObject> { + let provider_options = provider_options?; + self.read_keys() + .iter() + .find_map(|key| provider_options.get(*key)) + } + + /// [`read`](Self::read) for `CallOptions::provider_options`. + pub(crate) fn read_in( + self, + provider_options: Option<&SharedProviderOptions>, + ) -> Option<&JsonObject> { + self.read(provider_options) + } + + /// 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") + } +} + +/// The options under the public Gemini key only (`providerOptions.google`), +/// for the surfaces the AI SDK keys by `google` alone (embeddings, images, +/// files). +pub(crate) fn google_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(GOOGLE) +} + +/// `{ "google": payload }`. +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 239a6ab26..056376d8e 100644 --- a/aimux-providers/src/google/types.rs +++ b/aimux-providers/src/google/types.rs @@ -64,36 +64,37 @@ pub struct CandidateContent { // ── Usage metadata ─────────────────────────────────────────────────────────── /// `usageMetadata` block. Mirrors `convert-google-usage.ts`'s -/// `GoogleUsageMetadata`. +/// `GoogleUsageMetadata`. Fields the response did not carry stay absent when +/// the block is echoed in `providerMetadata`, as in the AI SDK. #[derive(Debug, Default, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] pub struct GoogleUsageMetadata { - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub prompt_token_count: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub candidates_token_count: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub total_token_count: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub cached_content_token_count: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub thoughts_token_count: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub traffic_type: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub service_tier: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub prompt_tokens_details: Option>, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub candidates_tokens_details: Option>, } #[derive(Debug, Deserialize, Serialize)] #[serde(rename_all = "camelCase")] pub struct TokenDetail { - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub modality: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub token_count: Option, } diff --git a/aimux-providers/src/google/video.rs b/aimux-providers/src/google/video.rs index dce85c6b0..70c66cffb 100644 --- a/aimux-providers/src/google/video.rs +++ b/aimux-providers/src/google/video.rs @@ -8,8 +8,6 @@ //! 2. GET `{base_url}/{operation_name}` — polled by Core via `do_status` until `done: true` //! 3. Return video URL(s) from operation result -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -19,9 +17,7 @@ use aimux_core::video_model::{ VideoResponse, VideoResult, }; -use aimux_provider_utils::HttpRequest; - -use super::GoogleConfig; +use crate::shared::EndpointConfig; /// A Google video generation model. /// @@ -29,49 +25,23 @@ use super::GoogleConfig; /// `Client` internally (RFC-0009 §4.1). pub struct GoogleVideoModel { model_id: String, - config: GoogleConfig, + config: EndpointConfig, } impl GoogleVideoModel { - #[must_use] - pub fn new(model_id: String, config: GoogleConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut h = HashMap::new(); - h.insert("x-goog-api-key".to_string(), self.config.api_key.clone()); - if let Some(extra) = extra { - for (k, v) in extra { - h.insert(k.clone(), v.clone()); - } - } - h - } - - fn predict_url(&self) -> String { - format!( - "{}/models/{}:predictLongRunning", - self.config.base_url, self.model_id - ) - } - - fn operation_url(&self, name: &str) -> String { - format!("{}/{}", self.config.base_url, name) - } } #[async_trait] impl VideoModel for GoogleVideoModel { fn provider(&self) -> &str { - "google" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } fn max_videos_per_call(&self) -> Option { Some(1) } @@ -88,6 +58,7 @@ impl VideoModel for GoogleVideoModel { } 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())); } @@ -109,25 +80,12 @@ impl VideoModel for GoogleVideoModel { "parameters": Value::Object(parameters), }); - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - - let url = self.predict_url(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let url = exchange.url(&format!("/models/{}:predictLongRunning", self.model_id)); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url, - headers: header_list.clone(), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, - body, + exchange.request(url, options), + exchange.transform_body(body), aimux_provider_utils::create_json_response_handler(), super::google_failed_response_handler(), ) @@ -139,9 +97,7 @@ impl VideoModel for GoogleVideoModel { .get("name") .and_then(|v| v.as_str()) .ok_or_else(|| { - AiMuxError::InvalidResponseData( - "Google video prediction missing operation name".to_string(), - ) + AiMuxError::InvalidResponseData("No operation name returned from API".to_string()) })? .to_string(); @@ -171,15 +127,10 @@ impl VideoModel for GoogleVideoModel { ) })?; - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - - let poll_url = self.operation_url(operation_name); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let poll_url = exchange.url(&format!("/{operation_name}")); let resp = aimux_provider_utils::get_from_api( - HttpRequest::new(poll_url.clone(), header_list, options), + exchange.request(poll_url.clone(), options), aimux_provider_utils::create_json_response_handler::(), super::google_failed_response_handler(), ) @@ -188,8 +139,9 @@ impl VideoModel for GoogleVideoModel { let response_headers = resp.response_headers; let response_body = resp.raw_value.as_ref().map(ToString::to_string); let raw_body: Value = resp.value; - // Check the in-band error first: a terminal response may carry both - // done:true and an error object (provider-declared failure). + 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") { let msg = err .get("message") @@ -205,40 +157,71 @@ impl VideoModel for GoogleVideoModel { ..ApiCallError::new(msg, poll_url, serde_json::json!({})) }))); } - if raw_body.get("done").and_then(serde_json::Value::as_bool) != Some(true) { - return Ok(VideoOperationStatus::Pending); - } - // Extract videos from response. - let videos: Vec = raw_body + let outputs = raw_body .get("response") - .and_then(|r| r.get("videos")) - .and_then(|v| v.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|v| { - v.get("gcsUri") - .or_else(|| v.get("url")) - .and_then(|u| u.as_str()) - .map(|url| VideoData::Url { - url: url.to_string(), - media_type: "video/mp4".to_string(), - }) - }) - .collect() - }) - .unwrap_or_default(); + .and_then(|r| r.get("generateVideoResponse")) + .and_then(|r| r.get("generatedSamples")) + .and_then(Value::as_array) + .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 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(), + }); + video_metadata.push(json!({"uri": uri})); + } + } if videos.is_empty() { return Err(AiMuxError::InvalidResponseData( - "Google video operation completed without any video output".to_string(), + "No valid videos in response".to_string(), )); } Ok(VideoOperationStatus::Completed(VideoResult { videos, warnings: Vec::new(), - provider_metadata: None, + provider_metadata: Some( + [( + crate::google::options::GOOGLE.to_string(), + [("videos".to_string(), json!(video_metadata))] + .into_iter() + .collect(), + )] + .into_iter() + .collect(), + ), response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), diff --git a/aimux-providers/src/google_pse.rs b/aimux-providers/src/google_pse.rs index cfe817bf9..807379a71 100644 --- a/aimux-providers/src/google_pse.rs +++ b/aimux-providers/src/google_pse.rs @@ -8,20 +8,26 @@ //! key (`GOOGLE_API_KEY`) and a search-engine ID / `cx` (`GOOGLE_CSE_ID`), //! both passed as query parameters. The `cx` may also be supplied at call //! time via `provider_options["google_pse"]["cx"]`. +//! +//! [`create_google_pse`] takes [`GooglePseProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`GooglePseProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `GOOGLE_API_KEY`. +//! [`google_pse()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::Value; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; /// Fixed model ID for the Google PSE search model. const MODEL_ID: &str = "google-pse-search"; @@ -47,102 +53,149 @@ fn google_pse_failed_response_handler() -> aimux_provider_utils::ResponseHandler }) } -/// Configuration for the Google PSE provider. -#[derive(Debug, Clone)] -pub struct GooglePseConfig { - pub api_key: String, - /// Search-engine ID (`cx`). May be `None` here and supplied per-call via - /// `provider_options["google_pse"]["cx"]`. - pub cx: Option, - pub base_url: String, -} +pub(crate) mod options; -impl GooglePseConfig { - /// Create from an API key (uses the default Google Custom Search base URL). - /// `cx` defaults to `None`. - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - cx: None, - base_url: "https://www.googleapis.com/customsearch/v1".to_string(), - } - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://www.googleapis.com/customsearch/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "GOOGLE_API_KEY"; +const CX_ENV_VAR: &str = "GOOGLE_CSE_ID"; +const DEFAULT_NAME: &str = "google_pse"; - /// Set the search-engine ID (`cx`). - #[must_use] - pub fn with_cx(mut self, cx: impl Into) -> Self { - self.cx = Some(cx.into()); - self - } +/// Settings of [`create_google_pse`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct GooglePseProviderSettings { + /// The endpoint URL, query parameters excluded. Default + /// `https://www.googleapis.com/customsearch/v1`; a trailing slash is + /// removed. + pub base_url: Option, + /// The API key, sent as the `key` query parameter. `None` loads + /// `GOOGLE_API_KEY` when a request is made and fails that request with + /// `AiMuxError::LoadApiKey` if it is unset. An explicit value is used as + /// given, `""` included: it never falls back to the environment. + pub api_key: Option>, + /// The search-engine ID (`cx`). `None` uses + /// `providerOptions.google_pse.cx` of the call, then `GOOGLE_CSE_ID`; + /// a call with none of them fails with `AiMuxError::InvalidArgument`. + pub cx: Option, + /// Extra headers on every request. A `None` value removes the header. + /// Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` string + /// (`"{name}.search"`). Default `"google_pse"`. The providerOptions key + /// stays `google_pse`. + 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, +} - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self +impl std::fmt::Debug for GooglePseProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GooglePseProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field("cx", &self.cx) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Create from the `GOOGLE_API_KEY` and (optional) `GOOGLE_CSE_ID` - /// environment variables. `GOOGLE_API_KEY` is required; `GOOGLE_CSE_ID` - /// may be omitted if `cx` is supplied per-call via `provider_options`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "GOOGLE_API_KEY", "Google PSE")?; - let cx = std::env::var("GOOGLE_CSE_ID").ok(); - Ok(Self::new(api_key).with_maybe_cx(cx)) - } +/// Create a Google PSE provider. +/// +/// # 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. +pub fn create_google_pse( + settings: GooglePseProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(GooglePseProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + credential: Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Google PSE"), + cx: settings.cx, + // The key travels in the query string, so the headers carry none. + headers: credential_headers( + Credential::None, + AuthScheme::Bearer, + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - fn with_maybe_cx(mut self, cx: Option) -> Self { - self.cx = cx; - self - } +/// The default provider: `create_google_pse` with default settings, created +/// on first use. Creating it reads nothing from the environment and cannot +/// fail; a missing key surfaces from the first request instead. +pub fn google_pse() -> &'static GooglePseProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_google_pse(GooglePseProviderSettings::default()) + .expect("default Google PSE settings are always valid") + }) } -/// Google PSE provider — creates `GooglePseSearchModel` instances. -/// -/// Google PSE is a search-only provider; it does not support language models. +/// A Google PSE provider. Search only; it holds no HTTP client. pub struct GooglePseProvider { - config: GooglePseConfig, + name: String, + base_url: String, + credential: Credential, + cx: Option, + headers: HeadersFn, + fetch: Option, } impl GooglePseProvider { - #[must_use] - pub fn new(config: GooglePseConfig) -> Self { - Self { config } - } - - /// Create a search model instance. + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> GooglePseSearchModel { - GooglePseSearchModel::new(self.config.clone()) + GooglePseSearchModel::from_config( + EndpointConfig::fixed( + format!("{}.search", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ), + self.credential.clone(), + self.cx.clone(), + ) } } -impl Provider for GooglePseProvider { - fn name(&self) -> &str { - "google_pse" - } -} +crate::impl_single_modality_provider!(GooglePseProvider, search_model, |p, _id| p.search_model()); -/// Resolve the `cx` (search-engine ID): prefer the config value, then -/// `provider_options["google_pse"]["cx"]`. +/// Resolve the `cx` (search-engine ID): the `cx` setting, then +/// `provider_options["google_pse"]["cx"]`, then `GOOGLE_CSE_ID`. fn resolve_cx( - config_cx: Option<&str>, + setting_cx: Option<&str>, provider_options: Option<&SharedProviderOptions>, ) -> Option { - if let Some(cx) = config_cx { + if let Some(cx) = setting_cx { return Some(cx.to_string()); } - if let Some(po) = provider_options - && let Some(google_pse) = po.get("google_pse") - && let Some(cx) = google_pse.get("cx").and_then(|v| v.as_str()) + if let Some(cx) = options::google_pse_options(provider_options) + .and_then(|options| options.get("cx")) + .and_then(Value::as_str) { return Some(cx.to_string()); } - None + std::env::var(CX_ENV_VAR).ok() } /// A single Google PSE result item. All fields are optional so unknown-but- @@ -180,26 +233,29 @@ fn map_results(entries: Vec) -> Vec { /// A Google PSE search model. pub struct GooglePseSearchModel { - config: GooglePseConfig, + config: EndpointConfig, + credential: Credential, + cx: Option, } impl GooglePseSearchModel { - #[must_use] - pub fn new(config: GooglePseConfig) -> Self { - Self { config } - } - - fn endpoint(&self) -> String { - // The base URL already includes the `/customsearch/v1` path; query - // parameters (key, cx, q, num) are appended by the caller. - self.config.base_url.clone() + pub(crate) fn from_config( + config: EndpointConfig, + credential: Credential, + cx: Option, + ) -> Self { + Self { + config, + credential, + cx, + } } } #[async_trait] impl SearchModel for GooglePseSearchModel { fn provider(&self) -> &str { - "google_pse" + &self.config.provider } fn model_id(&self) -> &str { @@ -207,31 +263,26 @@ impl SearchModel for GooglePseSearchModel { } async fn do_search(&self, options: &SearchCallOptions) -> Result { - let cx = resolve_cx(self.config.cx.as_deref(), options.provider_options.as_ref()) + let cx = resolve_cx(self.cx.as_deref(), options.provider_options.as_ref()) .ok_or_else(|| { AiMuxError::InvalidArgument( - "Google PSE requires a `cx` (search-engine ID). Set the `GOOGLE_CSE_ID` \ + "Google PSE requires a `cx` (search-engine ID). Set the `cx` setting, the `GOOGLE_CSE_ID` \ environment variable or pass it via `provider_options[\"google_pse\"][\"cx\"]`." .to_string(), ) })?; - // Google PSE authenticates via query parameters; only forward - // user-supplied extra headers. - let headers: Vec<(String, String)> = options - .headers - .as_ref() - .map(|extra: &HashMap| { - extra.iter().map(|(k, v)| (k.clone(), v.clone())).collect() - }) - .unwrap_or_default(); + // Google PSE authenticates via query parameters, so the provider headers + // carry no credential. + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let api_key = self.credential.secret().await?.unwrap_or_default(); - let mut url = url::Url::parse(&self.endpoint()).map_err(|e| { + let mut url = url::Url::parse(exchange.base_url()).map_err(|e| { AiMuxError::InvalidArgument(format!("invalid google_pse endpoint: {e}")) })?; { let mut qp = url.query_pairs_mut(); - qp.append_pair("key", &self.config.api_key) + qp.append_pair("key", &api_key) .append_pair("cx", &cx) .append_pair("q", &options.query); if let Some(num) = options.max_results { @@ -240,15 +291,7 @@ impl SearchModel for GooglePseSearchModel { } let resp = aimux_provider_utils::get_from_api( - HttpRequest { - url: url.to_string(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(url.to_string(), options), aimux_provider_utils::create_json_response_handler::(), google_pse_failed_response_handler(), ) diff --git a/aimux-providers/src/google_pse/options.rs b/aimux-providers/src/google_pse/options.rs new file mode 100644 index 000000000..8a14c4a0a --- /dev/null +++ b/aimux-providers/src/google_pse/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Google PSE. +//! +//! Google PSE reads `providerOptions.google_pse`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "google_pse"; + +/// The Google PSE options in a providerOptions map. +pub(crate) fn google_pse_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/groq/browser_search_models.rs b/aimux-providers/src/groq/browser_search_models.rs new file mode 100644 index 000000000..3495a264c --- /dev/null +++ b/aimux-providers/src/groq/browser_search_models.rs @@ -0,0 +1,18 @@ +//! Models that support Groq's browser search tool. +//! +//! Mirrors `groq-browser-search-models.ts`. Based on +//! . + +/// Models that support browser search functionality. +pub(crate) const BROWSER_SEARCH_SUPPORTED_MODELS: [&str; 2] = + ["openai/gpt-oss-20b", "openai/gpt-oss-120b"]; + +/// Whether a model supports browser search functionality. +pub(crate) fn is_browser_search_supported_model(model_id: &str) -> bool { + BROWSER_SEARCH_SUPPORTED_MODELS.contains(&model_id) +} + +/// A formatted list of the supported models for warning messages. +pub(crate) fn get_supported_models_string() -> String { + BROWSER_SEARCH_SUPPORTED_MODELS.join(", ") +} diff --git a/aimux-providers/src/groq/convert.rs b/aimux-providers/src/groq/convert.rs new file mode 100644 index 000000000..e8594f711 --- /dev/null +++ b/aimux-providers/src/groq/convert.rs @@ -0,0 +1,152 @@ +//! Mirrors `convert-to-groq-chat-messages.ts`. + +use base64::Engine; +use serde_json::{Map, Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultPart, UserPart, +}; +use aimux_core::shared::{FileBytes, FileData}; +use aimux_provider_utils::resolve_full_media_type; + +fn is_image(media_type: &str) -> bool { + media_type.split('/').next() == Some("image") +} + +fn image_url_part(media_type: &str, base64: &str) -> Value { + json!({ + "type": "image_url", + "image_url": { + "url": format!("data:{media_type};base64,{base64}"), + }, + }) +} + +fn non_image() -> AiMuxError { + AiMuxError::UnsupportedFunctionality("Non-image file content parts".to_string()) +} + +fn convert_user_part(part: &UserPart) -> Result { + Ok(match part { + UserPart::Text(TextPart { text, .. }) => json!({ "type": "text", "text": text }), + UserPart::File(file) => match &file.data { + FileData::Reference { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "file parts with provider references".to_string(), + )); + } + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "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::Data { data } => { + let data = match data { + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) + } + FileBytes::Base64(data) => data.clone(), + }; + image_url_part(&resolve_full_media_type(file)?, &data) + } + }, + }) +} + +fn convert_assistant_message(content: &[AssistantPart]) -> Value { + let mut text = String::new(); + let mut reasoning = String::new(); + let mut tool_calls = Vec::new(); + for part in content { + match part { + 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, + tool_name, + input, + .. + }) => tool_calls.push(json!({ + "id": tool_call_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": input.to_string(), + }, + })), + _ => {} + } + } + + let mut message = Map::new(); + message.insert("role".into(), json!("assistant")); + message.insert("content".into(), json!(text)); + if !reasoning.is_empty() { + message.insert("reasoning".into(), json!(reasoning)); + } + if !tool_calls.is_empty() { + message.insert("tool_calls".into(), Value::Array(tool_calls)); + } + Value::Object(message) +} + +/// The `messages` of the request body. +/// +/// # Errors +/// +/// `UnsupportedFunctionality` for a non-image file or an undetectable image subtype. +pub(crate) fn convert_to_groq_chat_messages( + prompt: &LanguageModelPrompt, +) -> Result, AiMuxError> { + let mut messages = Vec::new(); + + for message in prompt { + match message { + LanguageModelMessage::System { content, .. } => { + messages.push(json!({ "role": "system", "content": content })); + } + LanguageModelMessage::User { content, .. } => { + if let [UserPart::Text(TextPart { text, .. })] = content.as_slice() { + messages.push(json!({ "role": "user", "content": text })); + continue; + } + let parts = content + .iter() + .map(convert_user_part) + .collect::, _>>()?; + messages.push(json!({ "role": "user", "content": parts })); + } + LanguageModelMessage::Assistant { content, .. } => { + messages.push(convert_assistant_message(content)) + } + LanguageModelMessage::Tool { content, .. } => { + for part in content { + let ToolPart::ToolResult(ToolResultPart { + tool_call_id, + output, + .. + }) = part + else { + continue; + }; + messages.push(json!({ + "role": "tool", + "tool_call_id": tool_call_id, + "content": crate::openai::convert::tool_result_to_content(output), + })); + } + } + } + } + + Ok(messages) +} diff --git a/aimux-providers/src/groq/error.rs b/aimux-providers/src/groq/error.rs new file mode 100644 index 000000000..2baa7063e --- /dev/null +++ b/aimux-providers/src/groq/error.rs @@ -0,0 +1,40 @@ +//! Mirrors `groq-error.ts`: the `{ "error": { "message", "type" } }` shape of +//! a failed response. + +use serde::Deserialize; +use serde_json::Value; + +use aimux_core::error::AiMuxError; +use aimux_provider_utils::{ + ProviderErrorParts, ResponseHandler, create_json_error_response_handler, +}; + +/// `groqFailedResponseHandler` (groq@4.0.52 `groq-error.ts:4-16`): the +/// message is `error.message`; `error.type` is the provider code. +pub(crate) fn groq_failed_response_handler() -> ResponseHandler { + create_json_error_response_handler(|data| { + let error = data.get("error"); + let field = |name: &str| { + error + .and_then(|error| error.get(name)) + .and_then(Value::as_str) + .map(str::to_owned) + }; + ProviderErrorParts { + message: field("message").unwrap_or_default(), + provider_code: field("type"), + } + }) +} + +/// `groqErrorDataSchema`, also the shape of an error chunk of a stream. +#[derive(Debug, Deserialize)] +pub(crate) struct GroqErrorData { + pub error: GroqError, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqError { + pub message: String, + pub r#type: String, +} diff --git a/aimux-providers/src/groq/finish_reason.rs b/aimux-providers/src/groq/finish_reason.rs new file mode 100644 index 000000000..56b558ca5 --- /dev/null +++ b/aimux-providers/src/groq/finish_reason.rs @@ -0,0 +1,14 @@ +//! Mirrors `map-groq-finish-reason.ts`. + +use aimux_core::types::FinishReasonUnified; + +/// Map Groq's `finish_reason` to the unified finish reason. +pub(crate) fn map_groq_finish_reason(finish_reason: Option<&str>) -> FinishReasonUnified { + match finish_reason { + Some("stop") => FinishReasonUnified::Stop, + Some("length") => FinishReasonUnified::Length, + Some("content_filter") => FinishReasonUnified::ContentFilter, + Some("function_call" | "tool_calls") => FinishReasonUnified::ToolCalls, + _ => FinishReasonUnified::Other, + } +} diff --git a/aimux-providers/src/groq/mod.rs b/aimux-providers/src/groq/mod.rs new file mode 100644 index 000000000..1fbf49a51 --- /dev/null +++ b/aimux-providers/src/groq/mod.rs @@ -0,0 +1,179 @@ +//! Groq provider. +//! +//! [`create_groq`] is the Rust form of the AI SDK's `createGroq`: it takes +//! [`GroqProviderSettings`], validates the base URL, fixes the provider name +//! and returns a [`GroqProvider`] whose chat models are Groq's own +//! [`GroqChatLanguageModel`] (`"groq.chat"`). The API key is not read there; +//! it is loaded in the request headers of every call, from the setting or +//! from `GROQ_API_KEY`. [`groq()`] is the default instance. +//! +//! The files mirror the package's: `model.rs` is +//! `groq-chat-language-model.ts`, `convert.rs` is +//! `convert-to-groq-chat-messages.ts`, `prepare_tools.rs` is +//! `groq-prepare-tools.ts`, `usage.rs` is `convert-groq-usage.ts`, +//! `finish_reason.rs` is `map-groq-finish-reason.ts`, `options.rs` is +//! `groq-chat-language-model-options.ts`, `error.rs` is `groq-error.ts` and +//! `browser_search_models.rs` is `groq-browser-search-models.ts`. + +mod browser_search_models; +mod convert; +mod error; +mod finish_reason; +mod model; +mod options; +mod prepare_tools; +mod types; +mod usage; + +pub use model::GroqChatLanguageModel; + +use std::sync::{Arc, OnceLock}; + +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::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.groq.com/openai/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "GROQ_API_KEY"; +const DEFAULT_NAME: &str = "groq"; + +/// Settings of [`create_groq`] (the AI SDK's `GroqProviderSettings`). +#[derive(Clone, Default)] +pub struct GroqProviderSettings { + /// Base URL for the API calls. Default `https://api.groq.com/openai/v1`; + /// a trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `GROQ_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. + pub api_key: Option, + /// Extra headers on every request; a `None` value removes the header. + pub headers: Option, + /// The transport. `None` uses the process default, resolved per request. + pub fetch: Option, +} + +impl std::fmt::Debug for GroqProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GroqProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key.is_some()) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("fetch", &self.fetch.is_some()) + .finish() + } +} + +/// Create a Groq provider. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when `base_url` is not an `http(s)` +/// URL with a host. The key is not read here. +pub fn create_groq(settings: GroqProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(GroqProvider { + 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, + "Groq", + ), + Vec::new(), + settings.headers, + ), + "groq", + "4.0.52", + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_groq` with default settings, created on first +/// use. Creating it reads nothing from the environment and cannot fail. +pub fn groq() -> &'static GroqProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_groq(GroqProviderSettings::default()) + .expect("default Groq settings are always valid") + }) +} + +/// A Groq provider (the AI SDK's `GroqProvider`). Groq serves language models +/// only here: embedding and image models are `NoSuchModel`. +pub struct GroqProvider { + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, +} + +impl GroqProvider { + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) + } + + /// A chat model; `provider()` is `"{name}.chat"` (`"groq.chat"`). + #[must_use] + pub fn chat(&self, model_id: &str) -> GroqChatLanguageModel { + GroqChatLanguageModel::from_config(model_id.to_string(), self.model_config("chat")) + } + + /// The provider as a function: the default language model, the chat model. + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.chat(model_id)) + } +} + +impl Provider for GroqProvider { + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) + } + + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } +} + +impl ProviderDiscovery for GroqProvider { + /// `GET {base_url}/models`: one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config("models"); + Box::pin(async move { + crate::shared::list_data_models(&config, error::groq_failed_response_handler()).await + }) + } +} diff --git a/aimux-providers/src/groq/model.rs b/aimux-providers/src/groq/model.rs new file mode 100644 index 000000000..527b721fa --- /dev/null +++ b/aimux-providers/src/groq/model.rs @@ -0,0 +1,592 @@ +//! The Groq chat language model (`{name}.chat`). +//! +//! Mirrors `groq-chat-language-model.ts`: Groq's own implementation of the +//! chat completions endpoint. Streamed tool calls go through the shared +//! [`StreamingToolCallTracker`]. + +use aimux_core::tool::RawToolCall; +use std::collections::HashMap; + +use async_trait::async_trait; +use futures::StreamExt; +use serde_json::{Map, Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::options::{CallOptions, ResponseFormat}; +use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; +use aimux_core::shared::{RequestInfo, ResponseInfo, StreamResponseInfo}; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::{ + FinishReason, FinishReasonUnified, ReasoningEffort, ResponseMetadata, Warning, +}; +use aimux_provider_utils::{ + StreamingToolCallDelta, StreamingToolCallTracker, TypeValidation, generate_id, +}; + +use crate::shared::EndpointConfig; + +use super::convert::convert_to_groq_chat_messages; +use super::error::{GroqErrorData, groq_failed_response_handler}; +use super::finish_reason::map_groq_finish_reason; +use super::options::{self, GroqLanguageModelChatOptions, parse_groq_options}; +use super::prepare_tools::prepare_tools; +use super::types::{GroqChatChunk, GroqChatResponse}; +use super::usage::convert_groq_usage; + +/// A Groq chat language model. +pub struct GroqChatLanguageModel { + model_id: String, + config: EndpointConfig, +} + +/// `getResponseMetadata`. +fn response_metadata( + id: Option, + created: Option, + model_id: Option, +) -> ResponseMetadata { + ResponseMetadata { + id, + timestamp: created + .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) + .map(|dt| dt.to_rfc3339()), + model_id, + } +} + +/// `getGroqStreamErrorMetadata`: the HTTP status an error type stands for. +/// Whether the error is retryable follows from that status. +fn groq_stream_error_status(error_type: &str) -> Option { + match error_type { + "rate_limit_error" => Some(429), + "api_error" | "internal_server_error" | "server_error" => Some(500), + "overloaded_error" | "service_unavailable" => Some(503), + "timeout" | "timeout_error" => Some(504), + "authentication_error" | "invalid_api_key" => Some(401), + "permission_error" => Some(403), + "not_found_error" | "model_not_found" => Some(404), + "bad_request" | "context_length_exceeded" | "invalid_request_error" => Some(400), + _ => None, + } +} + +/// `createGroqStreamError`. +fn create_groq_stream_error( + data: &Value, + url: &str, + request_body_values: Value, + response_headers: HashMap, +) -> AiMuxError { + match serde_json::from_value::(data.clone()) { + Ok(GroqErrorData { error }) => aimux_provider_utils::stream_error_api_call( + error.message, + Some(error.r#type.clone()), + groq_stream_error_status(&error.r#type), + data, + url, + request_body_values, + response_headers, + ), + Err(error) => error.into(), + } +} + +impl GroqChatLanguageModel { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { + Self { model_id, config } + } + + /// The request body and the warnings raised while building it (`getArgs`). + fn get_args( + &self, + options: &CallOptions, + ) -> Result<(Map, Vec), AiMuxError> { + let mut warnings = Vec::new(); + + let groq_options: GroqLanguageModelChatOptions = parse_groq_options(options)?; + + let structured_outputs = groq_options.structured_outputs.unwrap_or(true); + let strict_json_schema = groq_options.strict_json_schema.unwrap_or(true); + + if options.top_k.is_some() { + warnings.push(Warning::Unsupported { + feature: "topK".to_string(), + details: None, + }); + } + + if let Some(ResponseFormat::Json { + schema: Some(_), .. + }) = &options.response_format + && !structured_outputs + { + warnings.push(Warning::Unsupported { + feature: "responseFormat".to_string(), + details: Some( + "JSON response format schema is only supported with structuredOutputs" + .to_string(), + ), + }); + } + + let prepared = prepare_tools(options, &self.model_id); + + let mut reasoning_effort = groq_options.reasoning_effort; + if reasoning_effort.is_none() + && let Some(reasoning) = options.reasoning.filter(|effort| effort.is_custom()) + { + reasoning_effort = match reasoning { + ReasoningEffort::None if self.model_id == "qwen/qwen3.6-27b" => { + Some(options::ReasoningEffort::None) + } + ReasoningEffort::None => { + warnings.push(Warning::Unsupported { + feature: "reasoning".to_string(), + details: Some(format!( + "reasoning \"{reasoning}\" is not supported by this model." + )), + }); + None + } + _ => { + // `mapReasoningToProviderEffort` with + // minimal -> low, low, medium, high, xhigh -> high. + let (mapped, name) = match reasoning { + ReasoningEffort::Minimal | ReasoningEffort::Low => { + (options::ReasoningEffort::Low, "low") + } + ReasoningEffort::Medium => (options::ReasoningEffort::Medium, "medium"), + _ => (options::ReasoningEffort::High, "high"), + }; + if reasoning.to_string() != name { + warnings.push(Warning::Compatibility { + feature: "reasoning".to_string(), + details: Some(format!( + "reasoning \"{reasoning}\" is not directly supported by this model. mapped to effort \"{name}\"." + )), + }); + } + Some(mapped) + } + }; + } + + let mut body = Map::new(); + // model id: + body.insert("model".into(), json!(self.model_id)); + + // model specific settings: + if let Some(user) = &groq_options.user { + body.insert("user".into(), json!(user)); + } + if let Some(parallel) = groq_options.parallel_tool_calls { + body.insert("parallel_tool_calls".into(), json!(parallel)); + } + + // standardized settings: + if let Some(max_tokens) = options.max_output_tokens { + body.insert("max_tokens".into(), json!(max_tokens)); + } + for (key, value) in [ + ("temperature", options.temperature), + ("top_p", options.top_p), + ("frequency_penalty", options.frequency_penalty), + ("presence_penalty", options.presence_penalty), + ] { + if let Some(value) = value { + body.insert(key.into(), json!(value)); + } + } + if let Some(stop) = &options.stop_sequences { + body.insert("stop".into(), json!(stop)); + } + if let Some(seed) = options.seed { + body.insert("seed".into(), json!(seed)); + } + + // response format: + if let Some(ResponseFormat::Json { + schema, + name, + description, + }) = &options.response_format + { + let response_format = match schema { + Some(schema) if structured_outputs => { + let mut json_schema = Map::new(); + json_schema.insert("schema".into(), schema.clone()); + json_schema.insert("strict".into(), json!(strict_json_schema)); + json_schema.insert( + "name".into(), + json!(name.clone().unwrap_or_else(|| "response".to_string())), + ); + if let Some(description) = description { + json_schema.insert("description".into(), json!(description)); + } + json!({ "type": "json_schema", "json_schema": json_schema }) + } + _ => json!({ "type": "json_object" }), + }; + body.insert("response_format".into(), response_format); + } + + // provider options: + if let Some(format) = groq_options.reasoning_format { + body.insert("reasoning_format".into(), json!(format)); + } + if let Some(effort) = reasoning_effort { + body.insert("reasoning_effort".into(), json!(effort)); + } + if let Some(tier) = groq_options.service_tier { + body.insert("service_tier".into(), json!(tier)); + } + + // messages: + body.insert( + "messages".into(), + Value::Array(convert_to_groq_chat_messages(&options.prompt)?), + ); + + // tools: + if let Some(tools) = prepared.tools { + body.insert("tools".into(), Value::Array(tools)); + } + if let Some(tool_choice) = prepared.tool_choice { + body.insert("tool_choice".into(), tool_choice); + } + + warnings.extend(prepared.tool_warnings); + Ok((body, warnings)) + } +} + +#[async_trait] +impl LanguageModel for GroqChatLanguageModel { + /// `"{name}.chat"`. + fn provider(&self) -> &str { + &self.config.provider + } + + fn model_id(&self) -> &str { + &self.model_id + } + + async fn supported_urls(&self) -> SupportedUrls { + let http = regex::Regex::new(r"^https?://.*$").expect("static pattern"); + SupportedUrls(std::iter::once(("image/*".to_string(), vec![http])).collect()) + } + + async fn do_generate(&self, options: &CallOptions) -> Result { + let (args, warnings) = self.get_args(options)?; + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let body = exchange.transform_body(Value::Object(args)); + + let resp = aimux_provider_utils::post_json_to_api( + exchange.request(exchange.url("/chat/completions"), options), + body.clone(), + aimux_provider_utils::create_json_response_handler::(), + groq_failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers; + let response = resp.value; + + let choice = response.choices.into_iter().next().ok_or_else(|| { + AiMuxError::InvalidResponseData("Response did not contain any choices.".to_string()) + })?; + + let mut content = Vec::new(); + + // text content: + if let Some(text) = choice.message.content.filter(|text| !text.is_empty()) { + content.push(GenerateContent::Text { + text, + provider_metadata: None, + }); + } + + // reasoning: + if let Some(text) = choice.message.reasoning.filter(|text| !text.is_empty()) { + content.push(GenerateContent::Reasoning(ReasoningOutput { + text, + provider_metadata: None, + })); + } + + // tool calls: + for tool_call in choice.message.tool_calls.into_iter().flatten() { + content.push(GenerateContent::ToolCall(RawToolCall { + tool_call_id: tool_call + .id + .filter(|id| !id.is_empty()) + .unwrap_or_else(generate_id), + tool_name: tool_call.function.name, + input: tool_call.function.arguments, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + } + + Ok(GenerateResult { + content, + finish_reason: FinishReason { + unified: map_groq_finish_reason(choice.finish_reason.as_deref()), + raw: choice.finish_reason, + }, + 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() + }), + }) + } + + async fn do_stream(&self, options: &CallOptions) -> Result { + let (mut args, warnings) = self.get_args(options)?; + args.insert("stream".into(), json!(true)); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let body = exchange.transform_body(Value::Object(args)); + let endpoint = exchange.url("/chat/completions"); + + let resp = aimux_provider_utils::post_json_to_api( + exchange.request(endpoint.clone(), options), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + groq_failed_response_handler(), + ) + .await?; + 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(); + let stream_response_headers = response_headers.clone(); + + let stream = async_stream::stream! { + yield Ok(StreamPart::StreamStart { warnings }); + + let mut tool_calls = StreamingToolCallTracker::new() + .with_generate_id(generate_id) + .with_type_validation(TypeValidation::Required); + let mut tool_parts = Vec::new(); + + let mut finish_reason = FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }; + let mut usage = None; + let mut is_first_chunk = true; + let mut is_active_text = false; + let mut is_active_reasoning = false; + + let mut event_iter = + futures::stream::iter(first_event.into_iter()).chain(sse_stream); + + while let Some(event) = event_iter.next().await { + let parsed = match event { + Ok(parsed) => parsed, + // a chunk that fails to parse is reported as an error part + // (`chunk.success === false`); a transport failure ends the stream. + Err(error) => { + if !error.is_recoverable_stream_error() { + yield Err(error); + return; + } + if emit_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: Value::Null }); + } + finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }; + yield Ok(StreamPart::Error { error }); + continue; + } + }; + + // Emit the raw chunk if requested (before anything else). + if emit_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); + } + + // handle error chunks: + if parsed.get("error").is_some() { + finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }; + yield Ok(StreamPart::Error { + error: create_groq_stream_error( + &parsed, + &stream_error_url, + stream_error_body.clone(), + stream_response_headers.clone(), + ), + }); + continue; + } + + // handle failed chunk parsing / validation: + let value: GroqChatChunk = match serde_json::from_value(parsed) { + Ok(value) => value, + Err(error) => { + finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }; + yield Ok(StreamPart::Error { error: error.into() }); + continue; + } + }; + + if is_first_chunk { + is_first_chunk = false; + let metadata = response_metadata(value.id, value.created, value.model); + yield Ok(StreamPart::ResponseMetadata(metadata)); + } + + if let Some(chunk_usage) = value.x_groq.and_then(|x_groq| x_groq.usage) { + usage = Some(chunk_usage); + } + + let Some(choice) = value.choices.into_iter().next() else { + continue; + }; + + if let Some(reason) = choice.finish_reason { + finish_reason = FinishReason { + unified: map_groq_finish_reason(Some(&reason)), + raw: Some(reason), + }; + } + + let Some(delta) = choice.delta else { + continue; + }; + + if let Some(reasoning) = delta.reasoning.filter(|text| !text.is_empty()) { + if !is_active_reasoning { + yield Ok(StreamPart::ReasoningStart { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + is_active_reasoning = true; + } + yield Ok(StreamPart::ReasoningDelta { + id: "reasoning-0".to_string(), + delta: reasoning, + provider_metadata: None, + }); + } + + if let Some(text) = delta.content.filter(|text| !text.is_empty()) { + // end the active reasoning block before text starts + if is_active_reasoning { + yield Ok(StreamPart::ReasoningEnd { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + is_active_reasoning = false; + } + if !is_active_text { + yield Ok(StreamPart::TextStart { + id: "txt-0".to_string(), + provider_metadata: None, + }); + is_active_text = true; + } + yield Ok(StreamPart::TextDelta { + id: "txt-0".to_string(), + delta: text, + provider_metadata: None, + }); + } + + if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { + // end the active reasoning block before tool calls start + if is_active_reasoning { + yield Ok(StreamPart::ReasoningEnd { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + is_active_reasoning = false; + } + for tool_call in &deltas { + let delta = StreamingToolCallDelta { + index: Some(tool_call.index), + id: tool_call.id.as_deref(), + r#type: tool_call.r#type.as_deref(), + name: tool_call.function.name.as_deref(), + arguments: tool_call.function.arguments.as_deref(), + provider_metadata: None, + }; + if let Err(error) = tool_calls.process(delta, &mut tool_parts) { + yield Err(error.into()); + return; + } + for part in tool_parts.drain(..) { + yield Ok(part); + } + } + } + } + + if is_active_reasoning { + yield Ok(StreamPart::ReasoningEnd { + id: "reasoning-0".to_string(), + provider_metadata: None, + }); + } + if is_active_text { + yield Ok(StreamPart::TextEnd { + id: "txt-0".to_string(), + provider_metadata: None, + }); + } + tool_calls.finish(&mut tool_parts); + for part in tool_parts.drain(..) { + yield Ok(part); + } + + yield Ok(StreamPart::Finish { + finish_reason, + usage: convert_groq_usage(usage.as_ref()), + provider_metadata: None, + }); + }; + + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} diff --git a/aimux-providers/src/groq/options.rs b/aimux-providers/src/groq/options.rs new file mode 100644 index 000000000..b86e731d7 --- /dev/null +++ b/aimux-providers/src/groq/options.rs @@ -0,0 +1,100 @@ +//! The providerOptions namespace and the chat options of Groq. +//! +//! Mirrors `groq-chat-language-model-options.ts`. `@ai-sdk/groq` reads +//! `providerOptions.groq`, whatever the provider is named; this module is the +//! one place that spells the key. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use aimux_core::error::AiMuxError; +use aimux_core::options::CallOptions; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "groq"; + +#[derive(Debug, Clone, Copy, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum ReasoningFormat { + Parsed, + Raw, + Hidden, +} + +/// Reasoning effort level for model inference. +/// See . +#[derive(Debug, Clone, Copy, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum ReasoningEffort { + None, + Default, + Low, + Medium, + High, +} + +/// Service tier for the request. +#[derive(Debug, Clone, Copy, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum ServiceTier { + OnDemand, + Performance, + Flex, + Auto, +} + +/// `groqLanguageModelChatOptions`. +#[derive(Debug, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct GroqLanguageModelChatOptions { + #[serde(default, deserialize_with = "deserialize_optional")] + pub reasoning_format: Option, + #[serde(default, deserialize_with = "deserialize_optional")] + pub reasoning_effort: Option, + /// Whether to enable parallel function calling during tool use. Default to true. + #[serde(default, deserialize_with = "deserialize_optional")] + pub parallel_tool_calls: Option, + /// A unique identifier representing the end-user. + #[serde(default, deserialize_with = "deserialize_optional")] + pub user: Option, + /// Whether to use structured outputs. Default true. + #[serde(default, deserialize_with = "deserialize_optional")] + pub structured_outputs: Option, + /// Whether to use strict JSON schema validation. Only used when structured + /// outputs are enabled and a schema is provided. Default true. + #[serde(default, deserialize_with = "deserialize_optional")] + pub strict_json_schema: Option, + #[serde(default, deserialize_with = "deserialize_optional")] + pub service_tier: Option, +} + +fn deserialize_optional<'de, D, T>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, + T: Deserialize<'de>, +{ + T::deserialize(deserializer).map(Some) +} + +/// Parse the Groq options of a call (`parseProviderOptions` with provider +/// `groq`): only the `groq` namespace is read. +/// +/// # Errors +/// +/// `InvalidArgument` when an option has the wrong type or value. +pub(crate) fn parse_groq_options( + options: &CallOptions, +) -> Result { + crate::openai::convert::parse_chat_provider_options( + options.provider_options.as_ref(), + NAMESPACE, + |groq| { + serde_json::from_value(Value::Object(groq.clone())).map_err(|error| { + AiMuxError::InvalidArgument(format!( + "invalid argument for parameter providerOptions: {error}" + )) + }) + }, + ) + .map(Option::unwrap_or_default) +} diff --git a/aimux-providers/src/groq/prepare_tools.rs b/aimux-providers/src/groq/prepare_tools.rs new file mode 100644 index 000000000..d36b97809 --- /dev/null +++ b/aimux-providers/src/groq/prepare_tools.rs @@ -0,0 +1,78 @@ +//! Mirrors `groq-prepare-tools.ts`. + +use serde_json::{Map, Value, json}; + +use aimux_core::options::CallOptions; +use aimux_core::tool::{LanguageModelV4ToolChoice, Tool}; +use aimux_core::types::Warning; + +use super::browser_search_models::{ + get_supported_models_string, is_browser_search_supported_model, +}; + +/// `tools` / `tool_choice` ready for the request body, with the warnings +/// raised while preparing them. +#[derive(Debug, Default)] +pub(crate) struct PreparedTools { + pub tools: Option>, + pub tool_choice: Option, + pub tool_warnings: Vec, +} + +pub(crate) fn prepare_tools(options: &CallOptions, model_id: &str) -> PreparedTools { + let mut prepared = PreparedTools::default(); + + // An empty tools array is the same as none, to prevent errors. + let Some(tools) = options.tools.as_ref().filter(|tools| !tools.is_empty()) else { + return prepared; + }; + + let mut groq_tools = Vec::new(); + for tool in tools { + match tool { + Tool::Provider(tool) if tool.id == "groq.browser_search" => { + if is_browser_search_supported_model(model_id) { + groq_tools.push(json!({ "type": "browser_search" })); + } else { + prepared.tool_warnings.push(Warning::Unsupported { + feature: format!("provider-defined tool {}", tool.id), + details: Some(format!( + "Browser search is only supported on the following models: {}. Current model: {model_id}", + get_supported_models_string() + )), + }); + } + } + Tool::Provider(tool) => prepared.tool_warnings.push(Warning::Unsupported { + feature: format!("provider-defined tool {}", tool.id), + details: None, + }), + Tool::Function(tool) => { + let mut function = Map::new(); + function.insert("name".into(), json!(tool.name)); + if let Some(description) = &tool.description { + function.insert("description".into(), json!(description)); + } + function.insert("parameters".into(), tool.input_schema.clone()); + if let Some(strict) = tool.strict { + function.insert("strict".into(), json!(strict)); + } + groq_tools.push(json!({ "type": "function", "function": function })); + } + } + } + prepared.tools = Some(groq_tools); + + prepared.tool_choice = options + .tool_choice + .as_ref() + .map(|tool_choice| match tool_choice { + LanguageModelV4ToolChoice::Auto => json!("auto"), + LanguageModelV4ToolChoice::None => json!("none"), + LanguageModelV4ToolChoice::Required => json!("required"), + LanguageModelV4ToolChoice::Tool { tool_name } => { + json!({ "type": "function", "function": { "name": tool_name } }) + } + }); + prepared +} diff --git a/aimux-providers/src/groq/types.rs b/aimux-providers/src/groq/types.rs new file mode 100644 index 000000000..76a2d597c --- /dev/null +++ b/aimux-providers/src/groq/types.rs @@ -0,0 +1,151 @@ +//! The response shapes the chat model reads. +//! +//! Mirrors the `groqChatResponseSchema` and `groqChatChunkSchema` of +//! `groq-chat-language-model.ts`, limited to what the implementation needs. + +use serde::{Deserialize, Serialize}; + +// ── Usage ─────────────────────────────────────────────────────────────────── + +#[derive(Debug, Default, Deserialize, Serialize)] +pub(crate) struct GroqUsage { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub completion_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub total_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt_tokens_details: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub completion_tokens_details: Option, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(crate) struct PromptTokensDetails { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cached_tokens: Option, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(crate) struct CompletionTokensDetails { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning_tokens: Option, +} + +// ── Non-streaming response ────────────────────────────────────────────────── + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqChatResponse { + #[serde(default)] + pub id: Option, + #[serde(default)] + pub created: Option, + #[serde(default)] + pub model: Option, + pub choices: Vec, + #[serde(default)] + pub usage: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqChatChoice { + pub message: GroqChatMessage, + #[serde(default)] + pub finish_reason: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqChatMessage { + #[serde(default)] + pub content: Option, + #[serde(default)] + pub reasoning: Option, + #[serde(default)] + pub tool_calls: Option>, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqToolCall { + #[serde(default)] + pub id: Option, + pub function: GroqToolCallFunction, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqToolCallFunction { + pub name: String, + pub arguments: String, +} + +// ── Streaming response ────────────────────────────────────────────────────── + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqChatChunk { + #[serde(default)] + pub id: Option, + #[serde(default)] + pub created: Option, + #[serde(default)] + pub model: Option, + pub choices: Vec, + #[serde(default)] + pub x_groq: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqChunkChoice { + #[serde(default)] + pub delta: Option, + #[serde(default)] + pub finish_reason: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqChunkDelta { + #[serde(default)] + pub content: Option, + #[serde(default)] + pub reasoning: Option, + #[serde(default)] + pub tool_calls: Option>, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqChunkToolCall { + pub index: usize, + #[serde(default)] + pub id: Option, + #[serde( + default, + rename = "type", + deserialize_with = "deserialize_tool_call_type" + )] + pub r#type: Option, + pub function: GroqChunkToolCallFunction, +} + +fn deserialize_tool_call_type<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let value = String::deserialize(deserializer)?; + if value != "function" { + return Err(serde::de::Error::custom("expected tool call type function")); + } + Ok(Some(value)) +} + +#[derive(Debug, Deserialize)] +pub(crate) struct GroqChunkToolCallFunction { + #[serde(default)] + pub name: Option, + #[serde(default)] + pub arguments: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct XGroq { + #[serde(default)] + pub usage: Option, +} diff --git a/aimux-providers/src/groq/usage.rs b/aimux-providers/src/groq/usage.rs new file mode 100644 index 000000000..c3995db22 --- /dev/null +++ b/aimux-providers/src/groq/usage.rs @@ -0,0 +1,48 @@ +//! Mirrors `convert-groq-usage.ts`. + +use aimux_core::types::{InputTokenUsage, OutputTokenUsage, Usage}; + +use super::types::GroqUsage; + +/// Convert Groq usage to the core `Usage`. A missing object is the null usage +/// (every count unknown). +pub(crate) fn convert_groq_usage(usage: Option<&GroqUsage>) -> Usage { + let Some(usage) = usage else { + return Usage::default(); + }; + + let prompt_tokens = usage.prompt_tokens.unwrap_or(0); + let cache_read_tokens = usage + .prompt_tokens_details + .as_ref() + .and_then(|details| details.cached_tokens); + let completion_tokens = usage.completion_tokens.unwrap_or(0); + let reasoning_tokens = usage + .completion_tokens_details + .as_ref() + .and_then(|details| details.reasoning_tokens); + let text_tokens = match reasoning_tokens { + Some(reasoning) => completion_tokens.saturating_sub(reasoning), + None => completion_tokens, + }; + + Usage { + input_tokens: InputTokenUsage { + total: Some(prompt_tokens), + no_cache: Some(match cache_read_tokens { + Some(cached) => prompt_tokens.saturating_sub(cached), + None => prompt_tokens, + }), + cache_read: cache_read_tokens, + cache_write: None, + }, + output_tokens: OutputTokenUsage { + total: Some(completion_tokens), + text: Some(text_tokens), + reasoning: reasoning_tokens, + }, + raw: serde_json::to_value(usage) + .ok() + .and_then(|value| value.as_object().cloned()), + } +} diff --git a/aimux-providers/src/huggingface.rs b/aimux-providers/src/huggingface.rs index 8cf265a47..63a2f88d5 100644 --- a/aimux-providers/src/huggingface.rs +++ b/aimux-providers/src/huggingface.rs @@ -1,132 +1,210 @@ -//! Hugging Face provider — a thin OpenAI-compatible wrapper. +//! Hugging Face provider. //! -//! Hugging Face exposes an OpenAI-compatible Chat Completions API through its -//! router at `https://router.huggingface.co/v1`. The TS SDK configures this base -//! URL and the `HUGGINGFACE_API_KEY` environment variable. The Rust -//! [`OpenAIProvider`](crate::openai::OpenAIProvider) appends `/chat/completions` -//! to the configured base URL, yielding -//! `https://router.huggingface.co/v1/chat/completions`. Everything else is -//! delegated to the shared `OpenAIProvider`. +//! [`create_huggingface`] is the Rust form of the AI SDK's `createHuggingFace`: +//! it takes [`HuggingFaceProviderSettings`], validates the base URL, fixes the +//! provider name and returns a [`HuggingFaceProvider`]. The API key is not read +//! there; it is loaded in the request headers of every call, from the setting +//! or from `HUGGINGFACE_API_KEY`. [`huggingface()`] is the default instance. //! -//! In addition to the Chat Completions API, Hugging Face also exposes a -//! Responses API (the lightest Responses implementation — function tools only, -//! no built-in tools). See [`responses::HuggingFaceResponsesModel`]. - +//! The AI SDK's package serves the Responses API only +//! ([`responses::HuggingFaceResponsesModel`], the lightest Responses +//! implementation: function tools only, no built-in tools), and +//! [`language_model`](Provider::language_model) returns it. There is no Chat +//! Completions model in this package; the router's OpenAI-compatible endpoint +//! is reachable through `create_openai_compatible` with +//! `https://router.huggingface.co/v1`. + +pub(crate) mod options; pub mod responses; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::load_api_key; - -use crate::openai::{OpenAIConfig, OpenAIModel}; +pub use crate::shared::TransformRequestBody; +pub use responses::HuggingFaceResponsesModel; -const DEFAULT_BASE_URL: &str = "https://router.huggingface.co/v1"; -const ENV_VAR: &str = "HUGGINGFACE_API_KEY"; -const PROVIDER_NAME: &str = "huggingface"; +use std::sync::{Arc, OnceLock}; -/// Configuration for the Hugging Face provider (wraps [`OpenAIConfig`]). -#[derive(Debug, Clone)] -pub struct HuggingFaceConfig(OpenAIConfig); +use futures::future::BoxFuture; -impl HuggingFaceConfig { - /// Create from an API key, using the default Hugging Face base URL. - pub fn new(api_key: impl Into) -> Self { - Self(OpenAIConfig::new(api_key).with_base_url(DEFAULT_BASE_URL)) - } - - /// Create from the `HUGGINGFACE_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `HUGGINGFACE_API_KEY` is not - /// set. - pub fn from_env() -> Result { - let key = load_api_key(None, ENV_VAR, "Hugging Face")?; - Ok(Self::new(key).with_api_key_source(Some("env:HUGGINGFACE_API_KEY"))) - } +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::LanguageModel; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::openai::config::OpenAIModelConfig; +use crate::shared::{Credential, EndpointConfig, provider_headers}; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://router.huggingface.co/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "HUGGINGFACE_API_KEY"; +const DEFAULT_NAME: &str = "huggingface"; + +/// Settings of [`create_huggingface`] (the AI SDK's +/// `HuggingFaceProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct HuggingFaceProviderSettings { + /// Base URL for the API calls. Default `https://router.huggingface.co/v1`; + /// a trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `HUGGINGFACE_API_KEY` when a request is made + /// and fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including `Authorization`. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings + /// (`"{name}.responses"`). Default `"huggingface"`. The providerOptions + /// key stays `huggingface`. + 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, + /// Rewrites every JSON request body once, after it is serialized and + /// before it is sent. + pub transform_request_body: Option, +} - /// Override the base URL (useful for tests / self-hosted endpoints). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self +impl std::fmt::Debug for HuggingFaceProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("HuggingFaceProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ) + .finish() } +} - /// 标注 api_key 来源(RFC-0023 回放重建用)。透传到内部 `OpenAIConfig`。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.0 = self.0.with_api_key_source(source); - self - } +/// Create a Hugging Face provider. +/// +/// # 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. +pub fn create_huggingface( + settings: HuggingFaceProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(HuggingFaceProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Hugging Face"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + transform_request_body: settings.transform_request_body, + }) +} - /// 内部 `OpenAIConfig` 引用(config_snapshot 复用 OpenAI helper 用,M2b)。 - pub(crate) fn openai_config(&self) -> &OpenAIConfig { - &self.0 - } +/// The default provider: `create_huggingface` with default settings, created +/// on first use. Creating it reads nothing from the environment and cannot +/// fail; a missing key surfaces from the first request instead. +pub fn huggingface() -> &'static HuggingFaceProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_huggingface(HuggingFaceProviderSettings::default()) + .expect("default Hugging Face settings are always valid") + }) } -/// Hugging Face provider — creates [`OpenAIModel`] (chat) and -/// [`responses::HuggingFaceResponsesModel`] (responses) instances pointed at HF. +/// A Hugging Face provider (the AI SDK's `HuggingFaceProvider`). Cheap to +/// clone the models out of; it holds no HTTP client. pub struct HuggingFaceProvider { - config: HuggingFaceConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, + transform_request_body: Option, } impl HuggingFaceProvider { - #[must_use] - pub fn new(config: HuggingFaceConfig) -> Self { - Self { config } + fn endpoint_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + self.transform_request_body.clone(), + ) } - /// Create a chat model instance for the given Hugging Face model id - /// (e.g. `"meta-llama/Llama-3.3-70B-Instruct"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - OpenAIModel::new(model_id.to_string(), self.config.0.clone()) + fn openai_config(&self, method: &str) -> OpenAIModelConfig { + OpenAIModelConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + self.transform_request_body.clone(), + ) } - /// Create a Responses model instance for the given Hugging Face model id. + /// A Responses model (e.g. `"deepseek-ai/DeepSeek-V3-0324"`); + /// `provider()` is `"{name}.responses"`. /// - /// The Hugging Face Responses API is the lightest Responses implementation: - /// it supports function tools only (no built-in tools), and uses the - /// `text.format` field for structured output. + /// The Hugging Face Responses API supports function tools only (no + /// built-in tools), and uses the `text.format` field for structured output. #[must_use] - pub fn responses_model(&self, model_id: &str) -> responses::HuggingFaceResponsesModel { - responses::HuggingFaceResponsesModel::new(model_id.to_string(), self.config.clone()) + pub fn responses(&self, model_id: &str) -> HuggingFaceResponsesModel { + HuggingFaceResponsesModel::from_config( + model_id.to_string(), + self.endpoint_config("responses"), + ) + } + + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; the same model as + /// [`responses`](Self::responses) and + /// [`language_model`](Provider::language_model). + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.responses(model_id)) } } impl Provider for HuggingFaceProvider { - fn name(&self) -> &str { - PROVIDER_NAME + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - fn list_models( - &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - // HuggingFaceProvider holds an OpenAIConfig (not an OpenAIProvider - // directly), so delegate via execute_list_models + catalogue resolve. - let config = self.config.0.clone(); - Box::pin(async move { - let headers = crate::openai::model::build_auth_headers(&config); - let runtime = crate::openai::model::execute_list_models( - &config.base_url, - &headers, - config.retry_config, - ) - .await?; - Ok(runtime) - }) + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } +} + +impl ProviderDiscovery for HuggingFaceProvider { + /// `GET {base_url}/models`: one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.openai_config("models"); + Box::pin(async move { crate::openai::model::list_models_once(&config).await }) } } diff --git a/aimux-providers/src/huggingface/options.rs b/aimux-providers/src/huggingface/options.rs new file mode 100644 index 000000000..da78f824c --- /dev/null +++ b/aimux-providers/src/huggingface/options.rs @@ -0,0 +1,24 @@ +//! The providerOptions and providerMetadata namespace of Hugging Face. +//! +//! `@ai-sdk/huggingface` reads `providerOptions.huggingface` and writes +//! response metadata under `huggingface`, 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; + +/// The providerOptions / providerMetadata key. +pub(crate) const NAMESPACE: &str = "huggingface"; + +/// The Hugging Face options in a providerOptions map. +pub(crate) fn huggingface_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} + +/// `{ "huggingface": payload }`. +pub(crate) fn huggingface_metadata(payload: Value) -> ProviderMetadata { + provider_namespace(NAMESPACE, payload).expect("provider metadata must be an object") +} diff --git a/aimux-providers/src/huggingface/responses.rs b/aimux-providers/src/huggingface/responses.rs index fdbcf4cb6..d9205aa54 100644 --- a/aimux-providers/src/huggingface/responses.rs +++ b/aimux-providers/src/huggingface/responses.rs @@ -31,18 +31,16 @@ use aimux_core::language_model_message::{ }; use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat}; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; -use aimux_core::shared::provider_namespace; use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::Tool; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage, Warning}; -use aimux_provider_utils::HttpRequest; - -use super::HuggingFaceConfig; -use crate::openai::responses::responses_convert::build_header_list; +use aimux_provider_utils::{ + MediaTypeData, detect_media_type, get_top_level_media_type, is_full_media_type, +}; -const PROVIDER_NAME: &str = "huggingface"; +use crate::shared::EndpointConfig; fn huggingface_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -129,74 +127,41 @@ fn huggingface_stream_error( /// A Hugging Face Responses API language model. /// -/// Created via [`super::HuggingFaceProvider::responses_model`]. +/// Created via [`super::HuggingFaceProvider::responses`]. pub struct HuggingFaceResponsesModel { model_id: String, - config: HuggingFaceConfig, + config: EndpointConfig, } impl HuggingFaceResponsesModel { - #[must_use] - pub fn new(model_id: String, config: HuggingFaceConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn endpoint(&self) -> String { - format!("{}/responses", self.config.0.base_url) - } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.0.api_key), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } } #[async_trait] impl LanguageModel for HuggingFaceResponsesModel { fn provider(&self) -> &str { - PROVIDER_NAME + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.openai_config().retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - // M2b: HuggingFace wraps OpenAIConfig — reuse the OpenAI snapshot helper. - crate::openai::config_snapshot_from_config( - self.provider(), - &self.model_id, - self.config.openai_config(), - ) - } - async fn do_generate(&self, options: &CallOptions) -> Result { + let exchange = self.config.exchange(options.headers.as_ref()).await?; let request = build_request_body_with_warnings(&self.model_id, options, false)?; - let body = request.body; - let headers = self.build_headers(options.headers.as_ref()); + let body = exchange.transform_body(request.body); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), build_header_list(&headers), options), + exchange.request(exchange.url("/responses"), options), body.clone(), huggingface_successful_response_handler(), huggingface_failed_response_handler(), ) .await?; - let response_body = resp.raw_value; let response_headers = resp.response_headers; let response: Value = resp.value; @@ -223,28 +188,28 @@ impl LanguageModel for HuggingFaceResponsesModel { finish_reason, usage, warnings: request.warnings, - provider_metadata: Some(provider_namespace( - "huggingface", + provider_metadata: Some(super::options::huggingface_metadata( json!({ "responseId": response_id }), - )?), + )), response: Some(aimux_core::shared::ResponseInfo { id: response_id, timestamp: format_timestamp(created_at), model_id: model, headers: Some(response_headers), - body: response_body, + body: Some(response), }), request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), }) } async fn do_stream(&self, options: &CallOptions) -> Result { + let exchange = self.config.exchange(options.headers.as_ref()).await?; let request = build_request_body_with_warnings(&self.model_id, options, true)?; - let body = request.body; - let headers = self.build_headers(options.headers.as_ref()); + let body = exchange.transform_body(request.body); + let endpoint = exchange.url("/responses"); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), build_header_list(&headers), options), + exchange.request(endpoint.clone(), options), body.clone(), aimux_provider_utils::create_event_source_response_handler::(), huggingface_failed_response_handler(), @@ -258,16 +223,12 @@ impl LanguageModel for HuggingFaceResponsesModel { first_event => first_event, }; if let Some(Ok(event)) = first_event.as_ref() - && let Some(error) = huggingface_stream_error( - event, - &self.endpoint(), - body.clone(), - response_headers.clone(), - ) + && let Some(error) = + huggingface_stream_error(event, &endpoint, body.clone(), response_headers.clone()) { return Err(error); } - let stream_error_url = self.endpoint(); + let stream_error_url = endpoint; let stream_request_body = body.clone(); let stream_response_headers = response_headers.clone(); @@ -344,7 +305,7 @@ impl LanguageModel for HuggingFaceResponsesModel { .to_string(); yield Ok(StreamPart::TextStart { id: id.clone(), - provider_metadata: Some(provider_namespace("huggingface", json!({ "itemId": id })).expect("provider metadata must be an object")), + provider_metadata: Some(super::options::huggingface_metadata(json!({ "itemId": id }))), }); } } @@ -376,7 +337,7 @@ impl LanguageModel for HuggingFaceResponsesModel { .to_string(); yield Ok(StreamPart::ReasoningStart { id: id.clone(), - provider_metadata: Some(provider_namespace("huggingface", json!({ "itemId": id })).expect("provider metadata must be an object")), + provider_metadata: Some(super::options::huggingface_metadata(json!({ "itemId": id }))), }); } _ => {} @@ -545,7 +506,7 @@ impl LanguageModel for HuggingFaceResponsesModel { yield Ok(StreamPart::Finish { finish_reason, usage, - provider_metadata: Some(provider_namespace("huggingface", json!({ "responseId": response_id })).expect("provider metadata must be an object")), + provider_metadata: Some(super::options::huggingface_metadata(json!({ "responseId": response_id }))), }); }; @@ -624,7 +585,7 @@ pub fn build_request_body_with_warnings( let hf_options = options .provider_options .as_ref() - .and_then(|m| m.get("huggingface")); + .and_then(|m| super::options::huggingface_options(Some(m))); let metadata = hf_options.and_then(|o| o.get("metadata")).cloned(); let instructions = hf_options .and_then(|o| o.get("instructions")) @@ -791,8 +752,11 @@ pub fn convert_to_huggingface_responses_messages( } // Tool calls and tool results are handled by the // Responses API — skip (no warning, matching TS). - AssistantPart::ToolCall(_) | AssistantPart::ToolResult(_) => {} - _ => {} + AssistantPart::ToolCall(_) + | AssistantPart::ToolResult(_) + | AssistantPart::File(_) + | AssistantPart::Custom(_) + | AssistantPart::ReasoningFile(_) => {} } } } @@ -847,28 +811,6 @@ fn convert_file_part_url(media_type: &str, url: &str) -> Result &str { - match media_type.find('/') { - Some(idx) => &media_type[..idx], - None => media_type, - } -} - -/// Returns `true` only when the media type has a non-empty, non-wildcard -/// subtype (i.e. matches `type/subtype` where `subtype` is not `*`). -fn is_full_media_type(media_type: &str) -> bool { - match media_type.find('/') { - Some(idx) => { - let subtype = &media_type[idx + 1..]; - !subtype.is_empty() && subtype != "*" - } - None => false, - } -} - /// Resolve a media type to its full `type/subtype` form. /// /// - If already a full media type, return as-is. @@ -879,8 +821,8 @@ fn resolve_full_media_type_base64(media_type: &str, data: &str) -> Result Result])] = &[ - ("image/gif", &[Some(0x47), Some(0x49), Some(0x46)]), - ( - "image/png", - &[Some(0x89), Some(0x50), Some(0x4E), Some(0x47)], - ), - ("image/jpeg", &[Some(0xFF), Some(0xD8)]), - ( - "image/webp", - &[ - Some(0x52), - Some(0x49), - Some(0x46), - Some(0x46), // RIFF - None, - None, - None, - None, // file size (variable) - Some(0x57), - Some(0x45), - Some(0x42), - Some(0x50), // WEBP - ], - ), - ("image/bmp", &[Some(0x42), Some(0x4D)]), - ( - "image/tiff", - &[Some(0x49), Some(0x49), Some(0x2A), Some(0x00)], - ), - ( - "image/tiff", - &[Some(0x4D), Some(0x4D), Some(0x00), Some(0x2A)], - ), - ( - "image/avif", - &[ - Some(0x00), - Some(0x00), - Some(0x00), - Some(0x20), - Some(0x66), - Some(0x74), - Some(0x79), - Some(0x70), - Some(0x61), - Some(0x76), - Some(0x69), - Some(0x66), - ], - ), - ( - "image/heic", - &[ - Some(0x00), - Some(0x00), - Some(0x00), - Some(0x20), - Some(0x66), - Some(0x74), - Some(0x79), - Some(0x70), - Some(0x68), - Some(0x65), - Some(0x69), - Some(0x63), - ], - ), -]; - -/// Detect the full media type from base64-encoded data, considering only -/// signatures for the given top-level type. -fn detect_media_type_from_base64(data: &str, top_level: &str) -> Option { - let signatures = match top_level { - "image" => IMAGE_SIGNATURES, - _ => return None, - }; - - // Decode enough bytes to cover the longest signature (12 bytes). - // 4 base64 chars → 3 bytes; ceil(12 / 3) * 4 = 16 chars. - let max_chars = 16.min(data.len()); - let bytes = base64::engine::general_purpose::STANDARD - .decode(&data[..max_chars]) - .ok()?; - - for (media_type, prefix) in signatures { - if bytes.len() >= prefix.len() - && prefix - .iter() - .enumerate() - .all(|(i, &b)| b.is_none() || bytes[i] == b.unwrap()) - { - return Some(media_type.to_string()); - } - } - - None -} - // ───────────────────────────────────────────────────────────────────────────── // Tool preparation // ───────────────────────────────────────────────────────────────────────────── @@ -1075,10 +915,9 @@ fn build_generate_content(response: &Value) -> Result, AiMu let item_id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); content.push(GenerateContent::Text { text: text.to_string(), - provider_metadata: Some( - provider_namespace("huggingface", json!({ "itemId": item_id })) - .expect("provider metadata must be an object"), - ), + provider_metadata: Some(super::options::huggingface_metadata( + json!({ "itemId": item_id }), + )), }); // Process annotations → source parts. @@ -1108,10 +947,9 @@ fn build_generate_content(response: &Value) -> Result, AiMu let item_id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); content.push(GenerateContent::Reasoning(ReasoningOutput { text: text.to_string(), - provider_metadata: Some( - provider_namespace("huggingface", json!({ "itemId": item_id })) - .expect("provider metadata must be an object"), - ), + provider_metadata: Some(super::options::huggingface_metadata( + json!({ "itemId": item_id }), + )), })); } } diff --git a/aimux-providers/src/hume.rs b/aimux-providers/src/hume.rs index 687b9b71c..728dc2a0d 100644 --- a/aimux-providers/src/hume.rs +++ b/aimux-providers/src/hume.rs @@ -14,8 +14,13 @@ //! //! `language` is not supported and emits a warning. The model ID is always an //! empty string (Hume does not use model IDs for TTS). +//! +//! [`create_hume`] takes [`HumeProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`HumeProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `HUME_API_KEY`. +//! [`hume()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -25,8 +30,9 @@ use aimux_core::shared::{SharedProviderOptions, Warning}; use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; -use aimux_provider_utils::{HttpRequest, load_api_key}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; /// Hume errors carry a top-level `message` with a `code` such as `"E0101"` /// (https://dev.hume.ai/docs/resources/errors); the nested `{error: @@ -63,72 +69,120 @@ fn hume_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, -} +pub(crate) mod options; -impl HumeConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.hume.ai".to_string(), - headers: None, - } - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.hume.ai"; +pub(crate) const API_KEY_ENV_VAR: &str = "HUME_API_KEY"; +const DEFAULT_NAME: &str = "hume"; - /// Override the base URL (for testing or proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = url.into().trim_end_matches('/').to_string(); - self - } +/// Settings of [`create_hume`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct HumeProviderSettings { + /// Base URL for the API calls. Default `https://api.hume.ai`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `HUME_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.speech"`). + /// Default `"hume"`. The providerOptions key stays `hume`. + 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, +} - /// Attach extra headers merged into every request. - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self +impl std::fmt::Debug for HumeProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("HumeProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Create from environment variable `HUME_API_KEY`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `HUME_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "HUME_API_KEY", "Hume")?; - Ok(Self::new(api_key)) - } +/// Create a Hume provider. +/// +/// # 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. +pub fn create_hume(settings: HumeProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(HumeProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Hume"), + AuthScheme::Header("X-Hume-Api-Key"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) } -// ── Provider ───────────────────────────────────────────────────────────────── +/// The default provider: `create_hume` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn hume() -> &'static HumeProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_hume(HumeProviderSettings::default()) + .expect("default Hume settings are always valid") + }) +} -/// Hume provider — creates `HumeSpeechModel` instances. +/// A Hume provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct HumeProvider { - config: HumeConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl HumeProvider { - #[must_use] - pub fn new(config: HumeConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a speech (TTS) model instance. Hume does not use model IDs for - /// TTS, so the model ID is always an empty string. + /// The speech (TTS) model. Hume does not use model ids for TTS, so the model id is always an empty string; `provider()` is `"{name}.speech"`. #[must_use] pub fn speech(&self) -> HumeSpeechModel { - HumeSpeechModel::new(self.config.clone()) + HumeSpeechModel::from_config(self.model_config("speech")) } } +crate::impl_single_modality_provider!(HumeProvider, speech_model, |p, _id| p.speech()); + // ── Speech model ───────────────────────────────────────────────────────────── /// The default voice ID used when no voice is provided. @@ -141,43 +195,22 @@ const SUPPORTED_OUTPUT_FORMATS: &[&str] = &["mp3", "pcm", "wav"]; pub struct HumeSpeechModel { /// Hume does not use model IDs for TTS; this is always an empty string. model_id: String, - config: HumeConfig, + config: EndpointConfig, } impl HumeSpeechModel { - #[must_use] - pub fn new(config: HumeConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { model_id: String::new(), config, } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("X-Hume-Api-Key".to_string(), self.config.api_key.clone()); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/v0/tts/file", self.config.base_url) - } } #[async_trait] impl SpeechModel for HumeSpeechModel { fn provider(&self) -> &str { - "hume.speech" + &self.config.provider } fn model_id(&self) -> &str { @@ -187,21 +220,10 @@ impl SpeechModel for HumeSpeechModel { async fn do_generate(&self, options: &SpeechCallOptions) -> Result { let (body, warnings) = build_request(options)?; - let headers: Vec<(String, String)> = self - .build_headers(options.headers.as_ref()) - .into_iter() - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(exchange.url("/v0/tts/file"), options), Value::Object(body.clone()), aimux_provider_utils::create_binary_response_handler(), hume_failed_response_handler(), @@ -318,7 +340,7 @@ struct HumeSpeechProviderOptions { fn parse_hume_provider_options( options: Option<&SharedProviderOptions>, ) -> Option { - let opts = options.and_then(|opts| opts.get("hume"))?; + let opts = options::hume_options(options)?; let context = opts.get("context").and_then(|c| c.as_object())?; // The context can be either { generationId: "..." } or { utterances: [...] }. diff --git a/aimux-providers/src/hume/options.rs b/aimux-providers/src/hume/options.rs new file mode 100644 index 000000000..137937efd --- /dev/null +++ b/aimux-providers/src/hume/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Hume. +//! +//! Hume reads `providerOptions.hume`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "hume"; + +/// The Hume options in a providerOptions map. +pub(crate) fn hume_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/jev.rs b/aimux-providers/src/jev.rs index 890bf584c..c8071ab48 100644 --- a/aimux-providers/src/jev.rs +++ b/aimux-providers/src/jev.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use async_trait::async_trait; -use serde_json::{Value, json}; +use serde_json::Value; use aimux_core::decision_model::*; use aimux_core::{AiMuxError, Provider}; @@ -11,6 +11,28 @@ use aimux_provider_utils::{HttpRequest, load_api_key}; use crate::systemone::{ResponseProfile, WireResponse}; +/// Failed-response handler of the Jev endpoint: `message`, and `type` +/// else `code` as the provider code, from the body's `error` object, or from +/// the top level when there is none. This is the mapping the endpoint used +/// with the removed shared handler; it is not an AI SDK schema. +fn jev_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(serde_json::Value::as_str) + .map(str::to_owned), + } + }) +} + const JEV_ROUNDING: DecisionRounding = DecisionRounding { probability_decimals: Some(2), score_decimals: Some(2), @@ -22,7 +44,6 @@ pub struct JevConfig { /// Full POST URL for the official API, or an explicitly configured proxy. pub endpoint: String, pub headers: HashMap, - pub retry_config: aimux_core::retry::RetryConfig, pub probability_source: DecisionProbabilitySource, } @@ -33,7 +54,6 @@ impl JevConfig { api_key: api_key.into(), endpoint: "https://api.typesafe.ai/v1/systemone".into(), headers: HashMap::new(), - retry_config: Default::default(), probability_source: DecisionProbabilitySource::Native, } } @@ -64,9 +84,9 @@ impl JevProvider { } impl Provider for JevProvider { - fn name(&self) -> &str { - "jev" - } + crate::__unsupported_required_model!(language_model, LanguageModel, "languageModel"); + crate::__unsupported_required_model!(embedding_model, EmbeddingModel, "embeddingModel"); + crate::__unsupported_required_model!(image_model, ImageModel, "imageModel"); fn decision_model(&self, model_id: &str) -> Result, AiMuxError> { if model_id.trim().is_empty() { @@ -94,23 +114,6 @@ impl DecisionModel for JevDecisionModel { fn model_id(&self) -> &str { &self.model_id } - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - let mut snapshot = - aimux_core::recording::ProviderRecord::minimal(self.provider(), self.model_id()); - snapshot.base_url = Some(self.config.endpoint.clone()); - snapshot.api_key_source = if self.config.api_key.is_empty() { - "none" - } else { - "explicit" - } - .into(); - snapshot.profile = Some(json!({"decision_capabilities": self.capabilities()})); - snapshot.provider_options = Some(json!({"headers": self.config.headers})); - snapshot - } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } fn capabilities(&self) -> DecisionCapabilities { DecisionCapabilities { supports_images: false, @@ -171,7 +174,7 @@ impl DecisionModel for JevDecisionModel { ), body.clone(), aimux_provider_utils::create_json_response_handler::(), - aimux_provider_utils::create_standard_json_error_response_handler(), + jev_failed_response_handler(), ) .await?; let raw = response.raw_value.unwrap_or(Value::Null); diff --git a/aimux-providers/src/jina_ai.rs b/aimux-providers/src/jina_ai.rs index 51ab3dd7e..e39bd999e 100644 --- a/aimux-providers/src/jina_ai.rs +++ b/aimux-providers/src/jina_ai.rs @@ -6,22 +6,28 @@ //! Jina AI is a modality-specific provider: it exposes a native rerank //! protocol (query + documents → relevance scores) and does not support //! language models. +//! +//! [`create_jina_ai`] takes [`JinaAiProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`JinaAiProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `JINA_AI_API_KEY`. +//! [`jina_ai()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::reranking_model::{ RerankingCallOptions, RerankingDocuments, RerankingModel, RerankingRank, RerankingResponse, RerankingResult, }; use aimux_core::shared::SharedProviderOptions; use aimux_core::types::Warning; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; /// Jina AI error response structure: `{ "detail": "...", "code": "..." }`. /// @@ -43,75 +49,119 @@ fn jina_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, -} +pub(crate) mod options; -impl JinaAiConfig { - /// Create from an API key (uses the default Jina AI base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.jina.ai".to_string(), - headers: None, - } - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.jina.ai"; +pub(crate) const API_KEY_ENV_VAR: &str = "JINA_AI_API_KEY"; +const DEFAULT_NAME: &str = "jina_ai"; - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - - /// Attach extra default headers sent with every request. - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } +/// Settings of [`create_jina_ai`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct JinaAiProviderSettings { + /// Base URL for the API calls. Default `https://api.jina.ai`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `JINA_AI_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.reranking"`). + /// Default `"jina_ai"`. The providerOptions key stays `jina`. + 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, +} - /// Create from the `JINA_AI_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `JINA_AI_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "JINA_AI_API_KEY", "Jina AI")?; - Ok(Self::new(api_key)) +impl std::fmt::Debug for JinaAiProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("JinaAiProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } -/// Jina AI provider — creates `JinaAiRerankingModel` instances. +/// Create a Jina AI provider. +/// +/// # Errors /// -/// Jina AI is a rerank-only provider; it does not support language models. +/// 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. +pub fn create_jina_ai(settings: JinaAiProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(JinaAiProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Jina AI"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_jina_ai` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn jina_ai() -> &'static JinaAiProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_jina_ai(JinaAiProviderSettings::default()) + .expect("default Jina AI settings are always valid") + }) +} + +/// A Jina AI provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct JinaAiProvider { - config: JinaAiConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl JinaAiProvider { - #[must_use] - pub fn new(config: JinaAiConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a reranking model instance for the given model name (e.g. - /// `"jina-reranker-v2-base-multilingual"` or `"jina-reranker-v3"`). + /// A reranking model (e.g. `"jina-reranker-v2-base-multilingual"`); `provider()` is `"{name}.reranking"`. #[must_use] pub fn reranking_model(&self, model_id: &str) -> JinaAiRerankingModel { - JinaAiRerankingModel::new(model_id.to_string(), self.config.clone()) + JinaAiRerankingModel::from_config(model_id.to_string(), self.model_config("reranking")) } } -impl Provider for JinaAiProvider { - fn name(&self) -> &str { - "jina_ai" - } -} +crate::impl_single_modality_provider!(JinaAiProvider, reranking_model, |p, id| p + .reranking_model(id)); /// Jina AI provider-specific reranking options. #[derive(Debug, Clone, Default)] @@ -123,8 +173,7 @@ fn parse_jina_reranking_options( provider_options: Option<&SharedProviderOptions>, ) -> JinaRerankingOptions { let mut opts = JinaRerankingOptions::default(); - if let Some(po) = provider_options - && let Some(jina) = po.get("jina") + if let Some(jina) = options::jina_options(provider_options) && let Some(v) = jina .get("returnDocuments") .and_then(serde_json::Value::as_bool) @@ -193,43 +242,19 @@ struct JinaRerankingResult { /// A Jina AI reranking model. pub struct JinaAiRerankingModel { model_id: String, - config: JinaAiConfig, + config: EndpointConfig, } impl JinaAiRerankingModel { - #[must_use] - pub fn new(model_id: String, config: JinaAiConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/v1/rerank", self.config.base_url) - } } #[async_trait] impl RerankingModel for JinaAiRerankingModel { fn provider(&self) -> &str { - "jina_ai" + &self.config.provider } fn model_id(&self) -> &str { @@ -245,21 +270,10 @@ impl RerankingModel for JinaAiRerankingModel { let mut warnings = Vec::new(); let body = build_request_body(&self.model_id, options, &jina_options, &mut warnings); - let headers: Vec<(String, String)> = self - .build_headers(options.headers.as_ref()) - .into_iter() - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(exchange.url("/v1/rerank"), options), body, aimux_provider_utils::create_json_response_handler::(), jina_failed_response_handler(), diff --git a/aimux-providers/src/jina_ai/options.rs b/aimux-providers/src/jina_ai/options.rs new file mode 100644 index 000000000..ba78e0681 --- /dev/null +++ b/aimux-providers/src/jina_ai/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Jina AI. +//! +//! Jina AI reads `providerOptions.jina`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "jina"; + +/// The Jina AI options in a providerOptions map. +pub(crate) fn jina_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/jlama.rs b/aimux-providers/src/jlama.rs deleted file mode 100644 index 5d9252a82..000000000 --- a/aimux-providers/src/jlama.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! Jlama provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `JLAMA_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "JLAMA_BASE_URL"; -const PROVIDER_NAME: &str = "jlama"; -const PLACEHOLDER_API_KEY: &str = "jlama"; - -pub struct JlamaConfig(OpenAIConfig); - -impl JlamaConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `JLAMA_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct JlamaProvider(OpenAIProvider); - -impl JlamaProvider { - #[must_use] - pub fn new(config: JlamaConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for JlamaProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/klingai.rs b/aimux-providers/src/klingai.rs index 7bb6739e2..cb7205a16 100644 --- a/aimux-providers/src/klingai.rs +++ b/aimux-providers/src/klingai.rs @@ -8,8 +8,13 @@ //! 2. GET `/v1/videos/text2video/{id}/{task_id}` — polled by Core via `do_status` //! until `succeeded` //! 3. Return the video URL from the result +//! +//! [`create_klingai`] takes [`KlingAIProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`KlingAIProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `KLINGAI_API_KEY`. +//! [`klingai()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; @@ -20,8 +25,9 @@ use aimux_core::video_model::{ VideoCallOptions, VideoData, VideoModel, VideoOperationStart, VideoOperationStatus, VideoResponse, VideoResult, }; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use crate::shared::{Credential, EndpointConfig, provider_headers}; fn klingai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -42,62 +48,119 @@ fn klingai_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, +pub(crate) mod options; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.klingai.com"; +pub(crate) const API_KEY_ENV_VAR: &str = "KLINGAI_API_KEY"; +const DEFAULT_NAME: &str = "klingai"; + +/// Settings of [`create_klingai`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct KlingAIProviderSettings { + /// Base URL for the API calls. Default `https://api.klingai.com`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `KLINGAI_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.video"`). + /// Default `"klingai"`. The providerOptions key stays `klingai`. + 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, } -impl KlingAIConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.klingai.com".to_string(), - headers: None, - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self +impl std::fmt::Debug for KlingAIProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("KlingAIProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } +/// Create a KlingAI provider. +/// +/// # 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. +pub fn create_klingai(settings: KlingAIProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(KlingAIProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "KlingAI"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `KLINGAI_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "KLINGAI_API_KEY", "KlingAI")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_klingai` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn klingai() -> &'static KlingAIProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_klingai(KlingAIProviderSettings::default()) + .expect("default KlingAI settings are always valid") + }) } +/// A KlingAI provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct KlingAIProvider { - config: KlingAIConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl KlingAIProvider { - #[must_use] - pub fn new(config: KlingAIConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// A video model (e.g. `"kling-v2.6-t2v"`); `provider()` is `"{name}.video"`. #[must_use] pub fn video(&self, model_id: &str) -> KlingAIVideoModel { - KlingAIVideoModel::new(model_id.to_string(), self.config.clone()) + KlingAIVideoModel::from_config(model_id.to_string(), self.model_config("video")) } } +crate::impl_single_modality_provider!(KlingAIProvider, video_model, |p, id| p.video(id)); + // ── Helpers ───────────────────────────────────────────────────────────────── /// Derive the KlingAI video mode from the model ID suffix. @@ -161,52 +224,22 @@ struct KlingAITaskVideos { pub struct KlingAIVideoModel { model_id: String, - config: KlingAIConfig, + config: EndpointConfig, } impl KlingAIVideoModel { - #[must_use] - pub fn new(model_id: String, config: KlingAIConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } +} - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn submit_endpoint(&self, mode: &str) -> String { - let path = match mode { - "i2v" => "/v1/videos/image2video", - "mi2v" => "/v1/videos/multi-image2video", - "motion-control" => "/v1/videos/motion-control", - _ => "/v1/videos/text2video", - }; - format!("{}{path}", self.config.base_url) - } - - fn poll_endpoint(&self, mode: &str, id: &str, task_id: &str) -> String { - let path = match mode { - "i2v" => "/v1/videos/image2video", - "mi2v" => "/v1/videos/multi-image2video", - "motion-control" => "/v1/videos/motion-control", - _ => "/v1/videos/text2video", - }; - format!("{}{path}/{id}/{task_id}", self.config.base_url) +/// The API path of a generation mode. +fn mode_path(mode: &str) -> &'static str { + match mode { + "i2v" => "/v1/videos/image2video", + "mi2v" => "/v1/videos/multi-image2video", + "motion-control" => "/v1/videos/motion-control", + _ => "/v1/videos/text2video", } } @@ -225,7 +258,7 @@ fn video_file_to_image_string(file: &aimux_core::video_model::VideoFile) -> Stri #[async_trait] impl VideoModel for KlingAIVideoModel { fn provider(&self) -> &str { - "klingai" + &self.config.provider } fn model_id(&self) -> &str { @@ -249,9 +282,7 @@ impl VideoModel for KlingAIVideoModel { if let Some(ref prompt) = options.prompt { body.insert("prompt".to_string(), json!(prompt)); } - if let Some(negative_prompt) = options - .provider_options - .get("klingai") + if let Some(negative_prompt) = options::klingai_options(Some(&options.provider_options)) .and_then(|v| v.get("negativePrompt")) .and_then(|v| v.as_str()) { @@ -286,14 +317,13 @@ impl VideoModel for KlingAIVideoModel { } } - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; // Submit task. - let submit_url = self.submit_endpoint(&mode); + let submit_url = exchange.url(mode_path(&mode)); let request_body = Value::Object(body); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(submit_url.clone(), header_list, options), + exchange.request(submit_url.clone(), options), request_body.clone(), aimux_provider_utils::create_json_response_handler::(), klingai_failed_response_handler(), @@ -364,12 +394,11 @@ impl VideoModel for KlingAIVideoModel { let id = field("id")?; let task_id = field("task_id")?; - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; - let poll_url = self.poll_endpoint(mode, id, task_id); + let poll_url = exchange.url(&format!("{}/{id}/{task_id}", mode_path(mode))); let resp = aimux_provider_utils::get_from_api( - HttpRequest::new(poll_url.clone(), header_list, options), + exchange.request(poll_url.clone(), options), aimux_provider_utils::create_json_response_handler::(), klingai_failed_response_handler(), ) diff --git a/aimux-providers/src/klingai/options.rs b/aimux-providers/src/klingai/options.rs new file mode 100644 index 000000000..03e74ec30 --- /dev/null +++ b/aimux-providers/src/klingai/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of KlingAI. +//! +//! KlingAI reads `providerOptions.klingai`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "klingai"; + +/// The KlingAI options in a providerOptions map. +pub(crate) fn klingai_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/lib.rs b/aimux-providers/src/lib.rs index 80f767ad9..7add1317b 100644 --- a/aimux-providers/src/lib.rs +++ b/aimux-providers/src/lib.rs @@ -4,43 +4,150 @@ //! //! Each provider implements the `LanguageModel` trait from `aimux-core`. -/// Delegate `Provider::list_models` to the inner `OpenAIProvider` (field `.0`). +/// One required `Provider` method that this vendor does not offer: returns +/// `NoSuchModel` with the AI SDK `modelType`. +#[doc(hidden)] +#[macro_export] +macro_rules! __unsupported_required_model { + ($method:ident, $trait:ident, $model_type:literal) => { + fn $method( + &self, + id: &str, + ) -> ::std::result::Result< + ::std::sync::Arc, + ::aimux_core::AiMuxError, + > { + ::std::result::Result::Err(::aimux_core::AiMuxError::no_such_model(id, $model_type)) + } + }; +} + +/// Emit `impl Provider for $ty` for a vendor that offers exactly one modality. /// -/// Used by newtype providers that wrap `OpenAIProvider` (e.g. `OllamaProvider`, -/// `VllmProvider`, `VertexAiOpenaiModelsProvider`, …). Inserts the full trait -/// method signature so the wrapping provider needs no extra imports. +/// The first argument after the type names the `Provider` method to wire +/// (`language_model`, `embedding_model`, `image_model`, `transcription_model`, +/// `speech_model`, `reranking_model`, `video_model` or `search_model`); the +/// closure-like tail binds the provider and the model id and evaluates to the +/// (infallible) concrete model. The two other required methods return +/// `NoSuchModel`; the other optional ones keep their `None` default. +/// +/// ```ignore +/// impl_single_modality_provider!(VllmProvider, language_model, |p, id| p.model(id)); +/// impl_single_modality_provider!(SerperProvider, search_model, |p, _id| p.search_model()); +/// ``` #[macro_export] -macro_rules! delegate_list_models { - () => { - fn list_models( +macro_rules! impl_single_modality_provider { + ($ty:ty, language_model, |$p:ident, $id:ident| $build:expr) => { + impl ::aimux_core::Provider for $ty { + fn language_model( + &self, + $id: &str, + ) -> ::std::result::Result< + ::std::sync::Arc, + ::aimux_core::AiMuxError, + > { + let $p = self; + ::std::result::Result::Ok(::std::sync::Arc::new($build)) + } + $crate::__unsupported_required_model!( + embedding_model, + EmbeddingModel, + "embeddingModel" + ); + $crate::__unsupported_required_model!(image_model, ImageModel, "imageModel"); + } + }; + ($ty:ty, embedding_model, |$p:ident, $id:ident| $build:expr) => { + impl ::aimux_core::Provider for $ty { + $crate::__unsupported_required_model!(language_model, LanguageModel, "languageModel"); + fn embedding_model( + &self, + $id: &str, + ) -> ::std::result::Result< + ::std::sync::Arc, + ::aimux_core::AiMuxError, + > { + let $p = self; + ::std::result::Result::Ok(::std::sync::Arc::new($build)) + } + $crate::__unsupported_required_model!(image_model, ImageModel, "imageModel"); + } + }; + ($ty:ty, image_model, |$p:ident, $id:ident| $build:expr) => { + impl ::aimux_core::Provider for $ty { + $crate::__unsupported_required_model!(language_model, LanguageModel, "languageModel"); + $crate::__unsupported_required_model!( + embedding_model, + EmbeddingModel, + "embeddingModel" + ); + fn image_model( + &self, + $id: &str, + ) -> ::std::result::Result< + ::std::sync::Arc, + ::aimux_core::AiMuxError, + > { + let $p = self; + ::std::result::Result::Ok(::std::sync::Arc::new($build)) + } + } + }; + ($ty:ty, $method:ident, |$p:ident, $id:ident| $build:expr) => { + impl ::aimux_core::Provider for $ty { + $crate::__unsupported_required_model!(language_model, LanguageModel, "languageModel"); + $crate::__unsupported_required_model!( + embedding_model, + EmbeddingModel, + "embeddingModel" + ); + $crate::__unsupported_required_model!(image_model, ImageModel, "imageModel"); + $crate::__optional_model!($method, |$p, $id| $build); + } + }; +} + +/// One optional `Provider` method wired to a concrete model (see +/// [`impl_single_modality_provider!`]). +#[doc(hidden)] +#[macro_export] +macro_rules! __optional_model { + (transcription_model, |$p:ident, $id:ident| $build:expr) => { + $crate::__optional_model!(@emit transcription_model, TranscriptionModel, |$p, $id| $build); + }; + (speech_model, |$p:ident, $id:ident| $build:expr) => { + $crate::__optional_model!(@emit speech_model, SpeechModel, |$p, $id| $build); + }; + (reranking_model, |$p:ident, $id:ident| $build:expr) => { + $crate::__optional_model!(@emit reranking_model, RerankingModel, |$p, $id| $build); + }; + (video_model, |$p:ident, $id:ident| $build:expr) => { + $crate::__optional_model!(@emit video_model, VideoModel, |$p, $id| $build); + }; + (search_model, |$p:ident, $id:ident| $build:expr) => { + $crate::__optional_model!(@emit search_model, SearchModel, |$p, $id| $build); + }; + (@emit $method:ident, $trait:ident, |$p:ident, $id:ident| $build:expr) => { + fn $method( &self, - ) -> ::std::pin::Pin< - ::std::boxed::Box< - dyn ::std::future::Future< - Output = ::std::result::Result< - ::std::vec::Vec, - aimux_core::AiMuxError, - >, - > + ::std::marker::Send - + '_, + $id: &str, + ) -> ::std::option::Option< + ::std::result::Result< + ::std::sync::Arc, + ::aimux_core::AiMuxError, >, > { - self.0.list_models() + let $p = self; + ::std::option::Option::Some(::std::result::Result::Ok(::std::sync::Arc::new($build))) } }; } -// Registry-backed provider construction (RFC-0017 phase 4): all built-in -// OpenAI-compatible providers are looked up by name from `provider_registry.json`. -// The 250 per-provider `XxxConfig`/`XxxProvider` shell types were retired in -// phase 4 — use [`provider`] / [`provider_from_env`] instead. +// Registry-backed provider construction: the embedded `provider_registry.json` +// is parsed once into `preset` descriptors and used by the by-name entry points. pub mod provider; pub mod replay; -pub use provider::{ - ExternalProviderEntry, ProviderOptions, ProviderProfile, is_external_provider, - load_providers_from_json, provider, provider_api_key_env, provider_from_env, provider_handle, - provider_names, provider_registry_entry, register_provider, -}; +pub use preset::{AuthMode, PresetDescriptor, PresetEntry, PresetSettings}; pub use replay::{rebuild_decision_provider, rebuild_provider}; pub mod catalogue; @@ -51,23 +158,28 @@ pub mod anthropic_aws; pub mod azure; pub mod bedrock; pub mod cohere; +pub mod deepseek; +mod default_providers; pub mod google; +pub mod groq; pub mod jev; +pub use jev::{JevConfig, JevDecisionModel, JevProvider}; pub mod mistral; pub mod openai; +pub mod openai_compatible; +pub use default_providers::{ + create_provider, default_providers, provider_api_key_env, provider_default_base_url, + provider_names, +}; +pub mod preset; +pub(crate) mod shared; pub mod vertex; pub mod voyage; pub mod codex; -pub mod openrouter; pub mod xai; -// OpenAI-compatible thin wrappers (second batch). pub mod huggingface; -pub mod llamafile; -pub mod lmstudio; -pub mod mistralrs; -pub mod ollama; // Speech-only providers (TTS). pub mod cartesia; @@ -98,158 +210,139 @@ pub mod open_responses; // Rust consumers and the FFI layer share one implementation. pub use aimux_provider_utils::logging::init_logging; -// Bulk-generated thin-wrapper providers. -pub mod cybertron; -pub mod docker_model_runner; -pub mod gaudi; -pub mod jlama; -pub mod litellm_proxy; -pub mod llamacpp; -pub mod local; -pub mod localai; -pub mod mlx; -pub mod omlx; -pub mod onnx; -pub mod oobabooba; -pub mod openvino; -pub mod sglang; -pub mod vllm; -pub mod xinference; - -pub use anthropic::{AnthropicConfig, AnthropicProvider}; -pub use anthropic_aws::{AnthropicAwsAuth, AnthropicAwsProvider, AnthropicAwsProviderConfig}; +pub use anthropic::{ + AnthropicMessagesModel, AnthropicProvider, AnthropicProviderSettings, create_anthropic, +}; +pub use anthropic_aws::{ + AnthropicAwsAuth, AnthropicAwsProvider, AnthropicAwsProviderSettings, create_anthropic_aws, +}; pub use azure::{ - AzureAuth, AzureConfig, AzureModel, AzureProvider, AzureResponsesModel, TokenProvider, + AzureChatModel, AzureOpenAIProvider, AzureOpenAIProviderSettings, AzureResponsesModel, azure, + create_azure, }; pub use bedrock::{ - BedrockAuth, BedrockEmbeddingModel, BedrockImageModel, BedrockProvider, BedrockProviderConfig, + AmazonBedrockProvider, AmazonBedrockProviderSettings, BedrockEmbeddingModel, BedrockImageModel, + BedrockModel, BedrockRerankingModel, amazon_bedrock, create_amazon_bedrock, +}; +pub use cohere::{ + CohereEmbeddingModel, CohereModel, CohereProvider, CohereProviderSettings, + CohereRerankingModel, cohere, create_cohere, }; -pub use cohere::{CohereConfig, CohereEmbeddingModel, CohereProvider}; pub use google::{ - GoogleConfig, GoogleEmbeddingModel, GoogleImageModel, GoogleImageSettings, GoogleProvider, - GoogleVideoModel, + GoogleEmbeddingModel, GoogleFiles, GoogleImageModel, GoogleImageSettings, GoogleModel, + GoogleProvider, GoogleProviderSettings, GoogleVideoModel, create_google, google, +}; +pub use mistral::{ + MistralEmbeddingModel, MistralModel, MistralProvider, MistralProviderSettings, create_mistral, + mistral, }; -pub use jev::{JevConfig, JevDecisionModel, JevProvider}; -pub use mistral::{MistralConfig, MistralEmbeddingModel, MistralProvider}; pub use openai::{ - OpenAIConfig, OpenAIDecisionModel, OpenAIEmbeddingModel, OpenAIImageModel, OpenAIProvider, - OpenAIResponsesModel, OpenAISpeechModel, OpenAITranscriptionModel, + OpenAIDecisionModel, OpenAIEmbeddingModel, OpenAIImageModel, OpenAIProvider, + OpenAIProviderSettings, OpenAIResponsesModel, OpenAISpeechModel, OpenAITranscriptionModel, + create_openai, }; pub use vertex::{ - VertexAuth, VertexEmbeddingModel, VertexImageModel, VertexProvider, VertexProviderConfig, - VertexTranscriptionModel, VertexVideoModel, + VertexAnthropicModel, VertexEmbeddingModel, VertexImageModel, VertexModel, VertexProvider, + VertexProviderSettings, VertexTranscriptionModel, VertexVideoModel, create_google_vertex, + google_vertex, +}; +pub use voyage::{ + VoyageEmbeddingModel, VoyageProvider, VoyageProviderSettings, VoyageRerankingModel, + create_voyage, voyage, }; -pub use voyage::{VoyageConfig, VoyageEmbeddingModel, VoyageProvider}; pub use codex::{ CODEX_API_BASE_URL, CODEX_API_KEY_ENV_VAR, CODEX_OAUTH_TOKEN_URL, CODEX_SUBSCRIPTION_BASE_URL, - CodexConfig, CodexMode, CodexModel, CodexProvider, CodexTokens, codex_refresh, - codex_refresh_at, + CodexMode, CodexModel, CodexProvider, CodexProviderSettings, CodexTokens, codex, codex_refresh, + codex_refresh_at, create_codex, }; -pub use openrouter::{OpenRouterConfig, OpenRouterProvider}; -pub use xai::{XAIConfig, XAIProvider}; +pub use xai::{XAIProvider, XAIProviderSettings, XaiResponsesModel, create_xai, xai}; pub use cartesia::{ - CartesiaConfig, CartesiaProvider, CartesiaSpeechModel, CartesiaTranscriptionModel, + CartesiaProvider, CartesiaProviderSettings, CartesiaSpeechModel, CartesiaTranscriptionModel, + cartesia, create_cartesia, }; pub use elevenlabs::{ - ElevenLabsConfig, ElevenLabsProvider, ElevenLabsSpeechModel, ElevenLabsTranscriptionModel, + ElevenLabsProvider, ElevenLabsProviderSettings, ElevenLabsSpeechModel, + ElevenLabsTranscriptionModel, create_elevenlabs, elevenlabs, }; -pub use huggingface::{HuggingFaceConfig, HuggingFaceProvider}; -pub use hume::{HumeConfig, HumeProvider, HumeSpeechModel}; -pub use llamafile::{LlamafileConfig, LlamafileProvider}; -pub use lmnt::{LMNTConfig, LMNTProvider, LMNTSpeechModel}; -pub use mistralrs::{MistralrsConfig, MistralrsProvider}; - -pub use lmstudio::{LmStudioConfig, LmStudioProvider}; -pub use ollama::{OllamaConfig, OllamaProvider}; +pub use huggingface::{ + HuggingFaceProvider, HuggingFaceProviderSettings, HuggingFaceResponsesModel, + create_huggingface, huggingface, +}; +pub use hume::{HumeProvider, HumeProviderSettings, HumeSpeechModel, create_hume, hume}; +pub use lmnt::{LMNTProvider, LMNTProviderSettings, LMNTSpeechModel, create_lmnt, lmnt}; -pub use assemblyai::{AssemblyAIConfig, AssemblyAIProvider, AssemblyAITranscriptionModel}; -pub use deepgram::{DeepgramConfig, DeepgramProvider, DeepgramTranscriptionModel}; -pub use fal::{FalConfig, FalImageModel, FalProvider, FalTranscriptionModel, FalVideoModel}; +pub use assemblyai::{ + AssemblyAIProvider, AssemblyAIProviderSettings, AssemblyAITranscriptionModel, assemblyai, + create_assemblyai, +}; +pub use deepgram::{ + DeepgramProvider, DeepgramProviderSettings, DeepgramTranscriptionModel, create_deepgram, + deepgram, +}; +pub use fal::{ + FalImageModel, FalProvider, FalProviderSettings, FalTranscriptionModel, FalVideoModel, + create_fal, fal, +}; // Image-only provider re-exports. pub use black_forest_labs::{ - BlackForestLabsConfig, BlackForestLabsImageModel, BlackForestLabsProvider, -}; -pub use gladia::{GladiaConfig, GladiaProvider, GladiaTranscriptionModel}; -pub use klingai::{KlingAIConfig, KlingAIProvider, KlingAIVideoModel}; -pub use luma::{LumaConfig, LumaImageModel, LumaProvider}; -pub use prodia::{ProdiaConfig, ProdiaImageModel, ProdiaProvider, ProdiaVideoModel}; -pub use replicate::{ReplicateConfig, ReplicateImageModel, ReplicateProvider, ReplicateVideoModel}; -pub use revai::{RevaiConfig, RevaiProvider, RevaiTranscriptionModel}; - -pub use open_responses::{OpenResponsesConfig, OpenResponsesModel, OpenResponsesProvider}; + BlackForestLabsImageModel, BlackForestLabsProvider, BlackForestLabsProviderSettings, + black_forest_labs, create_black_forest_labs, +}; +pub use gladia::{ + GladiaProvider, GladiaProviderSettings, GladiaTranscriptionModel, create_gladia, gladia, +}; +pub use klingai::{ + KlingAIProvider, KlingAIProviderSettings, KlingAIVideoModel, create_klingai, klingai, +}; +pub use luma::{LumaImageModel, LumaProvider, LumaProviderSettings, create_luma, luma}; +pub use prodia::{ + ProdiaImageModel, ProdiaProvider, ProdiaProviderSettings, ProdiaVideoModel, create_prodia, + prodia, +}; +pub use replicate::{ + ReplicateImageModel, ReplicateProvider, ReplicateProviderSettings, ReplicateVideoModel, + create_replicate, replicate, +}; +pub use revai::{ + RevaiProvider, RevaiProviderSettings, RevaiTranscriptionModel, create_revai, revai, +}; -// Bulk-generated provider re-exports. -pub use cybertron::{CybertronConfig, CybertronProvider}; -pub use docker_model_runner::{DockerModelRunnerConfig, DockerModelRunnerProvider}; -pub use gaudi::{GaudiConfig, GaudiProvider}; -pub use jlama::{JlamaConfig, JlamaProvider}; -pub use litellm_proxy::{LitellmProxyConfig, LitellmProxyProvider}; -pub use llamacpp::{LlamacppConfig, LlamacppProvider}; -pub use local::{LocalConfig, LocalProvider}; -pub use localai::{LocalaiConfig, LocalaiProvider}; -pub use mlx::{MlxConfig, MlxProvider}; -pub use omlx::{OmlxConfig, OmlxProvider}; -pub use onnx::{OnnxConfig, OnnxProvider}; -pub use oobabooba::{OobaboobaConfig, OobaboobaProvider}; -pub use openvino::{OpenvinoConfig, OpenvinoProvider}; -pub use sglang::{SglangConfig, SglangProvider}; -pub use vllm::{VllmConfig, VllmProvider}; -pub use xinference::{XinferenceConfig, XinferenceProvider}; +pub use open_responses::{ + OpenResponsesModel, OpenResponsesProvider, OpenResponsesProviderSettings, create_open_responses, +}; // Modality-specific providers (non-language, e.g. rerank-only). pub mod jina_ai; -pub use jina_ai::{JinaAiConfig, JinaAiProvider, JinaAiRerankingModel}; +pub use jina_ai::{ + JinaAiProvider, JinaAiProviderSettings, JinaAiRerankingModel, create_jina_ai, jina_ai, +}; // AWS Polly speech (TTS) provider — SigV4 authenticated, speech modality only. pub mod aws_polly; -pub use aws_polly::{AwsPollyConfig, AwsPollyProvider, AwsPollySpeechModel}; +pub use aws_polly::{ + AwsPollyProvider, AwsPollyProviderSettings, AwsPollySpeechModel, aws_polly, create_aws_polly, +}; // Recraft image provider (OpenAI Images-compatible + Recraft extension fields). pub mod recraft; -pub use recraft::{RecraftConfig, RecraftImageModel, RecraftProvider}; +pub use recraft::{ + RecraftImageModel, RecraftProvider, RecraftProviderSettings, create_recraft, recraft, +}; // Stability image provider (image modality only). pub mod stability; -pub use stability::{StabilityConfig, StabilityImageModel, StabilityProvider}; +pub use stability::{ + StabilityImageModel, StabilityProvider, StabilityProviderSettings, create_stability, stability, +}; // Video-only provider (runwayml). pub mod runwayml; -pub use runwayml::{RunwaymlConfig, RunwaymlProvider, RunwaymlVideoModel}; - -// P1 thin-wrapper providers (provider-research batch). -pub mod bedrock_mantle; - -pub use bedrock_mantle::{BedrockMantleConfig, BedrockMantleProvider}; - -// Vertex AI MaaS partner-model providers (OpenAI-compatible thin wrappers). -// Each wraps the shared OpenAIProvider against the Vertex AI MaaS OpenAPI -// endpoint, authenticating with a Google Cloud Bearer token. -pub mod vertex_ai_ai21_models; -pub mod vertex_ai_anthropic_models; -pub mod vertex_ai_deepseek_models; -pub mod vertex_ai_llama_models; -pub mod vertex_ai_minimax_models; -pub mod vertex_ai_mistral_models; -pub mod vertex_ai_moonshot_models; -pub mod vertex_ai_openai_models; -pub mod vertex_ai_qwen_models; -pub mod vertex_ai_zai_models; - -pub use vertex_ai_ai21_models::{VertexAiAi21ModelsConfig, VertexAiAi21ModelsProvider}; -pub use vertex_ai_anthropic_models::{ - VertexAiAnthropicModelsConfig, VertexAiAnthropicModelsProvider, -}; -pub use vertex_ai_deepseek_models::{VertexAiDeepseekModelsConfig, VertexAiDeepseekModelsProvider}; -pub use vertex_ai_llama_models::{VertexAiLlamaModelsConfig, VertexAiLlamaModelsProvider}; -pub use vertex_ai_minimax_models::{VertexAiMinimaxModelsConfig, VertexAiMinimaxModelsProvider}; -pub use vertex_ai_mistral_models::{VertexAiMistralModelsConfig, VertexAiMistralModelsProvider}; -pub use vertex_ai_moonshot_models::{VertexAiMoonshotModelsConfig, VertexAiMoonshotModelsProvider}; -pub use vertex_ai_openai_models::{VertexAiOpenaiModelsConfig, VertexAiOpenaiModelsProvider}; -pub use vertex_ai_qwen_models::{VertexAiQwenModelsConfig, VertexAiQwenModelsProvider}; -pub use vertex_ai_zai_models::{VertexAiZaiModelsConfig, VertexAiZaiModelsProvider}; +pub use runwayml::{ + RunwaymlProvider, RunwaymlProviderSettings, RunwaymlVideoModel, create_runwayml, runwayml, +}; // Search-only providers (web search modality). pub mod dataforseo; @@ -264,17 +357,40 @@ pub mod tavily; pub mod tinyfish; pub mod you_com; -pub use dataforseo::{DataforseoConfig, DataforseoProvider, DataforseoSearchModel}; -pub use exa_ai::{ExaAiConfig, ExaAiProvider, ExaAiSearchModel}; -pub use firecrawl::{FirecrawlConfig, FirecrawlProvider, FirecrawlSearchModel}; -pub use google_pse::{GooglePseConfig, GooglePseProvider, GooglePseSearchModel}; -pub use linkup::{LinkupConfig, LinkupProvider, LinkupSearchModel}; -pub use parallel_ai::{ParallelAiConfig, ParallelAiProvider, ParallelAiSearchModel}; -pub use searxng::{SearxngConfig, SearxngProvider, SearxngSearchModel}; -pub use serper::{SerperConfig, SerperProvider, SerperSearchModel}; -pub use tavily::{TavilyConfig, TavilyProvider, TavilySearchModel}; -pub use tinyfish::{TinyfishConfig, TinyfishProvider, TinyfishSearchModel}; -pub use you_com::{YouComConfig, YouComProvider, YouComSearchModel}; +pub use dataforseo::{ + DataforseoProvider, DataforseoProviderSettings, DataforseoSearchModel, create_dataforseo, + dataforseo, +}; +pub use exa_ai::{ExaAiProvider, ExaAiProviderSettings, ExaAiSearchModel, create_exa_ai, exa_ai}; +pub use firecrawl::{ + FirecrawlProvider, FirecrawlProviderSettings, FirecrawlSearchModel, create_firecrawl, firecrawl, +}; +pub use google_pse::{ + GooglePseProvider, GooglePseProviderSettings, GooglePseSearchModel, create_google_pse, + google_pse, +}; +pub use linkup::{ + LinkupProvider, LinkupProviderSettings, LinkupSearchModel, create_linkup, linkup, +}; +pub use parallel_ai::{ + ParallelAiProvider, ParallelAiProviderSettings, ParallelAiSearchModel, create_parallel_ai, + parallel_ai, +}; +pub use searxng::{ + SearxngProvider, SearxngProviderSettings, SearxngSearchModel, create_searxng, searxng, +}; +pub use serper::{ + SerperProvider, SerperProviderSettings, SerperSearchModel, create_serper, serper, +}; +pub use tavily::{ + TavilyProvider, TavilyProviderSettings, TavilySearchModel, create_tavily, tavily, +}; +pub use tinyfish::{ + TinyfishProvider, TinyfishProviderSettings, TinyfishSearchModel, create_tinyfish, tinyfish, +}; +pub use you_com::{ + YouComProvider, YouComProviderSettings, YouComSearchModel, create_you_com, you_com, +}; mod decision_scoring; mod decision_support; diff --git a/aimux-providers/src/linkup.rs b/aimux-providers/src/linkup.rs index 697931187..437a382e2 100644 --- a/aimux-providers/src/linkup.rs +++ b/aimux-providers/src/linkup.rs @@ -7,22 +7,50 @@ //! protocol (query → results) and does not support language models. The //! `depth` and `outputType` request fields are passed through via //! `provider_options` under the `"linkup"` key. +//! +//! [`create_linkup`] takes [`LinkupProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`LinkupProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `LINKUP_API_KEY`. +//! [`linkup()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; /// Fixed model ID for the Linkup search model. +/// Failed-response handler of the Linkup endpoint: `message`, and `type` +/// else `code` as the provider code, from the body's `error` object, or from +/// the top level when there is none. This is the mapping the endpoint used +/// with the removed shared handler; it is not an AI SDK schema. +fn linkup_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(serde_json::Value::as_str) + .map(str::to_owned), + } + }) +} + const MODEL_ID: &str = "linkup-search"; /// Linkup provider-specific search options (passed through via @@ -37,9 +65,7 @@ struct LinkupOptions { fn parse_linkup_options(provider_options: Option<&SharedProviderOptions>) -> LinkupOptions { let mut opts = LinkupOptions::default(); - if let Some(po) = provider_options - && let Some(linkup) = po.get("linkup") - { + if let Some(linkup) = options::linkup_options(provider_options) { if let Some(v) = linkup.get("depth").and_then(|v| v.as_str()) { opts.depth = Some(v.to_string()); } @@ -50,65 +76,118 @@ fn parse_linkup_options(provider_options: Option<&SharedProviderOptions>) -> Lin opts } -/// Configuration for the Linkup provider. -#[derive(Debug, Clone)] -pub struct LinkupConfig { - pub api_key: String, - pub base_url: String, -} +pub(crate) mod options; -impl LinkupConfig { - /// Create from an API key (uses the default Linkup base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.linkup.so".to_string(), - } - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.linkup.so"; +pub(crate) const API_KEY_ENV_VAR: &str = "LINKUP_API_KEY"; +const DEFAULT_NAME: &str = "linkup"; - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Settings of [`create_linkup`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct LinkupProviderSettings { + /// Base URL for the API calls. Default `https://api.linkup.so`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `LINKUP_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.search"`). + /// Default `"linkup"`. The providerOptions key stays `linkup`. + 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, +} - /// Create from the `LINKUP_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `LINKUP_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "LINKUP_API_KEY", "Linkup")?; - Ok(Self::new(api_key)) +impl std::fmt::Debug for LinkupProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LinkupProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } -/// Linkup provider — creates `LinkupSearchModel` instances. +/// Create a Linkup provider. +/// +/// # Errors /// -/// Linkup is a search-only provider; it does not support language models. +/// 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. +pub fn create_linkup(settings: LinkupProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(LinkupProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Linkup"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_linkup` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn linkup() -> &'static LinkupProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_linkup(LinkupProviderSettings::default()) + .expect("default Linkup settings are always valid") + }) +} + +/// A Linkup provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct LinkupProvider { - config: LinkupConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl LinkupProvider { - #[must_use] - pub fn new(config: LinkupConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a search model instance. + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> LinkupSearchModel { - LinkupSearchModel::new(self.config.clone()) + LinkupSearchModel::from_config(self.model_config("search")) } } -impl Provider for LinkupProvider { - fn name(&self) -> &str { - "linkup" - } -} +crate::impl_single_modality_provider!(LinkupProvider, search_model, |p, _id| p.search_model()); /// Build the Linkup `/v1/search` request body (pure function). fn build_request_body(options: &SearchCallOptions, linkup_options: &LinkupOptions) -> Value { @@ -175,38 +254,19 @@ fn map_results(entries: Vec) -> Vec { /// A Linkup search model. pub struct LinkupSearchModel { - config: LinkupConfig, + config: EndpointConfig, } impl LinkupSearchModel { - #[must_use] - pub fn new(config: LinkupConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/v1/search", self.config.base_url) - } } #[async_trait] impl SearchModel for LinkupSearchModel { fn provider(&self) -> &str { - "linkup" + &self.config.provider } fn model_id(&self) -> &str { @@ -217,24 +277,13 @@ impl SearchModel for LinkupSearchModel { let linkup_options = parse_linkup_options(options.provider_options.as_ref()); let body = build_request_body(options, &linkup_options); - let headers: Vec<(String, String)> = self - .build_headers(options.headers.as_ref()) - .into_iter() - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(exchange.url("/v1/search"), options), body, aimux_provider_utils::create_json_response_handler::(), - aimux_provider_utils::create_standard_json_error_response_handler(), + linkup_failed_response_handler(), ) .await?; diff --git a/aimux-providers/src/linkup/options.rs b/aimux-providers/src/linkup/options.rs new file mode 100644 index 000000000..5b576a60f --- /dev/null +++ b/aimux-providers/src/linkup/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Linkup. +//! +//! Linkup reads `providerOptions.linkup`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "linkup"; + +/// The Linkup options in a providerOptions map. +pub(crate) fn linkup_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/litellm_proxy.rs b/aimux-providers/src/litellm_proxy.rs deleted file mode 100644 index 2928fe951..000000000 --- a/aimux-providers/src/litellm_proxy.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! LiteLLM Proxy provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:4000/v1`. The `LITELLM_PROXY_API_KEY` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:4000/v1"; -const ENV_VAR: &str = "LITELLM_PROXY_API_KEY"; -const PROVIDER_NAME: &str = "litellm_proxy"; -const PLACEHOLDER_API_KEY: &str = "litellm_proxy"; - -pub struct LitellmProxyConfig(OpenAIConfig); - -impl LitellmProxyConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `LITELLM_PROXY_API_KEY` environment variable (which holds - /// a base URL), falling back to `http://127.0.0.1:4000/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct LitellmProxyProvider(OpenAIProvider); - -impl LitellmProxyProvider { - #[must_use] - pub fn new(config: LitellmProxyConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for LitellmProxyProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/llamacpp.rs b/aimux-providers/src/llamacpp.rs deleted file mode 100644 index 5cc96b621..000000000 --- a/aimux-providers/src/llamacpp.rs +++ /dev/null @@ -1,86 +0,0 @@ -//! llama.cpp provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `LLAMACPP_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "LLAMACPP_BASE_URL"; -const PROVIDER_NAME: &str = "llamacpp"; -const PLACEHOLDER_API_KEY: &str = "llamacpp"; - -pub struct LlamacppConfig(OpenAIConfig); - -impl LlamacppConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()), - ) - } - - /// Create from the `LLAMACPP_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let mut config = Self::new(PLACEHOLDER_API_KEY); - config.0 = config.0.with_api_key_source(Some("none")); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct LlamacppProvider(OpenAIProvider); - -impl LlamacppProvider { - #[must_use] - pub fn new(config: LlamacppConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for LlamacppProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - fn decision_model( - &self, - model_id: &str, - ) -> Result, AiMuxError> { - self.0.decision_model(model_id) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/llamafile.rs b/aimux-providers/src/llamafile.rs deleted file mode 100644 index 041107e7d..000000000 --- a/aimux-providers/src/llamafile.rs +++ /dev/null @@ -1,100 +0,0 @@ -//! Llamafile provider — a thin OpenAI-compatible wrapper for the local -//! [llamafile](https://github.com/Mozilla-Ocho/llamafile) inference server. -//! -//! llamafile exposes an OpenAI-compatible Chat Completions API at -//! `http://127.0.0.1:8080/v1` by default. The Rust -//! [`OpenAIProvider`] appends `/chat/completions` -//! to this base URL, yielding `http://127.0.0.1:8080/v1/chat/completions`. -//! -//! Unlike hosted providers, llamafile runs locally and does not require -//! authentication. Accordingly the `LLAMAFILE_BASE_URL` environment variable -//! holds a *base URL* (not an API key); [`LlamafileConfig::from_env`] reads it -//! and falls back to the default local endpoint when it is unset. A placeholder -//! API key is sent in the `Authorization` header — llamafile ignores it, but the -//! shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -/// Environment variable holding the llamafile base URL (not an API key). -const ENV_VAR: &str = "LLAMAFILE_BASE_URL"; -const PROVIDER_NAME: &str = "llamafile"; -/// Placeholder API key — llamafile does not authenticate, but the shared -/// `OpenAIConfig` requires a non-empty key string. -const PLACEHOLDER_API_KEY: &str = "llamafile"; - -/// Configuration for the llamafile provider (wraps [`OpenAIConfig`]). -pub struct LlamafileConfig(OpenAIConfig); - -impl LlamafileConfig { - /// Create from an API key, using the default local llamafile base URL. - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `LLAMAFILE_BASE_URL` environment variable. - /// - /// `LLAMAFILE_BASE_URL` holds a *base URL* (e.g. - /// `http://127.0.0.1:8080/v1`), not an API key — llamafile is a local - /// inference server that does not require authentication. When the variable - /// is unset (or empty), the default local endpoint - /// (`http://127.0.0.1:8080/v1`) is used. - /// - /// # Errors - /// - /// Never returns an error; an unset `LLAMAFILE_BASE_URL` falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - /// Override the base URL (useful for tests / non-default ports). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// llamafile provider — creates [`OpenAIModel`] instances pointed at llamafile. -pub struct LlamafileProvider(OpenAIProvider); - -impl LlamafileProvider { - #[must_use] - pub fn new(config: LlamafileConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given llamafile model id - /// (e.g. `"llama3.2:latest"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for LlamafileProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/lmnt.rs b/aimux-providers/src/lmnt.rs index 6fe407a01..d96340a8d 100644 --- a/aimux-providers/src/lmnt.rs +++ b/aimux-providers/src/lmnt.rs @@ -12,8 +12,13 @@ //! `speed`, and `language` in the request body, and returns raw binary audio. //! Provider options (`lmnt` key) support `conversational`, `length`, `seed`, //! `speed`, `temperature`, `topP`, `sampleRate`, and `format`. +//! +//! [`create_lmnt`] takes [`LMNTProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`LMNTProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `LMNT_API_KEY`. +//! [`lmnt()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -23,8 +28,9 @@ use aimux_core::shared::{SharedProviderOptions, Warning}; use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; -use aimux_provider_utils::{HttpRequest, load_api_key}; +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; fn lmnt_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -48,72 +54,120 @@ fn lmnt_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, -} +pub(crate) mod options; -impl LMNTConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.lmnt.com".to_string(), - headers: None, - } - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.lmnt.com"; +pub(crate) const API_KEY_ENV_VAR: &str = "LMNT_API_KEY"; +const DEFAULT_NAME: &str = "lmnt"; - /// Override the base URL (for testing or proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = url.into().trim_end_matches('/').to_string(); - self - } +/// Settings of [`create_lmnt`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct LMNTProviderSettings { + /// Base URL for the API calls. Default `https://api.lmnt.com`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `LMNT_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.speech"`). + /// Default `"lmnt"`. The providerOptions key stays `lmnt`. + 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, +} - /// Attach extra headers merged into every request. - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self +impl std::fmt::Debug for LMNTProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LMNTProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Create from environment variable `LMNT_API_KEY`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `LMNT_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "LMNT_API_KEY", "LMNT")?; - Ok(Self::new(api_key)) - } +/// Create a LMNT provider. +/// +/// # 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. +pub fn create_lmnt(settings: LMNTProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(LMNTProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "LMNT"), + AuthScheme::Header("x-api-key"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) } -// ── Provider ───────────────────────────────────────────────────────────────── +/// The default provider: `create_lmnt` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn lmnt() -> &'static LMNTProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_lmnt(LMNTProviderSettings::default()) + .expect("default LMNT settings are always valid") + }) +} -/// LMNT provider — creates `LMNTSpeechModel` instances. +/// A LMNT provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct LMNTProvider { - config: LMNTConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl LMNTProvider { - #[must_use] - pub fn new(config: LMNTConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a speech (TTS) model instance for the given model name (e.g. - /// `"aurora"`). + /// A speech (TTS) model (e.g. `"aurora"`); `provider()` is `"{name}.speech"`. #[must_use] pub fn speech(&self, model_id: &str) -> LMNTSpeechModel { - LMNTSpeechModel::new(model_id.to_string(), self.config.clone()) + LMNTSpeechModel::from_config(model_id.to_string(), self.model_config("speech")) } } +crate::impl_single_modality_provider!(LMNTProvider, speech_model, |p, id| p.speech(id)); + // ── Speech model ───────────────────────────────────────────────────────────── /// The default voice ID used when no voice is provided. @@ -125,40 +179,19 @@ const SUPPORTED_OUTPUT_FORMATS: &[&str] = &["mp3", "aac", "mulaw", "raw", "wav"] /// An LMNT speech (TTS) model. pub struct LMNTSpeechModel { model_id: String, - config: LMNTConfig, + config: EndpointConfig, } impl LMNTSpeechModel { - #[must_use] - pub fn new(model_id: String, config: LMNTConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("x-api-key".to_string(), self.config.api_key.clone()); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/v1/ai/speech/bytes", self.config.base_url) - } } #[async_trait] impl SpeechModel for LMNTSpeechModel { fn provider(&self) -> &str { - "lmnt.speech" + &self.config.provider } fn model_id(&self) -> &str { @@ -168,10 +201,10 @@ impl SpeechModel for LMNTSpeechModel { async fn do_generate(&self, options: &SpeechCallOptions) -> Result { let (body, warnings) = build_request(options, &self.model_id)?; - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), headers.into_iter().collect(), options), + exchange.request(exchange.url("/v1/ai/speech/bytes"), options), Value::Object(body.clone()), aimux_provider_utils::create_binary_response_handler(), lmnt_failed_response_handler(), @@ -298,7 +331,7 @@ struct LMNTSpeechProviderOptions { fn parse_lmnt_provider_options( options: Option<&SharedProviderOptions>, ) -> Option { - let opts = options.and_then(|opts| opts.get("lmnt"))?; + let opts = options::lmnt_options(options)?; Some(LMNTSpeechProviderOptions { conversational: opts diff --git a/aimux-providers/src/lmnt/options.rs b/aimux-providers/src/lmnt/options.rs new file mode 100644 index 000000000..1c5687ab1 --- /dev/null +++ b/aimux-providers/src/lmnt/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of LMNT. +//! +//! LMNT reads `providerOptions.lmnt`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "lmnt"; + +/// The LMNT options in a providerOptions map. +pub(crate) fn lmnt_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/lmstudio.rs b/aimux-providers/src/lmstudio.rs deleted file mode 100644 index c9decc8fa..000000000 --- a/aimux-providers/src/lmstudio.rs +++ /dev/null @@ -1,100 +0,0 @@ -//! LM Studio provider — a thin OpenAI-compatible wrapper for the local -//! [LM Studio](https://lmstudio.ai) inference server. -//! -//! LM Studio exposes an OpenAI-compatible Chat Completions API at -//! `http://127.0.0.1:1234/v1` by default. The Rust -//! [`OpenAIProvider`] appends `/chat/completions` -//! to this base URL, yielding `http://127.0.0.1:1234/v1/chat/completions`. -//! -//! Unlike hosted providers, LM Studio runs locally and does not require -//! authentication. Accordingly the `LMSTUDIO_BASE_URL` environment variable -//! holds a *base URL* (not an API key); [`LmStudioConfig::from_env`] reads it -//! and falls back to the default local endpoint when it is unset. A placeholder -//! API key is sent in the `Authorization` header — LM Studio ignores it, but -//! the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:1234/v1"; -/// Environment variable holding the LM Studio base URL (not an API key). -const ENV_VAR: &str = "LMSTUDIO_BASE_URL"; -const PROVIDER_NAME: &str = "lmstudio"; -/// Placeholder API key — LM Studio does not authenticate, but the shared -/// `OpenAIConfig` requires a non-empty key string. -const PLACEHOLDER_API_KEY: &str = "lmstudio"; - -/// Configuration for the LM Studio provider (wraps [`OpenAIConfig`]). -pub struct LmStudioConfig(OpenAIConfig); - -impl LmStudioConfig { - /// Create from an API key, using the default local LM Studio base URL. - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `LMSTUDIO_BASE_URL` environment variable. - /// - /// `LMSTUDIO_BASE_URL` holds a *base URL* (e.g. - /// `http://127.0.0.1:1234/v1`), not an API key — LM Studio is a local - /// inference server that does not require authentication. When the variable - /// is unset (or empty), the default local endpoint - /// (`http://127.0.0.1:1234/v1`) is used. - /// - /// # Errors - /// - /// Never returns an error; an unset `LMSTUDIO_BASE_URL` falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - /// Override the base URL (useful for tests / non-default ports). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// LM Studio provider — creates [`OpenAIModel`] instances pointed at LM Studio. -pub struct LmStudioProvider(OpenAIProvider); - -impl LmStudioProvider { - #[must_use] - pub fn new(config: LmStudioConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given LM Studio model id - /// (e.g. `"llama-3.2-3b-instruct"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for LmStudioProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/local.rs b/aimux-providers/src/local.rs deleted file mode 100644 index 528c748ca..000000000 --- a/aimux-providers/src/local.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! Local LLM provider — a thin OpenAI-compatible wrapper. -//! -//! See `localhost` for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `LOCAL_LLM_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "LOCAL_LLM_BASE_URL"; -const PROVIDER_NAME: &str = "local"; -const PLACEHOLDER_API_KEY: &str = "local"; - -pub struct LocalConfig(OpenAIConfig); - -impl LocalConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `LOCAL_LLM_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct LocalProvider(OpenAIProvider); - -impl LocalProvider { - #[must_use] - pub fn new(config: LocalConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for LocalProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/localai.rs b/aimux-providers/src/localai.rs deleted file mode 100644 index 249b87515..000000000 --- a/aimux-providers/src/localai.rs +++ /dev/null @@ -1,86 +0,0 @@ -//! LocalAI provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `LOCALAI_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "LOCALAI_BASE_URL"; -const PROVIDER_NAME: &str = "localai"; -const PLACEHOLDER_API_KEY: &str = "localai"; - -pub struct LocalaiConfig(OpenAIConfig); - -impl LocalaiConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()), - ) - } - - /// Create from the `LOCALAI_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let mut config = Self::new(PLACEHOLDER_API_KEY); - config.0 = config.0.with_api_key_source(Some("none")); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct LocalaiProvider(OpenAIProvider); - -impl LocalaiProvider { - #[must_use] - pub fn new(config: LocalaiConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for LocalaiProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - fn decision_model( - &self, - model_id: &str, - ) -> Result, AiMuxError> { - self.0.decision_model(model_id) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/luma.rs b/aimux-providers/src/luma.rs index f99e163a1..e905363e8 100644 --- a/aimux-providers/src/luma.rs +++ b/aimux-providers/src/luma.rs @@ -5,8 +5,14 @@ //! //! Uses an async submit + poll pattern: POST to create generation, then GET //! poll until state is "completed", then download the image. +//! +//! [`create_luma`] takes [`LumaProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`LumaProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `LUMA_API_KEY`. +//! [`luma()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::Duration; use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -16,9 +22,14 @@ use aimux_core::error::ApiCallError; use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; -use aimux_core::retry; use aimux_core::shared::Warning; -use aimux_provider_utils::{HttpRequest, load_api_key, sleep_or_abort, without_trailing_slash}; +use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{ + Credential, EndpointConfig, POLL_INTERVAL_MILLIS_KEY, PollStep, is_poll_control_key, + poll_interval_ms, poll_until, provider_headers, retry_download, +}; fn luma_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -37,109 +48,141 @@ fn luma_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.lumalabs.ai"; +pub(crate) const API_KEY_ENV_VAR: &str = "LUMA_API_KEY"; +const DEFAULT_NAME: &str = "luma"; + +/// Settings of [`create_luma`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct LumaProviderSettings { + /// Base URL for the API calls. Default `https://api.lumalabs.ai`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `LUMA_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.image"`). + /// Default `"luma"`. The providerOptions key stays `luma`. + 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, } -impl LumaConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.lumalabs.ai".to_string(), - headers: None, - } - } - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } - /// Create from the `LUMA_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "LUMA_API_KEY", "Luma")?; - Ok(Self::new(api_key)) +impl std::fmt::Debug for LumaProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LumaProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } +/// Create a Luma provider. +/// +/// # 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. +pub fn create_luma(settings: LumaProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(LumaProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Luma"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_luma` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn luma() -> &'static LumaProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_luma(LumaProviderSettings::default()) + .expect("default Luma settings are always valid") + }) +} + +/// A Luma provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct LumaProvider { - config: LumaConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } + impl LumaProvider { - #[must_use] - pub fn new(config: LumaConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + + /// An image model (e.g. `"photon-1"`); `provider()` is `"{name}.image"`. #[must_use] pub fn image(&self, model_id: &str) -> LumaImageModel { - LumaImageModel::new(model_id.to_string(), self.config.clone()) + LumaImageModel::from_config(model_id.to_string(), self.model_config("image")) } } +crate::impl_single_modality_provider!(LumaProvider, image_model, |p, id| p.image(id)); + /// A Luma image generation model. pub struct LumaImageModel { model_id: String, - config: LumaConfig, + config: EndpointConfig, } + impl LumaImageModel { - #[must_use] - pub fn new(model_id: String, config: LumaConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut h = HashMap::new(); - h.insert( - "Authorization".into(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - h.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - h.insert(k.clone(), v.clone()); - } - } - h - } - - fn generations_url(&self, generation_id: Option<&str>) -> String { - match generation_id { - Some(id) => format!( - "{}/dream-machine/v1/generations/{}", - self.config.base_url, id - ), - None => format!( - "{}/dream-machine/v1/generations/image", - self.config.base_url - ), - } - } } #[async_trait] impl ImageModel for LumaImageModel { fn provider(&self) -> &str { - "luma" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id @@ -167,17 +210,14 @@ impl ImageModel for LumaImageModel { }); } - let luma_opts = options.provider_options.get("luma"); + let luma_opts = options::luma_options(Some(&options.provider_options)); // Extract non-request options - let poll_interval = luma_opts - .and_then(|o| o.get("pollIntervalMillis")) - .and_then(serde_json::Value::as_u64) - .unwrap_or(DEFAULT_POLL_INTERVAL_MS); - let max_poll_attempts = luma_opts - .and_then(|o| o.get("maxPollAttempts")) - .and_then(serde_json::Value::as_u64) - .unwrap_or(DEFAULT_MAX_POLL_ATTEMPTS); + let poll_interval = Duration::from_millis(poll_interval_ms( + luma_opts, + POLL_INTERVAL_MILLIS_KEY, + POLL_INTERVAL_MS, + )); let reference_type = luma_opts .and_then(|o| o.get("referenceType")) .and_then(|v| v.as_str()) @@ -303,22 +343,23 @@ impl ImageModel for LumaImageModel { // Forward luma provider options (excluding non-request options) if let Some(luma) = luma_opts { for (k, v) in luma { - if matches!( - k.as_str(), - "pollIntervalMillis" | "maxPollAttempts" | "referenceType" | "images" - ) { + if is_poll_control_key(k) || matches!(k.as_str(), "referenceType" | "images") { continue; } body.insert(k.clone(), v.clone()); } } - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let generations_url = |generation_id: Option<&str>| match generation_id { + Some(id) => exchange.url(&format!("/dream-machine/v1/generations/{id}")), + None => exchange.url("/dream-machine/v1/generations/image"), + }; - // Submit + // Submit. This is the only request that creates a generation: nothing + // below sends it again. let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.generations_url(None), header_list.clone(), options), + exchange.request(generations_url(None), options), Value::Object(body), aimux_provider_utils::create_json_response_handler(), luma_failed_response_handler(), @@ -334,95 +375,72 @@ impl ImageModel for LumaImageModel { AiMuxError::InvalidResponseData("missing id in Luma response".to_string()) })? .to_string(); - let retries = retry::prepare_retries( - options.max_retries, - self.retry_config(), - options.abort_signal.clone(), - ); - - // Poll for completion - let mut image_url = None; - for _ in 0..max_poll_attempts { - let pr = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest::new( - self.generations_url(Some(&generation_id)), - header_list.clone(), - options, - ), - aimux_provider_utils::create_json_response_handler::(), - luma_failed_response_handler(), - ) - }) + + // Poll for completion. + let poll_url = generations_url(Some(&generation_id)); + let image_url = poll_until( + &format!("luma generation {generation_id}"), + options.abort_signal.as_ref(), + poll_interval, + MAX_POLL_ATTEMPTS, + || async { + let pr = aimux_provider_utils::get_from_api( + exchange.request(poll_url.clone(), options), + aimux_provider_utils::create_json_response_handler::(), + luma_failed_response_handler(), + ) .await?; - let response_body = pr.raw_value.as_ref().map(ToString::to_string); - let pv = pr.value; - - let state = pv.get("state").and_then(|v| v.as_str()).unwrap_or(""); - if state == "completed" { - image_url = pv - .get("assets") - .and_then(|a| a.get("image")) - .and_then(|v| v.as_str()) - .map(String::from); - if image_url.is_none() { - return Err(AiMuxError::InvalidResponseData(format!( - "Luma generation {generation_id} completed without assets.image" - ))); - } - break; - } - if state == "failed" { - return Err(AiMuxError::ApiCall(Box::new(ApiCallError { - status_code: Some(200), - provider_code: Some(state.to_string()), - message: "Image generation failed.".into(), - response_body, - ..ApiCallError::new( - "Image generation failed.", - self.generations_url(Some(&generation_id)), - serde_json::json!({}), - ) - }))); - } - sleep_or_abort( - std::time::Duration::from_millis(poll_interval), - options.abort_signal.as_ref(), - ) - .await?; - } + let response_body = pr.raw_value.as_ref().map(ToString::to_string); + let pv = pr.value; - let image_url = image_url.ok_or_else(|| { - AiMuxError::Timeout(format!( - "luma generation {generation_id} polling timed out after {max_poll_attempts} attempts ({}ms)", - max_poll_attempts * poll_interval - )) - })?; + let state = pv.get("state").and_then(|v| v.as_str()).unwrap_or(""); + if state == "completed" { + return pv + .get("assets") + .and_then(|a| a.get("image")) + .and_then(|v| v.as_str()) + .map(|url| PollStep::Ready(url.to_string())) + .ok_or_else(|| { + AiMuxError::InvalidResponseData(format!( + "Luma generation {generation_id} completed without assets.image" + )) + }); + } + if state == "failed" { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some(state.to_string()), + message: "Image generation failed.".into(), + response_body, + ..ApiCallError::new( + "Image generation failed.", + poll_url.clone(), + serde_json::json!({}), + ) + }))); + } + Ok(PollStep::Pending) + }, + ) + .await?; // Download image; assets.image is a URL from the poll response body, // so it goes through the SSRF download guard. - let ir = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: image_url.clone(), - headers: vec![], - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: true, - trusted_origin: Some(self.config.base_url.clone()), - credentialed_origin: Some(self.config.base_url.clone()), - }, - aimux_provider_utils::create_binary_response_handler(), - aimux_provider_utils::create_status_code_error_response_handler(), - ) - }) - .await?; + let ir = retry_download(options.abort_signal.as_ref(), poll_interval, || { + aimux_provider_utils::get_from_api( + HttpRequest { + url: image_url.clone(), + abort_signal: options.abort_signal.clone(), + validate_url: true, + trusted_origin: Some(exchange.base_url().to_string()), + credentialed_origin: Some(exchange.base_url().to_string()), + ..Default::default() + }, + aimux_provider_utils::create_binary_response_handler(), + aimux_provider_utils::create_status_code_error_response_handler(), + ) + }) + .await?; let image_bytes = ir.value.to_vec(); Ok(ImageResult { diff --git a/aimux-providers/src/luma/options.rs b/aimux-providers/src/luma/options.rs new file mode 100644 index 000000000..705780bd0 --- /dev/null +++ b/aimux-providers/src/luma/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Luma. +//! +//! Luma reads `providerOptions.luma`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "luma"; + +/// The Luma options in a providerOptions map. +pub(crate) fn luma_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/mistral/convert.rs b/aimux-providers/src/mistral/convert.rs index 5985bf014..0abad46fb 100644 --- a/aimux-providers/src/mistral/convert.rs +++ b/aimux-providers/src/mistral/convert.rs @@ -12,6 +12,7 @@ use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat use aimux_core::shared::{FileBytes, FileData}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified}; +use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; use serde_json::{Value, json}; // ── Prepared tools ────────────────────────────────────────────────────────── @@ -108,7 +109,7 @@ pub fn prepare_tools( /// does not carry the tool name on `ToolResult` parts). /// /// # Errors -/// Returns an error for unsupported text file data. +/// Returns an error for text file parts or an unresolved inline file media type. pub fn convert_prompt_to_mistral_messages( prompt: &LanguageModelPrompt, ) -> Result, AiMuxError> { @@ -231,21 +232,21 @@ fn convert_part_to_mistral(part: &UserPart) -> Result { UserPart::Text(TextPart { text, .. }) => json!({ "type": "text", "text": text }), UserPart::File(file) => { use base64::Engine; - let is_image = file.media_type.split('/').next() == Some("image"); + 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(), }; - let media_type = crate::google::convert::tool_file_media_type(file)?; - if !is_image && media_type != "application/pdf" { - return Err(AiMuxError::UnsupportedFunctionality( - "Only images and PDF file parts are supported".to_string(), - )); - } format!("data:{media_type};base64,{b64}") } FileData::Url { url, .. } => { @@ -258,12 +259,12 @@ fn convert_part_to_mistral(part: &UserPart) -> Result { } FileData::Reference { .. } => { return Err(AiMuxError::UnsupportedFunctionality( - "file parts with provider references".to_string(), + "file parts with provider references".into(), )); } FileData::Text { .. } => { return Err(AiMuxError::UnsupportedFunctionality( - "text file parts".to_string(), + "text file parts".into(), )); } }; @@ -281,7 +282,7 @@ fn convert_part_to_mistral(part: &UserPart) -> Result { /// Convert `CallOptions` to a Mistral request body. /// /// # Errors -/// Returns an error for unsupported text file data. +/// Returns an error for text file parts or an unresolved inline file media type. pub fn build_request_body( model_id: &str, options: &CallOptions, @@ -290,21 +291,25 @@ pub fn build_request_body( let provider_options = options .provider_options .as_ref() - .and_then(|namespaces| namespaces.get("mistral")) + .and_then(|namespaces| namespaces.get(crate::mistral::options::NAMESPACE)) .map(|namespace| { - crate::openai::convert::parse_option_fields(namespace, "mistral", |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(()) }) - }) + 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(); diff --git a/aimux-providers/src/mistral/embedding.rs b/aimux-providers/src/mistral/embedding.rs index 6e1411de3..e1dd790f1 100644 --- a/aimux-providers/src/mistral/embedding.rs +++ b/aimux-providers/src/mistral/embedding.rs @@ -5,8 +5,6 @@ //! //! Endpoint: `POST {base_url}/embeddings` -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -16,55 +14,30 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::HttpRequest; - -use super::MistralConfig; +use crate::shared::EndpointConfig; /// A Mistral embedding model (e.g. `"mistral-embed"`). pub struct MistralEmbeddingModel { model_id: String, - config: MistralConfig, + config: EndpointConfig, } impl MistralEmbeddingModel { - #[must_use] - pub fn new(model_id: String, config: MistralConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/embeddings", self.config.base_url) - } } #[async_trait] impl EmbeddingModel for MistralEmbeddingModel { fn provider(&self) -> &str { - "mistral" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - fn max_embeddings_per_call(&self) -> Option { Some(32) } @@ -93,17 +66,11 @@ impl EmbeddingModel for MistralEmbeddingModel { } body.insert("encoding_format".to_string(), json!("float")); - let headers = self.build_headers(options.headers.as_ref()); - - let mut header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - header_list.push(("Content-Type".to_string(), "application/json".to_string())); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), header_list, options), - Value::Object(body), + exchange.request(exchange.url("/embeddings"), options), + exchange.transform_body(Value::Object(body)), aimux_provider_utils::create_json_response_handler(), super::mistral_failed_response_handler(), ) @@ -163,7 +130,7 @@ struct MistralEmbeddingProviderOptions { fn parse_mistral_provider_options( options: Option<&SharedProviderOptions>, ) -> MistralEmbeddingProviderOptions { - let provider_opts = options.and_then(|opts| opts.get("mistral")); + let provider_opts = super::options::mistral_options(options); MistralEmbeddingProviderOptions { metadata: provider_opts.and_then(|o| o.get("metadata")).cloned(), output_dimension: provider_opts diff --git a/aimux-providers/src/mistral/mod.rs b/aimux-providers/src/mistral/mod.rs index 7b245b27b..9be840170 100644 --- a/aimux-providers/src/mistral/mod.rs +++ b/aimux-providers/src/mistral/mod.rs @@ -1,5 +1,11 @@ //! Mistral AI provider. //! +//! [`create_mistral`] is the Rust form of the AI SDK's `createMistral`: it +//! takes [`MistralProviderSettings`], validates the base URL, fixes the +//! provider name and returns a [`MistralProvider`]. The API key is not read +//! there; it is loaded in the request headers of every call, from the setting +//! or from `MISTRAL_API_KEY`. [`mistral()`] is the default instance. +//! //! OpenAI-compatible chat completions API with Mistral-specific differences: //! - Tool choice uses `"any"` instead of `"required"` //! - Content can be a string or an array of typed parts (text, thinking, image_url) @@ -9,16 +15,27 @@ pub mod convert; pub mod embedding; mod model; +pub(crate) mod options; mod types; pub use embedding::MistralEmbeddingModel; +pub use model::MistralModel; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::{RetryConfig, load_api_key, without_trailing_slash}; +use std::sync::{Arc, OnceLock}; + +use futures::future::BoxFuture; use serde_json::Value; +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; + pub(crate) fn mistral_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -75,156 +92,174 @@ pub(crate) fn mistral_stream_error( ) } -/// Configuration for the Mistral provider. -#[derive(Debug, Clone)] -pub struct MistralConfig { - pub api_key: String, - pub base_url: String, - /// api_key 来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 - pub api_key_source: Option, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, -} +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.mistral.ai/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "MISTRAL_API_KEY"; +const DEFAULT_NAME: &str = "mistral"; -impl MistralConfig { - /// Create from an API key (uses default Mistral base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.mistral.ai/v1".to_string(), - api_key_source: None, - retry_config: RetryConfig::default(), - } - } +/// The URL patterns the chat model fetches itself (`supportedUrls` of the AI +/// SDK's `MistralChatLanguageModel`): `https` PDFs. +fn chat_supported_urls() -> SupportedUrls { + let https = regex::Regex::new(r"^https://.*$").expect("static pattern"); + SupportedUrls(std::iter::once(("application/pdf".to_string(), vec![https])).collect()) +} - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Settings of [`create_mistral`] (the AI SDK's `MistralProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct MistralProviderSettings { + /// Base URL for the API calls. Default `https://api.mistral.ai/v1`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `MISTRAL_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. + pub api_key: Option, + /// Extra headers on every request. A `None` value removes the header, + /// including `Authorization`. Per-call headers win over these. + pub headers: Option, + /// The transport: a mock, a signing decorator, a proxy-aware client. + /// `None` uses the process default, resolved per request. + pub fetch: Option, + /// Generates unique IDs for chat output. + pub generate_id: Option String + Send + Sync>>, +} - /// 标注 api_key 来源(RFC-0023 回放重建用)。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self +impl std::fmt::Debug for MistralProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("MistralProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key.is_some()) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("fetch", &self.fetch.is_some()) + .field("generate_id", &self.generate_id.is_some()) + .finish() } +} - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self - } +/// Create a Mistral provider. +/// +/// # 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. +pub fn create_mistral(settings: MistralProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => 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, + ), + options::NAMESPACE, + "4.0.54", + ), + fetch: settings.fetch, + generate_id: settings.generate_id, + }) +} - /// Create from environment variable `MISTRAL_API_KEY`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `MISTRAL_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "MISTRAL_API_KEY", "Mistral")?; - Ok(Self::new(api_key).with_api_key_source(Some("env:MISTRAL_API_KEY"))) - } +/// The default provider: `create_mistral` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn mistral() -> &'static MistralProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_mistral(MistralProviderSettings::default()) + .expect("default Mistral settings are always valid") + }) } -/// Mistral provider — creates `MistralModel` instances. +/// A Mistral provider (the AI SDK's `MistralProvider`). Cheap to clone the +/// models out of; it holds no HTTP client. pub struct MistralProvider { - config: MistralConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, + generate_id: Option String + Send + Sync>>, } impl MistralProvider { + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) + } + + /// A chat model; `provider()` is `"{name}.chat"`. #[must_use] - pub fn new(config: MistralConfig) -> Self { - Self { config } + 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())), + ) + .with_generate_id(self.generate_id.clone()) } - /// Create a model instance for the given model name (e.g. `"mistral-small-latest"`). + /// An embedding model (e.g. `"mistral-embed"`); `provider()` is + /// `"{name}.embedding"`. #[must_use] - pub fn model(&self, model_id: &str) -> model::MistralModel { - model::MistralModel::new(model_id.to_string(), self.config.clone()) + pub fn embedding(&self, model_id: &str) -> MistralEmbeddingModel { + MistralEmbeddingModel::from_config(model_id.to_string(), self.model_config("embedding")) } - /// Create an embedding model instance for the given model name (e.g. - /// `"mistral-embed"`). + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; the same model as [`chat`](Self::chat) and + /// [`language_model`](Provider::language_model). #[must_use] - pub fn embedding_model(&self, model_id: &str) -> embedding::MistralEmbeddingModel { - embedding::MistralEmbeddingModel::new(model_id.to_string(), self.config.clone()) + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.chat(model_id)) } } impl Provider for MistralProvider { - fn name(&self) -> &str { - "mistral" + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) + } + + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) } - /// List models via `GET {base_url}/models` (OpenAI-compatible, RFC-0027). - fn list_models( - &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let config = self.config.clone(); + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } +} + +impl ProviderDiscovery for MistralProvider { + /// `GET {base_url}/models`: one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config("models"); Box::pin(async move { - let base = config.base_url.trim_end_matches('/'); - let url = format!("{base}/models"); - let headers = vec![ - ( - "Authorization".to_string(), - format!("Bearer {}", config.api_key), - ), - ("Content-Type".to_string(), "application/json".to_string()), - ]; - use aimux_provider_utils::HttpRequest; - // Retry rationale: see `openai::model::execute_list_models`. - let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: headers.clone(), - abort_signal: None, - call_id: None, - recording_context: None, - ..Default::default() - }, - aimux_provider_utils::create_json_response_handler(), - mistral_failed_response_handler(), - ) - }) - .await?; - #[derive(serde::Deserialize)] - struct Resp { - #[serde(default)] - data: Vec, - } - #[derive(serde::Deserialize)] - struct Entry { - id: String, - #[serde(default)] - owned_by: Option, - } - let parsed: Resp = resp.value; - let runtime: Vec = parsed - .data - .into_iter() - .map(|e| aimux_core::model_catalogue::RuntimeModel { - id: e.id, - owned_by: e.owned_by, - created: None, - }) - .collect(); - Ok(runtime) + crate::shared::list_data_models(&config, mistral_failed_response_handler()).await }) } } diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index b6441250d..c413602fa 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -5,14 +5,11 @@ //! - `content` in responses can be a string or an array of typed parts //! (text, thinking, image_url). Thinking parts are extracted as reasoning //! in streaming mode. -//! - Streaming tool calls arrive complete in a single chunk (no index-based -//! incremental accumulation). +//! - Streamed tool calls are assembled by the shared `StreamingToolCallTracker`. //! - Usage supports `num_cached_tokens` / `prompt_tokens_details.cached_tokens`. //! - Finish reasons include `model_length`. use aimux_core::tool::RawToolCall; -use std::collections::HashMap; - use async_trait::async_trait; use futures::StreamExt; use serde_json::Value; @@ -21,43 +18,40 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; -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; +use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker}; + +use crate::shared::EndpointConfig; -use super::MistralConfig; use super::convert::{build_request_body, parse_finish_reason}; use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; -/// An Mistral language model. +/// A Mistral language model. pub struct MistralModel { model_id: String, - config: MistralConfig, + config: EndpointConfig, + generate_id: std::sync::Arc String + Send + Sync>, } impl MistralModel { - pub fn new(model_id: String, config: MistralConfig) -> Self { - Self { model_id, config } - } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } + 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; } - headers + self } - fn endpoint(&self) -> String { - format!("{}/chat/completions", self.config.base_url) + 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), + } } } @@ -197,45 +191,23 @@ fn extract_reasoning_content(content: &Option) -> Option { #[async_trait] impl LanguageModel for MistralModel { fn provider(&self) -> &str { - "mistral" + &self.config.provider } - fn model_id(&self) -> &str { - &self.model_id + async fn supported_urls(&self) -> aimux_core::language_model::SupportedUrls { + (self.config.supported_urls)(&self.model_id) } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.base_url.clone()), - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: None, - } + fn model_id(&self) -> &str { + &self.model_id } async fn do_generate(&self, options: &CallOptions) -> Result { - let body = build_request_body(&self.model_id, options, false)?; - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let body = exchange.transform_body(build_request_body(&self.model_id, options, false)?); + let endpoint = exchange.url("/chat/completions"); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new( - self.endpoint(), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), + exchange.request(endpoint.clone(), options), body.clone(), aimux_provider_utils::create_json_response_handler(), super::mistral_failed_response_handler(), @@ -322,17 +294,11 @@ impl LanguageModel for MistralModel { } async fn do_stream(&self, options: &CallOptions) -> Result { - let body = build_request_body(&self.model_id, options, true)?; - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let body = exchange.transform_body(build_request_body(&self.model_id, options, true)?); + let endpoint = exchange.url("/chat/completions"); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new( - self.endpoint(), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), + exchange.request(endpoint.clone(), options), body.clone(), aimux_provider_utils::create_event_source_response_handler::(), super::mistral_failed_response_handler(), @@ -352,22 +318,26 @@ impl LanguageModel for MistralModel { { return Err(super::mistral_stream_error( err_obj, - &self.endpoint(), + &endpoint, body.clone(), response_headers.clone(), )); } - let stream_error_url = self.endpoint(); + 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 stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings: vec![] }); let text_id = 0usize; let mut text_started = false; let mut reasoning_started = false; + 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; let mut final_usage = Usage::default(); let mut final_finish_reason: Option = None; @@ -429,8 +399,8 @@ impl LanguageModel for MistralModel { final_usage = convert_usage(usage); } - // Process choices. - for choice in chunk.choices { + // 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) @@ -444,13 +414,7 @@ impl LanguageModel for MistralModel { }); text_started = false; } - let rid = format!( - "rc-{}", - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_nanos()) - .unwrap_or(0) - ); + let rid = generate_id(); reasoning_id = Some(rid.clone()); reasoning_started = true; yield Ok(StreamPart::ReasoningStart { id: rid, @@ -495,44 +459,26 @@ impl LanguageModel for MistralModel { }); } - // Tool calls (complete in a single chunk). + // Tool calls: assembled by the shared tracker. if let Some(tool_call_deltas) = choice.delta.tool_calls { - for dtc in tool_call_deltas { - let tool_id = dtc.id; - let tool_name = dtc.function.name; - let args = dtc.function.arguments; - - yield Ok(StreamPart::ToolInputStart { - id: tool_id.clone(), - tool_name: tool_name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - - if !args.is_empty() { - yield Ok(StreamPart::ToolInputDelta { - id: tool_id.clone(), - delta: args.clone(), - provider_metadata: None, - }); + for dtc in &tool_call_deltas { + let delta = StreamingToolCallDelta { + index: dtc.index, + id: dtc.id.as_deref(), + name: dtc.function.name.as_deref(), + arguments: dtc.function.arguments.as_deref(), + ..Default::default() + }; + // A new call without a function name is + // invalid response data, as in the AI SDK; + // the stream ends. + if let Err(error) = tool_calls.process(delta, &mut tool_parts) { + yield Err(error.into()); + return; + } + for part in tool_parts.drain(..) { + yield Ok(part); } - - yield Ok(StreamPart::ToolInputEnd { - id: tool_id.clone(), - provider_metadata: None, - }); - - let input = args; - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: tool_id, - tool_name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); } } @@ -584,6 +530,11 @@ impl LanguageModel for MistralModel { }); } + tool_calls.finish(&mut tool_parts); + for part in tool_parts.drain(..) { + yield Ok(part); + } + yield Ok(StreamPart::Finish { finish_reason: if stream_errored { FinishReason { @@ -601,7 +552,7 @@ impl LanguageModel for MistralModel { } else { final_usage }, - provider_metadata: Some(provider_namespace("mistral", serde_json::json!({})).expect("provider metadata must be an object")), + provider_metadata: Some(super::options::mistral_metadata(serde_json::json!({}))), }); }; @@ -614,24 +565,3 @@ impl LanguageModel for MistralModel { }) } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn config_snapshot_reports_config() { - let model = MistralModel::new( - "mistral-small-latest".to_string(), - MistralConfig::new("test-key"), - ); - let snapshot = model.config_snapshot(); - assert_eq!(snapshot.provider, "mistral"); - assert_eq!(snapshot.model_id, "mistral-small-latest"); - assert_eq!( - snapshot.base_url.as_deref(), - Some("https://api.mistral.ai/v1") - ); - assert_eq!(snapshot.api_key_source, "explicit"); - } -} diff --git a/aimux-providers/src/mistral/options.rs b/aimux-providers/src/mistral/options.rs new file mode 100644 index 000000000..8096e700b --- /dev/null +++ b/aimux-providers/src/mistral/options.rs @@ -0,0 +1,24 @@ +//! The providerOptions and providerMetadata namespace of Mistral. +//! +//! `@ai-sdk/mistral` reads `providerOptions.mistral` and writes response +//! 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; + +/// The providerOptions / providerMetadata key. +pub(crate) const NAMESPACE: &str = "mistral"; + +/// The Mistral options in a providerOptions map. +pub(crate) fn mistral_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + 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") +} diff --git a/aimux-providers/src/mistral/types.rs b/aimux-providers/src/mistral/types.rs index 6616e35e9..faa05b1e1 100644 --- a/aimux-providers/src/mistral/types.rs +++ b/aimux-providers/src/mistral/types.rs @@ -102,20 +102,23 @@ pub struct Delta { /// Content can be null, a string, or an array of typed parts. #[serde(default)] pub content: Option, - /// Mistral streams a complete tool call in a single chunk (no index-based - /// incremental accumulation like OpenAI). #[serde(default)] pub tool_calls: Option>, } #[derive(Debug, Deserialize)] pub struct DeltaToolCall { - pub id: String, + #[serde(default)] + pub index: Option, + #[serde(default)] + pub id: Option, pub function: DeltaFunction, } #[derive(Debug, Deserialize)] pub struct DeltaFunction { - pub name: String, - pub arguments: String, + #[serde(default)] + pub name: Option, + #[serde(default)] + pub arguments: Option, } diff --git a/aimux-providers/src/mistralrs.rs b/aimux-providers/src/mistralrs.rs deleted file mode 100644 index 22b6e8f42..000000000 --- a/aimux-providers/src/mistralrs.rs +++ /dev/null @@ -1,100 +0,0 @@ -//! Mistral.rs provider — a thin OpenAI-compatible wrapper for the local -//! [mistral.rs](https://github.com/EricLBuehler/mistral.rs) inference server. -//! -//! mistral.rs exposes an OpenAI-compatible Chat Completions API at -//! `http://127.0.0.1:8080/v1` by default. The Rust -//! [`OpenAIProvider`] appends `/chat/completions` -//! to this base URL, yielding `http://127.0.0.1:8080/v1/chat/completions`. -//! -//! Like llamafile, mistral.rs runs locally and does not require authentication. -//! Accordingly the `MISTRALRS_BASE_URL` environment variable holds a *base URL* -//! (not an API key); [`MistralrsConfig::from_env`] reads it and falls back to -//! the default local endpoint when it is unset. A placeholder API key is sent in -//! the `Authorization` header — mistral.rs ignores it, but the shared -//! `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -/// Environment variable holding the mistral.rs base URL (not an API key). -const ENV_VAR: &str = "MISTRALRS_BASE_URL"; -const PROVIDER_NAME: &str = "mistralrs"; -/// Placeholder API key — mistral.rs does not authenticate, but the shared -/// `OpenAIConfig` requires a non-empty key string. -const PLACEHOLDER_API_KEY: &str = "mistralrs"; - -/// Configuration for the mistral.rs provider (wraps [`OpenAIConfig`]). -pub struct MistralrsConfig(OpenAIConfig); - -impl MistralrsConfig { - /// Create from an API key, using the default local mistral.rs base URL. - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `MISTRALRS_BASE_URL` environment variable. - /// - /// `MISTRALRS_BASE_URL` holds a *base URL* (e.g. - /// `http://127.0.0.1:8080/v1`), not an API key — mistral.rs is a local - /// inference server that does not require authentication. When the variable - /// is unset (or empty), the default local endpoint - /// (`http://127.0.0.1:8080/v1`) is used. - /// - /// # Errors - /// - /// Never returns an error; an unset `MISTRALRS_BASE_URL` falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - /// Override the base URL (useful for tests / non-default ports). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// mistral.rs provider — creates [`OpenAIModel`] instances pointed at mistral.rs. -pub struct MistralrsProvider(OpenAIProvider); - -impl MistralrsProvider { - #[must_use] - pub fn new(config: MistralrsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given mistral.rs model id - /// (e.g. `"Qwen/Qwen3-4B"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for MistralrsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/mlx.rs b/aimux-providers/src/mlx.rs deleted file mode 100644 index cebacf5dd..000000000 --- a/aimux-providers/src/mlx.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! MLX (Apple Silicon) provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `MLX_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "MLX_BASE_URL"; -const PROVIDER_NAME: &str = "mlx"; -const PLACEHOLDER_API_KEY: &str = "mlx"; - -pub struct MlxConfig(OpenAIConfig); - -impl MlxConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `MLX_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct MlxProvider(OpenAIProvider); - -impl MlxProvider { - #[must_use] - pub fn new(config: MlxConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for MlxProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/ollama.rs b/aimux-providers/src/ollama.rs deleted file mode 100644 index 5ada0c676..000000000 --- a/aimux-providers/src/ollama.rs +++ /dev/null @@ -1,110 +0,0 @@ -//! Ollama provider — a thin OpenAI-compatible wrapper for the local -//! [Ollama](https://ollama.com) inference server. -//! -//! Ollama exposes an OpenAI-compatible Chat Completions API at -//! `http://127.0.0.1:11434/v1` by default. The Rust -//! [`OpenAIProvider`] appends `/chat/completions` -//! to this base URL, yielding `http://127.0.0.1:11434/v1/chat/completions`. -//! -//! Unlike hosted providers, Ollama runs locally and does not require -//! authentication. Accordingly the `OLLAMA_BASE_URL` environment variable -//! holds a *base URL* (not an API key); [`OllamaConfig::from_env`] reads it -//! and falls back to the default local endpoint when it is unset. A placeholder -//! API key is sent in the `Authorization` header — Ollama ignores it, but the -//! shared `OpenAIProvider` requires a non-empty key string. -//! -//! Note: Ollama also has a native API at `/api/chat` (NDJSON, not SSE). This -//! provider uses the OpenAI-compatible endpoint only. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:11434/v1"; -/// Environment variable holding the Ollama base URL (not an API key). -const ENV_VAR: &str = "OLLAMA_BASE_URL"; -const PROVIDER_NAME: &str = "ollama"; -/// Placeholder API key — Ollama does not authenticate, but the shared -/// `OpenAIConfig` requires a non-empty key string. -const PLACEHOLDER_API_KEY: &str = "ollama"; - -/// Configuration for the Ollama provider (wraps [`OpenAIConfig`]). -pub struct OllamaConfig(OpenAIConfig); - -impl OllamaConfig { - /// Create from an API key, using the default local Ollama base URL. - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()), - ) - } - - /// Create from the `OLLAMA_BASE_URL` environment variable. - /// - /// `OLLAMA_BASE_URL` holds a *base URL* (e.g. - /// `http://127.0.0.1:11434/v1`), not an API key — Ollama is a local - /// inference server that does not require authentication. When the variable - /// is unset (or empty), the default local endpoint - /// (`http://127.0.0.1:11434/v1`) is used. - /// - /// # Errors - /// - /// Never returns an error; an unset `OLLAMA_BASE_URL` falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let mut config = Self::new(PLACEHOLDER_API_KEY); - config.0 = config.0.with_api_key_source(Some("none")); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - /// Override the base URL (useful for tests / non-default ports). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// Ollama provider — creates [`OpenAIModel`] instances pointed at Ollama. -pub struct OllamaProvider(OpenAIProvider); - -impl OllamaProvider { - #[must_use] - pub fn new(config: OllamaConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Ollama model id - /// (e.g. `"llama3.2"` or `"qwen3:4b"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for OllamaProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - fn decision_model( - &self, - model_id: &str, - ) -> Result, AiMuxError> { - self.0.decision_model(model_id) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/omlx.rs b/aimux-providers/src/omlx.rs deleted file mode 100644 index 56b82ceff..000000000 --- a/aimux-providers/src/omlx.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! OMLX / MLX LM provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `OMLX_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "OMLX_BASE_URL"; -const PROVIDER_NAME: &str = "omlx"; -const PLACEHOLDER_API_KEY: &str = "omlx"; - -pub struct OmlxConfig(OpenAIConfig); - -impl OmlxConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `OMLX_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct OmlxProvider(OpenAIProvider); - -impl OmlxProvider { - #[must_use] - pub fn new(config: OmlxConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for OmlxProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/onnx.rs b/aimux-providers/src/onnx.rs deleted file mode 100644 index ed48f01f2..000000000 --- a/aimux-providers/src/onnx.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! ONNX Runtime provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `ONNX_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "ONNX_BASE_URL"; -const PROVIDER_NAME: &str = "onnx"; -const PLACEHOLDER_API_KEY: &str = "onnx"; - -pub struct OnnxConfig(OpenAIConfig); - -impl OnnxConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `ONNX_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct OnnxProvider(OpenAIProvider); - -impl OnnxProvider { - #[must_use] - pub fn new(config: OnnxConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for OnnxProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/oobabooba.rs b/aimux-providers/src/oobabooba.rs deleted file mode 100644 index b62b9ea09..000000000 --- a/aimux-providers/src/oobabooba.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! Oobabooga Text Generation WebUI provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:5000/v1`. The `OOBABOOBA_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:5000/v1"; -const ENV_VAR: &str = "OOBABOOBA_BASE_URL"; -const PROVIDER_NAME: &str = "oobabooba"; -const PLACEHOLDER_API_KEY: &str = "oobabooba"; - -pub struct OobaboobaConfig(OpenAIConfig); - -impl OobaboobaConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `OOBABOOBA_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:5000/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct OobaboobaProvider(OpenAIProvider); - -impl OobaboobaProvider { - #[must_use] - pub fn new(config: OobaboobaConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for OobaboobaProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs index c79ebcdb1..ac6270479 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -22,7 +22,6 @@ use aimux_core::language_model_message::{ ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, }; use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice}; -use aimux_core::provider::Provider; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; @@ -31,7 +30,14 @@ use aimux_core::types::{ FinishReason, FinishReasonUnified, ReasoningEffort, ResponseMetadata, Usage, Warning, }; -use aimux_provider_utils::HttpRequest; +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::image_model::ImageModel; +use aimux_core::provider::Provider; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, combine_headers, validate_base_url, +}; + +use crate::shared::{Credential, EndpointConfig, TransformRequestBody, provider_headers}; fn open_responses_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -84,125 +90,167 @@ fn open_responses_successful_response_handler() -> aimux_provider_utils::Respons }) } -// == Config == +// == Settings == -/// Configuration for the Open Responses provider. +/// Settings of [`create_open_responses`]. /// -/// Mirrors the TS `OpenResponsesConfig`. The `url` is the **full endpoint -/// URL** (e.g. `https://localhost:1234/v1/responses`), not a base URL. -pub struct OpenResponsesConfig { - /// Provider name reported by `LanguageModel::provider`. - pub provider: String, - /// Key used to look up provider-specific options in `providerOptions`. - pub provider_options_name: String, - /// Full endpoint URL for the Responses API. - pub url: String, - /// Optional header factory - called on every request and merged with - /// per-request headers (e.g. `Authorization`). - pub headers: Option HashMap + Send + Sync>>, - /// ID generator (retained for API parity with the TS config; not used - /// by the model itself). - pub generate_id: Arc String + Send + Sync>, - /// 凭证来源(RFC-0023):`None` = 未标注(回放时由 headers 闭包推断); - /// `Some("explicit")` / `Some("env:VAR")` = 调用方显式标注。Open Responses - /// 是通用包装,认证由调用方经 `headers` 闭包管理,故默认 `None`。 - pub api_key_source: Option, -} - -impl std::fmt::Debug for OpenResponsesConfig { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("OpenResponsesConfig") - .field("provider", &self.provider) - .field("provider_options_name", &self.provider_options_name) - .field("url", &self.url) - .field("headers", &self.headers.is_some()) - .field("generate_id", &"") - .field("api_key_source", &self.api_key_source) - .finish() - } +/// Everything but `api_key` and `headers` is fixed when the provider is +/// created; those two are evaluated on every request. +#[derive(Clone)] +pub struct OpenResponsesProviderSettings { + /// The provider name: `provider()` is `"{name}.responses"` and its first + /// dot-separated segment is the providerOptions key the model reads. + pub name: String, + /// Base URL of the server (e.g. `http://localhost:1234/v1`); requests go + /// to `{base_url}/responses`. A trailing slash is removed. + pub base_url: String, + /// The API key, sent as `Authorization: Bearer`. `None` sends no + /// credential (a local server) and reads no environment variable. An + /// explicit value is used as given, `""` included. + pub api_key: Option>, + /// Extra headers, resolved on every request (a `None` value removes a + /// header, including `Authorization`). Per-call headers win over these. + pub headers: Option>, + /// The transport: a mock, a signing decorator, a proxy-aware client. + /// `None` uses the process default, resolved per request. + pub fetch: Option, + /// Rewrites every JSON request body once, after it is serialized and + /// before it is sent. + pub transform_request_body: Option, } -impl OpenResponsesConfig { - /// Create a new config with the given provider name, provider-options - /// name, and endpoint URL. `headers` defaults to `None` and `generate_id` - /// defaults to a simple counter. - pub fn new( - provider: impl Into, - provider_options_name: impl Into, - url: impl Into, - ) -> Self { +impl OpenResponsesProviderSettings { + /// Settings with the two required fields; the rest unset. + #[must_use] + pub fn new(name: impl Into, base_url: impl Into) -> Self { Self { - provider: provider.into(), - provider_options_name: provider_options_name.into(), - url: url.into(), + name: name.into(), + base_url: base_url.into(), + api_key: None, headers: None, - generate_id: Arc::new(|| { - use std::sync::atomic::{AtomicU64, Ordering}; - static COUNTER: AtomicU64 = AtomicU64::new(0); - let n = COUNTER.fetch_add(1, Ordering::Relaxed); - format!("id-{n}") - }), - api_key_source: None, + fetch: None, + transform_request_body: None, } } +} - /// Set a header factory. - #[must_use] - pub fn with_headers(mut self, headers: F) -> Self - where - F: Fn() -> HashMap + Send + Sync + 'static, - { - self.headers = Some(Arc::new(headers)); - self +impl std::fmt::Debug for OpenResponsesProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OpenResponsesProviderSettings") + .field("name", &self.name) + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field("headers", &self.headers.is_some()) + .field("fetch", &self.fetch.is_some()) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ) + .finish() } +} - /// Set a generate-id factory. - #[must_use] - pub fn with_generate_id(mut self, generate_id: F) -> Self - where - F: Fn() -> String + Send + Sync + 'static, - { - self.generate_id = Arc::new(generate_id); - self - } +/// The first dot-separated segment of a provider name, trimmed: the +/// providerOptions key of the provider. +fn options_name_of(name: &str) -> String { + name.split('.') + .next() + .unwrap_or_default() + .trim() + .to_string() +} - /// 标注凭证来源(RFC-0023 回放重建用)。Open Responses 的认证由调用方经 - /// `headers` 闭包管理;此 setter 让调用方显式标注来源(如 `env:VAR`), - /// 覆盖 `config_snapshot` 对闭包的推断。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self +/// Create an Open Responses provider. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when `name` is empty or `base_url` is +/// not an `http(s)` URL with a host. Credentials are resolved per request, not +/// here. +pub fn create_open_responses( + settings: OpenResponsesProviderSettings, +) -> Result { + if settings.name.trim().is_empty() { + return Err(AiMuxError::InvalidArgument( + "Open Responses requires a non-empty `name`.".to_string(), + )); } + let base_url = validate_base_url(&settings.base_url)?; + let credential = match settings.api_key { + Some(key) => Credential::Explicit(key), + None => Credential::None, + }; + let provider_layer = provider_headers(credential, Vec::new(), None); + let user = settings.headers; + let headers: HeadersFn = Resolvable::from_async_fn(move || { + let provider_layer = provider_layer.clone(); + let user = user.clone(); + async move { + let layer = provider_layer.resolve().await?; + match &user { + Some(user) => Ok(combine_headers(&[&layer, &user.resolve().await?])), + None => Ok(layer), + } + } + }); + Ok(OpenResponsesProvider { + name: settings.name, + base_url, + headers, + fetch: settings.fetch, + transform_request_body: settings.transform_request_body, + }) } // == Provider == /// Open Responses provider - creates [`OpenResponsesModel`] instances. pub struct OpenResponsesProvider { - config: OpenResponsesConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, + transform_request_body: Option, } impl OpenResponsesProvider { + /// A Responses model; `provider()` is `"{name}.responses"`. #[must_use] - pub fn new(config: OpenResponsesConfig) -> Self { - Self { config } + pub fn responses(&self, model_id: &str) -> OpenResponsesModel { + OpenResponsesModel { + model_id: model_id.to_string(), + provider_options_name: options_name_of(&self.name), + config: EndpointConfig::fixed( + format!("{}.responses", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + self.transform_request_body.clone(), + ), + } } - /// Create a model instance for the given model id. + /// The provider as a function: the default language model for an id. The + /// same model as [`responses`](Self::responses) and + /// [`language_model`](Provider::language_model). #[must_use] - pub fn model(&self, model_id: &str) -> OpenResponsesModel { - OpenResponsesModel::new(model_id.to_string(), &self.config) + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.responses(model_id)) } } impl Provider for OpenResponsesProvider { - fn name(&self) -> &str { - &self.config.provider + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) } } @@ -211,75 +259,9 @@ impl Provider for OpenResponsesProvider { /// An Open Responses language model. pub struct OpenResponsesModel { model_id: String, - config: OpenResponsesConfig, -} - -impl OpenResponsesModel { - #[must_use] - pub fn new(model_id: String, config: &OpenResponsesConfig) -> Self { - Self { - model_id, - config: OpenResponsesConfig { - provider: config.provider.clone(), - provider_options_name: config.provider_options_name.clone(), - url: config.url.clone(), - headers: config.headers.clone(), - generate_id: config.generate_id.clone(), - api_key_source: config.api_key_source.clone(), - }, - } - } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = self - .config - .headers - .as_ref() - .map(|h| h()) - .unwrap_or_default(); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - /// Resolve the credential source for `config_snapshot` (M2b). - /// - /// Priority: - /// 1. An explicit `api_key_source` set via [`OpenResponsesConfig::with_api_key_source`] - /// (e.g. `"env:VAR"`) — lets callers mark env-sourced auth precisely. - /// 2. Otherwise, inspect the `headers` closure for a known auth header - /// *key* (`authorization` / `x-api-key` / `x-goog-api-key` / `api-key`). - /// Only key names are inspected — header *values* are never recorded, so - /// no secret leaks. If an auth header is present the source is - /// `"explicit"` (caller-managed); this is the same closure already - /// invoked by `build_headers` on every request, so calling it here is - /// no riskier. - /// 3. No headers closure (and no explicit source) → `"none"` (e.g. a local - /// LM Studio server with no auth). - fn resolve_api_key_source(&self) -> String { - if let Some(source) = &self.config.api_key_source { - return source.clone(); - } - let Some(headers_factory) = &self.config.headers else { - return "none".to_string(); - }; - let headers = headers_factory(); - let has_auth = headers.keys().any(|k| { - let lower = k.to_ascii_lowercase(); - matches!( - lower.as_str(), - "authorization" | "x-api-key" | "x-goog-api-key" | "api-key" - ) - }); - if has_auth { - "explicit".to_string() - } else { - "none".to_string() - } - } + /// The providerOptions key read (first segment of the provider name). + provider_options_name: String, + config: EndpointConfig, } #[async_trait] @@ -292,45 +274,21 @@ impl LanguageModel for OpenResponsesModel { &self.model_id } - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - // M2b: generic Responses wrapper — auth is caller-managed via the - // `headers` closure. Record the real credential source: an explicit - // marker if set, else inferred from the closure (auth header key - // present → "explicit"; no closure → "none"). Only key names are - // inspected — header values (secrets) are never serialized. - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.url.clone()), - api_key_source: self.resolve_api_key_source(), - profile: None, - provider_options: Some(serde_json::json!({ - "provider_options_name": self.config.provider_options_name, - "has_headers": self.config.headers.is_some(), - })), - } - } - async fn do_generate(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()); validate_tool_output_media(&options.prompt)?; + let exchange = self.config.exchange(options.headers.as_ref()).await?; let (body, warnings) = - build_request_body(&self.model_id, options, &self.config.provider_options_name); + build_request_body(&self.model_id, options, &self.provider_options_name); + let body = exchange.transform_body(body); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new( - self.config.url.clone(), - headers.into_iter().collect(), - options, - ), + exchange.request(exchange.url("/responses"), options), body.clone(), open_responses_successful_response_handler(), open_responses_failed_response_handler(), ) .await?; - let response_body = resp.raw_value; let response_headers = resp.response_headers; let raw: Value = resp.value; @@ -449,7 +407,7 @@ impl LanguageModel for OpenResponsesModel { provider_metadata: None, response: Some(aimux_core::shared::ResponseInfo { headers: Some(response_headers), - body: response_body, + body: Some(raw), ..response.into() }), request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), @@ -458,20 +416,21 @@ impl LanguageModel for OpenResponsesModel { async fn do_stream(&self, options: &CallOptions) -> Result { validate_tool_output_media(&options.prompt)?; - let headers = self.build_headers(options.headers.as_ref()); - let (mut stream_body, warnings) = - build_request_body(&self.model_id, options, &self.config.provider_options_name); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let (body, warnings) = + build_request_body(&self.model_id, options, &self.provider_options_name); - if let Some(obj) = stream_body.as_object_mut() { - obj.insert("stream".to_string(), json!(true)); - } + let stream_body = { + let mut b = body.clone(); + if let Some(obj) = b.as_object_mut() { + obj.insert("stream".to_string(), json!(true)); + } + exchange.transform_body(b) + }; + let endpoint = exchange.url("/responses"); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new( - self.config.url.clone(), - headers.into_iter().collect(), - options, - ), + exchange.request(endpoint.clone(), options), stream_body.clone(), aimux_provider_utils::create_event_source_response_handler::(), open_responses_failed_response_handler(), @@ -510,7 +469,7 @@ impl LanguageModel for OpenResponsesModel { provider_code, status_code, event, - self.config.url.clone(), + endpoint, stream_body, response_headers.clone(), )); @@ -824,9 +783,7 @@ impl LanguageModel for OpenResponsesModel { Ok(StreamResult { stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { - body: Some(stream_body), - }), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), response: Some(aimux_core::shared::StreamResponseInfo { headers: Some(response_headers), }), @@ -1294,7 +1251,6 @@ fn convert_user_content(content: &[UserPart], warnings: &mut Vec) -> Va } fn validate_tool_output_media(prompt: &LanguageModelPrompt) -> Result<(), AiMuxError> { - use base64::Engine; for message in prompt { let LanguageModelMessage::Tool { content, .. } = message else { continue; @@ -1308,27 +1264,11 @@ fn validate_tool_output_media(prompt: &LanguageModelPrompt) -> Result<(), AiMuxE continue; }; for item in value { - let ToolResultContent::File(FilePart { - data: FileData::Data { data }, - media_type, - .. - }) = item - else { - continue; - }; - let bytes = match data { - FileBytes::Binary(bytes) => std::borrow::Cow::Borrowed(bytes.as_slice()), - FileBytes::Base64(data) => std::borrow::Cow::Owned( - base64::engine::general_purpose::STANDARD - .decode(data) - .map_err(|error| { - AiMuxError::InvalidArgument(format!( - "invalid base64 file data: {error}" - )) - })?, - ), - }; - crate::anthropic::convert::resolve_full_media_type(media_type, &bytes)?; + if let ToolResultContent::File(part) = item + && matches!(part.data, FileData::Data { .. }) + { + aimux_provider_utils::resolve_full_media_type(part)?; + } } } } @@ -1356,16 +1296,7 @@ fn resolve_tool_result_output( let mut converted = match &part.data { FileData::Data { data } => { use base64::Engine; - let bytes = match data { - FileBytes::Binary(bytes) => bytes.clone(), - FileBytes::Base64(data) => base64::engine::general_purpose::STANDARD - .decode(data) - .unwrap_or_default(), - }; - let media_type = match crate::anthropic::convert::resolve_full_media_type( - &part.media_type, - &bytes, - ) { + let media_type = match aimux_provider_utils::resolve_full_media_type(part) { Ok(media_type) => media_type, Err(error) => { warnings.push(Warning::Other { @@ -1375,8 +1306,8 @@ fn resolve_tool_result_output( } }; let b64 = match data { - FileBytes::Binary(_) => { - base64::engine::general_purpose::STANDARD.encode(&bytes) + FileBytes::Binary(bytes) => { + base64::engine::general_purpose::STANDARD.encode(bytes) } FileBytes::Base64(data) => data.clone(), }; diff --git a/aimux-providers/src/openai/completion.rs b/aimux-providers/src/openai/completion.rs new file mode 100644 index 000000000..042aa3ee8 --- /dev/null +++ b/aimux-providers/src/openai/completion.rs @@ -0,0 +1,5 @@ +//! Native text-completion language model. + +/// The native provider's `/completions` model. +pub type OpenAICompletionModel = + crate::openai_compatible::completion::OpenAICompatibleCompletionModel; diff --git a/aimux-providers/src/openai/config.rs b/aimux-providers/src/openai/config.rs new file mode 100644 index 000000000..f9d29f127 --- /dev/null +++ b/aimux-providers/src/openai/config.rs @@ -0,0 +1,168 @@ +//! The private per-model configuration every OpenAI modality reads. +//! +//! This is the Rust shape of the object `createOpenAI` hands each model class +//! (`{ provider, url, headers, fetch }`): the model asks it for a URL, for the +//! request headers (evaluated per request, so a key is loaded when a call is +//! made and never when the provider is created) and for the transport. It has +//! no getters for the credential or the settings that produced it. The +//! credential origin is the one of each request's own URL, so the transport +//! refuses to send credentialed headers to any other origin after a redirect. + +use std::collections::HashMap; +use std::sync::Arc; + +use serde_json::Value; + +use aimux_core::AiMuxError; +use aimux_core::language_model::SupportedUrls; +use aimux_provider_utils::{ + ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, Resolvable, + combine_headers, normalize_headers, +}; + +use super::responses::ResponsesProfile; + +pub use crate::shared::TransformRequestBody; + +/// Maps an endpoint path (`"/chat/completions"`) to the full request URL. It +/// can fail because a host may only know the address when a call is made (an +/// Azure resource name read from the environment). +pub(crate) type UrlFn = Arc Result + Send + Sync>; + +/// What a model needs to talk to the API. Built by the provider that owns the +/// model: `OpenAIProvider`, Azure, Codex and the Hugging Face chat extension. +#[derive(Clone)] +pub(crate) struct OpenAIModelConfig { + /// The identity the model reports from `provider()`. + pub(crate) provider: String, + /// Endpoint path to full URL. + pub(crate) url: UrlFn, + /// Provider headers (credential, organization, project, user headers), + /// resolved on every request. + pub(crate) headers: HeadersFn, + pub(crate) token_provider: Option>, + /// Transport; `None` uses the process default, resolved per request. + pub(crate) fetch: Option, + /// URLs the model fetches itself. Empty: the caller downloads them. + pub(crate) supported_urls: SupportedUrls, + /// Provider-level request-body rewrite. + pub(crate) transform_request_body: Option, + /// How the Responses API model reads and writes providerOptions and file + /// ids for this host. + pub(crate) responses: ResponsesProfile, + pub(crate) chat_options: super::options::ChatOptionsParser, +} + +impl OpenAIModelConfig { + /// A config whose URLs are `base_url` followed by the endpoint path. + pub(crate) fn fixed( + provider: String, + base_url: String, + headers: HeadersFn, + fetch: Option, + transform_request_body: Option, + ) -> Self { + Self { + provider, + url: Arc::new(move |path| Ok(format!("{base_url}{path}"))), + headers, + token_provider: None, + fetch, + supported_urls: SupportedUrls::default(), + transform_request_body, + responses: ResponsesProfile::default(), + chat_options: super::options::parse_chat_options, + } + } + + /// The full URL of an endpoint path. + /// + /// # Errors + /// + /// Returns the host's setting error (`AiMuxError::LoadSetting` for an + /// unset Azure resource name). + pub(crate) fn url(&self, path: &str) -> Result { + (self.url)(path) + } + + /// The WebSocket form of an endpoint path: `https` becomes `wss` and + /// `http` becomes `ws`, as the AI SDK's `toWebSocketUrl` does. + /// + /// # Errors + /// + /// As [`url`](Self::url). + #[cfg(feature = "realtime")] + pub(crate) fn ws_url(&self, path: &str) -> Result { + let url = self.url(path)?; + Ok(if let Some(rest) = url.strip_prefix("https://") { + format!("wss://{rest}") + } else if let Some(rest) = url.strip_prefix("http://") { + format!("ws://{rest}") + } else { + url + }) + } + + /// Resolve the provider headers, layer the per-call headers over them + /// (case-insensitive, later wins, `None` removes) and return the list that + /// goes on the wire. + /// + /// # Errors + /// + /// Returns the error of the credential or header producer, for instance + /// `AiMuxError::LoadApiKey` when no key was given and the environment + /// variable is unset. + pub(crate) async fn request_headers( + &self, + call_headers: Option<&HashMap>, + ) -> Result, AiMuxError> { + let provider = self.headers.resolve().await?; + let call: HeaderMapOpt = call_headers + .map(|headers| { + headers + .iter() + .map(|(name, value)| (name.clone(), Some(value.clone()))) + .collect() + }) + .unwrap_or_default(); + let mut headers = combine_headers(&[&provider, &call]); + if let Some(token) = &self.token_provider + && !headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") && value.is_some()) + { + headers.insert( + "authorization".to_string(), + Some(format!("Bearer {}", token.resolve().await?)), + ); + } + Ok(normalize_headers(headers)) + } + + /// An exchange that inherits the operation's cancellation and recording + /// context and carries this config's transport and credential origin. + pub(crate) fn http_request( + &self, + url: String, + headers: Vec<(String, String)>, + options: &impl ExchangeContext, + ) -> HttpRequest { + self.with_transport(HttpRequest::new(url, headers, options)) + } + + /// Attach the transport and the credential origin (the request's own) to + /// a request. + pub(crate) fn with_transport(&self, mut request: HttpRequest) -> HttpRequest { + request.fetch = self.fetch.clone(); + request.credentialed_origin = Some(request.url.clone()); + request + } + + /// Run the provider-level body rewrite, when there is one. + pub(crate) fn transform_body(&self, body: Value) -> Value { + match &self.transform_request_body { + Some(transform) => transform(body), + None => body, + } + } +} diff --git a/aimux-providers/src/openai/convert.rs b/aimux-providers/src/openai/convert.rs index 092912c8a..71d246f6e 100644 --- a/aimux-providers/src/openai/convert.rs +++ b/aimux-providers/src/openai/convert.rs @@ -9,10 +9,10 @@ use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat use aimux_core::shared::{FileBytes, FileData, JsonObject, SharedProviderOptions}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Warning}; +use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; use serde::Serialize; use serde_json::{Value, json}; -use super::OpenAICompatProfile; /// Public capability enum used by the conversion helpers (moved to /// `convert_common` in M10; re-exported for API compatibility). pub use super::convert_common::SystemMessageMode; @@ -97,128 +97,8 @@ pub fn prepare_tools( } } -// ── Groq tool preparation ─────────────────────────────────────────────────── - -/// Models that support Groq's browser_search tool. -const GROQ_BROWSER_SEARCH_MODELS: &[&str] = &["openai/gpt-oss-20b", "openai/gpt-oss-120b"]; - -/// Prepare tools for Groq, supporting `groq.browser_search` provider tools -/// alongside standard function tools. Mirrors TS `prepareTools` in -/// `groq-prepare-tools.ts`. -fn prepare_tools_groq( - function_tools: &Option>, - all_tools: Option<&Vec>, - tool_choice: Option<&LanguageModelV4ToolChoice>, - model_id: &str, -) -> PreparedTools { - let non_empty_functions = function_tools.as_ref().filter(|&t| !t.is_empty()); - let has_any_tools = all_tools.as_ref().is_some_and(|t| !t.is_empty()); - - if !has_any_tools { - return PreparedTools { - tools: None, - tool_choice: None, - tool_warnings: vec![], - }; - } - - let mut groq_tools: Vec = Vec::new(); - let mut tool_warnings: Vec = Vec::new(); - - // Function tools - if let Some(tools) = non_empty_functions { - for t in tools { - let mut func = json!({ - "name": t.name, - "parameters": t.input_schema, - }); - if let Some(ref desc) = t.description { - func["description"] = json!(desc); - } - if let Some(strict) = t.strict { - func["strict"] = json!(strict); - } - groq_tools.push(json!({ "type": "function", "function": func })); - } - } - - // Provider-defined tools (browser_search) - if let Some(tools) = all_tools { - for t in tools { - if let Tool::Provider(pt) = t { - if pt.id == "groq.browser_search" { - if GROQ_BROWSER_SEARCH_MODELS.contains(&model_id) { - groq_tools.push(json!({ "type": "browser_search" })); - } else { - tool_warnings.push(ToolWarning { - warning_type: "unsupported".to_string(), - feature: format!("provider-defined tool {}", pt.id), - details: Some(format!( - "Browser search is only supported on the following models: {}. Current model: {}", - GROQ_BROWSER_SEARCH_MODELS.join(", "), - model_id - )), - }); - } - } else { - tool_warnings.push(ToolWarning { - warning_type: "unsupported".to_string(), - feature: format!("provider-defined tool {}", pt.id), - details: None, - }); - } - } - } - } - - let tools_opt = if groq_tools.is_empty() { - None - } else { - Some(groq_tools) - }; - - let tool_choice_opt = match (&tools_opt, tool_choice) { - (None, _) => None, - (Some(_), None) => None, - (Some(_), Some(tc)) => match tc { - LanguageModelV4ToolChoice::Auto => Some(json!("auto")), - LanguageModelV4ToolChoice::None => Some(json!("none")), - LanguageModelV4ToolChoice::Required => Some(json!("required")), - LanguageModelV4ToolChoice::Tool { tool_name } => { - Some(json!({ "type": "function", "function": { "name": tool_name } })) - } - }, - }; - - PreparedTools { - tools: tools_opt, - tool_choice: tool_choice_opt, - tool_warnings, - } -} - // ── Message conversion ────────────────────────────────────────────────────── -/// Convert a `LanguageModelPrompt` to OpenAI `messages` array. -/// -/// Panics on conversion failure. Production paths use the fallible variant -/// [`convert_prompt_to_openai_messages_with_mode_fallible`]; this panic -/// wrapper exists only for integration tests under `tests/`. It is -/// `#[doc(hidden)]` and `#[deprecated]` so it neither appears on the public -/// API surface nor can be pulled in by accident (release uses -/// `panic = "abort"`, so reaching a panic here via FFI would kill the host -/// process). -#[doc(hidden)] -#[deprecated( - since = "0.2.1", - note = "panics on failure; use convert_prompt_to_openai_messages_with_mode_fallible instead (issue #90 R1)" -)] -#[must_use] -pub fn convert_prompt_to_openai_messages(prompt: &LanguageModelPrompt) -> Vec { - convert_prompt_to_openai_messages_with_mode_fallible(prompt, SystemMessageMode::System) - .expect("convert_prompt_to_openai_messages: conversion failed") -} - /// Convert a `LanguageModelPrompt` to OpenAI `messages` array with a system /// message mode. /// @@ -230,64 +110,10 @@ pub fn convert_prompt_to_openai_messages(prompt: &LanguageModelPrompt) -> Vec Result, AiMuxError> { - convert_prompt_to_openai_messages_with_provider(prompt, system_message_mode, "openai") -} - -/// Convert a `LanguageModelPrompt` to OpenAI `messages` array with a system -/// message mode. -/// -/// Panics on conversion failure. Production paths use the fallible variant -/// [`convert_prompt_to_openai_messages_with_mode_fallible`]; this panic -/// wrapper exists only for integration tests under `tests/`. It is -/// `#[doc(hidden)]` and `#[deprecated]` so it neither appears on the public -/// API surface nor can be pulled in by accident (release uses -/// `panic = "abort"`, so reaching a panic here via FFI would kill the host -/// process). -#[doc(hidden)] -#[deprecated( - since = "0.2.1", - note = "panics on failure; use convert_prompt_to_openai_messages_with_mode_fallible instead (issue #90 R1)" -)] -#[must_use] -pub fn convert_prompt_to_openai_messages_with_mode( - prompt: &LanguageModelPrompt, - system_message_mode: SystemMessageMode, -) -> Vec { - convert_prompt_to_openai_messages_with_mode_fallible(prompt, system_message_mode) - .expect("convert_prompt_to_openai_messages_with_mode: conversion failed") -} - -/// Convert a `LanguageModelPrompt` to OpenAI `messages` array with a system -/// message mode and provider name (for provider-specific message conversion). -/// -/// # Errors -/// -/// Returns `AiMuxError::InvalidArgument` when a message part cannot be -/// converted (e.g. an unsupported file part shape or a missing provider -/// reference). -pub fn convert_prompt_to_openai_messages_with_provider( - prompt: &LanguageModelPrompt, - system_message_mode: SystemMessageMode, - provider: &str, -) -> Result, AiMuxError> { - convert_prompt_with_warnings(prompt, system_message_mode, provider, &mut Vec::new()) -} - -fn convert_prompt_with_warnings( - prompt: &LanguageModelPrompt, - system_message_mode: SystemMessageMode, - provider: &str, - warnings: &mut Vec, ) -> Result, AiMuxError> { let mut result = Vec::new(); for msg in prompt { - result.extend(convert_message_to_openai( - msg, - system_message_mode, - provider, - warnings, - )?); + result.extend(convert_message_to_openai(msg, system_message_mode)?); } Ok(result) } @@ -310,38 +136,6 @@ fn get_image_detail(provider_options: &Option) -> Option< .cloned() } -/// Get the top-level media type (e.g. "image" from "image/png"). -fn get_top_level_media_type(media_type: &str) -> &str { - media_type.split('/').next().unwrap_or("") -} - -/// Resolve a full media type from a top-level-only or wildcard media type. -/// For "image" or "image/*", detects "image/png" from the base64 data. -/// For "application", it stays as-is (cannot be resolved without full data). -fn resolve_full_media_type(media_type: &str, b64_data: &str) -> String { - let top_level = get_top_level_media_type(media_type); - if top_level == "image" && media_type != "image" && !media_type.ends_with("/*") { - return media_type.to_string(); - } - if top_level == "image" { - // Detect from base64 data - if b64_data.starts_with("iVBORw0KGgo") { - return "image/png".to_string(); - } - if b64_data.starts_with("/9j/") { - return "image/jpeg".to_string(); - } - if b64_data.starts_with("R0lGOD") { - return "image/gif".to_string(); - } - if b64_data.starts_with("UklGR") { - return "image/webp".to_string(); - } - return "image/png".to_string(); // default - } - media_type.to_string() -} - /// Resolve a provider reference, throwing if the provider is not found. fn resolve_provider_reference( reference: &aimux_core::shared::SharedProviderReference, @@ -359,36 +153,30 @@ fn resolve_provider_reference( } /// Convert a file part to the OpenAI format, handling images, audio, and PDF. -fn convert_file_part_to_openai( - file: &FilePart, - part_index: usize, - provider: &str, -) -> Result { +fn convert_file_part_to_openai(file: &FilePart, part_index: usize) -> Result { use base64::Engine; - let FilePart { data, media_type, filename, provider_options, } = file; - let prompt_cache_breakpoint = (provider == "openai") - .then(|| get_prompt_cache_breakpoint(provider_options)) - .flatten(); + let prompt_cache_breakpoint = get_prompt_cache_breakpoint(provider_options); let (data_b64, url) = match data { FileData::Reference { reference } => { - let file_id = resolve_provider_reference(reference, provider)?; - let mut part = if provider == "deepseek" { - json!({ "type": "file", "file_id": file_id }) - } else { - json!({ "type": "file", "file": { "file_id": file_id } }) - }; + let file_id = resolve_provider_reference(reference, "openai") + .map_err(AiMuxError::InvalidArgument)?; + let mut part = json!({ "type": "file", "file": { "file_id": file_id } }); if let Some(bpt) = prompt_cache_breakpoint { part["prompt_cache_breakpoint"] = bpt; } return Ok(part); } - FileData::Text { .. } => return Err("text file parts".to_string()), + FileData::Text { .. } => { + return Err(AiMuxError::UnsupportedFunctionality( + "text file parts".into(), + )); + } FileData::Url { url, .. } => (None, Some(url.as_str())), FileData::Data { data } => { let b64 = match data { @@ -408,10 +196,12 @@ fn convert_file_part_to_openai( let image_url = if let Some(url_str) = url { json!({ "url": url_str }) } else if let Some(b64) = data_b64 { - let full_mt = resolve_full_media_type(media_type, b64); + let full_mt = resolve_full_media_type(file)?; json!({ "url": format!("data:{};base64,{}", full_mt, b64) }) } else { - return Err("image part has no data or url".to_string()); + return Err(AiMuxError::InvalidArgument( + "image part has no data or url".into(), + )); }; let mut image_url_obj = image_url; @@ -432,14 +222,21 @@ fn convert_file_part_to_openai( // Audio if top_level == "audio" { if url.is_some() { - return Err("audio file parts with URLs".to_string()); + return Err(AiMuxError::UnsupportedFunctionality( + "audio file parts with URLs".into(), + )); } - let b64 = data_b64.ok_or("audio part has no data")?; - let full_mt = resolve_full_media_type(media_type, b64); + let b64 = + data_b64.ok_or_else(|| AiMuxError::InvalidArgument("audio part has no data".into()))?; + let full_mt = resolve_full_media_type(file)?; let format = match full_mt.as_str() { "audio/wav" => "wav", "audio/mp3" | "audio/mpeg" => "mp3", - _ => return Err(format!("audio content parts with media type {full_mt}")), + _ => { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "audio content parts with media type {full_mt}" + ))); + } }; let mut part = json!({ "type": "input_audio", @@ -452,21 +249,21 @@ fn convert_file_part_to_openai( } // PDF / application - let full_mt = if media_type == "application" { - return Err("media type \"application\".*is not passed as inline bytes.*".to_string()); - } else { - media_type.to_string() - }; + let full_mt = resolve_full_media_type(file)?; if full_mt != "application/pdf" { - return Err(format!("file part media type {full_mt}")); + return Err(AiMuxError::UnsupportedFunctionality(format!( + "file part media type {full_mt}" + ))); } if url.is_some() { - return Err("PDF file parts with URLs".to_string()); + return Err(AiMuxError::UnsupportedFunctionality( + "PDF file parts with URLs".into(), + )); } - let b64 = data_b64.ok_or("PDF part has no data")?; + let b64 = data_b64.ok_or_else(|| AiMuxError::InvalidArgument("PDF part has no data".into()))?; let fname = filename .map(std::string::ToString::to_string) .unwrap_or_else(|| format!("part-{part_index}.pdf")); @@ -487,37 +284,7 @@ fn convert_file_part_to_openai( fn convert_message_to_openai( msg: &LanguageModelMessage, system_message_mode: SystemMessageMode, - provider: &str, - warnings: &mut Vec, ) -> Result, AiMuxError> { - if provider != "openai" && provider != "deepseek" { - let has_reference = match msg { - LanguageModelMessage::User { content, .. } => content.iter().any(|part| { - matches!( - part, - UserPart::File(FilePart { - data: FileData::Reference { .. }, - .. - }) - ) - }), - LanguageModelMessage::Assistant { content, .. } => content.iter().any(|part| { - matches!( - part, - AssistantPart::File(FilePart { - data: FileData::Reference { .. }, - .. - }) - ) - }), - _ => false, - }; - if has_reference { - return Err(AiMuxError::UnsupportedFunctionality( - "file parts with provider references".into(), - )); - } - } let message = match msg { LanguageModelMessage::System { content, @@ -528,10 +295,7 @@ fn convert_message_to_openai( SystemMessageMode::Developer => "developer", SystemMessageMode::System => "system", }; - let content = match (provider == "openai") - .then(|| get_prompt_cache_breakpoint(provider_options)) - .flatten() - { + let content = match get_prompt_cache_breakpoint(provider_options) { None => json!(content), Some(bpt) => json!([{ "type": "text", "text": content, "prompt_cache_breakpoint": bpt, @@ -545,14 +309,9 @@ fn convert_message_to_openai( let ToolPart::ToolResult(result) = part else { continue; }; - let mut content = if provider == "deepseek" { - deepseek_tool_result_content(&result.output, warnings)? - } else { - tool_result_to_content(&result.output) - }; - if provider == "openai" - && let Some(breakpoint) = tool_result_cache_breakpoint(&result.output) - .or_else(|| get_prompt_cache_breakpoint(&result.provider_options)) + 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}]); } @@ -563,10 +322,11 @@ fn convert_message_to_openai( return Ok(messages); } LanguageModelMessage::User { content, .. } => { - let all_plain_text = content.len() == 1 && content.iter().all(|part| { - matches!(part, UserPart::Text(text) - if provider != "openai" || get_prompt_cache_breakpoint(&text.provider_options).is_none()) - }); + let all_plain_text = content.len() == 1 + && content.iter().all(|part| { + matches!(part, UserPart::Text(text) + if get_prompt_cache_breakpoint(&text.provider_options).is_none()) + }); let content = if all_plain_text { json!( content @@ -583,10 +343,8 @@ fn convert_message_to_openai( .iter() .enumerate() .map(|(index, part)| match part { - UserPart::Text(text) => Ok(convert_text_part_to_openai(text, provider)), - UserPart::File(file) => - convert_file_part_to_openai(file, index, provider) - .map_err(AiMuxError::InvalidArgument), + UserPart::Text(text) => Ok(convert_text_part_to_openai(text)), + UserPart::File(file) => convert_file_part_to_openai(file, index), }) .collect::, _>>()? ) @@ -595,84 +353,42 @@ fn convert_message_to_openai( } LanguageModelMessage::Assistant { content, .. } => { let mut text = String::new(); - let mut reasoning = String::new(); + let mut text_parts = Vec::new(); + let mut has_cache_breakpoint = false; let mut tool_calls = Vec::new(); for part in content { match part { - AssistantPart::Text(part) => text.push_str(&part.text), - AssistantPart::Reasoning(part) => reasoning.push_str(&part.text), + AssistantPart::Text(part) => { + text.push_str(&part.text); + has_cache_breakpoint |= + get_prompt_cache_breakpoint(&part.provider_options).is_some(); + text_parts.push(convert_text_part_to_openai(part)); + } AssistantPart::ToolCall(part) => { - let arguments = if provider == "openai" && !part.input.is_object() { - "{}".to_string() - } else { + let arguments = if part.input.is_object() { part.input.to_string() + } else { + "{}".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 }, })); } - AssistantPart::File(_) - | AssistantPart::ToolResult(_) - | AssistantPart::ReasoningFile(_) - | AssistantPart::Custom(_) => {} + _ => {} } } - let mut message = if provider == "groq" || !tool_calls.is_empty() { - let value = if provider != "groq" && text.is_empty() { - Value::Null - } else { - json!(text) - }; - let mut message = json!({ "role": "assistant", "content": value }); - if !tool_calls.is_empty() { - message["tool_calls"] = json!(tool_calls); - } - message + let content = if has_cache_breakpoint { + json!(text_parts) + } else if !tool_calls.is_empty() && text.is_empty() { + Value::Null } else { - let all_plain_text = content.iter().all(|part| { - matches!( - part, - AssistantPart::Text(TextPart { - provider_options: None, - .. - }) | AssistantPart::Reasoning(_) - | AssistantPart::ToolResult(_) - | AssistantPart::ReasoningFile(_) - | AssistantPart::Custom(_) - ) - }); - let value = if all_plain_text { - json!(text) - } else { - let mut parts = Vec::new(); - for part in content { - match part { - AssistantPart::Text(text) => { - parts.push(convert_text_part_to_openai(text, provider)) - } - AssistantPart::File(file) => parts.push( - convert_file_part_to_openai(file, parts.len(), provider) - .map_err(AiMuxError::InvalidArgument)?, - ), - AssistantPart::Reasoning(_) - | AssistantPart::ToolResult(_) - | AssistantPart::ReasoningFile(_) - | AssistantPart::Custom(_) - | AssistantPart::ToolCall(_) => {} - } - } - json!(parts) - }; - json!({ "role": "assistant", "content": value }) + json!(text) }; - if !reasoning.is_empty() { - let field = if provider == "groq" { - "reasoning" - } else { - "reasoning_content" - }; - message[field] = json!(reasoning); + let mut message = json!({ "role": "assistant", "content": content }); + if !tool_calls.is_empty() { + message["tool_calls"] = json!(tool_calls); } message } @@ -680,106 +396,8 @@ fn convert_message_to_openai( Ok(vec![message]) } -/// Serialize a tool-result value into the OpenAI tool message content string. -fn deepseek_tool_result_content( - output: &ToolResultOutput, - warnings: &mut Vec, -) -> Result { - let ToolResultOutput::Content { value } = output else { - return Ok(tool_result_to_content(output)); - }; - let is_image = |part: &FilePart| { - get_top_level_media_type(&part.media_type) == "image" - && !matches!(part.data, FileData::Text { .. }) - }; - if !value - .iter() - .any(|part| matches!(part, ToolResultContent::File(file) if is_image(file))) - { - return Ok(tool_result_to_content(output)); - } - let mut parts = Vec::new(); - for part in value { - match part { - ToolResultContent::Text(part) => parts.push(json!({"type":"text", "text":part.text})), - ToolResultContent::File(part) if is_image(part) => { - if let FileData::Reference { reference } = &part.data { - let file_id = resolve_provider_reference(reference, "deepseek") - .map_err(AiMuxError::InvalidArgument)?; - parts.push(json!({"type":"file", "file_id":file_id})); - continue; - } - let url = match &part.data { - FileData::Url { url, .. } => { - let media_type = resolve_full_media_type(&part.media_type, ""); - if !matches!( - media_type.as_str(), - "image/jpeg" | "image/jpg" | "image/png" | "image/gif" | "image/webp" - ) { - return Err(AiMuxError::UnsupportedFunctionality(format!( - "DeepSeek image media type {media_type}" - ))); - } - if url.len() > 8192 { - return Err(AiMuxError::InvalidPrompt( - "DeepSeek image URLs must not exceed 8192 characters.".into(), - )); - } - url.clone() - } - 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(), - }; - let media_type = resolve_full_media_type(&part.media_type, &b64); - let media_type = if media_type == "image/jpg" { - "image/jpeg" - } else { - media_type.as_str() - }; - if !matches!( - media_type, - "image/jpeg" | "image/png" | "image/gif" | "image/webp" - ) { - return Err(AiMuxError::UnsupportedFunctionality(format!( - "DeepSeek image media type {media_type}" - ))); - } - format!("data:{media_type};base64,{b64}") - } - _ => continue, - }; - let mut image_url = json!({"url":url}); - if let Some(detail) = part - .provider_options - .as_ref() - .and_then(|options| options.get("deepseek")) - .and_then(|options| options.get("imageDetail")) - { - image_url["detail"] = detail.clone(); - } - parts.push(json!({"type":"image_url", "image_url":image_url})); - } - part => warnings.push(Warning::Unsupported { - feature: format!( - "tool result content part type: {}", - if matches!(part, ToolResultContent::File(_)) { - "file" - } else { - "custom" - } - ), - details: None, - }), - } - } - Ok(json!(parts)) -} - +/// 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, .. } => { @@ -883,11 +501,10 @@ pub(crate) fn tool_result_cache_breakpoint(output: &ToolResultOutput) -> Option< } } -fn convert_text_part_to_openai(part: &TextPart, provider: &str) -> Value { +/// 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 provider == "openai" - && let Some(bpt) = get_prompt_cache_breakpoint(&part.provider_options) - { + if let Some(bpt) = get_prompt_cache_breakpoint(&part.provider_options) { value["prompt_cache_breakpoint"] = bpt; } value @@ -925,180 +542,16 @@ pub(crate) fn parse_option_fields( Ok(parsed) } -fn parse_chat_provider_options( - options: &Option, +/// Parse a namespace with the option schema supplied by its package. +pub(crate) fn parse_chat_provider_options( + options: Option<&SharedProviderOptions>, provider: &str, -) -> Result, AiMuxError> { - let provider = if provider == "azure" { - "openai" - } else { - provider - }; - if !matches!(provider, "openai" | "groq" | "deepseek") { - return Ok(options.clone()); - } - let Some(namespace) = options.as_ref().and_then(|options| options.get(provider)) else { - return Ok(None); - }; - let parsed = parse_option_fields(namespace, provider, |key, value| { - let one_of = |values: &[&str]| value.as_str().is_some_and(|s| values.contains(&s)); - let valid = match (provider, key) { - ("openai", "logitBias") => { - return Some(value.as_object().ok_or(()).and_then(|object| { - let mut parsed = JsonObject::new(); - for (key, value) in object { - if !value.is_number() { - return Err(()); - } - let key = key.trim_matches(|c| { - matches!(c, - '\u{0009}'..='\u{000d}' | '\u{0020}' | '\u{00a0}' | '\u{1680}' - | '\u{2000}'..='\u{200a}' | '\u{2028}' | '\u{2029}' | '\u{202f}' - | '\u{205f}' | '\u{3000}' | '\u{feff}' - ) - }); - let number = if key.is_empty() { - 0.0 - } else if let Some((digits, radix)) = key - .strip_prefix("0x") - .or_else(|| key.strip_prefix("0X")) - .map(|s| (s, 16)) - .or_else(|| { - key.strip_prefix("0o") - .or_else(|| key.strip_prefix("0O")) - .map(|s| (s, 8)) - }) - .or_else(|| { - key.strip_prefix("0b") - .or_else(|| key.strip_prefix("0B")) - .map(|s| (s, 2)) - }) - { - if digits.is_empty() { - return Err(()); - } - let bits_per_digit = match radix { - 16 => 4, - 8 => 3, - _ => 1, - }; - let mut significant_bits = 0usize; - let mut mantissa = 0u64; - let mut guard = false; - let mut sticky = false; - for digit in digits.chars() { - let digit = digit.to_digit(radix).ok_or(())?; - for shift in (0..bits_per_digit).rev() { - let bit = (digit >> shift) & 1; - if significant_bits == 0 && bit == 0 { - continue; - } - if significant_bits < 53 { - mantissa = (mantissa << 1) | u64::from(bit); - } else if significant_bits == 53 { - guard = bit != 0; - } else { - sticky |= bit != 0; - } - significant_bits += 1; - } - } - if guard && (sticky || mantissa & 1 != 0) { - mantissa += 1; - } - let exponent = significant_bits.saturating_sub(53); - if exponent > 1023 { - f64::INFINITY - } else { - mantissa as f64 * 2.0f64.powi(exponent as i32) - } - } else { - key.parse::().map_err(|_| ())? - }; - if !number.is_finite() { - return Err(()); - } - let key = if number == 0.0 { - "0".to_string() - } else if number.abs() >= 1e21 || number.abs() < 1e-6 { - let scientific = format!("{number:e}"); - let (mantissa, exponent) = scientific.split_once('e').ok_or(())?; - let exponent = exponent.parse::().map_err(|_| ())?; - format!("{mantissa}e{exponent:+}") - } else { - number.to_string() - }; - parsed.insert(key, value.clone()); - } - Ok(Value::Object(parsed)) - })); - } - ("openai", "logprobs") => value.is_boolean() || value.is_number(), - ("openai", "maxCompletionTokens") => value.is_number(), - ("openai", "metadata") => value.as_object().is_some_and(|object| { - object.iter().all(|(key, value)| { - key.encode_utf16().count() <= 64 - && value - .as_str() - .is_some_and(|s| s.encode_utf16().count() <= 512) - }) - }), - ("openai", "prediction") => value.is_object(), - ("openai", "reasoningEffort") => { - one_of(&["none", "minimal", "low", "medium", "high", "xhigh", "max"]) - } - ("openai", "serviceTier") => { - one_of(&["auto", "flex", "priority", "fast", "ultrafast", "default"]) - } - ("openai", "textVerbosity") => one_of(&["low", "medium", "high"]), - ("openai", "promptCacheRetention") => one_of(&["in_memory", "24h"]), - ("openai", "systemMessageMode") => one_of(&["system", "developer", "remove"]), - ("openai", "user" | "promptCacheKey" | "safetyIdentifier") | ("groq", "user") => { - value.is_string() - } - ("openai", "parallelToolCalls" | "store" | "strictJsonSchema" | "forceReasoning") - | ("groq", "parallelToolCalls" | "structuredOutputs" | "strictJsonSchema") - | ("deepseek", "logprobs" | "strictJsonSchema") => value.is_boolean(), - ("groq", "reasoningFormat") => one_of(&["parsed", "raw", "hidden"]), - ("groq", "reasoningEffort") => one_of(&["none", "default", "low", "medium", "high"]), - ("groq", "serviceTier") => one_of(&["on_demand", "performance", "flex", "auto"]), - ("deepseek", "reasoningEffort") => one_of(&["low", "medium", "high", "xhigh", "max"]), - ("deepseek", "topLogprobs") => value - .as_f64() - .is_some_and(|n| n.fract() == 0.0 && (0.0..=20.0).contains(&n)), - ("deepseek", "userId") => value.as_str().is_some_and(|s| { - !s.is_empty() - && s.len() <= 512 - && s.bytes() - .all(|c| c.is_ascii_alphanumeric() || matches!(c, b'_' | b'-')) - }), - ("openai", "promptCacheOptions") | ("deepseek", "thinking") => { - return Some(value.as_object().ok_or(()).and_then(|object| { - parse_option_fields(object, provider, |key, value| { - let valid = match (provider, key) { - ("openai", "mode") => value - .as_str() - .is_some_and(|s| matches!(s, "implicit" | "explicit")), - ("openai", "ttl") => value.as_str() == Some("30m"), - ("deepseek", "type") => value - .as_str() - .is_some_and(|s| matches!(s, "adaptive" | "enabled" | "disabled")), - _ => return None, - }; - Some(if valid { Ok(value.clone()) } else { Err(()) }) - }) - .map(Value::Object) - .map_err(|_| ()) - })); - } - _ => return None, - }; - Some(if valid { Ok(value.clone()) } else { Err(()) }) - })?; - Ok(Some(aimux_core::shared::NamespaceMap::from([( - provider.to_string(), - parsed, - )]))) + parse: impl FnOnce(&JsonObject) -> Result, +) -> Result, AiMuxError> { + options + .and_then(|options| options.get(provider)) + .map(parse) + .transpose() } /// Get a value from provider_options.openai.. @@ -1122,38 +575,7 @@ pub fn build_request_body( options: &CallOptions, stream: bool, ) -> Result { - build_request_body_with_warnings( - model_id, - options, - stream, - "openai", - &OpenAICompatProfile::full(), - ) - .map(|r| r.body) -} - -/// Convert `CallOptions` to an OpenAI request body, returning warnings. -/// `provider` controls provider-specific behaviour (e.g. groq reads provider -/// options from the `"groq"` key). -/// `profile` declares provider capability differences (top_k, tools, etc.). -/// -/// Conversion errors propagate to the caller (fail-fast, issue H2): the old -/// behaviour of silently returning `body: null` sent empty requests upstream -/// and made conversion failures invisible. -/// Look up a key from the selected vendor namespace. -fn provider_option( - provider_opts: &Option, - provider: &str, - key: &str, -) -> Option { - if matches!(provider, "groq" | "deepseek") { - return provider_opts - .as_ref() - .and_then(|m| m.get(provider)) - .and_then(|o| o.get(key)) - .cloned(); - } - openai_option(provider_opts, key) + build_request_body_with_warnings(model_id, options, stream).map(|r| r.body) } /// Resolve the effective reasoning effort — direct passthrough (v3: no built-in @@ -1161,10 +583,9 @@ fn provider_option( /// `reasoning`; custom top-level levels map verbatim to `reasoning_effort`. fn resolve_reasoning_effort( provider_opts: &Option, - provider: &str, reasoning: &Option, ) -> Option { - provider_option(provider_opts, provider, "reasoningEffort") + openai_option(provider_opts, "reasoningEffort") .map(|v| { v.as_str() .map(std::string::ToString::to_string) @@ -1181,21 +602,19 @@ fn resolve_reasoning_effort( fn resolve_is_reasoning_model( provider_opts: &Option, - provider: &str, caps: &ModelCapabilities, ) -> bool { - provider_option(provider_opts, provider, "forceReasoning") + openai_option(provider_opts, "forceReasoning") .map(|v| v.as_bool().unwrap_or(false)) .unwrap_or(caps.is_reasoning_model) } fn resolve_system_message_mode( provider_opts: &Option, - provider: &str, is_reasoning_model: bool, caps: &ModelCapabilities, ) -> SystemMessageMode { - provider_option(provider_opts, provider, "systemMessageMode") + openai_option(provider_opts, "systemMessageMode") .and_then(|v| v.as_str().map(std::string::ToString::to_string)) .map(|s| match s.as_str() { "developer" => SystemMessageMode::Developer, @@ -1209,40 +628,26 @@ fn resolve_system_message_mode( }) } -/// Insert `max_tokens` / `max_completion_tokens`. -/// -/// `max_tokens_key` 是内部数据（非用户概念），指定该厂商唯一认的 key： -/// - Some("max_tokens") → 只发 max_tokens（如 stepfun/siliconflow/perplexity 等） -/// - Some("max_completion_tokens") → 只发 max_completion_tokens（groq/heroku 等） -/// - None → 现状推断：推理模型发 mct，非推理发 max_tokens。 +/// Insert `max_tokens` / `max_completion_tokens`: reasoning models take the +/// latter, the rest the former; an explicit `maxCompletionTokens` option is +/// always sent as `max_completion_tokens`. fn apply_max_tokens( body: &mut Value, options: &CallOptions, provider_opts: &Option, - provider: &str, - profile: &OpenAICompatProfile, is_reasoning_model: bool, ) { - let max_completion_tokens_opt = provider_option(provider_opts, provider, "maxCompletionTokens"); - - let use_mct_key = match profile.max_tokens_key { - Some("max_tokens") => false, - Some("max_completion_tokens") => true, - _ => is_reasoning_model, - }; + let max_completion_tokens_opt = openai_option(provider_opts, "maxCompletionTokens"); if let Some(max_tokens) = options.max_output_tokens { - let key = if use_mct_key { + let key = if is_reasoning_model { "max_completion_tokens" } else { "max_tokens" }; body[key] = json!(max_tokens); } - // 显式 maxCompletionTokens 选项：只认 max_tokens 的厂商不发送 mct。 - if let Some(mct) = max_completion_tokens_opt - && profile.max_tokens_key != Some("max_tokens") - { + if let Some(mct) = max_completion_tokens_opt { body["max_completion_tokens"] = json!(mct); } } @@ -1348,30 +753,9 @@ fn insert_sampling_params(body: &mut Value, params: &SamplingParams, options: &C } } -/// `response_format` → body, respecting the profile's capability flag and -/// groq's structuredOutputs / strictJsonSchema semantics. -fn apply_response_format( - body: &mut Value, - options: &CallOptions, - provider_opts: &Option, - provider: &str, - profile: &OpenAICompatProfile, - warnings: &mut Vec, -) { - if !profile.supports_response_format { - // Provider does not support response_format: drop it and warn when the - // caller requested a (non-default) format. `Text` is the no-op default - // and needs no warning. - if let Some(ref rf) = options.response_format - && !matches!(rf, ResponseFormat::Text) - { - warnings.push(Warning::Unsupported { - feature: "responseFormat".to_string(), - details: Some("response_format is not supported by this provider".to_string()), - }); - } - return; - } +/// `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) { let Some(ref rf) = options.response_format else { return; }; @@ -1382,51 +766,22 @@ fn apply_response_format( name, description, } => { - // Groq: structuredOutputs defaults to true, strictJsonSchema - // defaults to true. When structuredOutputs is false and a schema - // is provided, emit a warning and use json_object. - let (structured_outputs, strict_json_schema) = if provider == "groq" { - ( - provider_option(provider_opts, provider, "structuredOutputs") - .and_then(|v| v.as_bool()) - .unwrap_or(true), - provider_option(provider_opts, provider, "strictJsonSchema") - .and_then(|v| v.as_bool()) - .unwrap_or(true), - ) - } else { - ( - true, - provider_option(provider_opts, provider, "strictJsonSchema") - .and_then(|v| v.as_bool()) - .unwrap_or(true), - ) - }; - - if schema.is_some() && structured_outputs { + if let Some(schema) = schema { let mut schema_obj = json!({}); - if let Some(s) = schema { - schema_obj["schema"] = s.clone(); - } + schema_obj["schema"] = schema.clone(); 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!(strict_json_schema); + schema_obj["strict"] = json!( + openai_option(&options.provider_options, "strictJsonSchema") + .and_then(|value| value.as_bool()) + .unwrap_or(true) + ); body["response_format"] = json!({ "type": "json_schema", "json_schema": schema_obj, }); - } else if schema.is_some() && !structured_outputs { - // Schema provided but structuredOutputs disabled → json_object + warning - body["response_format"] = json!({ "type": "json_object" }); - warnings.push(Warning::Unsupported { - feature: "responseFormat".to_string(), - details: Some( - "JSON response format schema is only supported with structuredOutputs" - .to_string(), - ), - }); } else { body["response_format"] = json!({ "type": "json_object" }); } @@ -1435,14 +790,13 @@ fn apply_response_format( } /// Pass through the simple provider-specific options that map 1:1 to a body -/// field (plus Groq's `reasoning_format`). +/// field. fn apply_provider_option_passthrough( body: &mut Value, provider_opts: &Option, - provider: &str, ) { let mut set = |key: &str, body_key: &str| { - if let Some(val) = provider_option(provider_opts, provider, key) { + if let Some(val) = openai_option(provider_opts, key) { body[body_key] = val; } }; @@ -1457,56 +811,26 @@ fn apply_provider_option_passthrough( set("promptCacheRetention", "prompt_cache_retention"); set("promptCacheOptions", "prompt_cache_options"); set("safetyIdentifier", "safety_identifier"); - if provider == "deepseek" { - set("userId", "user_id"); - } - if provider != "openai" { - set("topLogprobs", "top_logprobs"); - } - if let Some(logprobs) = provider_option(provider_opts, provider, "logprobs") { - if provider == "openai" { - if logprobs == json!(true) || logprobs.is_number() { - body["logprobs"] = json!(true); - body["top_logprobs"] = if logprobs.is_number() { - logprobs - } else { - json!(0) - }; - } - } else if provider == "deepseek" { - if logprobs == json!(true) || body.get("top_logprobs").is_some() { - body["logprobs"] = json!(true); - } - } else { - body["logprobs"] = logprobs; - } - } else if provider == "deepseek" && body.get("top_logprobs").is_some() { - body["logprobs"] = json!(true); - } - // Groq: reasoning_format provider option - if provider == "groq" - && let Some(val) = provider_option(provider_opts, provider, "reasoningFormat") + if let Some(logprobs) = openai_option(provider_opts, "logprobs") + && (logprobs == json!(true) || logprobs.is_number()) { - body["reasoning_format"] = val; + body["logprobs"] = json!(true); + body["top_logprobs"] = if logprobs.is_number() { + logprobs + } else { + json!(0) + }; } } -/// `service_tier` with model-capability validation (Groq passes it through -/// without validation). +/// `service_tier` with model-capability validation. fn apply_service_tier( body: &mut Value, provider_opts: &Option, - provider: &str, caps: &ModelCapabilities, warnings: &mut Vec, ) { - if provider == "groq" { - if let Some(val) = provider_option(provider_opts, provider, "serviceTier") { - body["service_tier"] = val; - } - return; - } - let service_tier = provider_option(provider_opts, provider, "serviceTier") + let service_tier = openai_option(provider_opts, "serviceTier") .and_then(|v| v.as_str().map(std::string::ToString::to_string)); if let Some(ref st) = service_tier { match st.as_str() { @@ -1542,17 +866,8 @@ fn apply_service_tier( } } -/// Function tools → `tools` / `tool_choice` (Groq also maps provider-defined -/// tools such as browser_search). Drops both when the profile declares no tool -/// support, warning when the caller supplied any. -fn apply_tools( - body: &mut Value, - options: &CallOptions, - provider: &str, - profile: &OpenAICompatProfile, - model_id: &str, - warnings: &mut Vec, -) { +/// Function tools → `tools` / `tool_choice`. +fn apply_tools(body: &mut Value, options: &CallOptions) { let function_tools: Option> = options.tools.as_ref().map(|tools| { tools .iter() @@ -1563,51 +878,16 @@ fn apply_tools( .collect() }); - let supports_tools = profile.supports_tools; - if !supports_tools && options.tools.as_ref().is_some_and(|t| !t.is_empty()) { - warnings.push(Warning::Unsupported { - feature: "tools".to_string(), - details: Some("tools are not supported by this provider".to_string()), - }); - } - - let prepared = if !supports_tools { - PreparedTools { - tools: None, - tool_choice: None, - tool_warnings: Vec::new(), - } - } else if provider == "groq" { - prepare_tools_groq( - &function_tools, - options.tools.as_ref(), - options.tool_choice.as_ref(), - model_id, - ) - } else { - prepare_tools(&function_tools, options.tool_choice.as_ref()) - }; - + 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; } } - - // Convert tool warnings to Warning type - for tw in prepared.tool_warnings { - warnings.push(Warning::Unsupported { - feature: tw.feature, - details: tw.details, - }); - } } /// Convert `CallOptions` to an OpenAI request body, returning warnings. -/// `provider` controls provider-specific behaviour (e.g. groq reads provider -/// options from the `"groq"` key). -/// `profile` declares provider capability differences (top_k, tools, etc.). /// /// Conversion errors propagate to the caller (fail-fast, issue H2): the old /// behaviour of silently returning `body: null` sent empty requests upstream @@ -1621,35 +901,47 @@ pub fn build_request_body_with_warnings( model_id: &str, options: &CallOptions, stream: bool, - provider: &str, - profile: &OpenAICompatProfile, +) -> Result { + build_request_body_with_chat_options( + model_id, + options, + stream, + super::options::parse_chat_options, + ) +} + +pub(crate) fn build_request_body_with_chat_options( + model_id: &str, + options: &CallOptions, + stream: bool, + parse: super::options::ChatOptionsParser, ) -> Result { let mut warnings: Vec = Vec::new(); let caps = get_model_capabilities(model_id); - let parsed_provider_opts = parse_chat_provider_options(&options.provider_options, provider)?; + let parsed_provider_opts = parse_chat_provider_options( + options.provider_options.as_ref(), + super::options::NAMESPACE, + parse, + )? + .map(|parsed| { + aimux_core::shared::NamespaceMap::from([(super::options::NAMESPACE.to_string(), parsed)]) + }); let provider_opts = &parsed_provider_opts; - let resolved_reasoning_effort = - resolve_reasoning_effort(provider_opts, provider, &options.reasoning); - let is_reasoning_model = resolve_is_reasoning_model(provider_opts, provider, &caps); - let system_message_mode = - resolve_system_message_mode(provider_opts, provider, is_reasoning_model, &caps); + let resolved_reasoning_effort = resolve_reasoning_effort(provider_opts, &options.reasoning); + 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); - // top_k: only send when the provider supports it. - let top_k_supported = profile.supports_top_k; - if options.top_k.is_some() && !top_k_supported { + // OpenAI has no top_k. + if options.top_k.is_some() { warnings.push(Warning::Unsupported { feature: "topK".to_string(), details: None, }); } - let messages = convert_prompt_with_warnings( - &options.prompt, - system_message_mode, - provider, - &mut warnings, - )?; + let messages = + convert_prompt_to_openai_messages_with_mode_fallible(&options.prompt, system_message_mode)?; let mut body = json!({ "model": model_id, @@ -1658,24 +950,10 @@ pub fn build_request_body_with_warnings( if stream { body["stream"] = json!(true); - // Groq does not send stream_options; other providers include usage. - if provider != "groq" { - body["stream_options"] = json!({ "include_usage": true }); - } + body["stream_options"] = json!({ "include_usage": true }); } - if let Some(tk) = options.top_k.filter(|_| top_k_supported) { - body["top_k"] = json!(tk); - } - - apply_max_tokens( - &mut body, - options, - provider_opts, - provider, - profile, - is_reasoning_model, - ); + apply_max_tokens(&mut body, options, provider_opts, is_reasoning_model); let sampling = strip_sampling_params( options, @@ -1687,62 +965,19 @@ pub fn build_request_body_with_warnings( ); insert_sampling_params(&mut body, &sampling, options); - apply_response_format( - &mut body, - options, - provider_opts, - provider, - profile, - &mut warnings, - ); - apply_provider_option_passthrough(&mut body, provider_opts, provider); + apply_response_format(&mut body, options); + apply_provider_option_passthrough(&mut body, provider_opts); - // Reasoning effort (v3 passthrough; 注：旧的"reasoning 无映射提示"warning - // 块已删除——v3 直传语义下该分支不可达：custom 值此时必已进 resolved)。 + // Reasoning effort (v3 passthrough: no built-in vendor normalization). if let Some(ref effort) = resolved_reasoning_effort { body["reasoning_effort"] = json!(effort); } - apply_service_tier(&mut body, provider_opts, provider, &caps, &mut warnings); - apply_tools( - &mut body, - options, - provider, - profile, - model_id, - &mut warnings, - ); - - // 厂商特化后处理已整体退役（RFC-0017 阶段 2）：不内置任何厂商映射， - // thinking 注入 / effort 重映射等差异由用户 bodyOverrides 定义。 - - // Per-call request body overrides (RFC-0017): deep-merge user-supplied - // JSON into the built body. `null` values delete the corresponding key. - // Applied last so users can override anything, including vendor-specific - // fields injected above. - if let Some(ref overrides) = options.body_overrides { - deep_merge_json(&mut body, overrides); - } + apply_service_tier(&mut body, provider_opts, &caps, &mut warnings); + apply_tools(&mut body, options); Ok(RequestBodyResult { body, warnings }) } -pub fn deep_merge_json(target: &mut Value, patch: &Value) { - match (target, patch) { - (Value::Object(t), Value::Object(p)) => { - for (k, v) in p { - match v { - Value::Null => { - t.remove(k); - } - _ => { - deep_merge_json(t.entry(k).or_insert(Value::Null), v); - } - } - } - } - (target, patch) => *target = patch.clone(), - } -} /// Parse OpenAI finish reason string into `FinishReason`. #[must_use] diff --git a/aimux-providers/src/openai/decision.rs b/aimux-providers/src/openai/decision.rs index 33eece648..91fa828ee 100644 --- a/aimux-providers/src/openai/decision.rs +++ b/aimux-providers/src/openai/decision.rs @@ -3,36 +3,36 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use aimux_core::AiMuxError; use aimux_core::decision_model::*; -use aimux_core::recording::ProviderRecord; use aimux_core::types::{InputTokenUsage, OutputTokenUsage, Usage}; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; -use super::OpenAIConfig; +use super::config::OpenAIModelConfig; +use crate::decision_support::DecisionConfig; +/// Native OpenAI Decisions model; [`OpenAIProvider::decision`](super::OpenAIProvider::decision) +/// creates it. pub struct OpenAIDecisionModel { - config: OpenAIConfig, + config: DecisionConfig, model_id: String, } impl OpenAIDecisionModel { - /// Create a native Decisions handle. The server validates model availability. + /// A native Decisions model. The server validates model availability. /// # Errors - /// Rejects an empty ID, another provider's config, and generation overrides. - pub fn new(model_id: &str, config: OpenAIConfig) -> Result { + /// Rejects an empty ID. + pub(crate) fn from_config( + model_id: &str, + config: OpenAIModelConfig, + ) -> Result { if model_id.trim().is_empty() { return Err(AiMuxError::InvalidArgument( "decision model ID cannot be empty".into(), )); } - if config.provider != "openai" || config.body_overrides.is_some() { - return Err(AiMuxError::UnsupportedFunctionality( - "native OpenAI decisions require an OpenAI config without body_overrides".into(), - )); - } Ok(Self { - config, + config: DecisionConfig::OpenAI(config), model_id: model_id.into(), }) } @@ -396,9 +396,6 @@ impl DecisionModel for OpenAIDecisionModel { fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } fn capabilities(&self) -> DecisionCapabilities { DecisionCapabilities { supports_images: true, @@ -416,17 +413,9 @@ impl DecisionModel for OpenAIDecisionModel { max_score_levels: None, } } - fn config_snapshot(&self) -> ProviderRecord { - let mut record = - super::config_snapshot_from_config(self.provider(), self.model_id(), &self.config); - record.profile = Some( - json!({"decision_protocol":"openai_decisions", "decision_capabilities":self.capabilities()}), - ); - record - } async fn do_decide(&self, options: &DecisionCallOptions) -> Result { let body = request_body(&self.model_id, options)?; - let url = format!("{}/decisions", self.config.base_url.trim_end_matches('/')); + let url = self.config.url("/decisions")?; crate::decision_support::post(&self.config, options, &url, body, |raw, headers| { let data = serde_json::from_value(raw.clone()) .map_err(|e| crate::decision_support::invalid(format!("OpenAI decisions: {e}")))?; diff --git a/aimux-providers/src/openai/embedding.rs b/aimux-providers/src/openai/embedding.rs index 4a7d0a7b0..bc3c77053 100644 --- a/aimux-providers/src/openai/embedding.rs +++ b/aimux-providers/src/openai/embedding.rs @@ -5,8 +5,6 @@ //! //! Endpoint: `POST {base_url}/embeddings` -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -16,9 +14,7 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::HttpRequest; - -use super::OpenAIConfig; +use super::config::OpenAIModelConfig; /// An OpenAI-compatible embedding model. /// @@ -26,44 +22,13 @@ use super::OpenAIConfig; /// HTTP client — the `aimux-provider-utils` API helpers use the shared `Client` internally (RFC-0009 §4.1). pub struct OpenAIEmbeddingModel { model_id: String, - config: OpenAIConfig, + config: OpenAIModelConfig, } impl OpenAIEmbeddingModel { - #[must_use] - pub fn new(model_id: String, config: OpenAIConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: OpenAIModelConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref org) = self.config.org_id { - headers.insert("OpenAI-Organization".to_string(), org.clone()); - } - if let Some(ref project) = self.config.project { - headers.insert("OpenAI-Project".to_string(), project.clone()); - } - // Config-level extra headers (lowest priority after auth/org/project). - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/embeddings", self.config.base_url) - } } #[async_trait] @@ -76,10 +41,6 @@ impl EmbeddingModel for OpenAIEmbeddingModel { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - fn max_embeddings_per_call(&self) -> Option { Some(2048) } @@ -105,17 +66,16 @@ impl EmbeddingModel for OpenAIEmbeddingModel { body.insert("user".to_string(), json!(user)); } - let headers = self.build_headers(options.headers.as_ref()); - - let mut header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - header_list.push(("Content-Type".to_string(), "application/json".to_string())); + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; + let body = self.config.transform_body(Value::Object(body)); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), header_list, options), - Value::Object(body), + self.config + .http_request(self.config.url("/embeddings")?, headers, options), + body, aimux_provider_utils::create_json_response_handler(), super::openai_failed_response_handler(), ) diff --git a/aimux-providers/src/openai/files.rs b/aimux-providers/src/openai/files.rs index 52560f0fa..6bb651827 100644 --- a/aimux-providers/src/openai/files.rs +++ b/aimux-providers/src/openai/files.rs @@ -4,8 +4,6 @@ //! Aligned with Vercel AI SDK `OpenAIFiles` //! (`reference/ai/packages/openai/src/files/openai-files.ts`). -use std::collections::HashMap; - use async_trait::async_trait; use serde::Deserialize; use serde_json::json; @@ -14,9 +12,9 @@ use aimux_core::error::AiMuxError; use aimux_core::files_model::{Files, UploadFileCallOptions, UploadFileData, UploadFileResult}; use aimux_core::shared::{FileBytes, SharedProviderOptions}; -use aimux_provider_utils::{HttpBody, HttpRequest}; +use aimux_provider_utils::HttpBody; -use super::OpenAIConfig; +use super::config::OpenAIModelConfig; /// OpenAI provider-specific file upload options. /// @@ -89,44 +87,19 @@ struct OpenAIFilesResponse { /// Aligned with TS `OpenAIFiles`. Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers /// uses the process-wide shared `Client` internally (RFC-0009 §4.1). pub struct OpenAIFiles { - config: OpenAIConfig, + config: OpenAIModelConfig, } impl OpenAIFiles { - #[must_use] - pub fn new(config: OpenAIConfig) -> Self { + pub(crate) fn from_config(config: OpenAIModelConfig) -> Self { Self { config } } - - fn build_headers(&self) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref org) = self.config.org_id { - headers.insert("OpenAI-Organization".to_string(), org.clone()); - } - if let Some(ref project) = self.config.project { - headers.insert("OpenAI-Project".to_string(), project.clone()); - } - if let Some(ref extra) = self.config.headers { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/files", self.config.base_url) - } } #[async_trait] impl Files for OpenAIFiles { fn provider(&self) -> &str { - "openai.files" + &self.config.provider } async fn upload_file( @@ -156,11 +129,8 @@ impl Files for OpenAIFiles { .map(|v| ("expires_after", v.to_string())), ); - let headers = self.build_headers(); - let header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + let header_list = self.config.request_headers(None).await?; + let url = self.config.url("/files")?; // `send()` returns Ok only for 2xx; non-2xx responses are mapped to an // error internally using the shared error structure. The multipart body @@ -168,25 +138,16 @@ impl Files for OpenAIFiles { // sets `Content-Type` from it, so it is intentionally not added to the // header list above. // - // Nothing above `upload_file` retries it (there is no Core - // `do_upload_file`), so the retry lives here — safe for the whole - // exchange because an upload is not billable and a failed - // create-a-file request returns no id to replay against (§9.4). - let retries = aimux_core::retry::prepare_retries( - None, - self.config.retry_config, - options.abort_signal.clone(), - ); - let resp = retries - .retry(|| { - aimux_provider_utils::post_to_api( - HttpRequest::new(self.endpoint(), header_list.clone(), options), - HttpBody::Bytes(body.clone(), content_type.clone()), - aimux_provider_utils::create_json_response_handler(), - super::openai_failed_response_handler(), - ) - }) - .await?; + // One attempt: like the AI SDK's `uploadFile`, an upload is never + // replayed implicitly. A caller that wants a retry makes it around + // `upload_file`. + let resp = aimux_provider_utils::post_to_api( + self.config.http_request(url, header_list, options), + HttpBody::Bytes(body, content_type), + aimux_provider_utils::create_json_response_handler(), + super::openai_failed_response_handler(), + ) + .await?; let data: OpenAIFilesResponse = resp.value; diff --git a/aimux-providers/src/openai/image.rs b/aimux-providers/src/openai/image.rs index 2e9fcf7a6..bf5cceef3 100644 --- a/aimux-providers/src/openai/image.rs +++ b/aimux-providers/src/openai/image.rs @@ -8,8 +8,6 @@ //! //! Both endpoints return the same JSON response shape (`openaiImageResponseSchema`). -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -20,9 +18,9 @@ use aimux_core::image_model::{ }; use aimux_core::shared::{SharedProviderMetadata, SharedProviderOptions, Warning}; -use aimux_provider_utils::{HttpBody, HttpRequest}; +use aimux_provider_utils::HttpBody; -use super::OpenAIConfig; +use super::config::OpenAIModelConfig; /// Models that default to `b64_json` response format and therefore do NOT need /// the explicit `response_format` field. Identified by prefix. @@ -65,47 +63,13 @@ fn get_max_images_per_call(model_id: &str) -> u32 { /// An OpenAI-compatible image generation/editing model. pub struct OpenAIImageModel { model_id: String, - config: OpenAIConfig, + config: OpenAIModelConfig, } impl OpenAIImageModel { - #[must_use] - pub fn new(model_id: String, config: OpenAIConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: OpenAIModelConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref org) = self.config.org_id { - headers.insert("OpenAI-Organization".to_string(), org.clone()); - } - if let Some(ref project) = self.config.project { - headers.insert("OpenAI-Project".to_string(), project.clone()); - } - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn generations_endpoint(&self) -> String { - format!("{}/images/generations", self.config.base_url) - } - - fn edits_endpoint(&self) -> String { - format!("{}/images/edits", self.config.base_url) - } } #[async_trait] @@ -118,10 +82,6 @@ impl ImageModel for OpenAIImageModel { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - fn max_images_per_call(&self) -> Option { Some(get_max_images_per_call(&self.model_id)) } @@ -146,7 +106,10 @@ impl ImageModel for OpenAIImageModel { } let timestamp = chrono::Utc::now().to_rfc3339(); - let headers = self.build_headers(options.headers.as_ref()); + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; let (body_value, response_headers, _response_body) = if options.files.is_some() { // ── Edit path: multipart form data ── @@ -155,7 +118,8 @@ impl ImageModel for OpenAIImageModel { build_edit_multipart(&self.model_id, options, &openai_options); let resp = aimux_provider_utils::post_to_api( - HttpRequest::new(self.edits_endpoint(), build_header_list(&headers), options), + self.config + .http_request(self.config.url("/images/edits")?, headers, options), HttpBody::Bytes(form_body, content_type), aimux_provider_utils::create_json_response_handler(), super::openai_failed_response_handler(), @@ -167,22 +131,25 @@ impl ImageModel for OpenAIImageModel { } else { // ── Generation path: JSON body ── let openai_options = parse_generation_provider_options(&options.provider_options); - let body = build_generation_body(&self.model_id, options, &openai_options); + let body = self + .config + .transform_body(Value::Object(build_generation_body( + &self.model_id, + options, + &openai_options, + ))); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new( - self.generations_endpoint(), - build_header_list(&headers), - options, - ), - Value::Object(body.clone()), + self.config + .http_request(self.config.url("/images/generations")?, headers, options), + body.clone(), aimux_provider_utils::create_json_response_handler(), super::openai_failed_response_handler(), ) .await?; let val: Value = resp.value; - (val, resp.response_headers, Some(Value::Object(body))) + (val, resp.response_headers, Some(body)) }; let images = extract_images(&body_value); @@ -205,19 +172,6 @@ impl ImageModel for OpenAIImageModel { // ── Helpers ───────────────────────────────────────────────────────────────── -/// Convert a `HashMap` into the `Vec<(String, String)>` header -/// list expected by [`aimux_provider_utils::HttpRequest`]. -/// -/// `Content-Type` is intentionally not added here: the http layer sets it -/// automatically from the [`HttpBody`] (`Json` → `application/json`, -/// `Bytes` → the supplied content-type). -fn build_header_list(headers: &HashMap) -> Vec<(String, String)> { - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect() -} - /// OpenAI generation provider options (camelCase → snake_case mapping). struct GenerationOptions { quality: Option, diff --git a/aimux-providers/src/openai/mod.rs b/aimux-providers/src/openai/mod.rs index 5f51ed721..5e2c3ce38 100644 --- a/aimux-providers/src/openai/mod.rs +++ b/aimux-providers/src/openai/mod.rs @@ -1,8 +1,17 @@ -//! OpenAI-compatible provider. +//! OpenAI provider. //! -//! Works with any OpenAI-compatible API endpoint (OpenAI, Azure OpenAI, -//! Together.ai, Groq, Fireworks, DeepSeek, etc.). +//! [`create_openai`] is the Rust form of the AI SDK's `createOpenAI`: it takes +//! [`OpenAIProviderSettings`], validates the base URL, fixes the provider name +//! and returns an [`OpenAIProvider`]. The API key is not read there; it is +//! loaded in the request headers of every call, from the setting or from +//! `OPENAI_API_KEY`. [`openai()`] is the default instance. +//! +//! This package is the native OpenAI API only. Servers that merely speak the +//! same wire format are served by [`crate::openai_compatible`] (and the +//! registry presets built on it). +pub mod completion; +pub(crate) mod config; pub mod convert; mod convert_common; pub mod decision; @@ -10,27 +19,42 @@ pub mod embedding; pub mod files; pub mod image; pub mod model; +pub(crate) mod options; pub mod responses; pub mod speech; 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; pub use image::OpenAIImageModel; pub use model::OpenAIModel; pub use responses::OpenAIResponsesModel; pub use speech::OpenAISpeechModel; pub use transcription::OpenAITranscriptionModel; -use std::collections::HashMap; +use std::sync::{Arc, OnceLock}; + +use futures::future::BoxFuture; +use serde_json::Value; +use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; +use aimux_core::files_model::Files; +use aimux_core::image_model::ImageModel; use aimux_core::language_model::LanguageModel; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_core::speech_model::SpeechModel; +use aimux_core::transcription_model::TranscriptionModel; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; -use aimux_core::provider::Provider; -use aimux_provider_utils::{RetryConfig, load_api_key, without_trailing_slash}; -use serde_json::Value; +use crate::shared::{Credential, provider_headers}; + +use config::OpenAIModelConfig; pub(crate) fn openai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -89,350 +113,269 @@ pub(crate) fn openai_stream_error( ) } -/// 描述 OpenAI 兼容厂商的差异。 -/// -/// 薄封装填这个结构，共享的请求构造和响应解析读它决定行为。 -/// 这样既保持 `dyn LanguageModel` 能跨厂商互换，又能表达差异。 -#[derive(Debug, Clone, Default)] -pub struct OpenAICompatProfile { - /// 是否支持 top_k 参数。Groq 等厂商不支持，设为 false 时请求体不发送 top_k。 - pub supports_top_k: bool, - /// 是否支持 tools。少数厂商不支持，设为 false 时请求体不发送 tools/tool_choice。 - pub supports_tools: bool, - /// 是否支持 response_format。 - pub supports_response_format: bool, - /// 流式 usage 的特殊 key。Groq 用 "x_groq"，大多数厂商用顶层 "usage"（留 None）。 - pub stream_usage_key: Option<&'static str>, - /// max-token 字段的 key（内部数据，非用户概念；RFC-0017 阶段 2）。 - /// - `Some("max_tokens")` → 只发 `max_tokens` - /// - `Some("max_completion_tokens")` → 只发 `max_completion_tokens` - /// - `None` → 按模型推断（推理模型 mct / 非推理 max_tokens） - pub max_tokens_key: Option<&'static str>, -} +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; -impl OpenAICompatProfile { - /// 默认 profile：支持全部能力，无特殊流式 usage key。 - /// 适用于 OpenAI 本身和大多数兼容厂商。 - #[must_use] - pub fn full() -> Self { - Self { - supports_top_k: true, - supports_tools: true, - supports_response_format: true, - stream_usage_key: None, - max_tokens_key: None, - } - } - - /// Groq profile：不支持 top_k，流式 usage 在 x_groq 字段； - /// `max_tokens` 已弃用，只发 `max_completion_tokens`（backlog B9）。 - #[must_use] - pub fn groq() -> Self { - Self { - supports_top_k: false, - supports_tools: true, - supports_response_format: true, - stream_usage_key: Some("x_groq"), - max_tokens_key: Some("max_completion_tokens"), - } - } - - /// 设置 `max_tokens_key`（内部数据，非用户概念）：`"max_tokens"` 或 - /// `"max_completion_tokens"`。注册表薄封装行用此构建差异 profile。 - #[must_use] - pub fn with_max_tokens_key(mut self, key: &'static str) -> Self { - self.max_tokens_key = Some(key); - self - } - - /// DeepSeek profile：特化已退役（RFC-0017 阶段 2），回归 `full()`—— - /// thinking / effort 映射等厂商差异由用户 bodyOverrides 定义。 - /// 保留此薄封装以维持注册表与调用方结构不变。 - #[must_use] - pub fn deepseek() -> Self { - Self::full() - } +/// 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 +/// on every request. +#[derive(Clone, Default)] +pub struct OpenAIProviderSettings { + /// Base URL for the API calls. `None` reads `OPENAI_BASE_URL` when the + /// provider is created, else `https://api.openai.com/v1`; a trailing + /// slash is removed. An invalid explicit value fails [`create_openai`]; + /// an invalid `OPENAI_BASE_URL` fails every request instead. + 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 + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. + pub api_key: Option, + /// Sent as `OpenAI-Organization`. + pub organization: Option, + /// Sent as `OpenAI-Project`. + pub project: Option, + /// Extra headers on every request. A `None` value removes the header, + /// including one of the fixed ones. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of every model's `provider()` string. + /// Default `"openai"`. + 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, + /// Custom socket connector for streaming transcription. + #[cfg(feature = "realtime")] + pub web_socket: Option>, } -/// RFC-0023:从 `OpenAIConfig` 构建 `ProviderRecord`(共享给 chat/responses 两个 model)。 -pub(crate) fn config_snapshot_from_config( - provider: &str, - model_id: &str, - config: &OpenAIConfig, -) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - ProviderRecord { - provider: provider.to_string(), - model_id: model_id.to_string(), - base_url: Some(config.base_url.clone()), - // 来源字段缺失时保守记为 explicit(显式 key 最常见;不泄露明文)。 - api_key_source: config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: Some(profile_to_json(&config.profile)), - provider_options: provider_options_to_json(config), +impl std::fmt::Debug for OpenAIProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OpenAIProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key.is_some()) + .field("organization", &self.organization) + .field("project", &self.project) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } -/// profile → JSON(`&'static str` 字段序列化为 String,回放重建时转回)。 -fn profile_to_json(p: &OpenAICompatProfile) -> serde_json::Value { - serde_json::json!({ - "supports_top_k": p.supports_top_k, - "supports_tools": p.supports_tools, - "supports_response_format": p.supports_response_format, - "stream_usage_key": p.stream_usage_key, - "max_tokens_key": p.max_tokens_key, +/// Create an OpenAI provider. +/// +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when the explicit `base_url` is not +/// an `http(s)` URL with a host. That is the only way this fails: the key is +/// loaded per request, and an invalid `OPENAI_BASE_URL` is reported by each +/// request, not here. +pub fn create_openai(settings: OpenAIProviderSettings) -> Result { + let base_url = crate::shared::base_url_setting( + settings.base_url.as_deref(), + "OPENAI_BASE_URL", + DEFAULT_BASE_URL, + validate_base_url, + )?; + let name = settings.name.unwrap_or_else(|| "openai".to_string()); + Ok(OpenAIProvider { + name, + 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, + "OpenAI", + ), + [ + ("OpenAI-Organization", settings.organization), + ("OpenAI-Project", settings.project), + ] + .into_iter() + .filter_map(|(name, value)| value.map(|value| (name.to_string(), value))) + .collect(), + settings.headers, + ), + "openai", + "4.0.80", + ), + fetch: settings.fetch, + #[cfg(feature = "realtime")] + web_socket: settings.web_socket, }) } -/// 可重建的 provider_options(与 `provider::ProviderOptions` 序列化形状一致, -/// rebuild_provider 直接反序列化)。base_url 已放 `ProviderRecord.base_url`,不重复。 -fn provider_options_to_json(config: &OpenAIConfig) -> Option { - let opts = crate::provider::ProviderOptions { - base_url: None, - headers: config.headers.clone(), - organization: config.org_id.clone(), - project: config.project.clone(), - max_retries: Some(config.retry_config.max_retries), - body_overrides: config.body_overrides.clone(), - }; - serde_json::to_value(opts).ok() +/// The default provider: `create_openai` with default settings, created on +/// first use. `OPENAI_BASE_URL` is read at creation (an invalid value fails +/// each request, not the creation); the key is loaded per request. +pub fn openai() -> &'static OpenAIProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_openai(OpenAIProviderSettings::default()) + .expect("default settings have no explicit base URL to reject") + }) } -/// Configuration for the OpenAI provider. -#[derive(Debug, Clone)] -pub struct OpenAIConfig { - pub api_key: String, - /// api_key 来源(RFC-0023 `ProviderRecord.api_key_source` 分类): - /// `Some("env:VAR")` = 来自环境变量;`Some("none")` = 本地无认证占位; - /// `None` = 显式传 key(config_snapshot 记为 "explicit")。不存明文之外 - /// 的信息,仅用于回放重建(S-1 建议的来源追踪字段)。 - pub api_key_source: Option, - pub base_url: String, - pub org_id: Option, - /// OpenAI project ID sent via the `OpenAI-Project` header. - pub project: Option, - /// Extra headers merged into every request. - pub headers: Option>, - /// Provider name — controls provider-specific behaviour in the shared - /// request builder (e.g. "groq" reads provider options from the "groq" - /// key and applies a reasoning-effort map). Defaults to "openai". - pub provider: String, - /// 厂商能力差异描述。默认 `full()`（支持全部能力）。 - /// 薄封装用 `with_profile()` 设置差异。 - pub profile: OpenAICompatProfile, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, - /// Provider 级请求体覆盖（RFC-0017）。在标准请求体 + 内置厂商 override - /// 之后 deep-merge。per-call 的 `CallOptions.body_overrides` 在此之后 - /// 再 merge（覆盖 provider 级）。 - pub body_overrides: Option, +/// An OpenAI provider (the AI SDK's `OpenAIProvider`). Cheap to clone the +/// models out of; it holds no HTTP client. +pub struct OpenAIProvider { + name: String, + /// The validated base URL, or the error of an invalid `OPENAI_BASE_URL`, + /// which every request reports. + base_url: Result, + headers: HeadersFn, + fetch: Option, + #[cfg(feature = "realtime")] + web_socket: Option>, } -impl OpenAIConfig { - /// Create from an API key (uses default OpenAI base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - api_key_source: None, - base_url: "https://api.openai.com/v1".to_string(), - org_id: None, - project: None, - headers: None, - provider: "openai".to_string(), - profile: OpenAICompatProfile::full(), - retry_config: RetryConfig::default(), - body_overrides: None, +impl OpenAIProvider { + fn model_config(&self, method: &str) -> OpenAIModelConfig { + let mut config = OpenAIModelConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone().unwrap_or_default(), + self.headers.clone(), + self.fetch.clone(), + None, + ); + if let Err(error) = &self.base_url { + let error = error.clone(); + config.url = Arc::new(move |_| Err(error.clone())); } + config } - /// Use a custom base URL (for Azure, Groq, etc.). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - - #[must_use] - pub fn with_org_id(mut self, org_id: impl Into) -> Self { - self.org_id = Some(org_id.into()); - self - } - - /// Set the OpenAI project ID (sent via the `OpenAI-Project` header). + /// A chat-completions model; `provider()` is `"{name}.chat"`. #[must_use] - pub fn with_project(mut self, project: impl Into) -> Self { - self.project = Some(project.into()); - self + pub fn chat(&self, model_id: &str) -> OpenAIModel { + OpenAIModel::from_config(model_id.to_string(), self.model_config("chat")) } - /// Attach extra headers merged into every request. + /// A text-completion model; `provider()` is `"{name}.completion"`. #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self + pub fn completion(&self, model_id: &str) -> OpenAICompletionModel { + OpenAICompletionModel::from_native_config( + model_id.to_string(), + self.model_config("completion"), + ) } - /// Set the provider name (e.g. "groq") for provider-specific behaviour. + /// A Responses API model; `provider()` is `"{name}.responses"`. #[must_use] - pub fn with_provider(mut self, provider: impl Into) -> Self { - self.provider = provider.into(); - self + pub fn responses(&self, model_id: &str) -> OpenAIResponsesModel { + let mut config = self.model_config("responses"); + config.responses.file_id_prefixes = vec!["file-"]; + OpenAIResponsesModel::from_config(model_id.to_string(), config) } - /// 设置厂商能力差异描述。 + /// An embedding model; `provider()` is `"{name}.embedding"`. #[must_use] - pub fn with_profile(mut self, profile: OpenAICompatProfile) -> Self { - self.profile = profile; - self + pub fn embedding(&self, model_id: &str) -> OpenAIEmbeddingModel { + OpenAIEmbeddingModel::from_config(model_id.to_string(), self.model_config("embedding")) } - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self - } - - /// 设置 provider 级请求体覆盖（RFC-0017）。 - #[must_use] - pub fn with_body_overrides(mut self, overrides: Value) -> Self { - self.body_overrides = Some(overrides); - self - } - - /// Create from environment variable `OPENAI_API_KEY`. + /// A native Decisions model (an aimux extension; the AI SDK has no + /// decision model). Its `provider()` is `"openai"`. /// /// # Errors /// - /// Returns `AiMuxError::InvalidArgument` when `OPENAI_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "OPENAI_API_KEY", "OpenAI")?; - Ok(Self::new(api_key).with_api_key_source(Some("env:OPENAI_API_KEY"))) + /// `InvalidArgument` for an empty model id. + pub fn decision(&self, model_id: &str) -> Result { + OpenAIDecisionModel::from_config(model_id, self.model_config("decision")) } - /// 标注 api_key 来源(RFC-0023 回放重建用)。 + /// An image model; `provider()` is `"{name}.image"`. #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self + pub fn image(&self, model_id: &str) -> OpenAIImageModel { + OpenAIImageModel::from_config(model_id.to_string(), self.model_config("image")) } -} -/// OpenAI provider — creates `OpenAIModel` instances. -/// -/// 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 OpenAIProvider { - config: OpenAIConfig, -} - -impl OpenAIProvider { + /// A speech (TTS) model; `provider()` is `"{name}.speech"`. #[must_use] - pub fn new(config: OpenAIConfig) -> Self { - Self { config } + pub fn speech(&self, model_id: &str) -> OpenAISpeechModel { + OpenAISpeechModel::from_config(model_id.to_string(), self.model_config("speech")) } - /// Create a model instance for the given model name (e.g. `"gpt-4o"`). + /// A transcription (STT) model; `provider()` is `"{name}.transcription"`. #[must_use] - pub fn model(&self, model_id: &str) -> model::OpenAIModel { - model::OpenAIModel::new(model_id.to_string(), self.config.clone()) + pub fn transcription(&self, model_id: &str) -> OpenAITranscriptionModel { + let model = OpenAITranscriptionModel::from_config( + model_id.to_string(), + self.model_config("transcription"), + ); + #[cfg(feature = "realtime")] + let model = model.with_web_socket(self.web_socket.clone()); + model } - /// Create a Responses API model instance for the given model name (e.g. - /// `"gpt-4o"`). Uses the `/v1/responses` endpoint instead of - /// `/v1/chat/completions`. + /// The files interface; `provider()` is `"{name}.files"`. #[must_use] - pub fn responses_model(&self, model_id: &str) -> responses::OpenAIResponsesModel { - responses::OpenAIResponsesModel::new(model_id.to_string(), self.config.clone()) + pub fn files(&self) -> OpenAIFiles { + OpenAIFiles::from_config(self.model_config("files")) } - /// Create a Files interface for uploading files to OpenAI. + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; it returns the Responses model, the same + /// model as [`responses`](Self::responses) and + /// [`language_model`](Provider::language_model). #[must_use] - pub fn files(&self) -> files::OpenAIFiles { - files::OpenAIFiles::new(self.config.clone()) - } - - /// Create an embedding model instance for the given model name (e.g. - /// `"text-embedding-3-large"`). - #[must_use] - pub fn embedding_model(&self, model_id: &str) -> embedding::OpenAIEmbeddingModel { - embedding::OpenAIEmbeddingModel::new(model_id.to_string(), self.config.clone()) + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.responses(model_id)) } +} - /// Create a speech (TTS) model instance for the given model name (e.g. - /// `"tts-1"`). Uses the `/audio/speech` endpoint. - #[must_use] - pub fn speech(&self, model_id: &str) -> speech::OpenAISpeechModel { - speech::OpenAISpeechModel::new(model_id.to_string(), self.config.clone()) +impl Provider for OpenAIProvider { + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) } - /// Create an image generation model instance for the given model name - /// (e.g. `"dall-e-3"` or `"gpt-image-1"`). Uses the `/images/generations` - /// endpoint for generation and `/images/edits` for editing. - #[must_use] - pub fn image(&self, model_id: &str) -> image::OpenAIImageModel { - image::OpenAIImageModel::new(model_id.to_string(), self.config.clone()) + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - /// Create a transcription (STT) model instance for the given model name - /// (e.g. `"whisper-1"` or `"gpt-4o-transcribe"`). Uses the - /// `/audio/transcriptions` endpoint. - #[must_use] - pub fn transcription(&self, model_id: &str) -> transcription::OpenAITranscriptionModel { - transcription::OpenAITranscriptionModel::new(model_id.to_string(), self.config.clone()) + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) } -} -impl Provider for OpenAIProvider { - fn name(&self) -> &str { - "openai" + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.image(model_id))) } fn decision_model( &self, model_id: &str, ) -> Result, AiMuxError> { - if self.config.provider == "openai" { - Ok(Box::new(OpenAIDecisionModel::new( - model_id, - self.config.clone(), - )?)) - } else { - Ok(Box::new( - crate::runtime_decision::RuntimeDecisionModel::new(model_id, self.config.clone())?, - )) - } + Ok(Box::new(self.decision(model_id)?)) + } + + fn transcription_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.transcription(model_id)))) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn speech_model(&self, model_id: &str) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.speech(model_id)))) } - /// List models via `GET {base_url}/models` (OpenAI-compatible), enriched - /// with the community catalogue portrait when available (RFC-0027). - /// - /// The provider name used for catalogue lookup is `config.provider` (the - /// registry entry name, e.g. `"deepseek"`), not the hardcoded `"openai"` - /// returned by [`name`](Provider::name) — that lets the same shared impl - /// attach the right portrait per registry-backed provider. - fn list_models( - &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let config = self.config.clone(); - Box::pin(async move { - let headers = model::build_auth_headers(&config); - let runtime = - model::execute_list_models(&config.base_url, &headers, config.retry_config).await?; - Ok(runtime) - }) + fn files(&self) -> Option> { + Some(Arc::new(self.files())) + } +} + +impl ProviderDiscovery for OpenAIProvider { + /// `GET {base_url}/models`: one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config("models"); + Box::pin(async move { model::list_models_once(&config).await }) } } diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index 9ab26b76d..72c374e43 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -1,15 +1,13 @@ //! OpenAI language model — implements `LanguageModel` trait. //! -//! The HTTP request/response handling lives in the free functions -//! [`execute_generate`] and [`execute_stream`], which take an endpoint URL, a -//! header map and a model id. They call the `aimux-provider-utils` API helpers — +//! The HTTP request/response handling lives in the crate-internal functions +//! `execute_generate` and `execute_stream`, which take a prepared request (URL, +//! headers, transport) and a model id. They call the `aimux-provider-utils` API helpers — //! **no `reqwest` types cross this boundary**. This lets other providers that //! speak the OpenAI chat-completions wire format (notably Azure OpenAI) reuse //! the conversion + streaming logic while supplying their own URL and auth. use aimux_core::tool::RawToolCall; -use std::collections::HashMap; - use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; @@ -24,88 +22,25 @@ use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usa use aimux_provider_utils::{HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker}; -use super::OpenAIConfig; -use super::convert::{RequestBodyResult, build_request_body_with_warnings, parse_finish_reason}; +use super::config::{OpenAIModelConfig, TransformRequestBody}; +use super::convert::{ + RequestBodyResult, build_request_body_with_chat_options, parse_finish_reason, +}; use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; -/// An OpenAI-compatible language model. +/// An OpenAI chat-completions language model. /// /// 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 OpenAIModel { model_id: String, - config: OpenAIConfig, + config: OpenAIModelConfig, } impl OpenAIModel { - #[must_use] - pub fn new(model_id: String, config: OpenAIConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: OpenAIModelConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = build_auth_headers(&self.config); - // Per-call headers (from CallOptions.headers), overriding provider-level. - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/chat/completions", self.config.base_url) - } -} - -/// Build the auth + provider-level headers from a config (no per-call headers). -/// -/// Shared by `OpenAIModel::build_headers` and the model-listing path -/// ([`execute_list_models`]) so both use identical auth wiring. -#[must_use] -pub fn build_auth_headers(config: &super::OpenAIConfig) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", config.api_key), - ); - if let Some(ref org) = config.org_id { - headers.insert("OpenAI-Organization".to_string(), org.clone()); - } - if let Some(ref project) = config.project { - headers.insert("OpenAI-Project".to_string(), project.clone()); - } - if let Some(ref config_headers) = config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - headers -} - -/// Merge provider-level `body_overrides` into the per-call options (RFC-0017). -/// -/// Provider-level overrides are applied first (lower priority); per-call -/// `body_overrides` from `CallOptions` are merged on top. If neither is -/// present, the options are returned unchanged (cheap clone). -fn merge_body_overrides(options: &CallOptions, provider_overrides: &Option) -> CallOptions { - match (provider_overrides, &options.body_overrides) { - (Some(provider), Some(call)) => { - // Merge: provider first, then call (call wins). - let mut merged = provider.clone(); - crate::openai::convert::deep_merge_json(&mut merged, call); - let mut opts = options.clone(); - opts.body_overrides = Some(merged); - opts - } - (Some(provider), None) => { - let mut opts = options.clone(); - opts.body_overrides = Some(provider.clone()); - opts - } - (None, _) => options.clone(), - } } // ── Usage conversion ───────────────────────────────────────────────────────── @@ -120,7 +55,7 @@ fn merge_body_overrides(options: &CallOptions, provider_overrides: &Option) -> Usage { let prompt_tokens = usage.prompt_tokens.unwrap_or(0); @@ -170,7 +105,7 @@ fn convert_usage(usage: &UsageResponse, usage_raw: Option<&Value>) -> Usage { reasoning: Some(reasoning_tokens), }, // M10 (RFC-0016): keep the provider's original usage JSON verbatim — - // vendor-specific fields (e.g. Moonshot `cached_tokens`, DeepSeek + // vendor-specific fields (e.g. Moonshot `cached_tokens`, prompt-cache // `prompt_cache_hit_tokens`) are otherwise lost for audit/billing. raw: usage_raw.and_then(|value| value.as_object().cloned()), } @@ -178,11 +113,8 @@ fn convert_usage(usage: &UsageResponse, usage_raw: Option<&Value>) -> Usage { #[async_trait] impl LanguageModel for OpenAIModel { - /// Provider identity for recording/routing. Uses `config.provider` (the - /// registry entry name, e.g. `"deepseek"`/`"groq"`) rather than a hardcoded - /// `"openai"`, so OpenAI-compatible providers keep their real identity — - /// mirroring the Responses path (`OpenAIResponsesModel::provider`). - /// Direct `OpenAIProvider` use defaults to `"openai"`. + /// Provider identity for recording/routing: `"{name}.chat"` (`"openai.chat"` + /// by default). fn provider(&self) -> &str { &self.config.provider } @@ -191,38 +123,42 @@ impl LanguageModel for OpenAIModel { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - super::config_snapshot_from_config(&self.config.provider, &self.model_id, &self.config) + async fn supported_urls(&self) -> aimux_core::language_model::SupportedUrls { + self.config.supported_urls.clone() } async fn do_generate(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()); - let options = merge_body_overrides(options, &self.config.body_overrides); + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; + let http = + self.config + .http_request(self.config.url("/chat/completions")?, headers, options); execute_generate( - &self.endpoint(), - &headers, + http, &self.model_id, - &options, - &self.config.provider, - &self.config.profile, + options, + self.config.transform_request_body.as_ref(), + self.config.chat_options, ) .await } async fn do_stream(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()); - let options = merge_body_overrides(options, &self.config.body_overrides); + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; + let http = + self.config + .http_request(self.config.url("/chat/completions")?, headers, options); execute_stream( - &self.endpoint(), - &headers, + http, &self.model_id, - &options, - &self.config.provider, - &self.config.profile, + options, + self.config.transform_request_body.as_ref(), + self.config.chat_options, ) .await } @@ -231,46 +167,37 @@ impl LanguageModel for OpenAIModel { // ── Shared OpenAI chat-completions execution ───────────────────────────────── // // These free functions contain the actual HTTP + response-parsing logic. They -// are `pub` so that providers speaking the OpenAI wire format (Azure OpenAI) -// can reuse them with their own endpoint URL and auth headers. - -/// Build the header list for a JSON POST: auth/extra headers + `Content-Type`. -/// -/// Returns a `Vec<(String, String)>` for `HttpRequest` — no reqwest types. -fn build_header_list(headers: &HashMap) -> Vec<(String, String)> { - let mut list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - list.push(("Content-Type".to_string(), "application/json".to_string())); - list -} +// are shared with providers that speak the OpenAI chat wire format (Azure +// OpenAI), which supply their own request: URL, auth headers and transport. /// Execute a non-streaming OpenAI chat-completion request. /// -/// `endpoint` is the full chat-completions URL; `headers` carries the auth -/// headers (and any extra/request headers); `model_id` is placed in the -/// request body's `model` field. +/// `http` carries the full chat-completions URL, the auth and request headers +/// and the transport; `model_id` is placed in the request body's `model` +/// field. `transform_request_body`, when set, rewrites the finished body once +/// before it is sent (and the rewritten body is what the result reports). /// /// # Errors /// /// Returns request-build conversion errors, `ApiCall` for HTTP/transport /// failures, `JsonParse` for a malformed body, and `InvalidResponseData` when /// `choices` is empty. -pub async fn execute_generate( - endpoint: &str, - headers: &HashMap, +pub(crate) async fn execute_generate( + http: HttpRequest, model_id: &str, options: &CallOptions, - provider: &str, - profile: &super::OpenAICompatProfile, + transform_request_body: Option<&TransformRequestBody>, + parse_chat_options: super::options::ChatOptionsParser, ) -> Result { let request_result = - build_request_body_with_warnings(model_id, options, false, provider, profile)?; - let body = request_result.body; + build_request_body_with_chat_options(model_id, options, false, parse_chat_options)?; + let body = match transform_request_body { + Some(transform) => transform(request_result.body), + None => request_result.body, + }; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.to_string(), build_header_list(headers), options), + http, body.clone(), aimux_provider_utils::create_json_response_handler::(), super::openai_failed_response_handler(), @@ -300,7 +227,7 @@ pub async fn execute_generate( provider_metadata: None, }); } - // Reasoning: prefer reasoning_content over reasoning (DeepSeek/阿里通义). + // Reasoning: prefer reasoning_content over reasoning. let reasoning_text = choice .message .reasoning_content @@ -402,30 +329,34 @@ pub async fn execute_generate( /// Execute a streaming OpenAI chat-completion request. /// -/// `endpoint` is the full chat-completions URL; `headers` carries the auth -/// headers (and any extra/request headers); `model_id` is placed in the -/// request body's `model` field. +/// `http` carries the full chat-completions URL, the auth and request headers +/// and the transport; `model_id` is placed in the request body's `model` +/// field. `transform_request_body` is applied as in [`execute_generate`]. /// /// # Errors /// /// Returns request-build conversion errors and `ApiCall` when establishing the /// stream fails; transport errors surface as `Err` items in the stream. -pub async fn execute_stream( - endpoint: &str, - headers: &HashMap, +pub(crate) async fn execute_stream( + http: HttpRequest, model_id: &str, options: &CallOptions, - provider: &str, - profile: &super::OpenAICompatProfile, + transform_request_body: Option<&TransformRequestBody>, + parse_chat_options: super::options::ChatOptionsParser, ) -> Result { let request_result = - build_request_body_with_warnings(model_id, options, true, provider, profile)?; + build_request_body_with_chat_options(model_id, options, true, parse_chat_options)?; // M9 (RFC-0016): keep the warnings computed while building the body — // they are emitted in `StreamStart` below instead of being dropped. let RequestBodyResult { body, warnings } = request_result; + let body = match transform_request_body { + Some(transform) => transform(body), + None => body, + }; + let endpoint = http.url.clone(); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.to_string(), build_header_list(headers), options), + http, body.clone(), aimux_provider_utils::create_event_source_response_handler::(), super::openai_failed_response_handler(), @@ -450,22 +381,16 @@ pub async fn execute_stream( { return Err(super::openai_stream_error( err_obj, - endpoint, + &endpoint, body.clone(), response_headers.clone(), )); } - // Capture the provider's stream usage key before entering the async stream - // block (the borrowed `profile` cannot be moved into the generator). - // Some(key) → read streaming usage from `chunk[key].usage`; - // None → read from the top-level `usage`. - let stream_usage_key = profile.stream_usage_key; - // M2 (RFC-0016): capture whether raw chunks should be emitted — the // borrowed `options` cannot be moved into the generator. let emit_raw_chunks = options.include_raw_chunks == Some(true); - let stream_error_url = endpoint.to_owned(); + let stream_error_url = endpoint; let stream_error_body = body.clone(); let stream_response_headers = response_headers.clone(); @@ -525,17 +450,10 @@ pub async fn execute_stream( break; } - // Extract streaming usage based on profile.stream_usage_key - // (captured before the stream block). Some(key) reads usage - // from the provider-specific sub-object `key` (e.g. Groq's - // "x_groq"); None reads the top-level "usage". This is done - // from the raw JSON chunk before it is consumed below. - // `chunk_usage_raw` keeps the provider's original object for - // `Usage.raw` (M10). - let chunk_usage_raw: Option = match stream_usage_key { - Some(key) => parsed.get(key).and_then(|v| v.get("usage")).cloned(), - None => parsed.get("usage").cloned(), - }; + // The chunk's top-level `usage`, taken from the raw JSON before + // the chunk is consumed below. `chunk_usage_raw` keeps the + // provider's original object for `Usage.raw` (M10). + let chunk_usage_raw: Option = parsed.get("usage").cloned(); let chunk_usage: Option = chunk_usage_raw .as_ref() .and_then(|u| serde_json::from_value(u.clone()).ok()); @@ -564,7 +482,7 @@ pub async fn execute_stream( })); } - // Update usage based on profile.stream_usage_key. + // Update usage from the chunk that carries it. if let Some(usage) = &chunk_usage { final_usage = convert_usage(usage, chunk_usage_raw.as_ref()); final_usage_raw = Some(usage.clone()); @@ -803,47 +721,29 @@ struct ModelEntry { created: Option, } -/// Execute a `GET {base_url}/models` request (OpenAI-compatible) and return the -/// sparse runtime model list. Used by `OpenAIProvider::list_models`. -/// -/// `headers` carries the auth headers (Bearer key etc.); `base_url` is the -/// provider's API base (e.g. `https://api.openai.com/v1`). +/// One `GET {base_url}/models` exchange (OpenAI-compatible): no retry, no +/// recording. Discovery is not a Core operation and the AI SDK has no +/// equivalent, so a failure is reported to the caller as it happened. /// /// # Errors /// -/// Returns `ApiCall` for HTTP/transport failures and `JsonParse` when the -/// body does not deserialize into the models list. -pub async fn execute_list_models( - base_url: &str, - headers: &HashMap, - retry_config: aimux_core::retry::RetryConfig, +/// Returns the header-resolution error (a missing key is `LoadApiKey`), +/// `ApiCall` for HTTP/transport failures and `JsonParse` when the body does +/// not deserialize into the models list. +pub(crate) async fn list_models_once( + config: &OpenAIModelConfig, ) -> Result, AiMuxError> { - // Strip a trailing slash so we don't get `//models`. - let base = base_url.trim_end_matches('/'); - let url = format!("{base}/models"); - - // `list_models` is not a Core user operation, so nothing above this call - // retries it — apply the Core retry primitive here (a catalogue GET is - // idempotent) so a transient 429/503/transport failure behaves like every - // other exchange instead of failing on the first hiccup. - let retries = aimux_core::retry::prepare_retries(None, retry_config, None); - let resp = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: build_header_list(headers), - abort_signal: None, - call_id: None, - recording_context: None, - ..Default::default() - }, - aimux_provider_utils::create_json_response_handler(), - super::openai_failed_response_handler(), - ) - }) - .await?; - + let headers = config.request_headers(None).await?; + let resp = aimux_provider_utils::get_from_api( + config.with_transport(HttpRequest { + url: config.url("/models")?, + headers, + ..Default::default() + }), + aimux_provider_utils::create_json_response_handler(), + super::openai_failed_response_handler(), + ) + .await?; let parsed: ModelsListResponse = resp.value; Ok(parsed diff --git a/aimux-providers/src/openai/options.rs b/aimux-providers/src/openai/options.rs new file mode 100644 index 000000000..9de515682 --- /dev/null +++ b/aimux-providers/src/openai/options.rs @@ -0,0 +1,148 @@ +//! Chat option schema mirrored from `openai-chat-language-model-options.ts`. + +use aimux_core::error::AiMuxError; +use aimux_core::shared::JsonObject; +use serde_json::Value; + +use super::convert::parse_option_fields; + +pub(crate) const NAMESPACE: &str = "openai"; +pub(crate) type ChatOptionsParser = fn(&JsonObject) -> Result; + +pub(crate) fn parse_chat_options(options: &JsonObject) -> Result { + parse_option_fields(options, NAMESPACE, |key, value| { + let one_of = |values: &[&str]| value.as_str().is_some_and(|s| values.contains(&s)); + let valid = match key { + "logitBias" => { + return Some(value.as_object().ok_or(()).and_then(|object| { + let mut parsed = JsonObject::new(); + for (key, value) in object { + if !value.is_number() { + return Err(()); + } + let key = key.trim_matches(|c| { + matches!(c, + '\u{0009}'..='\u{000d}' | '\u{0020}' | '\u{00a0}' | '\u{1680}' + | '\u{2000}'..='\u{200a}' | '\u{2028}' | '\u{2029}' | '\u{202f}' + | '\u{205f}' | '\u{3000}' | '\u{feff}' + ) + }); + let number = if key.is_empty() { + 0.0 + } else if let Some((digits, radix)) = key + .strip_prefix("0x") + .or_else(|| key.strip_prefix("0X")) + .map(|s| (s, 16)) + .or_else(|| { + key.strip_prefix("0o") + .or_else(|| key.strip_prefix("0O")) + .map(|s| (s, 8)) + }) + .or_else(|| { + key.strip_prefix("0b") + .or_else(|| key.strip_prefix("0B")) + .map(|s| (s, 2)) + }) + { + if digits.is_empty() { + return Err(()); + } + let bits_per_digit = match radix { + 16 => 4, + 8 => 3, + _ => 1, + }; + let mut significant_bits = 0usize; + let mut mantissa = 0u64; + let mut guard = false; + let mut sticky = false; + for digit in digits.chars() { + let digit = digit.to_digit(radix).ok_or(())?; + for shift in (0..bits_per_digit).rev() { + let bit = (digit >> shift) & 1; + if significant_bits == 0 && bit == 0 { + continue; + } + if significant_bits < 53 { + mantissa = (mantissa << 1) | u64::from(bit); + } else if significant_bits == 53 { + guard = bit != 0; + } else { + sticky |= bit != 0; + } + significant_bits += 1; + } + } + if guard && (sticky || mantissa & 1 != 0) { + mantissa += 1; + } + let exponent = significant_bits.saturating_sub(53); + if exponent > 1023 { + f64::INFINITY + } else { + mantissa as f64 * 2.0f64.powi(exponent as i32) + } + } else { + key.parse::().map_err(|_| ())? + }; + if !number.is_finite() { + return Err(()); + } + let key = if number == 0.0 { + "0".to_string() + } else if number.abs() >= 1e21 || number.abs() < 1e-6 { + let scientific = format!("{number:e}"); + let (mantissa, exponent) = scientific.split_once('e').ok_or(())?; + let exponent = exponent.parse::().map_err(|_| ())?; + format!("{mantissa}e{exponent:+}") + } else { + number.to_string() + }; + parsed.insert(key, value.clone()); + } + Ok(Value::Object(parsed)) + })); + } + "logprobs" => value.is_boolean() || value.is_number(), + "maxCompletionTokens" => value.is_number(), + "metadata" => value.as_object().is_some_and(|object| { + object.iter().all(|(key, value)| { + key.encode_utf16().count() <= 64 + && value + .as_str() + .is_some_and(|s| s.encode_utf16().count() <= 512) + }) + }), + "prediction" => value.is_object(), + "reasoningEffort" => { + one_of(&["none", "minimal", "low", "medium", "high", "xhigh", "max"]) + } + "serviceTier" => one_of(&["auto", "flex", "priority", "fast", "ultrafast", "default"]), + "textVerbosity" => one_of(&["low", "medium", "high"]), + "promptCacheRetention" => one_of(&["in_memory", "24h"]), + "systemMessageMode" => one_of(&["system", "developer", "remove"]), + "user" | "promptCacheKey" | "safetyIdentifier" => value.is_string(), + "parallelToolCalls" | "store" | "strictJsonSchema" | "forceReasoning" => { + value.is_boolean() + } + "promptCacheOptions" => { + return Some(value.as_object().ok_or(()).and_then(|object| { + parse_option_fields(object, NAMESPACE, |key, value| { + let valid = match key { + "mode" => value + .as_str() + .is_some_and(|s| matches!(s, "implicit" | "explicit")), + "ttl" => value.as_str() == Some("30m"), + _ => return None, + }; + Some(if valid { Ok(value.clone()) } else { Err(()) }) + }) + .map(Value::Object) + .map_err(|_| ()) + })); + } + _ => return None, + }; + Some(if valid { Ok(value.clone()) } else { Err(()) }) + }) +} diff --git a/aimux-providers/src/openai/responses/convert.rs b/aimux-providers/src/openai/responses/convert.rs index 992095870..30b90d6ef 100644 --- a/aimux-providers/src/openai/responses/convert.rs +++ b/aimux-providers/src/openai/responses/convert.rs @@ -14,13 +14,14 @@ use serde_json::{Value, json}; use aimux_core::error::AiMuxError; use aimux_core::language_model_message::{ - AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + AssistantPart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, }; use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat}; -use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; +use aimux_core::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::types::ResponsesUsage; @@ -31,20 +32,73 @@ use super::types::ResponsesUsage; // converter (issue M10). use crate::openai::convert_common::{ModelCapabilities, SystemMessageMode, get_model_capabilities}; -/// Compatibility alias: the Responses-specific enum name was merged into the -/// shared [`SystemMessageMode`] during M10. Keeping the alias lets existing -/// import paths (`openai::responses::convert::ResponsesSystemMessageMode`) -/// keep compiling while the module signature uses the shared type. -#[doc(hidden)] -pub use crate::openai::convert_common::SystemMessageMode as ResponsesSystemMessageMode; - // -- Provider options helper ------------------------------------------------- -/// Get a value from `provider_options.openai.`. -fn openai_option(options: &Option, key: &str) -> Option { +/// The providerOptions keys a Responses model reads, in order of precedence +/// (the first one present wins as a whole), and the key it writes response +/// metadata under. The AI SDK's `OpenAIResponsesLanguageModel` picks +/// `providerOptionsName` from the host: `openai` for the OpenAI API, `azure` +/// for Azure (which still falls back to `openai` when no `azure` options were +/// given). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ResponsesNamespace { + read: &'static [&'static str], + write: &'static str, +} + +impl ResponsesNamespace { + /// The OpenAI API: reads and writes `openai`. + pub const OPENAI: Self = Self::new(&["openai"], "openai"); + + /// A host namespace. `read` is non-empty and in order of precedence; + /// `write` is the metadata key. + #[must_use] + pub const fn new(read: &'static [&'static str], write: &'static str) -> Self { + Self { read, write } + } + + /// The key response metadata is written under. + #[must_use] + pub fn write_key(self) -> &'static str { + self.write + } + + /// The first present options object among the read keys. + fn find(self, provider_options: &SharedProviderOptions) -> Option<&JsonObject> { + self.read.iter().find_map(|key| provider_options.get(*key)) + } + + /// [`find`](Self::find) for `CallOptions::provider_options`. + pub(crate) fn find_in(self, provider_options: &SharedProviderOptions) -> Option<&JsonObject> { + self.read.iter().find_map(|key| provider_options.get(*key)) + } +} + +impl Default for ResponsesNamespace { + fn default() -> Self { + Self::OPENAI + } +} + +/// How a host differs in the Responses API model: its namespace, and the +/// prefixes that mark a file part's data as an uploaded file id instead of +/// content to inline (`fileIdPrefixes` in the AI SDK). +#[derive(Debug, Clone, Default)] +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>, +} + +/// Get a value from the host namespace's options (`openai.` for OpenAI). +fn openai_option( + ns: ResponsesNamespace, + options: &Option, + key: &str, +) -> Option { options .as_ref() - .and_then(|m| m.get("openai")) + .and_then(|m| ns.find_in(m)) .and_then(|o| o.get(key)) .cloned() } @@ -73,8 +127,9 @@ pub struct ResponsesInputResult { /// response chain). /// /// # Errors -/// Returns an unsupported-functionality error for text file parts. +/// Returns an error for text file parts or an unresolved inline image media type. pub fn convert_to_responses_input( + ns: ResponsesNamespace, prompt: &LanguageModelPrompt, system_message_mode: SystemMessageMode, store: bool, @@ -109,18 +164,23 @@ pub fn convert_to_responses_input( LanguageModelMessage::User { content, .. } => { let content: Vec = content .iter() - .map(convert_user_part) + .map(|part| convert_user_part(ns, part)) .collect::>()?; input.push(json!({ "role": "user", "content": content })); } LanguageModelMessage::Assistant { content, .. } => { + // Reasoning item id → its index in `input`: the parts of one + // reasoning item become one input item (per assistant + // message, as `reasoningMessages` upstream). + let mut reasoning_items: std::collections::HashMap> = + std::collections::HashMap::new(); for part in content { match part { AssistantPart::Text(TextPart { text, provider_options, }) => { - let id = item_id(provider_options); + let id = item_id(ns, provider_options); if (has_previous_response_id || has_conversation) && id.is_some() { continue; } @@ -128,7 +188,7 @@ pub fn convert_to_responses_input( input.push(json!({ "type": "item_reference", "id": id })); continue; } - let phase = phase_from_provider_options(provider_options); + let phase = phase_from_provider_options(ns, provider_options); let mut item = json!({ "role": "assistant", "content": [{ "type": "output_text", "text": text }], @@ -148,19 +208,19 @@ pub fn convert_to_responses_input( provider_options, .. }) => { - let id = item_id(provider_options); + let id = item_id(ns, provider_options); if (has_previous_response_id || has_conversation) && id.is_some() { continue; } - let namespace = namespace_from_provider_options(provider_options); + let namespace = namespace_from_provider_options(ns, provider_options); let mut item = json!({ "type": "function_call", "call_id": tool_call_id, "name": tool_name, "arguments": serialize_arguments(tool_input), }); - if let Some(ref ns) = namespace { - item["namespace"] = json!(ns); + if let Some(ref namespace) = namespace { + item["namespace"] = json!(namespace); } input.push(item); } @@ -169,56 +229,76 @@ pub fn convert_to_responses_input( provider_options, .. }) => { - let reasoning_id = openai_sub_option(provider_options, "itemId"); + // convert-to-openai-responses-input.ts:1151-1250 (openai@4.0.83). + let reasoning_id = openai_sub_option(ns, provider_options, "itemId") + .and_then(|id| id.as_str().map(str::to_owned)); + let encrypted = openai_sub_option( + ns, + provider_options, + "reasoningEncryptedContent", + ) + .filter(|value| !value.is_null()); if (has_previous_response_id || has_conversation) && reasoning_id.is_some() { continue; } - if let Some(ref rid) = reasoning_id { + let summary_part = (!text.is_empty()) + .then(|| json!({ "type": "summary_text", "text": text })); + if let Some(rid) = reasoning_id { if store { - input.push(json!({ "type": "item_reference", "id": rid })); - } else { - let encrypted = openai_sub_option( - provider_options, - "reasoningEncryptedContent", - ); - let mut summary: Vec = Vec::new(); - if !text.is_empty() { - summary - .push(json!({ "type": "summary_text", "text": text })); + // One item reference, at the first part. + reasoning_items.entry(rid).or_insert_with_key(|rid| { + input.push(json!({ "type": "item_reference", "id": rid })); + None + }); + } else if let Some(&Some(index)) = reasoning_items.get(&rid) { + // Later parts of the same item append their + // summary; the last encrypted content wins. + let item = &mut input[index]; + match summary_part { + Some(summary_part) => { + if let Some(summary) = + item["summary"].as_array_mut() + { + summary.push(summary_part); + } + } + None => warnings.push(Warning::Other { + message: format!( + "Cannot append empty reasoning part to existing reasoning sequence. Skipping reasoning part: {}.", + serde_json::to_string(part).unwrap_or_default() + ), + }), + } + if let Some(enc) = encrypted { + item["encrypted_content"] = enc; } + } else { let mut item = json!({ "type": "reasoning", "id": rid, - "summary": summary, + "summary": summary_part.into_iter().collect::>(), }); if let Some(enc) = encrypted { - item["encrypted_content"] = json!(enc); + item["encrypted_content"] = enc; } + reasoning_items.insert(rid, Some(input.len())); input.push(item); } + } else if let Some(enc) = encrypted { + input.push(json!({ + "type": "reasoning", + "encrypted_content": enc, + "summary": summary_part.into_iter().collect::>(), + })); } else { - let encrypted = openai_sub_option( - provider_options, - "reasoningEncryptedContent", - ); - if let Some(enc) = encrypted { - let mut summary: Vec = Vec::new(); - if !text.is_empty() { - summary - .push(json!({ "type": "summary_text", "text": text })); - } - input.push(json!({ - "type": "reasoning", - "encrypted_content": enc, - "summary": summary, - })); - } else { - warnings.push(Warning::Other { - message: "Non-OpenAI reasoning parts are not supported. Skipping reasoning part.".to_string(), - }); - } + warnings.push(Warning::Other { + message: format!( + "Non-OpenAI reasoning parts are not supported. Skipping reasoning part: {}.", + serde_json::to_string(part).unwrap_or_default() + ), + }); } } AssistantPart::ToolResult(part) => { @@ -229,7 +309,7 @@ pub fn convert_to_responses_input( continue; } if store { - let id = item_id(&part.provider_options) + 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 { @@ -239,7 +319,7 @@ pub fn convert_to_responses_input( } } AssistantPart::Custom(part) if part.kind == "openai.compaction" => { - let id = item_id(&part.provider_options); + let id = item_id(ns, &part.provider_options); if has_conversation && id.is_some() { continue; } @@ -247,7 +327,7 @@ pub fn convert_to_responses_input( if store { input.push(json!({"type":"item_reference", "id":id})); } else { - input.push(json!({"type":"compaction", "id":id, "encrypted_content":openai_sub_option(&part.provider_options, "encryptedContent")})); + input.push(json!({"type":"compaction", "id":id, "encrypted_content":openai_sub_option(ns, &part.provider_options, "encryptedContent")})); } } } @@ -279,18 +359,35 @@ pub fn convert_to_responses_input( if let ToolResultOutput::ExecutionDenied { provider_options, .. } = output - && openai_sub_option(provider_options, "approvalId") + && openai_sub_option(ns, provider_options, "approvalId") .and_then(|value| value.as_str().map(|id| !id.is_empty())) == Some(true) { continue; } - let mut content_value = convert_tool_result_output(output, &mut warnings)?; + let mut content_value = convert_tool_result_output(ns, output, &mut warnings)?; if !matches!(output, ToolResultOutput::Content { .. }) - && let Some(breakpoint) = - crate::openai::convert::tool_result_cache_breakpoint(output).or_else( - || openai_sub_option(provider_options, "promptCacheBreakpoint"), - ) + && 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") + }) { content_value = json!([{"type":"input_text", "text":content_value, "prompt_cache_breakpoint":breakpoint}]); } @@ -326,17 +423,11 @@ pub fn convert_to_responses_input( } /// Convert a single user-message content part into the Responses input shape. -fn convert_user_part(part: &UserPart) -> Result { +fn convert_user_part(ns: ResponsesNamespace, part: &UserPart) -> Result { Ok(match part { - UserPart::Text(TextPart { text, .. }) => json!({ "type": "input_text", "text": text }), - UserPart::File(FilePart { - data, - media_type, - filename, - provider_options, - }) => { - let is_image = media_type.split('/').next() == Some("image"); - let mut item = match data { + 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) => { @@ -345,40 +436,53 @@ fn convert_user_part(part: &UserPart) -> Result { } FileBytes::Base64(data) => data.clone(), }; - let data_url = format!("data:{media_type};base64,{b64}"); - if is_image { - json!({ "type": "input_image", "image_url": data_url }) + 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": filename.as_deref().unwrap_or("part.pdf"), - "file_data": data_url, + "filename": fname, + "file_data": format!("data:{};base64,{}", file.media_type, b64), }) } } FileData::Url { url, .. } => { - if is_image { + 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 }) } } - FileData::Reference { reference } => { - let file_id = reference.get("openai").ok_or_else(|| { - AiMuxError::InvalidArgument( - "missing file reference for provider openai".into(), - ) - })?; - json!({ "type": if is_image { "input_image" } else { "input_file" }, "file_id": file_id }) - } 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(provider_options, "imageDetail") + && let Some(detail) = openai_sub_option(ns, &file.provider_options, "imageDetail") { item["detail"] = detail; } @@ -388,6 +492,7 @@ fn convert_user_part(part: &UserPart) -> Result { } fn convert_tool_result_output( + ns: ResponsesNamespace, output: &ToolResultOutput, warnings: &mut Vec, ) -> Result { @@ -409,12 +514,13 @@ fn convert_tool_result_output( }); continue; } - let mut converted = convert_user_part(&UserPart::File(part.clone()))?; + let mut converted = convert_user_part(ns, &UserPart::File(part.clone()))?; if converted["type"] == "input_file" && converted.get("file_data").is_some() { converted["filename"] = json!(part.filename.as_deref().unwrap_or("data")); } if converted["type"] == "input_image" - && let Some(detail) = openai_sub_option(&part.provider_options, "imageDetail") + && let Some(detail) = + openai_sub_option(ns, &part.provider_options, "imageDetail") { converted["detail"] = detail; } @@ -427,7 +533,7 @@ fn convert_tool_result_output( continue; } }; - if let Some(breakpoint) = openai_sub_option(provider_options, "promptCacheBreakpoint") { + if let Some(breakpoint) = openai_sub_option(ns, provider_options, "promptCacheBreakpoint") { converted["prompt_cache_breakpoint"] = breakpoint; } parts.push(converted); @@ -441,20 +547,26 @@ fn serialize_arguments(input: &Value) -> String { } /// Read the `itemId` from a content part's `providerOptions.openai.itemId`. -fn item_id(provider_options: &Option) -> Option { +fn item_id( + ns: ResponsesNamespace, + provider_options: &Option, +) -> Option { provider_options .as_ref() - .and_then(|v| v.get("openai")) + .and_then(|v| ns.find(v)) .and_then(|o| o.get("itemId")) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string) } /// Read the `phase` from a content part's `providerOptions.openai.phase`. -fn phase_from_provider_options(provider_options: &Option) -> Option { +fn phase_from_provider_options( + ns: ResponsesNamespace, + provider_options: &Option, +) -> Option { provider_options .as_ref() - .and_then(|v| v.get("openai")) + .and_then(|v| ns.find(v)) .and_then(|o| o.get("phase")) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string) @@ -462,21 +574,26 @@ fn phase_from_provider_options(provider_options: &Option) /// Read the `namespace` from a content part's `providerOptions.openai.namespace`. fn namespace_from_provider_options( + ns: ResponsesNamespace, provider_options: &Option, ) -> Option { provider_options .as_ref() - .and_then(|v| v.get("openai")) + .and_then(|v| ns.find(v)) .and_then(|o| o.get("namespace")) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string) } /// Read a sub-key from `providerOptions.openai.` on a content part. -fn openai_sub_option(provider_options: &Option, key: &str) -> Option { +fn openai_sub_option( + ns: ResponsesNamespace, + provider_options: &Option, + key: &str, +) -> Option { provider_options .as_ref() - .and_then(|v| v.get("openai")) + .and_then(|v| ns.find(v)) .and_then(|o| o.get(key)) .cloned() } @@ -610,26 +727,28 @@ fn push_unsupported_call_option_warnings(options: &CallOptions, warnings: &mut V /// wins over top-level `reasoning`), `reasoningSummary` (defaults to "detailed" /// when an effort other than "none" applies), and whether the model reasons. fn resolve_responses_reasoning( + ns: ResponsesNamespace, provider_opts: &Option, options: &CallOptions, caps: &ModelCapabilities, ) -> (Option, Option, bool) { - let resolved_reasoning_effort: Option = openai_option(provider_opts, "reasoningEffort") - .map(|v| { - v.as_str() - .map(std::string::ToString::to_string) - .unwrap_or_else(|| v.to_string()) - }) - .or_else(|| { - if options.reasoning.is_some_and(ReasoningEffort::is_custom) { - options.reasoning.map(|r| r.to_string()) - } else { - None - } - }); + let resolved_reasoning_effort: Option = + openai_option(ns, provider_opts, "reasoningEffort") + .map(|v| { + v.as_str() + .map(std::string::ToString::to_string) + .unwrap_or_else(|| v.to_string()) + }) + .or_else(|| { + if options.reasoning.is_some_and(ReasoningEffort::is_custom) { + options.reasoning.map(|r| r.to_string()) + } else { + None + } + }); let resolved_reasoning_summary: Option = - openai_option(provider_opts, "reasoningSummary") + openai_option(ns, provider_opts, "reasoningSummary") .map(|v| { v.as_str() .map(std::string::ToString::to_string) @@ -646,7 +765,7 @@ fn resolve_responses_reasoning( } }); - let is_reasoning_model = openai_option(provider_opts, "forceReasoning") + let is_reasoning_model = openai_option(ns, provider_opts, "forceReasoning") .map(|v| v.as_bool().unwrap_or(false)) .unwrap_or(caps.is_reasoning_model); @@ -659,11 +778,12 @@ fn resolve_responses_reasoning( /// Warn when `conversation` and `previousResponseId` are both set. fn warn_conversation_conflict( + ns: ResponsesNamespace, provider_opts: &Option, warnings: &mut Vec, ) { - let has_conversation = openai_option(provider_opts, "conversation").is_some(); - let has_previous_response_id = openai_option(provider_opts, "previousResponseId").is_some(); + let has_conversation = openai_option(ns, provider_opts, "conversation").is_some(); + let has_previous_response_id = openai_option(ns, provider_opts, "previousResponseId").is_some(); if has_conversation && has_previous_response_id { warnings.push(Warning::Unsupported { feature: "conversation".to_string(), @@ -675,11 +795,12 @@ fn warn_conversation_conflict( } fn resolve_responses_system_message_mode( + ns: ResponsesNamespace, provider_opts: &Option, is_reasoning_model: bool, caps: &ModelCapabilities, ) -> SystemMessageMode { - openai_option(provider_opts, "systemMessageMode") + openai_option(ns, provider_opts, "systemMessageMode") .and_then(|v| v.as_str().map(std::string::ToString::to_string)) .map(|s| match s.as_str() { "developer" => SystemMessageMode::Developer, @@ -695,14 +816,15 @@ fn resolve_responses_system_message_mode( /// temperature / top_p, subject to reasoning-model restrictions. fn apply_responses_sampling( + ns: ResponsesNamespace, body: &mut Value, options: &CallOptions, - provider_opts: &Option, caps: &ModelCapabilities, is_reasoning_model: bool, resolved_reasoning_effort: &Option, warnings: &mut Vec, ) { + let provider_opts = &options.provider_options; let mut temperature = options.temperature; let mut top_p = options.top_p; @@ -732,7 +854,7 @@ fn apply_responses_sampling( "reasoningMode", "reasoningContext", ] { - if openai_option(provider_opts, key).is_some() { + if openai_option(ns, provider_opts, key).is_some() { warnings.push(Warning::Unsupported { feature: key.to_string(), details: Some(format!("{key} is not supported for non-reasoning models")), @@ -751,6 +873,7 @@ fn apply_responses_sampling( /// `text.format` (json_schema / json_object) plus `verbosity`. fn apply_responses_text_format( + ns: ResponsesNamespace, body: &mut Value, options: &CallOptions, provider_opts: &Option, @@ -766,7 +889,7 @@ fn apply_responses_text_format( let mut text = json!({}); match schema { Some(schema) => { - let strict_json = openai_option(provider_opts, "strictJsonSchema") + let strict_json = openai_option(ns, provider_opts, "strictJsonSchema") .and_then(|v| v.as_bool()) .unwrap_or(true); text["format"] = json!({ @@ -786,7 +909,7 @@ fn apply_responses_text_format( } } - if let Some(verbosity) = openai_option(provider_opts, "textVerbosity") { + if let Some(verbosity) = openai_option(ns, provider_opts, "textVerbosity") { let text = body.get_mut("text").and_then(|t| t.as_object_mut()); match text { Some(obj) => { @@ -802,11 +925,12 @@ fn apply_responses_text_format( /// The computed `include` list (store=false on reasoning models adds /// `reasoning.encrypted_content`). fn resolve_responses_include( + ns: ResponsesNamespace, provider_opts: &Option, is_reasoning_model: bool, ) -> Option> { let mut include: Option> = - openai_option(provider_opts, "include").and_then(|v| v.as_array().cloned()); + openai_option(ns, provider_opts, "include").and_then(|v| v.as_array().cloned()); let add_include = |key: &str, inc: &mut Option>| { let already = inc @@ -821,7 +945,7 @@ fn resolve_responses_include( }; // store defaults to true; only the explicit `false` triggers encrypted_content. - let store_explicit = openai_option(provider_opts, "store").and_then(|v| v.as_bool()); + let store_explicit = openai_option(ns, provider_opts, "store").and_then(|v| v.as_bool()); if store_explicit == Some(false) && is_reasoning_model { add_include("reasoning.encrypted_content", &mut include); } @@ -832,11 +956,12 @@ fn resolve_responses_include( /// Pass-through of the remaining Responses provider options (only sent when /// set). fn apply_responses_provider_options( + ns: ResponsesNamespace, body: &mut Value, provider_opts: &Option, ) { let mut set = |key: &str, body_key: &str| { - if let Some(v) = openai_option(provider_opts, key) { + if let Some(v) = openai_option(ns, provider_opts, key) { body[body_key] = v; } }; @@ -851,7 +976,7 @@ fn apply_responses_provider_options( set("promptCacheRetention", "prompt_cache_retention"); set("safetyIdentifier", "safety_identifier"); set("truncation", "truncation"); - if let Some(Value::Object(options)) = openai_option(provider_opts, "promptCacheOptions") { + if let Some(Value::Object(options)) = openai_option(ns, provider_opts, "promptCacheOptions") { body["prompt_cache_options"] = Value::Object( options .into_iter() @@ -863,12 +988,13 @@ fn apply_responses_provider_options( /// `service_tier` with model-capability validation. fn apply_responses_service_tier( + ns: ResponsesNamespace, body: &mut Value, provider_opts: &Option, caps: &ModelCapabilities, warnings: &mut Vec, ) { - if let Some(st) = openai_option(provider_opts, "serviceTier") + if let Some(st) = openai_option(ns, provider_opts, "serviceTier") .and_then(|v| v.as_str().map(std::string::ToString::to_string)) { match st.as_str() { @@ -896,6 +1022,7 @@ fn apply_responses_service_tier( /// `reasoning` block for reasoning models. fn apply_responses_reasoning_block( + ns: ResponsesNamespace, body: &mut Value, provider_opts: &Option, is_reasoning_model: bool, @@ -907,9 +1034,9 @@ fn apply_responses_reasoning_block( } let effort = resolved_reasoning_effort.as_ref(); let summary = resolved_reasoning_summary.as_ref(); - let mode = openai_option(provider_opts, "reasoningMode") + let mode = openai_option(ns, provider_opts, "reasoningMode") .and_then(|v| v.as_str().map(std::string::ToString::to_string)); - let context = openai_option(provider_opts, "reasoningContext") + let context = openai_option(ns, provider_opts, "reasoningContext") .and_then(|v| v.as_str().map(std::string::ToString::to_string)); if effort.is_some() || summary.is_some() || mode.is_some() || context.is_some() { @@ -936,11 +1063,25 @@ fn apply_responses_reasoning_block( /// behavior is unchanged. /// /// # Errors -/// Returns an unsupported-functionality error for text file parts. +/// Returns an error for text file parts or an unresolved inline image media type. pub fn build_responses_request_body( model_id: &str, options: &CallOptions, stream: bool, +) -> Result { + build_responses_request_body_for(ResponsesNamespace::OPENAI, model_id, options, stream) +} + +/// [`build_responses_request_body`] for a host with its own providerOptions +/// namespace. +/// +/// # Errors +/// Returns an error for text file parts or an unresolved inline image media type. +pub fn build_responses_request_body_for( + ns: ResponsesNamespace, + model_id: &str, + options: &CallOptions, + stream: bool, ) -> Result { let mut warnings: Vec = Vec::new(); let caps = get_model_capabilities(model_id); @@ -951,26 +1092,28 @@ pub fn build_responses_request_body( // -- Reasoning resolution -- let (resolved_reasoning_effort, resolved_reasoning_summary, is_reasoning_model) = - resolve_responses_reasoning(provider_opts, options, &caps); + resolve_responses_reasoning(ns, provider_opts, options, &caps); // -- conversation + previousResponseId conflict -- - warn_conversation_conflict(provider_opts, &mut warnings); + warn_conversation_conflict(ns, provider_opts, &mut warnings); // -- System message mode -- let system_message_mode = - resolve_responses_system_message_mode(provider_opts, is_reasoning_model, &caps); + resolve_responses_system_message_mode(ns, provider_opts, is_reasoning_model, &caps); // -- Input conversion -- - let store_bool = openai_option(provider_opts, "store") + let store_bool = openai_option(ns, provider_opts, "store") .and_then(|v| v.as_bool()) .unwrap_or(true); - let has_previous_response_id = openai_option(provider_opts, "previousResponseId").is_some(); + 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( + ns, &options.prompt, system_message_mode, store_bool, has_previous_response_id, - openai_option(provider_opts, "conversation").is_some(), + has_conversation, )?; warnings.extend(input_result.warnings); @@ -990,9 +1133,9 @@ pub fn build_responses_request_body( // temperature / top_p (subject to reasoning-model restrictions) apply_responses_sampling( + ns, &mut body, options, - provider_opts, &caps, is_reasoning_model, &resolved_reasoning_effort, @@ -1000,27 +1143,28 @@ pub fn build_responses_request_body( ); // -- Response format (text.format) + verbosity -- - apply_responses_text_format(&mut body, options, provider_opts); + apply_responses_text_format(ns, &mut body, options, provider_opts); // -- include (computed) -- - let include = resolve_responses_include(provider_opts, is_reasoning_model); + let include = resolve_responses_include(ns, provider_opts, is_reasoning_model); if let Some(inc) = include { body["include"] = json!(inc); } // -- store (only sent when explicitly set) -- - if let Some(s) = openai_option(provider_opts, "store").and_then(|v| v.as_bool()) { + if let Some(s) = openai_option(ns, provider_opts, "store").and_then(|v| v.as_bool()) { body["store"] = json!(s); } // -- Other provider options (only sent when set) -- - apply_responses_provider_options(&mut body, provider_opts); + apply_responses_provider_options(ns, &mut body, provider_opts); // -- service_tier (with capability validation) -- - apply_responses_service_tier(&mut body, provider_opts, &caps, &mut warnings); + apply_responses_service_tier(ns, &mut body, provider_opts, &caps, &mut warnings); // -- reasoning block (reasoning models only) -- apply_responses_reasoning_block( + ns, &mut body, provider_opts, is_reasoning_model, @@ -1043,6 +1187,87 @@ pub fn build_responses_request_body( Ok(ResponsesRequestBodyResult { body, warnings }) } +// -- File id prefixes -------------------------------------------------------- + +/// Apply the host's file-id prefixes (`fileIdPrefixes` in the AI SDK) to a +/// built request body: when a file content part's base64 data starts with one +/// of `prefixes`, it is an uploaded file's id and the part carries a `file_id` +/// field instead of `image_url` / `file_data`. A no-op for no prefixes. +pub(crate) fn apply_file_id_prefixes(body: &mut Value, prefixes: &[&str]) { + if prefixes.is_empty() { + return; + } + if let Some(input) = body.get_mut("input").and_then(|v| v.as_array_mut()) { + for msg in input.iter_mut() { + if let Some(content) = msg.get_mut("content").and_then(|v| v.as_array_mut()) { + for part in content.iter_mut() { + apply_prefix_to_part(part, prefixes); + } + } + } + } +} + +/// Check if a content part's `image_url` or `file_data` contains a data URL +/// whose payload starts with a file-id prefix and, if so, replace it with a +/// `file_id`. +/// +/// For `input_file` parts whose media type is an image, the type is also +/// changed to `input_image`, mirroring the AI SDK, which sends image files as +/// `input_image` with a `file_id`. +fn apply_prefix_to_part(part: &mut Value, prefixes: &[&str]) { + let is_file_id = |data: &str| prefixes.iter().any(|prefix| data.starts_with(prefix)); + // input_image: { type: "input_image", image_url: "data:;base64," } + if part.get("type").and_then(|v| v.as_str()) == Some("input_image") { + let file_id = part + .get("image_url") + .and_then(|v| v.as_str()) + .and_then(extract_base64_data) + .filter(|data| is_file_id(data)) + .map(std::string::ToString::to_string); + if let Some(file_id) = file_id + && let Some(obj) = part.as_object_mut() + { + obj.remove("image_url"); + obj.insert("file_id".to_string(), json!(file_id)); + } + } + + // input_file: { type: "input_file", file_data: "data:;base64," } + if part.get("type").and_then(|v| v.as_str()) == Some("input_file") { + let file_data = part + .get("file_data") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + if let Some(file_data) = file_data { + let media_type = extract_media_type(&file_data).unwrap_or(""); + if let Some(data) = extract_base64_data(&file_data) + && is_file_id(data) + && let Some(obj) = part.as_object_mut() + { + obj.remove("file_data"); + obj.remove("filename"); + obj.insert("file_id".to_string(), json!(data)); + if get_top_level_media_type(media_type) == "image" { + obj.insert("type".to_string(), json!("input_image")); + } + } + } + } +} + +/// Extract the MIME type from a `data:;base64,` URL. +fn extract_media_type(data_url: &str) -> Option<&str> { + let prefix = data_url.strip_prefix("data:")?; + let end = prefix.find(";base64,")?; + Some(&prefix[..end]) +} + +/// Extract the base64 payload from a `data:;base64,` URL. +fn extract_base64_data(data_url: &str) -> Option<&str> { + data_url.split_once(";base64,").map(|(_, rest)| rest) +} + // -- Usage conversion -------------------------------------------------------- /// Convert a Responses API usage into the core `Usage` type. @@ -1149,3 +1374,125 @@ pub fn map_responses_finish_reason( raw: finish_reason.map(std::string::ToString::to_string), } } + +#[cfg(test)] +mod tests { + use super::*; + + fn reasoning(text: &str, options: Value) -> AssistantPart { + AssistantPart::Reasoning(ReasoningPart { + text: text.to_string(), + provider_options: Some(serde_json::from_value(json!({ "openai": options })).unwrap()), + }) + } + + fn convert(store: bool, content: Vec) -> ResponsesInputResult { + let prompt = vec![LanguageModelMessage::Assistant { + content, + provider_options: None, + }]; + convert_to_responses_input( + ResponsesNamespace::OPENAI, + &prompt, + SystemMessageMode::System, + store, + false, + false, + ) + .unwrap() + } + + /// Parts of one reasoning item become one input item: the summaries are + /// appended and the last encrypted content wins (openai@4.0.83 + /// `convert-to-openai-responses-input.ts:1167-1220`). + #[test] + fn reasoning_parts_with_one_item_id_are_merged_without_store() { + let result = convert( + false, + vec![ + reasoning( + "first", + json!({ "itemId": "rs_1", "reasoningEncryptedContent": "enc-1" }), + ), + reasoning( + "second", + json!({ "itemId": "rs_1", "reasoningEncryptedContent": "enc-2" }), + ), + reasoning( + "other", + json!({ "itemId": "rs_2", "reasoningEncryptedContent": "enc-3" }), + ), + ], + ); + assert_eq!( + result.input, + vec![ + json!({ + "type": "reasoning", + "id": "rs_1", + "encrypted_content": "enc-2", + "summary": [ + { "type": "summary_text", "text": "first" }, + { "type": "summary_text", "text": "second" }, + ], + }), + json!({ + "type": "reasoning", + "id": "rs_2", + "encrypted_content": "enc-3", + "summary": [{ "type": "summary_text", "text": "other" }], + }), + ] + ); + assert!(result.warnings.is_empty(), "{:?}", result.warnings); + } + + /// An empty part cannot extend an item that already has parts. + #[test] + fn an_empty_reasoning_part_is_not_appended_to_an_item() { + let result = convert( + false, + vec![ + reasoning( + "first", + json!({ "itemId": "rs_1", "reasoningEncryptedContent": "enc-1" }), + ), + reasoning("", json!({ "itemId": "rs_1" })), + ], + ); + assert_eq!( + result.input, + vec![json!({ + "type": "reasoning", + "id": "rs_1", + "encrypted_content": "enc-1", + "summary": [{ "type": "summary_text", "text": "first" }], + })] + ); + assert!(matches!( + result.warnings.as_slice(), + [Warning::Other { message }] + if message.starts_with("Cannot append empty reasoning part to existing reasoning sequence. Skipping reasoning part: ") + )); + } + + /// With `store`, one item reference stands for every part of the item. + #[test] + fn reasoning_parts_with_one_item_id_are_one_reference_with_store() { + let result = convert( + true, + vec![ + reasoning("first", json!({ "itemId": "rs_1" })), + reasoning("second", json!({ "itemId": "rs_1" })), + reasoning("other", json!({ "itemId": "rs_2" })), + ], + ); + assert_eq!( + result.input, + vec![ + json!({ "type": "item_reference", "id": "rs_1" }), + json!({ "type": "item_reference", "id": "rs_2" }), + ] + ); + } +} diff --git a/aimux-providers/src/openai/responses/mod.rs b/aimux-providers/src/openai/responses/mod.rs index eb703841e..d3a588045 100644 --- a/aimux-providers/src/openai/responses/mod.rs +++ b/aimux-providers/src/openai/responses/mod.rs @@ -26,10 +26,11 @@ pub mod convert; pub mod responses_convert; pub mod types; +pub(crate) use convert::ResponsesProfile; pub use convert::{ - ResponsesInputResult, ResponsesRequestBodyResult, build_responses_request_body, - convert_responses_usage, convert_to_responses_input, map_responses_finish_reason, - prepare_responses_tools, + ResponsesInputResult, ResponsesNamespace, ResponsesRequestBodyResult, + build_responses_request_body, build_responses_request_body_for, convert_responses_usage, + convert_to_responses_input, map_responses_finish_reason, prepare_responses_tools, }; use std::collections::HashMap; @@ -43,61 +44,56 @@ use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateResult, StreamResult}; -use aimux_provider_utils::HttpRequest; - -use super::OpenAIConfig; -use responses_convert::build_header_list; +use super::config::OpenAIModelConfig; /// An OpenAI Responses API language model. /// /// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use /// the process-wide shared `Client` internally (RFC-0009 §4.1). /// -/// Created via `OpenAIResponsesProvider` or -/// directly with [`OpenAIResponsesModel::new`]. +/// Created with +/// [`OpenAIProvider::responses`](crate::openai::OpenAIProvider::responses). pub struct OpenAIResponsesModel { model_id: String, - config: OpenAIConfig, + config: OpenAIModelConfig, } impl OpenAIResponsesModel { - #[must_use] - pub fn new(model_id: String, config: OpenAIConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: OpenAIModelConfig) -> Self { Self { model_id, config } } - pub(crate) fn build_headers( + /// The request headers for one call: provider headers resolved now, with + /// the per-call headers layered over them. + pub(crate) async fn request_headers( &self, - extra: Option<&HashMap>, - ) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref org) = self.config.org_id { - headers.insert("OpenAI-Organization".to_string(), org.clone()); - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers + call_headers: Option<&HashMap>, + ) -> Result, AiMuxError> { + self.config.request_headers(call_headers).await } - fn endpoint(&self) -> String { - format!("{}/responses", self.config.base_url) + fn endpoint(&self) -> Result { + self.config.url("/responses") } - /// The provider-metadata key: `"azure"` when the provider string contains - /// `"azure"`, otherwise `"openai"`. Mirrors the TS `providerOptionsName`. + /// The request body for one call: built with the host's providerOptions + /// namespace, file-id prefixes applied, then the provider-level rewrite. + fn request_body( + &self, + options: &CallOptions, + stream: bool, + ) -> Result { + let profile = &self.config.responses; + 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); + result.body = self.config.transform_body(result.body); + Ok(result) + } + + /// The provider-metadata key of the host. fn provider_options_name(&self) -> &str { - if self.config.provider.contains("azure") { - "azure" - } else { - "openai" - } + self.config.responses.namespace.write_key() } } @@ -107,27 +103,23 @@ impl LanguageModel for OpenAIResponsesModel { &self.config.provider } - fn model_id(&self) -> &str { - &self.model_id - } - - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config + async fn supported_urls(&self) -> aimux_core::language_model::SupportedUrls { + self.config.supported_urls.clone() } - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - super::config_snapshot_from_config(&self.config.provider, &self.model_id, &self.config) + fn model_id(&self) -> &str { + &self.model_id } async fn do_generate(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()); - let request_result = build_responses_request_body(&self.model_id, options, false)?; + let headers = self.request_headers(options.headers.as_ref()).await?; + let request_result = self.request_body(options, false)?; let body = request_result.body; let provider_key = self.provider_options_name().to_string(); - let endpoint = self.endpoint(); + let endpoint = self.endpoint()?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.clone(), build_header_list(&headers), options), + self.config.http_request(endpoint.clone(), headers, options), body.clone(), aimux_provider_utils::create_json_response_handler::(), super::openai_failed_response_handler(), @@ -154,8 +146,8 @@ impl LanguageModel for OpenAIResponsesModel { } async fn do_stream(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()); - let request_result = build_responses_request_body(&self.model_id, options, false)?; + let headers = self.request_headers(options.headers.as_ref()).await?; + let request_result = self.request_body(options, false)?; let body = request_result.body; let warnings = request_result.warnings; let provider_key = self.provider_options_name().to_string(); @@ -165,16 +157,16 @@ impl LanguageModel for OpenAIResponsesModel { let store_flag = options .provider_options .as_ref() - .and_then(|m| m.get("openai")) + .and_then(|m| self.config.responses.namespace.find_in(m)) .and_then(|o| o.get("store")) .and_then(serde_json::Value::as_bool) == Some(true); - let endpoint = self.endpoint(); + let endpoint = self.endpoint()?; let mut stream_body = body.clone(); stream_body["stream"] = Value::Bool(true); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.clone(), build_header_list(&headers), options), + self.config.http_request(endpoint.clone(), headers, options), stream_body, aimux_provider_utils::create_event_source_response_handler::(), super::openai_failed_response_handler(), diff --git a/aimux-providers/src/openai/speech.rs b/aimux-providers/src/openai/speech.rs index d1d302415..144c2b175 100644 --- a/aimux-providers/src/openai/speech.rs +++ b/aimux-providers/src/openai/speech.rs @@ -10,8 +10,6 @@ //! bytes in the response body. The `language` option is not supported and produces //! an `unsupported` warning. -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -21,9 +19,7 @@ use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::HttpRequest; - -use super::OpenAIConfig; +use super::config::OpenAIModelConfig; /// The output formats accepted by the OpenAI TTS API. const SUPPORTED_OUTPUT_FORMATS: &[&str] = &["mp3", "opus", "aac", "flac", "wav", "pcm"]; @@ -33,44 +29,13 @@ const SUPPORTED_OUTPUT_FORMATS: &[&str] = &["mp3", "opus", "aac", "flac", "wav", /// Works with any OpenAI-compatible `/audio/speech` endpoint. pub struct OpenAISpeechModel { model_id: String, - config: OpenAIConfig, + config: OpenAIModelConfig, } impl OpenAISpeechModel { - #[must_use] - pub fn new(model_id: String, config: OpenAIConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: OpenAIModelConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref org) = self.config.org_id { - headers.insert("OpenAI-Organization".to_string(), org.clone()); - } - if let Some(ref project) = self.config.project { - headers.insert("OpenAI-Project".to_string(), project.clone()); - } - // Config-level extra headers (lowest priority after auth/org/project). - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/audio/speech", self.config.base_url) - } } #[async_trait] @@ -83,23 +48,19 @@ impl SpeechModel for OpenAISpeechModel { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - async fn do_generate(&self, options: &SpeechCallOptions) -> Result { let (body, warnings) = build_request_body_and_warnings(options, &self.model_id)?; - let headers = self.build_headers(options.headers.as_ref()); - - let header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; + let body = self.config.transform_body(Value::Object(body)); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), header_list, options), - Value::Object(body.clone()), + self.config + .http_request(self.config.url("/audio/speech")?, headers, options), + body.clone(), aimux_provider_utils::create_binary_response_handler(), super::openai_failed_response_handler(), ) @@ -116,9 +77,7 @@ impl SpeechModel for OpenAISpeechModel { Ok(SpeechResult { audio: AudioData::Binary(audio_bytes), warnings, - request: Some(SpeechRequest { - body: Some(Value::Object(body)), - }), + request: Some(SpeechRequest { body: Some(body) }), response: SpeechResponse { timestamp: Some(timestamp), model_id: Some(self.model_id.clone()), diff --git a/aimux-providers/src/openai/transcription.rs b/aimux-providers/src/openai/transcription.rs index cf96ceea7..70ed0c080 100644 --- a/aimux-providers/src/openai/transcription.rs +++ b/aimux-providers/src/openai/transcription.rs @@ -17,8 +17,6 @@ //! (RFC-0028, requires the `realtime` feature) is realtime-models-only — the //! inverse gating. -use std::collections::HashMap; - use async_trait::async_trait; use serde::Deserialize; use serde_json::Value; @@ -30,9 +28,9 @@ use aimux_core::transcription_model::{ TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::{HttpBody, HttpRequest, MultipartForm, media_type_to_extension}; +use aimux_provider_utils::{HttpBody, MultipartForm, media_type_to_extension}; -use super::OpenAIConfig; +use super::config::OpenAIModelConfig; // ── Language map ──────────────────────────────────────────────────────────── @@ -209,6 +207,7 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid base64: {e}"))) } + AudioInput::Url(_) => Err(AudioInput::url_error()), } } @@ -217,42 +216,27 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { /// Works with any OpenAI-compatible `/audio/transcriptions` endpoint. pub struct OpenAITranscriptionModel { model_id: String, - config: OpenAIConfig, + config: OpenAIModelConfig, + #[cfg(feature = "realtime")] + web_socket: Option>, } impl OpenAITranscriptionModel { - #[must_use] - pub fn new(model_id: String, config: OpenAIConfig) -> Self { - Self { model_id, config } - } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref org) = self.config.org_id { - headers.insert("OpenAI-Organization".to_string(), org.clone()); - } - if let Some(ref project) = self.config.project { - headers.insert("OpenAI-Project".to_string(), project.clone()); - } - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } + pub(crate) fn from_config(model_id: String, config: OpenAIModelConfig) -> Self { + Self { + model_id, + config, + #[cfg(feature = "realtime")] + web_socket: None, } - headers } - - fn endpoint(&self) -> String { - format!("{}/audio/transcriptions", self.config.base_url) + #[cfg(feature = "realtime")] + pub(crate) fn with_web_socket( + mut self, + web_socket: Option>, + ) -> Self { + self.web_socket = web_socket; + self } } @@ -266,10 +250,6 @@ impl TranscriptionModel for OpenAITranscriptionModel { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - async fn do_generate( &self, options: &TranscriptionCallOptions, @@ -340,11 +320,14 @@ impl TranscriptionModel for OpenAITranscriptionModel { let (body_bytes, content_type) = form.finish(); - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; let resp = aimux_provider_utils::post_to_api( - HttpRequest::new(self.endpoint(), header_list, options), + self.config + .http_request(self.config.url("/audio/transcriptions")?, headers, options), HttpBody::Bytes(body_bytes, content_type), aimux_provider_utils::create_json_response_handler::(), super::openai_failed_response_handler(), @@ -431,16 +414,7 @@ impl TranscriptionModel for OpenAITranscriptionModel { // wss URL from the REST base URL (https→wss / http→ws), matching the // AI SDK's toWebSocketUrl. - let ws_url = { - let base = self.config.base_url.trim_end_matches('/'); - if let Some(rest) = base.strip_prefix("https://") { - format!("wss://{rest}/realtime?intent=transcription") - } else if let Some(rest) = base.strip_prefix("http://") { - format!("ws://{rest}/realtime?intent=transcription") - } else { - format!("{base}/realtime?intent=transcription") - } - }; + let ws_url = self.config.ws_url("/realtime?intent=transcription")?; // session.update — exact upstream wire shape (RFC-0028 §3.2 B1): the // model rides in session.audio.input.transcription.model, NOT in the @@ -473,10 +447,10 @@ impl TranscriptionModel for OpenAITranscriptionModel { session_update["session"]["include"] = serde_json::json!(include); } - let mut header_list: Vec<(String, String)> = self - .build_headers(options.headers.as_ref()) - .into_iter() - .collect(); + let mut header_list = self + .config + .request_headers(options.headers.as_ref()) + .await?; let abort = options.abort_signal.clone(); let include_raw = options.include_raw_chunks; @@ -492,6 +466,7 @@ impl TranscriptionModel for OpenAITranscriptionModel { subprotocols: Vec::new(), abort_signal: abort.clone(), timeout: options.timeout, + connector: self.web_socket.clone(), }; let mut ws = ws_connect(&req).await?; diff --git a/aimux-providers/src/openai/types.rs b/aimux-providers/src/openai/types.rs index 1a8c83400..c5d9d2e42 100644 --- a/aimux-providers/src/openai/types.rs +++ b/aimux-providers/src/openai/types.rs @@ -31,10 +31,10 @@ pub struct Choice { pub struct MessageResponse { pub role: String, pub content: Option, - /// Reasoning text (Groq / OpenAI o-series models). + /// Reasoning text (`reasoning` field of OpenAI-shaped servers). #[serde(default)] pub reasoning: Option, - /// DeepSeek / 阿里通义等厂商的 reasoning_content 字段。 + /// `reasoning_content` field of OpenAI-shaped servers (DashScope and others). /// 优先于 `reasoning`（当两者同时存在时）。 #[serde(default)] pub reasoning_content: Option, @@ -113,16 +113,6 @@ pub struct StreamChunk { pub choices: Vec, #[serde(default)] pub usage: Option, - /// Groq-specific extension: usage is sent in `x_groq.usage` during streaming. - #[serde(default)] - pub x_groq: Option, -} - -/// Groq streaming extension (`x_groq` field in SSE chunks). -#[derive(Debug, Deserialize, Default)] -pub struct XGroq { - #[serde(default)] - pub usage: Option, } #[derive(Debug, Deserialize, Default)] @@ -140,10 +130,10 @@ pub struct StreamChoice { pub struct Delta { #[serde(default)] pub content: Option, - /// Reasoning text delta (Groq / OpenAI o-series). + /// Reasoning text delta (`reasoning` field of OpenAI-shaped servers). #[serde(default)] pub reasoning: Option, - /// DeepSeek / 阿里通义等厂商的 reasoning_content delta。 + /// `reasoning_content` delta of OpenAI-shaped servers (DashScope and others). /// 优先于 `reasoning`。 #[serde(default)] pub reasoning_content: Option, diff --git a/aimux-providers/src/openai_compatible/chat.rs b/aimux-providers/src/openai_compatible/chat.rs new file mode 100644 index 000000000..0912761c3 --- /dev/null +++ b/aimux-providers/src/openai_compatible/chat.rs @@ -0,0 +1,881 @@ +//! The OpenAI-compatible chat language model (`{name}.chat`). +//! +//! The Rust form of `OpenAICompatibleChatLanguageModel`. Request construction +//! lives in [`super::convert`]; vendor differences arrive through the +//! [`ChatDialect`](super::config::ChatDialect) of the model's config. + +use aimux_core::tool::RawToolCall; +use std::collections::{HashMap, HashSet}; + +use async_trait::async_trait; +use futures::StreamExt; +use serde_json::{Map, Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::options::CallOptions; +use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; +use aimux_core::shared::provider_namespace; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::{ + FinishReason, FinishReasonUnified, InputTokenUsage, OutputTokenUsage, ResponseMetadata, Usage, +}; +use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker, generate_id}; + +use super::config::CompatModelConfig; +use super::convert::{ChatBodySpec, RequestBodyResult, build_request_body, parse_finish_reason}; +use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; + +/// An OpenAI-compatible chat model. Holds no HTTP client: the request helpers +/// of `aimux-provider-utils` use the injected or process-default transport. +pub struct OpenAICompatibleChatModel { + model_id: String, + config: CompatModelConfig, +} + +impl OpenAICompatibleChatModel { + pub(crate) fn from_config(model_id: String, config: CompatModelConfig) -> Self { + Self { model_id, config } + } + + /// The JSON body a call would send, before the provider's + /// `transform_request_body`, with the warnings raised while building it. + /// Nothing is read from the environment and nothing is sent: a key is not + /// resolved here. + /// + /// # Errors + /// + /// `InvalidArgument` for a malformed provider option or an unconvertible + /// prompt part. + pub(crate) fn request_body( + &self, + options: &CallOptions, + stream: bool, + ) -> Result { + build_request_body( + &self.model_id, + options, + stream, + &ChatBodySpec { + provider_options_name: self.config.provider_options_name(), + include_usage: self.config.chat.include_usage, + supports_structured_outputs: self.config.chat.supports_structured_outputs, + supports_multi_part_tool_content: self.config.chat.supports_multi_part_tool_content, + dialect: &self.config.chat.dialect, + }, + ) + } +} + +// ── Usage ──────────────────────────────────────────────────────────────────── + +/// 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. +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 + .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 reasoning_tokens = usage + .completion_tokens_details + .as_ref() + .and_then(|d| d.reasoning_tokens) + .unwrap_or(0); + + Usage { + input_tokens: InputTokenUsage { + total: Some(prompt_tokens), + no_cache: Some(no_cache), + cache_read: Some(cached), + cache_write, + }, + 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()), + } +} + +/// Usage from a raw `usage` value. +pub(crate) fn usage_from_raw(raw: Option<&Value>) -> Usage { + let Some(raw) = raw.filter(|value| !value.is_null()) else { + return Usage::default(); + }; + match serde_json::from_value::(raw.clone()) { + Ok(parsed) => convert_usage(&parsed, Some(raw)), + Err(_) => Usage { + raw: raw.as_object().cloned(), + ..Usage::default() + }, + } +} + +fn prediction_tokens(usage: Option<&UsageResponse>, metadata: &mut Map) { + if let Some(details) = usage.and_then(|u| u.completion_tokens_details.as_ref()) { + if let Some(accepted) = details.accepted_prediction_tokens { + metadata.insert("acceptedPredictionTokens".into(), json!(accepted)); + } + if let Some(rejected) = details.rejected_prediction_tokens { + metadata.insert("rejectedPredictionTokens".into(), json!(rejected)); + } + } +} + +fn thought_signature(extra_content: Option<&Value>) -> Option { + extra_content + .and_then(|extra| extra.pointer("/google/thought_signature")) + .and_then(Value::as_str) + .filter(|signature| !signature.is_empty()) + .map(str::to_string) +} + +fn stream_error( + name: &str, + error: &Value, + url: &str, + request_body_values: Value, + response_headers: HashMap, +) -> AiMuxError { + let message = error + .get("message") + .and_then(Value::as_str) + .map(str::to_owned) + .unwrap_or_else(|| format!("{name} stream failed before any output was generated")); + let code = error.get("code").or_else(|| error.get("type")); + let provider_code = 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 + // 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( + message, + provider_code, + status_code, + error, + url, + request_body_values, + response_headers, + ) +} + +// ── LanguageModel ──────────────────────────────────────────────────────────── + +#[async_trait] +impl LanguageModel for OpenAICompatibleChatModel { + /// `"{name}.chat"`. + fn provider(&self) -> &str { + &self.config.provider + } + + fn model_id(&self) -> &str { + &self.model_id + } + + async fn supported_urls(&self) -> SupportedUrls { + match &self.config.chat.supported_urls { + Some(urls) => urls().await, + None => SupportedUrls::default(), + } + } + + async fn do_generate(&self, options: &CallOptions) -> Result { + let built = self.request_body(options, false)?; + let metadata_key = built.metadata_key; + let body = self.config.transform_body(built.body); + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; + let http = self + .config + .http_request("/chat/completions", headers, options)?; + + let resp = aimux_provider_utils::post_json_to_api( + http, + body.clone(), + aimux_provider_utils::create_json_response_handler::(), + self.config.failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers; + let raw = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; + + let choice = data.choices.into_iter().next().ok_or_else(|| { + AiMuxError::InvalidResponseData("Response did not contain any choices.".to_string()) + })?; + + let mut content = Vec::new(); + for text in content_texts(choice.message.content.as_ref()) { + content.push(GenerateContent::Text { + text, + provider_metadata: None, + }); + } + let reasoning = choice + .message + .reasoning_content + .or(choice.message.reasoning) + .filter(|text| !text.is_empty()); + if let Some(text) = reasoning { + content.push(GenerateContent::Reasoning(ReasoningOutput { + text, + provider_metadata: None, + })); + } + for call in choice.message.tool_calls.into_iter().flatten() { + let signature = thought_signature(call.extra_content.as_ref()); + content.push(GenerateContent::ToolCall(RawToolCall { + tool_call_id: call + .id + .filter(|id| !id.is_empty()) + .unwrap_or_else(generate_id), + tool_name: call.function.name, + input: call.function.arguments.unwrap_or_default(), + provider_executed: None, + dynamic: None, + provider_metadata: signature.as_ref().map(|signature| { + provider_namespace(&metadata_key, json!({ "thoughtSignature": signature })) + .expect("provider metadata must be an object") + }), + })); + } + for (i, annotation) in choice.message.annotations.iter().flatten().enumerate() { + if annotation.get("type").and_then(Value::as_str) == Some("url_citation") + && let Some(citation) = annotation.get("url_citation") + { + content.push(GenerateContent::Source(Source::Url { + id: format!("annotation-{i}"), + url: citation + .get("url") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(), + title: citation + .get("title") + .and_then(Value::as_str) + .map(str::to_string), + provider_metadata: None, + })); + } + } + + let finish_reason = choice + .finish_reason + .as_deref() + .map(parse_finish_reason) + .unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }); + + 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 mut namespace = Map::new(); + prediction_tokens(data.usage.as_ref(), &mut namespace); + let mut metadata = aimux_core::shared::SharedProviderMetadata::new(); + metadata.insert(metadata_key, namespace); + if let Some(extractor) = &dialect.metadata_extractor + && let Some(extra) = extractor.extract_metadata(&raw) + { + for (key, value) in extra { + metadata.insert(key, value); + } + } + + Ok(GenerateResult { + content, + finish_reason, + usage, + warnings: built.warnings, + provider_metadata: Some(metadata), + response: Some(aimux_core::shared::ResponseInfo { + id: data.id, + timestamp: data + .created + .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) }), + }) + } + + async fn do_stream(&self, options: &CallOptions) -> Result { + let built = self.request_body(options, true)?; + let metadata_key = built.metadata_key; + let warnings = built.warnings; + let body = self.config.transform_body(built.body); + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; + let http = self + .config + .http_request("/chat/completions", headers, options)?; + let endpoint = http.url.clone(); + + let resp = aimux_provider_utils::post_json_to_api( + http, + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + self.config.failed_response_handler(), + ) + .await?; + 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(); + let stream_response_headers = response_headers.clone(); + + let stream = async_stream::stream! { + yield Ok(StreamPart::StreamStart { warnings }); + + let text_id = 0usize; + let mut text_started = false; + let reasoning_id = "reasoning-0".to_string(); + let mut reasoning_started = false; + let mut final_usage_raw: Option = None; + let mut final_usage_parsed: Option = None; + let mut final_finish_reason: Option = None; + let mut response_metadata_emitted = false; + + let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(generate_id); + let mut tool_parts = Vec::new(); + let signature_metadata = |extra: Option<&Value>| { + thought_signature(extra).map(|signature| { + provider_namespace(&metadata_key, json!({ "thoughtSignature": signature })) + .expect("provider metadata must be an object") + }) + }; + + // Some compatible servers send the first delta of a call without + // `function.name`; buffer by index until the name is known. + let mut pending: HashMap = HashMap::new(); + let mut forwarded: HashSet = HashSet::new(); + + 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; + } + match event { + Ok(parsed) => { + if emit_raw_chunks { + yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); + } + if let Some(extractor) = metadata_extractor.as_mut() { + extractor.process_chunk(&parsed); + } + if let Some(error) = parsed.get("error") { + yield Ok(StreamPart::Error { + error: stream_error( + &metadata_key, + error, + &stream_error_url, + stream_error_body.clone(), + stream_response_headers.clone(), + ), + }); + stream_errored = true; + break; + } + + let chunk_usage_raw: Option = match &dialect.stream_usage_key { + Some(key) => parsed.get(key).and_then(|v| v.get("usage")).cloned(), + None => parsed.get("usage").cloned(), + } + .filter(|usage| !usage.is_null()); + + let chunk: StreamChunk = match serde_json::from_value(parsed) { + Ok(chunk) => chunk, + Err(e) => { + yield Err(e.into()); + continue; + } + }; + + if !response_metadata_emitted + && (chunk.id.is_some() || chunk.model.is_some()) + { + response_metadata_emitted = true; + yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: chunk + .created + .and_then(|secs| chrono::DateTime::from_timestamp(secs as i64, 0)) + .map(|dt| dt.to_rfc3339()), + model_id: chunk.model.clone(), + })); + } + + if let Some(raw_usage) = &chunk_usage_raw { + final_usage_raw = Some(raw_usage.clone()); + final_usage_parsed = serde_json::from_value(raw_usage.clone()).ok(); + } + + for choice in chunk.choices { + 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, + 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, + }); + } + yield Ok(StreamPart::TextDelta { + id: format!("{text_id}"), + delta: content, + provider_metadata: None, + }); + } + + if let Some(deltas) = choice.delta.tool_calls { + if reasoning_started { + yield Ok(StreamPart::ReasoningEnd { + id: reasoning_id.clone(), + provider_metadata: None, + }); + reasoning_started = false; + } + for dtc in &deltas { + let function = dtc.function.as_ref(); + let name = function + .and_then(|f| f.name.as_deref()) + .filter(|n| !n.trim().is_empty()); + let arguments = function.and_then(|f| f.arguments.as_deref()); + let buffered; + let delta = match dtc.index { + Some(index) if !forwarded.contains(&index) => { + let entry = pending.entry(index).or_default(); + if entry.id.is_none() { + entry.id.clone_from(&dtc.id); + } + if entry.extra.is_null() + && let Some(extra) = &dtc.extra_content + { + entry.extra = extra.clone(); + } + if let Some(arguments) = arguments { + entry.arguments.push_str(arguments); + } + let Some(name) = name else { continue }; + buffered = pending.remove(&index).unwrap_or_default(); + forwarded.insert(index); + StreamingToolCallDelta { + index: Some(index), + id: buffered.id.as_deref(), + name: Some(name), + arguments: Some(&buffered.arguments), + provider_metadata: signature_metadata(Some(&buffered.extra)), + ..Default::default() + } + } + _ => StreamingToolCallDelta { + index: dtc.index, + id: dtc.id.as_deref(), + name: function.and_then(|f| f.name.as_deref()), + arguments, + provider_metadata: signature_metadata(dtc.extra_content.as_ref()), + ..Default::default() + }, + }; + if let Err(error) = tool_calls.process(delta, &mut tool_parts) { + yield Err(error.into()); + return; + } + for part in tool_parts.drain(..) { + yield Ok(part); + } + } + } + + if let Some(annotations) = choice.delta.annotations { + for (i, annotation) in annotations.iter().enumerate() { + if annotation.get("type").and_then(Value::as_str) + == Some("url_citation") + && let Some(citation) = annotation.get("url_citation") + { + yield Ok(StreamPart::Source(Source::Url { + id: format!("annotation-{i}"), + url: citation + .get("url") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(), + title: citation + .get("title") + .and_then(Value::as_str) + .map(str::to_string), + provider_metadata: None, + })); + } + } + } + + if let Some(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; + } + } + } + } + + if reasoning_started { + yield Ok(StreamPart::ReasoningEnd { + id: reasoning_id.clone(), + provider_metadata: None, + }); + } + if text_started { + yield Ok(StreamPart::TextEnd { + id: format!("{text_id}"), + provider_metadata: None, + }); + } + // A parsable argument buffer can still be a prefix of a longer + // input: finalize only when the stream flushes. + tool_calls.finish(&mut tool_parts); + for part in tool_parts.drain(..) { + yield Ok(part); + } + + let mut namespace = Map::new(); + prediction_tokens(final_usage_parsed.as_ref(), &mut namespace); + let mut metadata = aimux_core::shared::SharedProviderMetadata::new(); + metadata.insert(metadata_key.clone(), namespace); + if let Some(extractor) = &metadata_extractor + && let Some(extra) = extractor.build_metadata() + { + for (key, value) in extra { + metadata.insert(key, value); + } + } + + 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())) + }, + provider_metadata: Some(metadata), + }); + }; + + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} + +#[derive(Default)] +struct PendingToolCall { + id: Option, + arguments: String, + 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()], + 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) + .collect(), + _ => Vec::new(), + } +} + +/// One `GET {base_url}/models` exchange for discovery: no retry, no +/// recording. +pub(crate) async fn list_models_once( + config: &CompatModelConfig, +) -> Result, AiMuxError> { + #[derive(serde::Deserialize)] + struct ModelsList { + #[serde(default)] + data: Vec, + } + #[derive(serde::Deserialize)] + struct ModelEntry { + id: String, + #[serde(default)] + owned_by: Option, + #[serde(default)] + created: Option, + } + + let headers = config.request_headers(None).await?; + let (url, origin) = config.url("/models")?; + let resp = aimux_provider_utils::get_from_api( + config.with_transport( + aimux_provider_utils::HttpRequest { + url, + headers, + ..Default::default() + }, + origin, + ), + aimux_provider_utils::create_json_response_handler::(), + config.failed_response_handler(), + ) + .await?; + Ok(resp + .value + .data + .into_iter() + .map(|m| aimux_core::model_catalogue::RuntimeModel { + id: m.id, + owned_by: m.owned_by, + created: m.created, + }) + .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/completion.rs b/aimux-providers/src/openai_compatible/completion.rs new file mode 100644 index 000000000..186bd88f7 --- /dev/null +++ b/aimux-providers/src/openai_compatible/completion.rs @@ -0,0 +1,558 @@ +//! Text-completion endpoint shared by the native and compatible packages. + +use async_trait::async_trait; +use futures::StreamExt; +use serde::Deserialize; +use serde_json::{Map, Value, json}; + +use aimux_core::AiMuxError; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::language_model_message::{AssistantPart, LanguageModelMessage, UserPart}; +use aimux_core::options::{CallOptions, ResponseFormat}; +use aimux_core::result::StreamResult; +use aimux_core::result::{GenerateContent, GenerateResult}; +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, + Warning, +}; +use aimux_provider_utils::HttpRequest; + +use super::config::CompatModelConfig; +use super::convert::{parse_finish_reason, to_camel_case}; +use crate::openai::config::OpenAIModelConfig; + +#[derive(Clone)] +enum CompletionConfig { + Compatible(CompatModelConfig), + Native(OpenAIModelConfig), +} + +/// A language model for the `/completions` endpoint. +#[derive(Clone)] +pub struct OpenAICompatibleCompletionModel { + model_id: String, + config: CompletionConfig, +} + +impl OpenAICompatibleCompletionModel { + pub(crate) fn from_config(model_id: String, config: CompatModelConfig) -> Self { + Self { + model_id, + config: CompletionConfig::Compatible(config), + } + } + + pub(crate) fn from_native_config(model_id: String, config: OpenAIModelConfig) -> Self { + Self { + model_id, + config: CompletionConfig::Native(config), + } + } + + fn native(&self) -> bool { + matches!(self.config, CompletionConfig::Native(_)) + } + + fn arguments(&self, options: &CallOptions) -> Result<(Value, Vec), AiMuxError> { + let mut warnings = Vec::new(); + for (present, feature) in [ + (options.top_k.is_some(), "topK"), + ( + options + .tools + .as_ref() + .is_some_and(|tools| !tools.is_empty()), + "tools", + ), + (options.tool_choice.is_some(), "toolChoice"), + ] { + if present { + warnings.push(Warning::Unsupported { + feature: feature.into(), + details: None, + }); + } + } + if matches!(options.response_format, Some(ResponseFormat::Json { .. })) { + warnings.push(Warning::Unsupported { + feature: "responseFormat".into(), + details: Some("JSON response format is not supported.".into()), + }); + } + let name = self.provider().split('.').next().unwrap_or_default().trim(); + let camel = to_camel_case(name); + let namespaces = if self.native() { + vec!["openai", name] + } else { + vec![name, camel.as_str()] + }; + let mut provider_options = Map::new(); + if !self.native() + && 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."), + }); + } + for namespace in namespaces { + if let Some(values) = options + .provider_options + .as_ref() + .and_then(|all| all.get(namespace)) + { + for (key, value) in values { + let valid = match key.as_str() { + "echo" => value.is_boolean(), + "user" | "suffix" => value.is_string(), + "logitBias" => value + .as_object() + .is_some_and(|map| map.values().all(Value::is_number)), + "logprobs" if self.native() => value.is_boolean() || value.is_number(), + _ => true, + }; + if !valid { + return Err(AiMuxError::InvalidArgument(format!( + "invalid {namespace} provider option {key}" + ))); + } + } + provider_options.extend(values.clone()); + } + } + let mut body = Map::new(); + body.insert("model".into(), json!(self.model_id)); + for (source, target) in [ + ("echo", "echo"), + ("logitBias", "logit_bias"), + ("suffix", "suffix"), + ("user", "user"), + ] { + if let Some(value) = provider_options.get(source) { + body.insert(target.into(), value.clone()); + } + } + if self.native() { + match provider_options.get("logprobs") { + Some(Value::Bool(true)) => { + body.insert("logprobs".into(), json!(0)); + } + Some(value) if value.is_number() => { + body.insert("logprobs".into(), value.clone()); + } + _ => {} + } + } + for (key, value) in [ + ("max_tokens", json!(options.max_output_tokens)), + ("temperature", json!(options.temperature)), + ("top_p", json!(options.top_p)), + ("frequency_penalty", json!(options.frequency_penalty)), + ("presence_penalty", json!(options.presence_penalty)), + ("seed", json!(options.seed)), + ] { + if !value.is_null() { + body.insert(key.into(), value); + } + } + if !self.native() { + body.extend(provider_options); + } + body.insert("prompt".into(), json!(completion_prompt(options)?)); + let mut stop = vec!["\nuser:".to_string()]; + stop.extend(options.stop_sequences.iter().flatten().cloned()); + body.insert("stop".into(), json!(stop)); + Ok((Value::Object(body), warnings)) + } + + async fn request(&self, options: &CallOptions) -> Result { + match &self.config { + CompletionConfig::Compatible(config) => { + let headers = config.request_headers(options.headers.as_ref()).await?; + config.http_request("/completions", headers, options) + } + CompletionConfig::Native(config) => { + let headers = config.request_headers(options.headers.as_ref()).await?; + Ok(config.http_request(config.url("/completions")?, headers, options)) + } + } + } + + fn transform_body(&self, body: Value) -> Value { + match &self.config { + CompletionConfig::Compatible(_) => body, + CompletionConfig::Native(config) => match &config.transform_request_body { + Some(transform) => transform(body), + None => body, + }, + } + } + + fn failed_response_handler(&self) -> aimux_provider_utils::ResponseHandler { + match &self.config { + CompletionConfig::Compatible(config) => config.failed_response_handler(), + CompletionConfig::Native(_) => crate::openai::openai_failed_response_handler(), + } + } +} + +fn completion_prompt(options: &CallOptions) -> Result { + let mut text = String::new(); + for (index, message) in options.prompt.iter().enumerate() { + match message { + LanguageModelMessage::System { content, .. } if index == 0 => { + text.push_str(content); + text.push_str("\n\n"); + } + LanguageModelMessage::System { .. } => { + return Err(AiMuxError::InvalidPrompt( + "Unexpected system message in prompt".into(), + )); + } + LanguageModelMessage::User { content, .. } => { + text.push_str("user:\n"); + for part in content { + if let UserPart::Text(part) = part { + text.push_str(&part.text); + } + } + text.push_str("\n\n"); + } + LanguageModelMessage::Assistant { content, .. } => { + text.push_str("assistant:\n"); + for part in content { + match part { + AssistantPart::Text(part) => text.push_str(&part.text), + AssistantPart::ToolCall(_) => { + return Err(AiMuxError::UnsupportedFunctionality( + "tool-call messages".into(), + )); + } + _ => {} + } + } + text.push_str("\n\n"); + } + LanguageModelMessage::Tool { .. } => { + return Err(AiMuxError::UnsupportedFunctionality("tool messages".into())); + } + } + } + text.push_str("assistant:\n"); + Ok(text) +} + +#[derive(Deserialize)] +struct CompletionResponse { + id: Option, + created: Option, + model: Option, + choices: Vec, + usage: Option>, +} + +#[derive(Deserialize)] +struct CompletionChoice { + text: String, + finish_reason: Option, + logprobs: Option, +} + +fn parse_response( + value: Value, + native: bool, + streaming: bool, +) -> Result { + let invalid = || AiMuxError::InvalidResponseData("invalid completion response".into()); + if let Some(usage) = value.get("usage").filter(|value| !value.is_null()) + && !["prompt_tokens", "completion_tokens", "total_tokens"] + .iter() + .all(|key| usage.get(key).is_some_and(Value::is_number)) + { + return Err(invalid()); + } + let choices = value + .get("choices") + .and_then(Value::as_array) + .ok_or_else(invalid)?; + for choice in choices { + let valid_finish = if streaming { + choice.get("index").is_some_and(Value::is_number) + } else { + choice.get("finish_reason").is_some_and(Value::is_string) + }; + if !valid_finish { + return Err(invalid()); + } + if native && let Some(logprobs) = choice.get("logprobs").filter(|value| !value.is_null()) { + if !logprobs + .get("tokens") + .and_then(Value::as_array) + .is_some_and(|values| values.iter().all(Value::is_string)) + || !logprobs + .get("token_logprobs") + .and_then(Value::as_array) + .is_some_and(|values| values.iter().all(Value::is_number)) + { + return Err(invalid()); + } + if let Some(top) = logprobs + .get("top_logprobs") + .filter(|value| !value.is_null()) + && !top.as_array().is_some_and(|values| { + values.iter().all(|value| { + value + .as_object() + .is_some_and(|object| object.values().all(Value::is_number)) + }) + }) + { + return Err(invalid()); + } + } + } + Ok(serde_json::from_value(value)?) +} + +fn metadata(response: &CompletionResponse) -> ResponseMetadata { + ResponseMetadata { + id: response.id.clone(), + timestamp: response + .created + .and_then(|secs| chrono::DateTime::from_timestamp(secs, 0)) + .map(|dt| dt.to_rfc3339()), + model_id: response.model.clone(), + } +} + +fn usage(raw: Option>, native: bool) -> Usage { + let Some(mut raw) = raw else { + return Usage::default(); + }; + if native { + raw.retain(|key, _| { + matches!( + key.as_str(), + "prompt_tokens" | "completion_tokens" | "total_tokens" + ) + }); + } + let input = raw + .get("prompt_tokens") + .and_then(Value::as_u64) + .map(|n| n as u32); + let output = raw + .get("completion_tokens") + .and_then(Value::as_u64) + .map(|n| n as u32); + Usage { + input_tokens: InputTokenUsage { + total: if native { + input + } else { + Some(input.unwrap_or(0)) + }, + no_cache: Some(input.unwrap_or(0)), + ..Default::default() + }, + output_tokens: OutputTokenUsage { + total: if native { + output + } else { + Some(output.unwrap_or(0)) + }, + text: Some(output.unwrap_or(0)), + ..Default::default() + }, + raw: Some(raw), + } +} + +#[async_trait] +impl LanguageModel for OpenAICompatibleCompletionModel { + fn provider(&self) -> &str { + match &self.config { + CompletionConfig::Compatible(config) => &config.provider, + CompletionConfig::Native(config) => &config.provider, + } + } + fn model_id(&self) -> &str { + &self.model_id + } + async fn supported_urls(&self) -> SupportedUrls { + match &self.config { + CompletionConfig::Compatible(_) => SupportedUrls::default(), + CompletionConfig::Native(_) => SupportedUrls::default(), + } + } + + async fn do_generate(&self, options: &CallOptions) -> Result { + let (body, warnings) = self.arguments(options)?; + let body = self.transform_body(body); + let response = aimux_provider_utils::post_json_to_api( + self.request(options).await?, + body.clone(), + aimux_provider_utils::create_json_response_handler::(), + self.failed_response_handler(), + ) + .await?; + let mut parsed = parse_response(response.value, self.native(), false)?; + let metadata = metadata(&parsed); + let choice = parsed + .choices + .drain(..) + .next() + .ok_or_else(|| AiMuxError::InvalidResponseData("no choices in response".into()))?; + let provider_metadata = self.native().then(|| { + provider_namespace( + "openai", + match choice.logprobs { + Some(logprobs) => json!({ "logprobs": logprobs }), + None => json!({}), + }, + ) + .expect("provider metadata must be an object") + }); + let finish_reason = choice + .finish_reason + .as_deref() + .map(parse_finish_reason) + .unwrap_or(FinishReason { + unified: FinishReasonUnified::Other, + raw: None, + }); + Ok(GenerateResult { + content: if self.native() || !choice.text.is_empty() { + vec![GenerateContent::Text { + text: choice.text, + provider_metadata: None, + }] + } else { + vec![] + }, + finish_reason, + usage: usage(parsed.usage, self.native()), + warnings, + provider_metadata, + request: Some(RequestInfo { body: Some(body) }), + response: Some(ResponseInfo { + id: metadata.id, + timestamp: metadata.timestamp, + model_id: metadata.model_id, + headers: Some(response.response_headers), + body: response.raw_value, + }), + }) + } + + async fn do_stream(&self, options: &CallOptions) -> Result { + let (mut body, warnings) = self.arguments(options)?; + body["stream"] = json!(true); + let include_usage = match &self.config { + CompletionConfig::Compatible(config) => config.chat.include_usage, + CompletionConfig::Native(_) => true, + }; + if include_usage { + body["stream_options"] = json!({ "include_usage": true }); + } + let body = self.transform_body(body); + let request = self.request(options).await?; + let error_url = request.url.clone(); + let response = aimux_provider_utils::post_json_to_api( + request, + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + self.failed_response_handler(), + ) + .await?; + let mut events = response.value; + let response_headers = response.response_headers; + let stream_response_headers = response_headers.clone(); + let native = self.native(); + let raw_chunks = options.include_raw_chunks == Some(true); + let mut buffered = Vec::new(); + if native { + while let Some(event) = events.next().await { + if let Ok(parsed_event) = &event { + if let Some(error) = parsed_event.get("error") { + return Err(crate::openai::openai_stream_error( + error, + &error_url, + body.clone(), + response_headers.clone(), + )); + } + let output = parsed_event + .get("choices") + .and_then(Value::as_array) + .is_some_and(|choices| { + choices.iter().any(|choice| { + choice + .get("text") + .and_then(Value::as_str) + .is_some_and(|text| !text.is_empty()) + }) + }); + buffered.push(event); + if output { + break; + } + } else { + buffered.push(event); + break; + } + } + } + let stream_body = body.clone(); + let stream = async_stream::stream! { + yield Ok(StreamPart::StreamStart { warnings }); + let mut events = futures::stream::iter(buffered).chain(events); + let mut started = false; + let mut finish_reason = FinishReason { unified: FinishReasonUnified::Other, raw: None }; + let mut final_usage = None; + let mut logprobs = None; + while let Some(event) = events.next().await { + let event = match event { + Ok(event) => event, + Err(error) => { finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; yield Ok(StreamPart::Error { error }); continue; } + }; + if raw_chunks { yield Ok(StreamPart::Raw { raw_value: event.clone() }); } + if let Some(error) = event.get("error") { + finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; + yield Ok(StreamPart::Error { error: crate::openai::openai_stream_error(error, &error_url, stream_body.clone(), response_headers.clone()) }); + continue; + } + let parsed: CompletionResponse = match parse_response(event, native, true) { + Ok(parsed) => parsed, + Err(error) => { finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; yield Ok(StreamPart::Error { error }); continue; } + }; + if !started { + started = true; + yield Ok(StreamPart::ResponseMetadata(metadata(&parsed))); + yield Ok(StreamPart::TextStart { id: "0".into(), provider_metadata: None }); + } + if parsed.usage.is_some() { final_usage = parsed.usage; } + if let Some(choice) = parsed.choices.into_iter().next() { + if let Some(reason) = choice.finish_reason { finish_reason = parse_finish_reason(&reason); } + if native && choice.logprobs.is_some() { logprobs = choice.logprobs; } + if !native || !choice.text.is_empty() { yield Ok(StreamPart::TextDelta { id: "0".into(), delta: choice.text, provider_metadata: None }); } + } + } + if started { yield Ok(StreamPart::TextEnd { id: "0".into(), provider_metadata: None }); } + yield Ok(StreamPart::Finish { finish_reason, usage: usage(final_usage, native), provider_metadata: native.then(|| provider_namespace("openai", match logprobs { Some(logprobs) => json!({ "logprobs": logprobs }), None => json!({}) }).expect("provider metadata must be an object")) }); + }; + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { + headers: Some(stream_response_headers), + }), + }) + } +} diff --git a/aimux-providers/src/openai_compatible/config.rs b/aimux-providers/src/openai_compatible/config.rs new file mode 100644 index 000000000..4bce71c3f --- /dev/null +++ b/aimux-providers/src/openai_compatible/config.rs @@ -0,0 +1,239 @@ +//! The private per-model configuration of the OpenAI-compatible package. +//! +//! Like the native OpenAI config this is the Rust shape of the object +//! `createOpenAICompatible` hands each model class (`{ provider, url, headers, +//! fetch, ... }`). It has no getters for the credential or the settings that +//! produced it. Compatible endpoint behavior is data (flags and a usage-key +//! name); the shared model never compares a provider name. + +use futures::future::BoxFuture; +use std::collections::HashMap; +use std::sync::Arc; + +use aimux_core::types::{ProviderMetadata, Usage}; +use serde_json::Value; + +use aimux_core::AiMuxError; +use aimux_core::language_model::SupportedUrls; +use aimux_provider_utils::{ + ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, ProviderErrorParts, + combine_headers, normalize_headers, +}; + +pub use crate::shared::TransformRequestBody; + +/// Where the base URL comes from. +/// +/// A plain provider has a [`Fixed`](Self::Fixed) URL. A preset whose URL is +/// read from the environment or expanded from template parameters is +/// [`Lazy`](Self::Lazy): it is evaluated per request, so creating the provider +/// (and the preset default instances) never reads the environment and never +/// fails. +#[derive(Clone)] +pub(crate) enum BaseUrl { + Fixed(String), + Lazy(Arc Result + Send + Sync>), +} + +impl BaseUrl { + pub(crate) fn resolve(&self) -> Result { + match self { + Self::Fixed(url) => Ok(url.clone()), + Self::Lazy(produce) => produce(), + } + } +} + +/// Maps a provider error payload to the message and code of the API error. +pub(crate) type ErrorStructure = Arc ProviderErrorParts + Send + Sync>; + +/// Captures vendor-specific metadata from responses (the AI SDK's +/// `MetadataExtractor`). The returned map is merged into `provider_metadata` +/// next to the namespace entry. +pub trait MetadataExtractor: Send + Sync { + /// Metadata of a non-streaming response body. + fn extract_metadata(&self, parsed_body: &Value) -> Option; + + /// A fresh extractor for one streaming response. + fn create_stream_extractor(&self) -> Box; +} + +/// The per-stream half of [`MetadataExtractor`]. +pub trait StreamMetadataExtractor: Send { + /// Called with every parsed chunk (including the one carrying usage). + fn process_chunk(&mut self, parsed_chunk: &Value); + + /// Metadata merged into the finish part's `provider_metadata`. + fn build_metadata(&self) -> Option; +} + +/// Everything that distinguishes one compatible vendor's chat endpoint from +/// the baseline, other than the public provider settings. +#[derive(Clone)] +pub(crate) struct ChatDialect { + /// Send `top_k`. The AI SDK baseline never does (it warns and drops it). + pub supports_top_k: bool, + /// Send `tools` / `tool_choice`. + pub supports_tools: bool, + /// Send `response_format`. + pub supports_response_format: bool, + /// The only `max-tokens` key the vendor accepts: `"max_tokens"` or + /// `"max_completion_tokens"`. Preset capability data (RFC-0036 section 5). + pub max_tokens_key: Option<&'static str>, + /// Streaming usage rides in `chunk[key].usage` instead of `chunk.usage`. + pub stream_usage_key: Option, + pub metadata_extractor: Option>, + pub error_structure: ErrorStructure, +} + +impl ChatDialect { + /// The AI SDK baseline. + pub(crate) fn baseline() -> Self { + Self { + supports_top_k: false, + supports_tools: true, + supports_response_format: true, + max_tokens_key: None, + stream_usage_key: None, + metadata_extractor: None, + error_structure: Arc::new(default_error_structure), + } + } +} + +/// `{ error: { message, type?, param?, code? } }`, the structure nearly every +/// compatible server answers errors with. +pub(crate) fn default_error_structure(data: &Value) -> ProviderErrorParts { + let error = data.get("error").unwrap_or(data); + ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error.get("code").or_else(|| error.get("type")).and_then( + |value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }, + ), + } +} + +/// Converts the upstream token usage object into language-model usage. +pub type ConvertUsage = Arc Usage + Send + Sync>; +/// Returns the URLs supported by chat models, read each time a chat model's +/// supported URLs are: `supportedUrls?: () => LanguageModelV4['supportedUrls']` +/// (openai-compatible@3.0.62 `openai-compatible-chat-language-model.ts:84`, +/// `:161-163`), whose value may be a `PromiseLike`. +pub type SupportedUrlsFn = Arc BoxFuture<'static, SupportedUrls> + Send + Sync>; + +/// Chat-only settings of a provider. +#[derive(Clone)] +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 dialect: Arc, +} + +/// What a model needs to talk to the API. +#[derive(Clone)] +pub(crate) struct CompatModelConfig { + /// The identity the model reports: `"{name}.{method}"`. + pub provider: String, + pub base_url: BaseUrl, + /// Appended to every request URL, in this order. + pub query_params: Option>>, + /// Provider headers (credential, user headers), resolved on every request. + pub headers: HeadersFn, + /// Transport; `None` uses the process default, resolved per request. + pub fetch: Option, + pub transform_request_body: Option, + pub chat: ChatSettings, +} + +impl CompatModelConfig { + /// The providerOptions namespace of this provider: the first segment of + /// the provider string. + pub(crate) fn provider_options_name(&self) -> &str { + self.provider.split('.').next().unwrap_or_default().trim() + } + + /// The full URL of an endpoint path (with the query parameters) and the + /// origin credentialed headers may be sent to. + pub(crate) fn url(&self, path: &str) -> Result<(String, String), AiMuxError> { + let base = self.base_url.resolve()?; + let mut url = url::Url::parse(&format!("{base}{path}")).map_err(|e| { + AiMuxError::InvalidArgument(format!("invalid request URL for {}: {e}", self.provider)) + })?; + if let Some(params) = &self.query_params { + // As the AI SDK does: the configured parameters replace whatever + // query the base URL carried. + url.set_query(None); + if !params.is_empty() { + url.query_pairs_mut() + .extend_pairs(params.iter().map(|(k, v)| (k.as_str(), v.as_str()))); + } + } + Ok((url.to_string(), base)) + } + + /// Resolve the provider headers, layer the per-call headers over them + /// (case-insensitive, later wins, `None` removes) and return the list that + /// goes on the wire. + pub(crate) async fn request_headers( + &self, + call_headers: Option<&HashMap>, + ) -> Result, AiMuxError> { + let provider = self.headers.resolve().await?; + let call: HeaderMapOpt = call_headers + .map(|headers| { + headers + .iter() + .map(|(name, value)| (name.clone(), Some(value.clone()))) + .collect() + }) + .unwrap_or_default(); + Ok(normalize_headers(combine_headers(&[&provider, &call]))) + } + + /// An exchange for `path` that inherits the operation's cancellation and + /// recording context and carries this config's transport and credential + /// origin. + pub(crate) fn http_request( + &self, + path: &str, + headers: Vec<(String, String)>, + options: &impl ExchangeContext, + ) -> Result { + let (url, origin) = self.url(path)?; + Ok(self.with_transport(HttpRequest::new(url, headers, options), origin)) + } + + /// Attach the transport and the credential origin to a request. + pub(crate) fn with_transport(&self, mut request: HttpRequest, origin: String) -> HttpRequest { + request.fetch = self.fetch.clone(); + request.credentialed_origin = Some(origin); + request + } + + /// Run the provider-level body rewrite, when there is one. + pub(crate) fn transform_body(&self, body: Value) -> Value { + match &self.transform_request_body { + Some(transform) => transform(body), + None => body, + } + } + + /// The failed-response handler of this provider's error structure. + pub(crate) fn failed_response_handler( + &self, + ) -> aimux_provider_utils::ResponseHandler { + let structure = self.chat.dialect.error_structure.clone(); + aimux_provider_utils::create_json_error_response_handler(move |data| structure(data)) + } +} diff --git a/aimux-providers/src/openai_compatible/convert.rs b/aimux-providers/src/openai_compatible/convert.rs new file mode 100644 index 000000000..41abba1ce --- /dev/null +++ b/aimux-providers/src/openai_compatible/convert.rs @@ -0,0 +1,736 @@ +//! Call options and prompt to the OpenAI-compatible chat-completions body. +//! +//! The Rust form of `getArgs` and `convertToOpenAICompatibleChatMessages` of +//! `@ai-sdk/openai-compatible`. The body is built in the AI SDK's order (model, +//! `user`, sampling settings, `response_format`, `stop`, `seed`, the +//! pass-through fields of the provider's own options namespace, then +//! `reasoning_effort`, `verbosity`, `messages`, `tools`, `tool_choice`). +//! Compatible endpoint capabilities arrive as [`ChatDialect`] data. + +use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; +use base64::Engine; +use serde_json::{Map, Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model_message::{ + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultPart, UserPart, +}; +use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat}; +use aimux_core::tool::Tool; +use aimux_core::types::{FinishReason, FinishReasonUnified, Warning}; +use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; + +use super::config::ChatDialect; + +/// Option keys of the generic compatible schema +/// (`openaiCompatibleLanguageModelChatOptions`): consumed by the model, never +/// passed through. +const SCHEMA_KEYS: [&str; 4] = [ + "user", + "reasoningEffort", + "textVerbosity", + "strictJsonSchema", +]; + +/// `tools` / `tool_choice` ready for the body, with the warnings raised while +/// preparing them. +#[derive(Debug, Clone, Default)] +pub(crate) struct PreparedTools { + pub tools: Option>, + pub tool_choice: Option, + pub warnings: Vec, +} + +/// The wire form of a function tool. +pub(crate) fn function_tool_value(tool: &aimux_core::tool::FunctionTool) -> Value { + let mut function = Map::new(); + function.insert("name".into(), json!(tool.name)); + if let Some(description) = &tool.description { + function.insert("description".into(), json!(description)); + } + function.insert("parameters".into(), tool.input_schema.clone()); + if let Some(strict) = tool.strict { + function.insert("strict".into(), json!(strict)); + } + json!({ "type": "function", "function": function }) +} + +/// The wire form of the tool choice. +pub(crate) fn tool_choice_value(choice: &LanguageModelV4ToolChoice) -> Value { + match choice { + LanguageModelV4ToolChoice::Auto => json!("auto"), + LanguageModelV4ToolChoice::None => json!("none"), + LanguageModelV4ToolChoice::Required => json!("required"), + LanguageModelV4ToolChoice::Tool { tool_name } => { + json!({ "type": "function", "function": { "name": tool_name } }) + } + } +} + +/// Function tools only; a provider-defined tool is unsupported by the generic +/// endpoint and warned about (the AI SDK's `prepareTools`). +pub(crate) fn prepare_function_tools(options: &CallOptions) -> PreparedTools { + let mut prepared = PreparedTools::default(); + let Some(tools) = options.tools.as_ref().filter(|tools| !tools.is_empty()) else { + return prepared; + }; + let mut wire = Vec::new(); + for tool in tools { + match tool { + Tool::Function(function) => wire.push(function_tool_value(function)), + Tool::Provider(provider) => prepared.warnings.push(Warning::Unsupported { + feature: format!("provider-defined tool {}", provider.id), + details: None, + }), + } + } + if !wire.is_empty() { + prepared.tools = Some(wire); + prepared.tool_choice = options.tool_choice.as_ref().map(tool_choice_value); + } + prepared +} + +/// `snake_case` / `kebab-case` provider name to `camelCase` (the AI SDK's +/// `toCamelCase`). +#[must_use] +pub(crate) fn to_camel_case(name: &str) -> String { + let mut out = String::with_capacity(name.len()); + let mut chars = name.chars().peekable(); + while let Some(ch) = chars.next() { + match chars.peek() { + Some(next) if (ch == '_' || ch == '-') && next.is_ascii_lowercase() => { + out.push(next.to_ascii_uppercase()); + chars.next(); + } + _ => out.push(ch), + } + } + out +} + +/// What the converter needs to know about the provider it builds for. +pub(crate) struct ChatBodySpec<'a> { + /// First segment of the provider string: the providerOptions namespace. + pub provider_options_name: &'a str, + pub include_usage: bool, + pub supports_structured_outputs: bool, + pub supports_multi_part_tool_content: bool, + pub dialect: &'a ChatDialect, +} + +/// A request body and what building it produced. +#[derive(Debug, Clone)] +pub(crate) struct RequestBodyResult { + /// The JSON body, before the provider's `transform_request_body`. + pub body: Value, + /// Warnings raised while building it. + pub warnings: Vec, + /// The key provider metadata is reported under. + pub metadata_key: String, +} + +fn namespace<'a>(options: &'a CallOptions, key: &str) -> Option<&'a Map> { + options + .provider_options + .as_ref() + .and_then(|all| all.get(key)) +} + +/// The key provider metadata is reported under (`resolveProviderOptionsKey`): +/// the provider name, unless the call used only its camelCase form. +fn resolve_metadata_key(name: &str, options: &CallOptions) -> String { + let camel = to_camel_case(name); + if camel != name && namespace(options, name).is_none() && namespace(options, &camel).is_some() { + camel + } else { + name.to_string() + } +} + +fn option_string( + merged: &Map, + namespace_name: &str, + key: &str, +) -> Result, AiMuxError> { + match merged.get(key) { + None | Some(Value::Null) => 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" + ))), + } +} + +/// Build the request body (`getArgs`). +/// +/// # Errors +/// +/// `InvalidArgument` for a malformed provider option or a prompt part that has +/// no wire form. +pub(crate) fn build_request_body( + model_id: &str, + options: &CallOptions, + stream: bool, + spec: &ChatBodySpec<'_>, +) -> Result { + let name = spec.provider_options_name; + let camel = to_camel_case(name); + let dialect = spec.dialect; + let mut warnings = Vec::new(); + + // Resolved provider options: the generic namespace, then the provider's + // own, then its camelCase form; later wins per key. + let mut merged = Map::new(); + if let Some(deprecated) = namespace(options, "openai-compatible") { + warnings.push(Warning::Deprecated { + setting: "providerOptions key 'openai-compatible'".to_string(), + message: "Use 'openaiCompatible' instead.".to_string(), + }); + merged.extend(deprecated.iter().map(|(k, v)| (k.clone(), v.clone()))); + } + for key in ["openaiCompatible", name, camel.as_str()] { + if let Some(object) = namespace(options, key) { + merged.extend(object.iter().map(|(k, v)| (k.clone(), v.clone()))); + } + } + if camel != name && namespace(options, name).is_some() { + warnings.push(Warning::Deprecated { + setting: format!("providerOptions key '{name}'"), + message: format!("Use '{camel}' instead."), + }); + } + + let user = option_string(&merged, name, "user")?; + let reasoning_effort = option_string(&merged, name, "reasoningEffort")?.or_else(|| { + options + .reasoning + .filter(|effort| effort.is_custom()) + .map(|effort| effort.to_string()) + }); + let verbosity = option_string(&merged, name, "textVerbosity")?; + let strict_json_schema = merged + .get("strictJsonSchema") + .and_then(Value::as_bool) + .unwrap_or(true); + + let metadata_key = resolve_metadata_key(name, options); + let mut body = Map::new(); + body.insert("model".into(), json!(model_id)); + if let Some(user) = user { + body.insert("user".into(), json!(user)); + } + if let Some(max_tokens) = options.max_output_tokens { + body.insert( + dialect.max_tokens_key.unwrap_or("max_tokens").into(), + json!(max_tokens), + ); + } + for (key, value) in [ + ("temperature", options.temperature), + ("top_p", options.top_p), + ("frequency_penalty", options.frequency_penalty), + ("presence_penalty", options.presence_penalty), + ] { + if let Some(value) = value { + body.insert(key.into(), json!(value)); + } + } + if let Some(top_k) = options.top_k { + if dialect.supports_top_k { + body.insert("top_k".into(), json!(top_k)); + } else { + warnings.push(Warning::Unsupported { + feature: "topK".to_string(), + details: None, + }); + } + } + + match &options.response_format { + Some(ResponseFormat::Json { + schema, + name: format_name, + description, + }) if dialect.supports_response_format => match schema { + Some(schema) if spec.supports_structured_outputs => { + let mut json_schema = Map::new(); + json_schema.insert("schema".into(), schema.clone()); + json_schema.insert("strict".into(), json!(strict_json_schema)); + json_schema.insert( + "name".into(), + json!(format_name.clone().unwrap_or_else(|| "response".into())), + ); + if let Some(description) = description { + json_schema.insert("description".into(), json!(description)); + } + body.insert( + "response_format".into(), + json!({ "type": "json_schema", "json_schema": json_schema }), + ); + } + Some(_) => { + warnings.push(Warning::Unsupported { + feature: "responseFormat".to_string(), + details: Some( + "JSON response format schema is only supported with structuredOutputs" + .to_string(), + ), + }); + body.insert("response_format".into(), json!({ "type": "json_object" })); + } + None => { + body.insert("response_format".into(), json!({ "type": "json_object" })); + } + }, + Some(ResponseFormat::Json { .. }) => warnings.push(Warning::Unsupported { + feature: "responseFormat".to_string(), + details: Some("response_format is not supported by this provider".to_string()), + }), + Some(ResponseFormat::Text) | None => {} + } + + if let Some(stop) = &options.stop_sequences { + body.insert("stop".into(), json!(stop)); + } + if let Some(seed) = options.seed { + body.insert("seed".into(), json!(seed)); + } + + // Fields of the provider's own namespace that the schema does not know go + // to the body as given (the generic `openaiCompatible` namespace is + // schema-only: unknown fields there are dropped). + for key in [name, camel.as_str()] { + if let Some(object) = namespace(options, key) { + for (field, value) in object { + if !SCHEMA_KEYS.contains(&field.as_str()) { + body.insert(field.clone(), value.clone()); + } + } + } + } + + if let Some(effort) = reasoning_effort { + body.insert("reasoning_effort".into(), json!(effort)); + } + if let Some(verbosity) = verbosity { + body.insert("verbosity".into(), json!(verbosity)); + } + + let messages = convert_messages( + &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)); + + if dialect.supports_tools { + let prepared = prepare_function_tools(options); + warnings.extend(prepared.warnings); + if let Some(tools) = prepared.tools { + body.insert("tools".into(), Value::Array(tools)); + } + if let Some(choice) = prepared.tool_choice { + body.insert("tool_choice".into(), choice); + } + } else if options + .tools + .as_ref() + .is_some_and(|tools| !tools.is_empty()) + { + warnings.push(Warning::Unsupported { + feature: "tools".to_string(), + details: Some("tools are not supported by this provider".to_string()), + }); + } + + if stream { + body.insert("stream".into(), json!(true)); + if spec.include_usage { + body.insert("stream_options".into(), json!({ "include_usage": true })); + } + } + + Ok(RequestBodyResult { + body: Value::Object(body), + warnings, + metadata_key, + }) +} + +// ── Messages ──────────────────────────────────────────────────────────────── + +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 +} + +fn with_metadata(mut object: Value, metadata: Map) -> Value { + if let Some(map) = object.as_object_mut() { + map.extend(metadata); + } + object +} + +/// `messages` for the body. +/// +/// # Errors +/// +/// `InvalidArgument` when a part has no wire form (an unsupported media type, +/// an audio or PDF part given by URL). +pub(crate) fn convert_messages( + prompt: &LanguageModelPrompt, + spec: &MessageSpec<'_>, +) -> Result, AiMuxError> { + let mut out = Vec::new(); + for message in prompt { + out.extend(convert_message(message, spec)?); + } + Ok(out) +} + +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( + json!({ "role": "system", "content": content }), + wire_metadata(provider_options.as_ref(), key), + )]), + LanguageModelMessage::User { content, provider_options } => Ok(vec![convert_user_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(), key), + )]), + LanguageModelMessage::Tool { content, .. } => Ok(content + .iter() + .filter_map(|part| { + let ToolPart::ToolResult(ToolResultPart { + 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) + }, + }), + wire_metadata(provider_options.as_ref(), key), + )) + }) + .collect()), + } +} + +fn convert_user_message( + content: &[UserPart], + spec: &MessageSpec<'_>, + message_metadata: Map, +) -> Result { + let key = spec.metadata_key; + if let [ + UserPart::Text(TextPart { + text, + provider_options, + }), + ] = content + { + let mut metadata = message_metadata; + metadata.extend(wire_metadata(provider_options.as_ref(), key)); + return Ok(with_metadata( + json!({ "role": "user", "content": text }), + metadata, + )); + } + 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>>()?; + 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"), + }; + let media_type = part.media_type.as_str(); + 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 kind = format!("{kind}_url"); + Ok(json!({ "type": kind, kind: { "url": url } })) + } + "audio" => { + let Some(b64) = base64_data.as_deref() else { + return Err(unsupported("audio file parts with URLs".to_string())); + }; + let full_media_type = resolve_full_media_type(part)?; + let format = match full_media_type.as_str() { + "audio/wav" => "wav", + "audio/mp3" | "audio/mpeg" => "mp3", + other => { + return Err(unsupported(format!( + "audio content parts with media type {other}" + ))); + } + }; + Ok(json!({ "type": "input_audio", "input_audio": { "data": b64, "format": format } })) + } + "application" => { + let Some(b64) = base64_data.as_deref() else { + return Err(unsupported("PDF file parts with URLs".to_string())); + }; + let full_media_type = resolve_full_media_type(part)?; + if full_media_type != "application/pdf" { + return Err(unsupported(format!( + "file part media type {full_media_type}" + ))); + } + let filename = part.filename.as_deref().unwrap_or("document.pdf"); + Ok(json!({ + "type": "file", + "file": { + "filename": filename, + "file_data": format!("data:application/pdf;base64,{b64}"), + } + })) + } + "text" => { + let Some(b64) = base64_data.as_deref() else { + return Err(unsupported("text file parts with URLs".to_string())); + }; + 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()))?; + Ok(json!({ "type": "text", "text": text })) + } + _ => Err(unsupported(format!("file part media type {media_type}"))), + } +} + +fn convert_assistant_message( + content: &[AssistantPart], + spec: &MessageSpec<'_>, + message_metadata: Map, +) -> Value { + let key = spec.metadata_key; + let mut text = String::new(); + let mut reasoning = String::new(); + let mut tool_calls = Vec::new(); + let mut 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::Reasoning(ReasoningPart { text: t, .. }) => reasoning.push_str(t), + AssistantPart::ToolCall(ToolCallPart { + tool_call_id, + tool_name, + input, + provider_options, + .. + }) => { + let arguments = input.to_string(); + let mut call = json!({ + "id": tool_call_id, + "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["extra_content"] = json!({ "google": { "thought_signature": signature } }); + } + tool_calls.push(with_metadata( + call, + wire_metadata(provider_options.as_ref(), key), + )); + } + // Provider-executed results and media have no assistant-content + // form on this wire. + _ => {} + } + } + let mut wire = Map::new(); + wire.insert("role".into(), json!("assistant")); + wire.insert( + "content".into(), + if tool_calls.is_empty() || !text.is_empty() { + json!(text) + } else { + Value::Null + }, + ); + if !reasoning.is_empty() { + wire.insert("reasoning_content".into(), json!(reasoning)); + } + if !tool_calls.is_empty() { + wire.insert("tool_calls".into(), Value::Array(tool_calls)); + } + wire.extend(metadata); + Value::Object(wire) +} + +/// Map an OpenAI finish reason string to the unified one, keeping the raw. +#[must_use] +pub(crate) fn parse_finish_reason(raw: &str) -> FinishReason { + let unified = match raw { + "stop" => FinishReasonUnified::Stop, + "length" => FinishReasonUnified::Length, + "tool_calls" | "function_call" => FinishReasonUnified::ToolCalls, + "content_filter" => FinishReasonUnified::ContentFilter, + _ => FinishReasonUnified::Other, + }; + FinishReason { + unified, + raw: Some(raw.to_string()), + } +} + +// String(value), as used by the compatible package for thought signatures. +fn signature_string(value: &Value) -> String { + match value { + Value::String(value) => value.clone(), + Value::Object(_) => "[object Object]".into(), + Value::Array(values) => values + .iter() + .map(|value| { + if value.is_null() { + String::new() + } else { + signature_string(value) + } + }) + .collect::>() + .join(","), + Value::Number(value) => value.as_f64().map_or_else( + || value.to_string(), + |value| { + if value == 0.0 { + "0".into() + } else { + value.to_string() + } + }, + ), + _ => value.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn camel_case_follows_the_ai_sdk() { + assert_eq!(to_camel_case("acme"), "acme"); + assert_eq!(to_camel_case("my_provider"), "myProvider"); + assert_eq!(to_camel_case("my-provider"), "myProvider"); + assert_eq!( + to_camel_case("vertex_ai_openai_models"), + "vertexAiOpenaiModels" + ); + assert_eq!(to_camel_case("a_1"), "a_1"); + } +} diff --git a/aimux-providers/src/openai_compatible/embedding.rs b/aimux-providers/src/openai_compatible/embedding.rs new file mode 100644 index 000000000..767a674d5 --- /dev/null +++ b/aimux-providers/src/openai_compatible/embedding.rs @@ -0,0 +1,142 @@ +//! The OpenAI-compatible embedding model (`{name}.embedding`). +//! +//! 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. + +use async_trait::async_trait; +use serde_json::{Map, Value, json}; + +use aimux_core::embedding_model::{ + EmbeddingCallOptions, EmbeddingModel, EmbeddingResponse, EmbeddingResult, EmbeddingUsage, +}; +use aimux_core::error::AiMuxError; + +use super::config::CompatModelConfig; +use super::convert::to_camel_case; + +/// An OpenAI-compatible embedding model. +pub struct OpenAICompatibleEmbeddingModel { + model_id: String, + config: CompatModelConfig, +} + +impl OpenAICompatibleEmbeddingModel { + pub(crate) fn from_config(model_id: String, config: CompatModelConfig) -> Self { + Self { model_id, config } + } +} + +/// `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()] { + let Some(object) = options + .provider_options + .as_ref() + .and_then(|all| all.get(key)) + 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()); + } + } + (dimensions, user) +} + +#[async_trait] +impl EmbeddingModel for OpenAICompatibleEmbeddingModel { + fn provider(&self) -> &str { + &self.config.provider + } + + fn model_id(&self) -> &str { + &self.model_id + } + + fn max_embeddings_per_call(&self) -> Option { + Some(2048) + } + + fn supports_parallel_calls(&self) -> bool { + true + } + + async fn do_embed( + &self, + options: &EmbeddingCallOptions, + ) -> Result { + let (dimensions, user) = embedding_options(options, self.config.provider_options_name()); + + 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)); + } + let body = Value::Object(body); + + let headers = self + .config + .request_headers(options.headers.as_ref()) + .await?; + let resp = aimux_provider_utils::post_json_to_api( + self.config.http_request("/embeddings", headers, options)?, + body, + aimux_provider_utils::create_json_response_handler(), + self.config.failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers; + let raw_value: Value = resp.value; + + let embeddings: Vec> = raw_value + .get("data") + .and_then(Value::as_array) + .map(|items| { + items + .iter() + .filter_map(|item| { + item.get("embedding")?.as_array().map(|numbers| { + numbers + .iter() + .filter_map(|v| v.as_f64().map(|f| f as f32)) + .collect() + }) + }) + .collect() + }) + .unwrap_or_default(); + + let usage = raw_value + .get("usage") + .and_then(|u| u.get("prompt_tokens")) + .and_then(Value::as_u64) + .map(|tokens| EmbeddingUsage { + tokens: tokens as u32, + }); + + Ok(EmbeddingResult { + embeddings, + usage, + provider_metadata: None, + response: Some(EmbeddingResponse { + headers: Some(response_headers), + body: Some(raw_value), + }), + warnings: Vec::new(), + }) + } +} diff --git a/aimux-providers/src/openai_compatible/image.rs b/aimux-providers/src/openai_compatible/image.rs new file mode 100644 index 000000000..008b99cb6 --- /dev/null +++ b/aimux-providers/src/openai_compatible/image.rs @@ -0,0 +1,195 @@ +//! The OpenAI-compatible image model (`{name}.image`). +//! +//! 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. +//! +//! The provider options of the `openaiCompatible` namespace, then the +//! provider's own, are forwarded to the endpoint as given. + +use async_trait::async_trait; +use serde_json::{Map, Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::image_model::{ + ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, ImageResponse, + ImageResult, +}; +use aimux_core::shared::Warning; +use aimux_provider_utils::{HttpBody, MultipartForm}; + +use super::config::CompatModelConfig; +use super::convert::to_camel_case; + +/// An OpenAI-compatible image generation / editing model. +pub struct OpenAICompatibleImageModel { + model_id: String, + config: CompatModelConfig, +} + +impl OpenAICompatibleImageModel { + pub(crate) fn from_config(model_id: String, config: CompatModelConfig) -> Self { + 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 { + 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 let Some(object) = options.provider_options.get(key) { + out.extend(object.iter().map(|(k, v)| (k.clone(), v.clone()))); + } + } + out + } +} + +fn file_bytes(file: &ImageFile) -> Result<(Vec, String), AiMuxError> { + match file { + ImageFile::File { media_type, data } => { + let bytes = match data { + ImageFileData::Binary(bytes) => bytes.clone(), + ImageFileData::Base64(b64) => { + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) + .map_err(|e| AiMuxError::InvalidArgument(format!("invalid base64: {e}")))? + } + }; + Ok((bytes, media_type.clone())) + } + ImageFile::Url { .. } => Err(AiMuxError::InvalidArgument( + "image edit files must be inline data, not URLs".to_string(), + )), + } +} + +fn form_value(value: &Value) -> String { + match value { + Value::String(text) => text.clone(), + other => other.to_string(), + } +} + +#[async_trait] +impl ImageModel for OpenAICompatibleImageModel { + fn provider(&self) -> &str { + &self.config.provider + } + + fn model_id(&self) -> &str { + &self.model_id + } + + fn max_images_per_call(&self) -> Option { + Some(10) + } + + async fn do_generate(&self, options: &ImageCallOptions) -> Result { + let mut warnings = Vec::new(); + if options.aspect_ratio.is_some() { + warnings.push(Warning::Unsupported { + feature: "aspectRatio".to_string(), + details: Some( + "This model does not support aspect ratio. Use `size` instead.".to_string(), + ), + }); + } + if options.seed.is_some() { + warnings.push(Warning::Unsupported { + feature: "seed".to_string(), + details: None, + }); + } + + let timestamp = chrono::Utc::now().to_rfc3339(); + let headers = self + .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 images: Vec = value + .get("data") + .and_then(Value::as_array) + .map(|items| { + items + .iter() + .filter_map(|item| item.get("b64_json").and_then(Value::as_str)) + .map(str::to_string) + .collect() + }) + .unwrap_or_default(); + + Ok(ImageResult { + images: ImageOutputs::Base64(images), + warnings, + provider_metadata: None, + response: ImageResponse { + timestamp: Some(timestamp), + model_id: Some(self.model_id.clone()), + headers: Some(response_headers), + }, + usage: None, + }) + } +} diff --git a/aimux-providers/src/openai_compatible/mod.rs b/aimux-providers/src/openai_compatible/mod.rs new file mode 100644 index 000000000..285f3f16b --- /dev/null +++ b/aimux-providers/src/openai_compatible/mod.rs @@ -0,0 +1,342 @@ +//! OpenAI-compatible provider. +//! +//! [`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`] +//! whose models speak the chat-completions, embeddings and image endpoints of +//! any OpenAI-compatible server. +//! +//! What is fixed when the provider is created, and what is not: +//! +//! - The explicit settings (`name`, `base_url`, `headers`, `query_params`, +//! `fetch`, the capability flags, `transform_request_body`) are fixed in the +//! factory. +//! - A non-empty `api_key` is used as given. `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: +//! options are read from `openaiCompatible`, then the name, then its camelCase +//! form (later wins), fields of the name's own namespace that the generic +//! schema does not know are passed through to the body, and provider metadata +//! is reported under the name. +//! +//! Vendor packages ([`groq`](crate::groq), [`deepseek`](crate::deepseek)) and the +//! registry presets ([`preset`](crate::preset)) are built on the same +//! internals; the shared model never branches on a vendor name. + +pub(crate) mod chat; +pub mod completion; +pub(crate) mod config; +pub(crate) mod convert; +pub mod embedding; +pub mod image; +mod types; + +pub use chat::OpenAICompatibleChatModel; +pub use completion::OpenAICompatibleCompletionModel; +pub use config::{ + ConvertUsage, MetadataExtractor, StreamMetadataExtractor, SupportedUrlsFn, TransformRequestBody, +}; +pub use embedding::OpenAICompatibleEmbeddingModel; +pub use image::OpenAICompatibleImageModel; + +use std::collections::HashMap; +use std::sync::Arc; + +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::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, provider_headers}; +use config::{BaseUrl, ChatDialect, ChatSettings, CompatModelConfig}; + +/// Settings of [`create_openai_compatible`] (the AI SDK's +/// `OpenAICompatibleProviderSettings`). +#[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 `.`. + pub name: String, + /// Base URL for the API calls. Required, `http(s)` with a host; a trailing + /// slash is removed. + pub base_url: String, + /// The API key. `None` or `""` sends no `Authorization` header; + /// a non-empty value is sent as `Bearer `. + pub api_key: Option, + /// Extra headers on every request. A `None` value removes the header, + /// including the `Authorization` one. Per-call headers win over these. + pub headers: Option, + /// Query parameters appended to every request URL (sorted by name). + pub query_params: Option>, + /// The transport: a mock, a signing decorator, a proxy-aware client. + /// `None` uses the process default, resolved per request. + pub fetch: Option, + /// Ask streaming responses to include usage (`stream_options.include_usage`). + pub include_usage: Option, + /// 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 + /// before it is sent. + pub transform_request_body: Option, + /// Extract provider-specific metadata from generated and streamed responses. + pub metadata_extractor: Option>, + /// Returns the supported URLs for chat models. + pub supported_urls: Option, + /// Converts chat response token usage. + pub convert_usage: Option, +} + +impl std::fmt::Debug for OpenAICompatibleProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OpenAICompatibleProviderSettings") + .field("name", &self.name) + .field("base_url", &self.base_url) + .field("api_key", &self.api_key.is_some()) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("query_params", &self.query_params) + .field("fetch", &self.fetch.is_some()) + .field("include_usage", &self.include_usage) + .field( + "supports_structured_outputs", + &self.supports_structured_outputs, + ) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ) + .field("metadata_extractor", &self.metadata_extractor.is_some()) + .field("supported_urls", &self.supported_urls.is_some()) + .field("convert_usage", &self.convert_usage.is_some()) + .finish() + } +} + +/// Create an OpenAI-compatible provider. +/// +/// # 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 +/// 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 { + name: settings.name, + base_url: BaseUrl::Fixed(base_url), + credential: match settings.api_key { + Some(key) => Credential::Optional(Resolvable::Value(key)), + None => Credential::None, + }, + fixed_headers: Vec::new(), + headers: settings.headers, + query_params: settings.query_params, + fetch: settings.fetch, + transform_request_body: settings.transform_request_body, + 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, + }, + }) +} + +/// The chat-endpoint behavior of a compatible vendor: the capability flags the +/// public settings expose, and the [`ChatDialect`] they cannot. A vendor +/// package or preset supplies its own; the public factory uses the AI SDK +/// baseline. +#[derive(Clone)] +pub(crate) struct ChatProfile { + pub include_usage: bool, + pub supports_structured_outputs: bool, + pub 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 +/// from the settings; vendor packages and presets fill it themselves. +pub(crate) struct Assembly { + pub name: String, + pub base_url: BaseUrl, + pub credential: Credential, + /// Headers sent before the user's (organization, project, ...). + pub fixed_headers: Vec<(String, String)>, + pub headers: Option, + pub query_params: Option>, + pub fetch: Option, + pub transform_request_body: Option, + pub profile: ChatProfile, +} + +/// An OpenAI-compatible provider (the AI SDK's `OpenAICompatibleProvider`). +/// Cheap to take models out of; it holds no HTTP client. +pub struct OpenAICompatibleProvider { + name: String, + base_url: BaseUrl, + query_params: Option>>, + headers: HeadersFn, + fetch: Option, + transform_request_body: Option, + chat: ChatSettings, +} + +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 query_params = assembly.query_params.map(|params| { + let mut pairs: Vec<(String, String)> = params.into_iter().collect(); + pairs.sort(); + Arc::new(pairs) + }); + Ok(Self { + name, + base_url: assembly.base_url, + query_params, + headers: aimux_provider_utils::headers::with_user_agent_suffix_fn( + provider_headers( + assembly.credential, + assembly.fixed_headers, + assembly.headers, + ), + "openai-compatible", + "3.0.59", + ), + fetch: assembly.fetch, + transform_request_body: assembly.transform_request_body, + chat: ChatSettings { + include_usage: assembly.profile.include_usage, + supports_structured_outputs: assembly.profile.supports_structured_outputs, + supports_multi_part_tool_content: assembly.profile.supports_multi_part_tool_content, + supported_urls: assembly.profile.supported_urls, + convert_usage: assembly.profile.convert_usage, + dialect: Arc::new(assembly.profile.dialect), + }, + }) + } + + fn model_config(&self, method: &str) -> CompatModelConfig { + CompatModelConfig { + provider: format!("{}.{method}", self.name), + base_url: self.base_url.clone(), + query_params: self.query_params.clone(), + headers: self.headers.clone(), + fetch: self.fetch.clone(), + transform_request_body: if method == "chat" { + self.transform_request_body.clone() + } else { + None + }, + chat: self.chat.clone(), + } + } + + /// A chat-completions model; `provider()` is `"{name}.chat"`. + #[must_use] + pub fn chat(&self, model_id: &str) -> OpenAICompatibleChatModel { + OpenAICompatibleChatModel::from_config(model_id.to_string(), self.model_config("chat")) + } + + /// A text-completion model; `provider()` is `"{name}.completion"`. + #[must_use] + pub fn completion(&self, model_id: &str) -> OpenAICompatibleCompletionModel { + OpenAICompatibleCompletionModel::from_config( + model_id.to_string(), + self.model_config("completion"), + ) + } + + /// An embedding model; `provider()` is `"{name}.embedding"`. + #[must_use] + pub fn embedding(&self, model_id: &str) -> OpenAICompatibleEmbeddingModel { + OpenAICompatibleEmbeddingModel::from_config( + model_id.to_string(), + self.model_config("embedding"), + ) + } + + /// An image model; `provider()` is `"{name}.image"`. + #[must_use] + pub fn image(&self, model_id: &str) -> OpenAICompatibleImageModel { + OpenAICompatibleImageModel::from_config(model_id.to_string(), self.model_config("image")) + } + + /// The provider as a function: the default language model for an id, the + /// chat model. The AI SDK's callable provider. + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.chat(model_id)) + } +} + +impl Provider for OpenAICompatibleProvider { + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) + } + + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.image(model_id))) + } + + /// Native decisions of the official runtimes (Ollama, llama.cpp, LocalAI, + /// Laya, vLLM, SGLang, Cloudflare); any other preset has none. + fn decision_model( + &self, + model_id: &str, + ) -> Result, AiMuxError> { + Ok(Box::new( + crate::runtime_decision::RuntimeDecisionModel::new( + &self.name, + model_id, + self.model_config("decision"), + )?, + )) + } +} + +impl ProviderDiscovery for OpenAICompatibleProvider { + /// `GET {base_url}/models`: one exchange, no retry. Without a key no + /// `Authorization` header is sent. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config("models"); + Box::pin(async move { chat::list_models_once(&config).await }) + } +} diff --git a/aimux-providers/src/openai_compatible/types.rs b/aimux-providers/src/openai_compatible/types.rs new file mode 100644 index 000000000..4a6bdfa3c --- /dev/null +++ b/aimux-providers/src/openai_compatible/types.rs @@ -0,0 +1,160 @@ +//! Wire types of the OpenAI-compatible chat-completions API. +//! +//! Deliberately loose: compatible servers omit fields the OpenAI API always +//! sends (`usage`, tool-call ids), so everything but the essentials is +//! optional, as in the AI SDK's `z.looseObject` schemas. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +// ── Non-streaming response ── + +#[derive(Debug, Deserialize)] +pub(crate) struct ChatCompletionResponse { + #[serde(default)] + pub id: Option, + #[serde(default)] + pub model: Option, + #[serde(default)] + pub created: Option, + #[serde(default)] + pub choices: Vec, + #[serde(default)] + pub usage: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct Choice { + pub message: MessageResponse, + #[serde(default)] + pub finish_reason: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct MessageResponse { + #[serde(default)] + pub content: Option, + /// Reasoning text (`reasoning`, Groq and others). + #[serde(default)] + pub reasoning: Option, + /// Reasoning text (`reasoning_content`, DeepSeek and others); wins over + /// `reasoning` when both are present. + #[serde(default)] + pub reasoning_content: Option, + #[serde(default)] + pub tool_calls: Option>, + /// Annotations (URL citations), kept as raw JSON. + #[serde(default)] + pub annotations: Option>, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct ToolCallResponse { + #[serde(default)] + pub id: Option, + pub function: FunctionCallResponse, + /// Google-hosted compatible endpoints carry the thought signature here. + #[serde(default)] + pub extra_content: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct FunctionCallResponse { + pub name: String, + #[serde(default)] + pub arguments: Option, +} + +#[derive(Debug, Deserialize, Clone, Serialize, Default)] +pub(crate) struct UsageResponse { + #[serde(default)] + pub prompt_tokens: Option, + #[serde(default)] + 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)] + pub completion_tokens_details: Option, +} + +#[derive(Debug, Deserialize, Clone, Serialize)] +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)] +pub(crate) struct CompletionTokensDetails { + #[serde(default)] + pub reasoning_tokens: Option, + #[serde(default)] + pub accepted_prediction_tokens: Option, + #[serde(default)] + pub rejected_prediction_tokens: Option, +} + +// ── Streaming response ── + +#[derive(Debug, Deserialize)] +pub(crate) struct StreamChunk { + #[serde(default)] + pub id: Option, + #[serde(default)] + pub model: Option, + #[serde(default)] + pub created: Option, + #[serde(default)] + pub choices: Vec, +} + +#[derive(Debug, Deserialize, Default)] +pub(crate) struct StreamChoice { + #[serde(default)] + pub delta: Delta, + #[serde(default)] + pub finish_reason: Option, +} + +#[derive(Debug, Deserialize, Default)] +pub(crate) struct Delta { + #[serde(default)] + pub content: Option, + #[serde(default)] + pub reasoning: Option, + #[serde(default)] + pub reasoning_content: Option, + #[serde(default)] + pub tool_calls: Option>, + #[serde(default)] + pub annotations: Option>, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct DeltaToolCall { + /// Some compatible providers omit it (`index: z.number().nullish()`); the + /// tracker falls back to id / name. + #[serde(default)] + pub index: Option, + #[serde(default)] + pub id: Option, + #[serde(default)] + pub function: Option, + #[serde(default)] + pub extra_content: Option, +} + +#[derive(Debug, Deserialize, Default)] +pub(crate) struct DeltaFunction { + #[serde(default)] + pub name: Option, + #[serde(default)] + pub arguments: Option, +} diff --git a/aimux-providers/src/openrouter.rs b/aimux-providers/src/openrouter.rs deleted file mode 100644 index a82a0ca06..000000000 --- a/aimux-providers/src/openrouter.rs +++ /dev/null @@ -1,90 +0,0 @@ -//! OpenRouter provider — a thin OpenAI-compatible wrapper. -//! -//! OpenRouter exposes an OpenAI-compatible Chat Completions API at -//! `https://openrouter.ai/api/v1`. The only provider-specific detail is the -//! base URL and the `OPENROUTER_API_KEY` environment variable; everything else -//! is delegated to the shared [`OpenAIProvider`]. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::load_api_key; - -use crate::openai::{ - OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider, OpenAIResponsesModel, -}; - -const DEFAULT_BASE_URL: &str = "https://openrouter.ai/api/v1"; -const ENV_VAR: &str = "OPENROUTER_API_KEY"; -const PROVIDER_NAME: &str = "openrouter"; - -/// Configuration for the OpenRouter provider (wraps [`OpenAIConfig`]). -pub struct OpenRouterConfig(OpenAIConfig); - -impl OpenRouterConfig { - /// Create from an API key, using the default OpenRouter base URL. - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()), - ) - } - - /// Create from the `OPENROUTER_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `OPENROUTER_API_KEY` is not - /// set. - pub fn from_env() -> Result { - let key = load_api_key(None, ENV_VAR, "OpenRouter")?; - Ok(Self::new(key)) - } - - /// Override the base URL (useful for tests / self-hosted endpoints). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// OpenRouter provider — creates [`OpenAIModel`] instances pointed at OpenRouter. -pub struct OpenRouterProvider(OpenAIProvider); - -impl OpenRouterProvider { - #[must_use] - pub fn new(config: OpenRouterConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given OpenRouter model id - /// (e.g. `"openai/gpt-4o-mini"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } - - /// Create a Responses-API model instance for the given OpenRouter model id. - /// - /// OpenRouter exposes an OpenAI-compatible Responses API at `/v1/responses`. - /// This delegates to the underlying [`OpenAIProvider::responses_model`]. - #[must_use] - pub fn responses_model(&self, model_id: &str) -> OpenAIResponsesModel { - self.0.responses_model(model_id) - } -} - -impl Provider for OpenRouterProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/openvino.rs b/aimux-providers/src/openvino.rs deleted file mode 100644 index 41cc94d3a..000000000 --- a/aimux-providers/src/openvino.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! OpenVINO provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8080/v1`. The `OPENVINO_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8080/v1"; -const ENV_VAR: &str = "OPENVINO_BASE_URL"; -const PROVIDER_NAME: &str = "openvino"; -const PLACEHOLDER_API_KEY: &str = "openvino"; - -pub struct OpenvinoConfig(OpenAIConfig); - -impl OpenvinoConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `OPENVINO_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8080/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct OpenvinoProvider(OpenAIProvider); - -impl OpenvinoProvider { - #[must_use] - pub fn new(config: OpenvinoConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for OpenvinoProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/parallel_ai.rs b/aimux-providers/src/parallel_ai.rs index 41c1251e7..6b360e7d9 100644 --- a/aimux-providers/src/parallel_ai.rs +++ b/aimux-providers/src/parallel_ai.rs @@ -7,82 +7,164 @@ //! search protocol (objective + search queries → results) and does not //! support language models. The single `query` is mapped to the //! `search_queries` array (and used as the `objective`). +//! +//! [`create_parallel_ai`] takes [`ParallelAiProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`ParallelAiProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `PARALLEL_API_KEY`. +//! [`parallel_ai()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; /// Fixed model ID for the Parallel AI search model. +/// Failed-response handler of the Parallel AI endpoint: `message`, and `type` +/// else `code` as the provider code, from the body's `error` object, or from +/// the top level when there is none. This is the mapping the endpoint used +/// with the removed shared handler; it is not an AI SDK schema. +fn parallel_ai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(serde_json::Value::as_str) + .map(str::to_owned), + } + }) +} + const MODEL_ID: &str = "parallel-search"; -/// Configuration for the Parallel AI provider. -#[derive(Debug, Clone)] -pub struct ParallelAiConfig { - pub api_key: String, - pub base_url: String, +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.parallel.ai"; +pub(crate) const API_KEY_ENV_VAR: &str = "PARALLEL_API_KEY"; +const DEFAULT_NAME: &str = "parallel_ai"; + +/// Settings of [`create_parallel_ai`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct ParallelAiProviderSettings { + /// Base URL for the API calls. Default `https://api.parallel.ai`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `PARALLEL_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.search"`). + /// Default `"parallel_ai"`. The providerOptions key stays `parallel_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, } -impl ParallelAiConfig { - /// Create from an API key (uses the default Parallel AI base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.parallel.ai".to_string(), - } +impl std::fmt::Debug for ParallelAiProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ParallelAiProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Create a Parallel AI provider. +/// +/// # 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. +pub fn create_parallel_ai( + settings: ParallelAiProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(ParallelAiProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Parallel AI"), + AuthScheme::Header("x-api-key"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `PARALLEL_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `PARALLEL_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "PARALLEL_API_KEY", "Parallel AI")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_parallel_ai` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn parallel_ai() -> &'static ParallelAiProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_parallel_ai(ParallelAiProviderSettings::default()) + .expect("default Parallel AI settings are always valid") + }) } -/// Parallel AI provider — creates `ParallelAiSearchModel` instances. -/// -/// Parallel AI is a search-only provider; it does not support language models. +/// A Parallel AI provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct ParallelAiProvider { - config: ParallelAiConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl ParallelAiProvider { - #[must_use] - pub fn new(config: ParallelAiConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a search model instance. + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> ParallelAiSearchModel { - ParallelAiSearchModel::new(self.config.clone()) + ParallelAiSearchModel::from_config(self.model_config("search")) } } -impl Provider for ParallelAiProvider { - fn name(&self) -> &str { - "parallel_ai" - } -} +crate::impl_single_modality_provider!(ParallelAiProvider, search_model, |p, _id| p.search_model()); /// Build the Parallel AI `/v1/search` request body (pure function). /// @@ -137,35 +219,19 @@ fn map_results(entries: Vec) -> Vec { /// A Parallel AI search model. pub struct ParallelAiSearchModel { - config: ParallelAiConfig, + config: EndpointConfig, } impl ParallelAiSearchModel { - #[must_use] - pub fn new(config: ParallelAiConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("x-api-key".to_string(), self.config.api_key.clone()); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/v1/search", self.config.base_url) - } } #[async_trait] impl SearchModel for ParallelAiSearchModel { fn provider(&self) -> &str { - "parallel_ai" + &self.config.provider } fn model_id(&self) -> &str { @@ -175,24 +241,13 @@ impl SearchModel for ParallelAiSearchModel { async fn do_search(&self, options: &SearchCallOptions) -> Result { let body = build_request_body(options); - let headers: Vec<(String, String)> = self - .build_headers(options.headers.as_ref()) - .into_iter() - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(exchange.url("/v1/search"), options), body, aimux_provider_utils::create_json_response_handler::(), - aimux_provider_utils::create_standard_json_error_response_handler(), + parallel_ai_failed_response_handler(), ) .await?; diff --git a/aimux-providers/src/preset.rs b/aimux-providers/src/preset.rs new file mode 100644 index 000000000..79c14f163 --- /dev/null +++ b/aimux-providers/src/preset.rs @@ -0,0 +1,687 @@ +//! Presets: a runtime table parsed once from the embedded registry. +//! +//! A preset is an OpenAI-compatible vendor described by data: a name, a +//! default base URL, the environment variable of its key (or no key at all), +//! an optional environment variable for the base URL and optional template +//! parameters. [`lookup`] finds a descriptor by name; [`create`] assembles an +//! [`OpenAICompatibleProvider`] from it. A vendor that has its own package is +//! not a preset: [`crate::create_provider`] builds it through that package. +//! +//! What is evaluated when, following the AI SDK's `loadApiKey` / +//! `loadSetting` timing: +//! +//! - **Creating** a preset validates the explicit settings only: a base URL +//! must be `http(s)` with a host, the `params` keys must be declared by the +//! descriptor and their values must be plain host or path segments. It reads +//! no environment variable; creating any valid row with default settings +//! cannot fail. +//! - **Requests** evaluate the rest, on every request: the API key (the +//! explicit one, or the descriptor's environment variable; unset fails the +//! request with `LoadApiKey`) and the base URL (explicit, then the +//! descriptor's base-URL variable, then the template expanded from +//! explicit parameters, environment and defaults; a missing parameter fails +//! the request with `InvalidArgument` naming it). +//! - A preset with [`AuthMode::None`] (a local server) needs no key and injects +//! no placeholder. It sends `Authorization` only with an explicit key or, +//! when the row names one, a set and nonempty key variable (a runtime +//! started with `--api-key`). + +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, OnceLock}; + +use serde::Deserialize; + +use aimux_core::error::AiMuxError; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; + +use crate::openai_compatible::config::{BaseUrl, ChatDialect}; +use crate::openai_compatible::{ + Assembly, ChatProfile, OpenAICompatibleProvider, TransformRequestBody, +}; +use crate::shared::Credential; + +/// How a preset authenticates. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AuthMode { + /// `Authorization: Bearer `; the key comes from the settings or from + /// the descriptor's environment variable. + ApiKey, + /// No credential (a local server): no key is resolved and no + /// `Authorization` header is sent. + None, +} + +/// A template parameter whose value is derived from another parameter: looked +/// up in `map`, else `otherwise` with `{from}` replaced. +#[derive(Debug, Clone)] +pub struct DeriveSpec { + pub from: String, + pub map: Vec<(String, String)>, + pub otherwise: String, +} + +/// A parameter of a templated base URL (`{name}` in the template). +#[derive(Debug, Clone)] +pub struct ParamSpec { + pub name: String, + /// Environment variables tried in order when the parameter is not given. + pub env: Vec, + pub default: Option, + /// Set for a derived parameter, which the caller cannot give. + pub derive: Option, +} + +/// One row of the embedded provider registry. +#[derive(Debug, Clone)] +pub struct PresetDescriptor { + pub name: String, + pub display: String, + /// Default base URL; a template when `params` is not empty. + pub base_url: String, + /// Environment variable of the API key. Optional for [`AuthMode::None`], + /// where an unset or empty variable sends no credential. + pub env_var: String, + pub auth: AuthMode, + /// Environment variable holding the base URL (local servers). + pub base_url_env: Option, + /// The only max-token key the vendor accepts (`"max_tokens"` or + /// `"max_completion_tokens"`); `None` sends `max_tokens`. + pub max_tokens_key: Option, + pub params: Vec, +} + +/// Settings of a preset factory. Every field is optional. +#[derive(Clone, Default)] +pub struct PresetSettings { + /// The API key. `None` loads the descriptor's environment variable when a + /// request is made (nothing at all for [`AuthMode::None`]). An explicit + /// value is used as given, `""` included, and never falls back to the + /// environment. + pub api_key: Option>, + /// Base URL for the API calls; wins over the descriptor's variable and + /// template. A trailing slash is removed. + pub base_url: Option, + /// Extra headers on every request; a `None` value removes the header. + pub headers: Option, + /// The transport. `None` uses the process default, resolved per request. + pub fetch: Option, + /// Values of the descriptor's template parameters. Only declared, + /// non-derived parameters are accepted; a value must be a plain segment + /// (`A-Z a-z 0-9 . _ -`). + pub params: HashMap, + /// Rewrites every JSON request body once, before it is sent. + pub transform_request_body: Option, +} + +impl std::fmt::Debug for PresetSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PresetSettings") + .field("api_key", &self.api_key) + .field("base_url", &self.base_url) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("fetch", &self.fetch.is_some()) + .field("params", &self.params) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ) + .finish() + } +} + +/// A registry descriptor found by the by-name entry points. +#[derive(Clone, Copy)] +pub struct PresetEntry { + pub descriptor: &'static PresetDescriptor, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RegistryRow { + name: String, + display: String, + base_url: String, + #[serde(default, deserialize_with = "registry_optional")] + env_var: Option, + #[serde(default, deserialize_with = "registry_optional")] + auth: Option, + #[serde(default, deserialize_with = "registry_optional")] + base_url_env: Option, + #[serde(default)] + params: Vec, + #[serde(default)] + profile: HashMap, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RegistryParam { + name: String, + #[serde(default, deserialize_with = "registry_optional")] + env: Option>, + #[serde(default, deserialize_with = "registry_optional")] + default: Option, + #[serde(default, deserialize_with = "registry_optional")] + derive: Option, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RegistryDerive { + from: String, + map: HashMap, + otherwise: String, +} + +// Missing optional fields are allowed; explicit null is not a valid value. +fn registry_optional<'de, D: serde::Deserializer<'de>, T: Deserialize<'de>>( + d: D, +) -> Result, D::Error> { + T::deserialize(d).map(Some) +} + +/// Parse and check the registry rows in `json`. Any row that breaks the +/// schema panics with the row's name and the rule it breaks. +fn parse_registry(json: &str) -> Vec { + let json: Vec = + serde_json::from_str(json).expect("invalid registry: expected an array of rows"); + let mut rows: Vec = json + .iter() + .map(|row| { + RegistryRow::deserialize(row) + .unwrap_or_else(|error| panic!("registry row {}: {error}", row["name"])) + }) + .collect(); + rows.sort_by(|a, b| a.name.cmp(&b.name)); + let identifier = regex::Regex::new(r"^[a-z][a-z0-9_]*$").unwrap(); + let env_name = regex::Regex::new(r"^[A-Z][A-Z0-9_]*$").unwrap(); + let placeholder = regex::Regex::new(r"\{([^{}]*)\}").unwrap(); + let mut seen = HashSet::new(); + let mut descriptors = Vec::with_capacity(rows.len()); + for row in rows { + let name = row.name.clone(); + let check = |valid: bool, message: &str| assert!(valid, "registry row '{name}': {message}"); + check(identifier.is_match(&row.name), "invalid name"); + check(seen.insert(row.name.clone()), "duplicate registry name"); + check(!row.display.is_empty(), "missing display"); + check( + row.base_url.starts_with("http://") || row.base_url.starts_with("https://"), + "base_url must start with http(s)://", + ); + check(row.base_url.is_ascii(), "base_url must be ASCII"); + check( + !row.base_url.contains(['<', '>', '$']), + "base_url contains <, > or $", + ); + let auth = match row.auth.as_deref().unwrap_or("api_key") { + "api_key" => AuthMode::ApiKey, + "none" => AuthMode::None, + value => panic!("registry row '{name}': unknown auth {value:?}"), + }; + check( + match auth { + AuthMode::ApiKey => row.env_var.as_ref().is_some_and(|var| !var.is_empty()), + AuthMode::None => row + .env_var + .as_ref() + .is_none_or(|var| env_name.is_match(var)), + }, + "env_var must be nonempty for api_key and a valid name if present for none", + ); + check( + row.base_url_env + .as_ref() + .is_none_or(|var| env_name.is_match(var)), + "invalid base_url_env name", + ); + check( + row.profile.keys().all(|key| key == "max_tokens_key"), + "unknown profile key", + ); + let max_tokens_key = row.profile.get("max_tokens_key").cloned(); + check( + max_tokens_key + .as_deref() + .is_none_or(|key| matches!(key, "max_tokens" | "max_completion_tokens")), + "invalid max_tokens_key", + ); + let mut params = Vec::with_capacity(row.params.len()); + for param in row.params { + check(identifier.is_match(¶m.name), "bad param name"); + let derive = param.derive.map(|derive| { + check( + param.env.is_none() && param.default.is_none(), + "derived param has env/default", + ); + DeriveSpec { + from: derive.from, + map: derive.map.into_iter().collect(), + otherwise: derive.otherwise, + } + }); + params.push(ParamSpec { + name: param.name, + env: param.env.unwrap_or_default(), + default: param.default, + derive, + }); + } + let names: HashSet<&str> = params.iter().map(|param| param.name.as_str()).collect(); + check(names.len() == params.len(), "duplicate param names"); + for param in ¶ms { + if let Some(derive) = ¶m.derive { + check( + params + .iter() + .any(|source| source.name == derive.from && source.derive.is_none()), + "derive.from must name a plain param", + ); + } + } + let placeholders: HashSet<&str> = placeholder + .captures_iter(&row.base_url) + .map(|c| c.get(1).unwrap().as_str()) + .collect(); + check(placeholders == names, "URL placeholders differ from params"); + descriptors.push(PresetDescriptor { + name: row.name, + display: row.display, + base_url: row.base_url, + env_var: row.env_var.unwrap_or_default(), + auth, + base_url_env: row.base_url_env, + max_tokens_key, + params, + }); + } + descriptors +} + +/// Create the provider of the registry row `name`: the row's base URL, key +/// variable and dialect, with `settings` overriding individual fields. The +/// result is an ordinary [`OpenAICompatibleProvider`]; nothing is read from +/// the environment here. +/// +/// # Errors +/// +/// [`AiMuxError::NoSuchProvider`] for a name that is not in the registry, +/// `InvalidArgument` for an unusable explicit `base_url`, an undeclared (or +/// derived) `params` key, or a `params` value that is not a plain segment. +pub fn create( + name: &str, + settings: PresetSettings, +) -> Result { + let entry = lookup(name).ok_or_else(|| AiMuxError::NoSuchProvider { + provider_id: name.to_string(), + model_id: name.to_string(), + model_type: String::new(), + available_providers: entries() + .map(|entry| entry.descriptor.name.to_string()) + .collect(), + })?; + assemble(entry.descriptor, settings) +} + +/// Every preset, in name order. Invalid registry data panics on first use. +/// +/// The registry is parsed once into a process-wide table; the entries borrow +/// from it. +pub fn entries() -> impl Iterator { + static DESCRIPTORS: OnceLock> = OnceLock::new(); + static TABLE: OnceLock> = OnceLock::new(); + TABLE + .get_or_init(|| { + DESCRIPTORS + .get_or_init(|| parse_registry(include_str!("provider_registry.json"))) + .iter() + .map(|descriptor| PresetEntry { descriptor }) + .collect() + }) + .iter() +} + +/// The names of every preset, in order. +pub fn names() -> impl Iterator { + entries().map(|entry| entry.descriptor.name.as_str()) +} + +/// The preset called `name`, or `None`; unknown names never fall back. +#[must_use] +pub fn lookup(name: &str) -> Option<&'static PresetEntry> { + entries().find(|entry| entry.descriptor.name == name) +} + +fn assemble( + descriptor: &'static PresetDescriptor, + settings: PresetSettings, +) -> Result { + for (name, value) in &settings.params { + let declared = descriptor + .params + .iter() + .find(|spec| &spec.name == name && spec.derive.is_none()); + if declared.is_none() { + return Err(AiMuxError::InvalidArgument(format!( + "preset '{}' has no template parameter `{name}`; declared parameters: [{}]", + descriptor.name, + settable_params(descriptor).join(", ") + ))); + } + check_param_value(descriptor, name, value)?; + } + + let base_url = match settings.base_url.as_deref() { + Some(url) => BaseUrl::Fixed(validate_base_url(url)?), + None if descriptor.params.is_empty() && descriptor.base_url_env.is_none() => { + BaseUrl::Fixed(validate_base_url(&descriptor.base_url)?) + } + None => { + let explicit = settings.params; + BaseUrl::Lazy(Arc::new(move || resolve_base_url(descriptor, &explicit))) + } + }; + + let credential = match (descriptor.auth, settings.api_key) { + (_, Some(key)) => Credential::Explicit(key), + (AuthMode::ApiKey, None) => Credential::Env { + var: descriptor.env_var.clone(), + description: descriptor.display.clone(), + }, + (AuthMode::None, None) if descriptor.env_var.is_empty() => Credential::None, + (AuthMode::None, None) => Credential::OptionalEnv(descriptor.env_var.clone()), + }; + + let mut dialect = ChatDialect::baseline(); + dialect.supports_top_k = true; + dialect.max_tokens_key = descriptor.max_tokens_key.as_deref(); + 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 { + name: descriptor.name.clone(), + base_url, + credential, + fixed_headers: Vec::new(), + headers: settings.headers, + query_params: None, + fetch: settings.fetch, + transform_request_body: settings.transform_request_body, + profile, + }) +} + +// ── Base URL resolution ────────────────────────────────────────────────────── + +fn settable_params(descriptor: &PresetDescriptor) -> Vec<&str> { + descriptor + .params + .iter() + .filter(|spec| spec.derive.is_none()) + .map(|spec| spec.name.as_str()) + .collect() +} + +/// A template parameter value must be a plain host or path segment: letters, +/// digits, `.`, `_` and `-`. That keeps a value from adding a path, user +/// information, a port or a query to the URL it is expanded into. +fn check_param_value( + descriptor: &PresetDescriptor, + name: &str, + value: &str, +) -> Result<(), AiMuxError> { + let plain = !value.is_empty() + && value.len() <= 253 + && value.chars().any(|c| c != '.') + && value + .chars() + .all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-')); + if plain { + Ok(()) + } else { + Err(AiMuxError::InvalidArgument(format!( + "preset '{}': template parameter `{name}` must be a plain segment \ + (letters, digits, `.`, `_`, `-`; no `/`, `@`, `:`, `?`, `#` or spaces), got {value:?}", + descriptor.name + ))) + } +} + +fn env_value(names: &[impl AsRef]) -> Option { + names.iter().find_map(|name| { + std::env::var(name.as_ref()) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + }) +} + +/// The base URL a request goes to: the descriptor's base-URL variable, else +/// the template expanded from `explicit`, the environment and the defaults. +fn resolve_base_url( + descriptor: &PresetDescriptor, + explicit: &HashMap, +) -> Result { + if let Some(var) = descriptor.base_url_env.as_deref() + && let Some(url) = env_value(&[var]) + { + return validate_base_url(&url) + .map_err(|e| AiMuxError::InvalidArgument(format!("{var}: {e}"))); + } + + let mut values: HashMap<&str, String> = HashMap::new(); + for spec in descriptor + .params + .iter() + .filter(|spec| spec.derive.is_none()) + { + let value = explicit + .get(&spec.name) + .cloned() + .or_else(|| env_value(&spec.env)) + .or_else(|| spec.default.clone()) + .ok_or_else(|| { + let source = if spec.env.is_empty() { + String::new() + } else { + format!(" or set {}", spec.env.join(" / ")) + }; + AiMuxError::InvalidArgument(format!( + "preset '{}' needs the template parameter `{}`: pass it in `params`{source}, \ + or pass a concrete `base_url`", + descriptor.name, spec.name + )) + })?; + check_param_value(descriptor, &spec.name, &value)?; + values.insert(&spec.name, value); + } + for spec in descriptor.params.iter() { + let Some(derive) = &spec.derive else { continue }; + let from = values.get(derive.from.as_str()).cloned().ok_or_else(|| { + AiMuxError::InvalidArgument(format!( + "preset '{}': derived parameter `{}` depends on undeclared `{}`", + descriptor.name, spec.name, derive.from + )) + })?; + let derived = derive + .map + .iter() + .find(|(key, _)| *key == from) + .map(|(_, host)| host.clone()) + .unwrap_or_else(|| { + derive + .otherwise + .replace(&format!("{{{}}}", derive.from), &from) + }); + values.insert(&spec.name, derived); + } + + let mut url = descriptor.base_url.clone(); + for (name, value) in &values { + url = url.replace(&format!("{{{name}}}"), value); + } + if url.chars().any(|c| matches!(c, '{' | '}' | '<' | '>')) { + return Err(AiMuxError::InvalidArgument(format!( + "preset '{}': base URL {:?} still has an unexpanded placeholder after expanding its \ + parameters", + descriptor.name, url + ))); + } + validate_base_url(&url) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The registry is parsed once: every lookup hands out the same + /// descriptor, owned by the process-wide table (nothing is leaked per + /// call). + #[test] + fn the_table_is_parsed_once_and_shared() { + let first = lookup("abacus").expect("abacus is a registry row"); + let again = lookup("abacus").expect("abacus is a registry row"); + assert!(std::ptr::eq(first.descriptor, again.descriptor)); + let from_entries = entries() + .find(|entry| entry.descriptor.name == "abacus") + .expect("abacus is listed"); + assert!(std::ptr::eq(first.descriptor, from_entries.descriptor)); + assert_eq!(names().count(), entries().count()); + } + + fn registry(rows: serde_json::Value) -> Vec { + parse_registry(&rows.to_string()) + } + + fn row(fields: serde_json::Value) -> serde_json::Value { + let mut row = serde_json::json!({ + "name": "vendor", + "display": "Vendor", + "base_url": "https://api.vendor.example/v1", + "env_var": "VENDOR_API_KEY", + }); + row.as_object_mut() + .unwrap() + .extend(fields.as_object().cloned().unwrap()); + row + } + + #[test] + fn a_valid_row_parses_into_a_descriptor() { + let descriptors = registry(serde_json::json!([ + row(serde_json::json!({})), + row(serde_json::json!({ + "name": "local", + "base_url": "http://localhost:{port}/v1", + "auth": "none", + "env_var": "LOCAL_API_KEY", + "base_url_env": "LOCAL_BASE_URL", + "params": [{ "name": "port", "default": "8080" }], + "profile": { "max_tokens_key": "max_completion_tokens" }, + })), + ])); + assert_eq!(descriptors.len(), 2); + let local = &descriptors[0]; + assert_eq!(local.name, "local"); + assert_eq!(local.auth, AuthMode::None); + assert_eq!(local.base_url_env.as_deref(), Some("LOCAL_BASE_URL")); + assert_eq!( + local.max_tokens_key.as_deref(), + Some("max_completion_tokens") + ); + assert_eq!(local.params[0].default.as_deref(), Some("8080")); + let vendor = &descriptors[1]; + assert_eq!(vendor.auth, AuthMode::ApiKey); + assert_eq!(vendor.env_var, "VENDOR_API_KEY"); + } + + #[test] + #[should_panic(expected = "registry row 'vendor': duplicate registry name")] + fn a_name_is_listed_once() { + registry(serde_json::json!([ + row(serde_json::json!({})), + row(serde_json::json!({})) + ])); + } + + #[test] + #[should_panic(expected = "registry row 'Vendor-1': invalid name")] + fn a_name_is_a_lowercase_identifier() { + registry(serde_json::json!([row( + serde_json::json!({ "name": "Vendor-1" }) + )])); + } + + #[test] + #[should_panic(expected = "base_url must start with http(s)://")] + fn a_base_url_is_http() { + registry(serde_json::json!([row( + serde_json::json!({ "base_url": "ftp://vendor" }) + )])); + } + + #[test] + #[should_panic(expected = "unknown auth \"oauth\"")] + fn auth_is_api_key_or_none() { + registry(serde_json::json!([row( + serde_json::json!({ "auth": "oauth" }) + )])); + } + + #[test] + #[should_panic(expected = "env_var must be nonempty for api_key")] + fn an_api_key_row_names_its_variable() { + registry(serde_json::json!([row( + serde_json::json!({ "env_var": "" }) + )])); + } + + #[test] + #[should_panic(expected = "URL placeholders differ from params")] + fn every_placeholder_is_a_param() { + registry(serde_json::json!([row(serde_json::json!({ + "base_url": "https://{region}.vendor.example/v1", + }))])); + } + + #[test] + #[should_panic(expected = "derived param has env/default")] + fn a_derived_param_has_no_env_or_default() { + registry(serde_json::json!([row(serde_json::json!({ + "base_url": "https://{region}.{host}/v1", + "params": [ + { "name": "region", "default": "eu" }, + { + "name": "host", + "default": "x", + "derive": { "from": "region", "map": {}, "otherwise": "vendor.example" }, + }, + ], + }))])); + } + + #[test] + #[should_panic(expected = "registry row \"vendor\"")] + fn an_explicit_null_is_not_a_value() { + registry(serde_json::json!([row( + serde_json::json!({ "env_var": null }) + )])); + } + + #[test] + #[should_panic(expected = "registry row \"vendor\"")] + fn an_unknown_field_is_refused() { + registry(serde_json::json!([row( + serde_json::json!({ "region": "eu" }) + )])); + } +} diff --git a/aimux-providers/src/prodia.rs b/aimux-providers/src/prodia.rs index b4abeb16e..c66f35e3f 100644 --- a/aimux-providers/src/prodia.rs +++ b/aimux-providers/src/prodia.rs @@ -5,8 +5,13 @@ //! //! POST to `/job?price=true`, returns a multipart response with a JSON "job" //! part and a binary "output" image part. +//! +//! [`create_prodia`] takes [`ProdiaProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`ProdiaProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `PRODIA_API_KEY`. +//! [`prodia()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -16,7 +21,9 @@ use aimux_core::image_model::{ ImageCallOptions, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; use aimux_core::shared::Warning; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; + +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; fn prodia_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -47,94 +54,176 @@ fn prodia_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, +pub(crate) mod options; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://prodia.com/api"; +pub(crate) const API_KEY_ENV_VAR: &str = "PRODIA_API_KEY"; +const DEFAULT_NAME: &str = "prodia"; + +/// Settings of [`create_prodia`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct ProdiaProviderSettings { + /// Base URL for the API calls. Default `https://prodia.com/api`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `PRODIA_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.image"`, `"{name}.video"`). + /// Default `"prodia"`. The providerOptions key stays `prodia`. + 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, } -impl ProdiaConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://prodia.com/api".to_string(), - headers: None, - } - } - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } - /// Create from the `PRODIA_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "PRODIA_API_KEY", "Prodia")?; - Ok(Self::new(api_key)) +impl std::fmt::Debug for ProdiaProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ProdiaProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } +/// Create a Prodia provider. +/// +/// # 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. +pub fn create_prodia(settings: ProdiaProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(ProdiaProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + credential: Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Prodia"), + user_headers: settings.headers, + fetch: settings.fetch, + }) +} + +/// The default provider: `create_prodia` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn prodia() -> &'static ProdiaProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_prodia(ProdiaProviderSettings::default()) + .expect("default Prodia settings are always valid") + }) +} + +/// A Prodia provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct ProdiaProvider { - config: ProdiaConfig, + name: String, + base_url: String, + credential: Credential, + user_headers: Option, + fetch: Option, } + impl ProdiaProvider { - #[must_use] - pub fn new(config: ProdiaConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + match method { + "image" => credential_headers( + self.credential.clone(), + AuthScheme::Header("X-Prodia-Key"), + vec![( + "Accept".to_string(), + "multipart/form-data; image/png".to_string(), + )], + self.user_headers.clone(), + ), + _ => credential_headers( + self.credential.clone(), + AuthScheme::Header("X-Prodia-Key"), + Vec::new(), + self.user_headers.clone(), + ), + }, + self.fetch.clone(), + None, + ) } + + /// An image model (e.g. `"inference.flux-fast.schnell.txt2img.v2"`); `provider()` is `"{name}.image"`. #[must_use] pub fn image(&self, model_id: &str) -> ProdiaImageModel { - ProdiaImageModel::new(model_id.to_string(), self.config.clone()) + ProdiaImageModel::from_config(model_id.to_string(), self.model_config("image")) } - /// Create a video generation model instance. + + /// A video model; `provider()` is `"{name}.video"`. #[must_use] pub fn video(&self, model_id: &str) -> ProdiaVideoModel { - ProdiaVideoModel::new(model_id.to_string(), self.config.clone()) + ProdiaVideoModel::from_config(model_id.to_string(), self.model_config("video")) + } +} + +impl ::aimux_core::Provider for ProdiaProvider { + fn language_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "languageModel")) + } + + fn embedding_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Ok(::std::sync::Arc::new(self.image(model_id))) + } + + fn video_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(::std::sync::Arc::new(self.video(model_id)))) } } -/// A Prodia image generation model. pub struct ProdiaImageModel { model_id: String, - config: ProdiaConfig, + config: EndpointConfig, } + impl ProdiaImageModel { - #[must_use] - pub fn new(model_id: String, config: ProdiaConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut h = HashMap::new(); - h.insert("X-Prodia-Key".into(), self.config.api_key.clone()); - h.insert("Accept".into(), "multipart/form-data; image/png".into()); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - h.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - h.insert(k.clone(), v.clone()); - } - } - h - } - - fn endpoint(&self) -> String { - format!("{}/job?price=true", self.config.base_url) - } } /// Parse a multipart body into parts (name, content_type, body). @@ -204,7 +293,7 @@ fn parse_multipart(body: &[u8], boundary: &str) -> Vec<(String, String, Vec) #[async_trait] impl ImageModel for ProdiaImageModel { fn provider(&self) -> &str { - "prodia" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id @@ -216,7 +305,7 @@ impl ImageModel for ProdiaImageModel { async fn do_generate(&self, options: &ImageCallOptions) -> Result { let warnings: Vec = Vec::new(); - let prodia_opts = options.provider_options.get("prodia"); + let prodia_opts = options::prodia_options(Some(&options.provider_options)); // Build job config let mut job_config = Map::new(); @@ -254,11 +343,10 @@ impl ImageModel for ProdiaImageModel { let body = json!({ "type": self.model_id, "config": job_config }); - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), header_list, options), + exchange.request(exchange.url("/job?price=true"), options), body, aimux_provider_utils::create_binary_response_handler(), prodia_failed_response_handler(), @@ -310,7 +398,7 @@ impl ImageModel for ProdiaImageModel { let mut metadata = aimux_core::shared::NamespaceMap::new(); let mut prodia_meta = Map::new(); prodia_meta.insert("images".into(), json!([job_result])); - metadata.insert("prodia".into(), prodia_meta); + metadata.insert(options::NAMESPACE.into(), prodia_meta); Ok(ImageResult { images: ImageOutputs::Binary(vec![image_bytes]), @@ -341,36 +429,19 @@ use aimux_core::video_model::{ /// (`reference/ai/packages/prodia/src/prodia-video-model.ts`). pub struct ProdiaVideoModel { model_id: String, - config: ProdiaConfig, + config: EndpointConfig, } impl ProdiaVideoModel { - #[must_use] - pub fn new(model_id: String, config: ProdiaConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut h = HashMap::new(); - h.insert("X-Prodia-Key".to_string(), self.config.api_key.clone()); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - h.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - h.insert(k.clone(), v.clone()); - } - } - h - } } #[async_trait] impl VideoModel for ProdiaVideoModel { fn provider(&self) -> &str { - "prodia" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id @@ -395,16 +466,11 @@ impl VideoModel for ProdiaVideoModel { let body = json!({"type": self.model_id, "config": Value::Object(config_obj)}); - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; // Submit job. let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new( - format!("{}/job", self.config.base_url), - header_list, - options, - ), + exchange.request(exchange.url("/job"), options), body, aimux_provider_utils::create_json_response_handler(), prodia_failed_response_handler(), @@ -449,12 +515,11 @@ impl VideoModel for ProdiaVideoModel { ) })?; - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); - let poll_url = format!("{}/job/{}", self.config.base_url, job_id); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let poll_url = exchange.url(&format!("/job/{job_id}")); let resp = aimux_provider_utils::get_from_api( - HttpRequest::new(poll_url.clone(), header_list, options), + exchange.request(poll_url.clone(), options), aimux_provider_utils::create_json_response_handler::(), prodia_failed_response_handler(), ) diff --git a/aimux-providers/src/prodia/options.rs b/aimux-providers/src/prodia/options.rs new file mode 100644 index 000000000..2ea5b4e67 --- /dev/null +++ b/aimux-providers/src/prodia/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Prodia. +//! +//! Prodia reads `providerOptions.prodia`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "prodia"; + +/// The Prodia options in a providerOptions map. +pub(crate) fn prodia_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/provider.rs b/aimux-providers/src/provider.rs index a9eefdd5c..d3a446239 100644 --- a/aimux-providers/src/provider.rs +++ b/aimux-providers/src/provider.rs @@ -1,23 +1,16 @@ -//! Registry-backed provider construction (RFC-0017 phase 4). +//! By-name creation for the bindings. //! -//! Single source of truth: [`provider_registry.json`](provider_registry.json), -//! edited by hand (one row per provider). All built-in -//! OpenAI-compatible providers are looked up by name at runtime — there are no -//! per-provider `XxxConfig`/`XxxProvider` types anymore (retired in phase 4). -//! -//! ``` -//! use aimux_providers::{provider, ProviderOptions}; -//! -//! # fn smoke() -> Result<(), aimux_core::error::AiMuxError> { -//! // Key from env var (GROQ_API_KEY), base URL & profile from the registry. -//! let model = provider("groq", None, "llama-3.3-70b", None)?; -//! # let _ = model; -//! # Ok(()) -//! # } -//! ``` +//! The Rust surface for creating providers is the factories, +//! [`create_provider`](crate::create_provider) and the registry. The FFI and +//! the Node / Python bindings still expose an older shape (a provider name +//! plus a JSON options object, and providers registered at runtime from a +//! JSON document); this module is that shape, implemented on +//! `create_provider`. It is not re-exported from the crate root and is to be +//! removed in #265 (the ops layer), where the C ABI and the bindings stop +//! calling it. use std::collections::HashMap; -use std::sync::{OnceLock, RwLock}; +use std::sync::{Arc, OnceLock, RwLock}; use serde::Deserialize; use serde_json::Value; @@ -25,27 +18,17 @@ use serde_json::Value; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; +use aimux_provider_utils::{Resolvable, load_api_key, validate_base_url}; -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIProvider}; +use crate::openai_compatible::config::{BaseUrl, ChatDialect}; +use crate::openai_compatible::{Assembly, ChatProfile, OpenAICompatibleProvider}; +use crate::preset::PresetSettings; +use crate::shared::Credential; -/// One entry of `provider-registry.json` (registry slice). -#[derive(Debug, Clone, Deserialize)] -struct RegistryEntry { - name: String, - display: String, - base_url: String, - env_var: String, - #[serde(default)] - profile: ProviderProfile, -} - -/// Provider capability profile, expressed as data (non-default fields only -/// in the registry JSON). -/// -/// Shared between the built-in registry and the runtime overlay layer -/// (RFC-0020). The `Default` impl and the per-field -/// `#[serde(default = "default_true")]` both yield `true` for the three -/// `supports_*` flags, matching the OpenAI-compatible baseline — so a +/// Chat capabilities of an external provider entry, expressed as data +/// (RFC-0020). The `Default` impl and the per-field `#[serde(default = +/// "default_true")]` both yield `true` for the three `supports_*` flags, +/// matching the OpenAI-compatible baseline of the registry presets - so a /// profile that is omitted entirely (or present but with individual fields /// missing) stays at full capability. #[derive(Debug, Clone, serde::Deserialize, serde::Serialize)] @@ -56,7 +39,10 @@ pub struct ProviderProfile { pub supports_tools: bool, #[serde(default = "default_true")] pub supports_response_format: bool, + /// Streaming usage rides in `chunk[key].usage` instead of `chunk.usage`. pub stream_usage_key: Option, + /// `"max_tokens"` or `"max_completion_tokens"`: the only max-token key the + /// endpoint accepts. Any other value is rejected at registration. pub max_tokens_key: Option, } @@ -76,49 +62,46 @@ fn default_true() -> bool { true } -static REGISTRY: OnceLock> = OnceLock::new(); - -/// Load + validate the embedded registry exactly once. -fn registry() -> &'static [RegistryEntry] { - REGISTRY.get_or_init(|| { - let raw = include_str!("provider_registry.json"); - let entries: Vec = serde_json::from_str(raw) - .unwrap_or_else(|e| panic!("provider_registry.json is invalid: {e}")); - for entry in &entries { - assert!(!entry.name.is_empty(), "registry entry missing name"); - assert!( - !entry.display.is_empty(), - "registry entry '{}' missing display", - entry.name - ); - assert!( - !entry.base_url.is_empty(), - "registry entry '{}' missing base_url", - entry.name - ); - assert!( - !entry.env_var.is_empty(), - "registry entry '{}' missing env_var", - entry.name - ); - } - entries - }) -} - -/// Detect unexpanded placeholder syntax in a registry `base_url`. -/// -/// Templated entries (cloudflare, neon, snowflake, oci, ...) carry account- or -/// region-scoped placeholders that the caller must fill in via -/// [`ProviderOptions::base_url`]. Recognized forms: `{VAR}` (covers `${var}` -/// too) and ``. If a new placeholder form is introduced in the -/// registry, extend this function. -fn base_url_has_placeholder(base_url: &str) -> bool { - base_url.contains('{') || base_url.contains('<') +impl ProviderProfile { + /// The chat behavior this profile describes. + fn chat_profile(&self) -> Result { + let max_tokens_key = match self.max_tokens_key.as_deref() { + None => None, + Some("max_tokens") => Some("max_tokens"), + Some("max_completion_tokens") => Some("max_completion_tokens"), + Some(other) => { + return Err(AiMuxError::InvalidArgument(format!( + "profile.max_tokens_key must be \"max_tokens\" or \"max_completion_tokens\", \ + got {other:?}" + ))); + } + }; + let mut dialect = ChatDialect::baseline(); + dialect.supports_top_k = self.supports_top_k; + dialect.supports_tools = self.supports_tools; + dialect.supports_response_format = self.supports_response_format; + dialect.stream_usage_key = self.stream_usage_key.clone(); + dialect.max_tokens_key = max_tokens_key; + Ok(ChatProfile { + include_usage: true, + supports_structured_outputs: true, + supports_multi_part_tool_content: false, + dialect, + supported_urls: None, + convert_usage: None, + }) + } } /// Per-call construction options for [`provider`] (overrides the registry entry). +/// +/// Deserializing a JSON object that carries `max_retries` or `body_overrides` +/// fails: retry is a call-level option (`CallOptions::max_retries`) and +/// request-body overrides are gone, so neither is silently ignored here. The +/// FFI `config_json` and the Node/Python configs reach this check and report it +/// as `InvalidArgument`. #[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] +#[serde(try_from = "ProviderOptionsWire")] pub struct ProviderOptions { /// Override the registry base URL. pub base_url: Option, @@ -128,36 +111,84 @@ pub struct ProviderOptions { pub organization: Option, /// OpenAI project ID (`OpenAI-Project` header). pub project: Option, - /// Retry count override; `Some(0)` disables retries. - pub max_retries: Option, - /// Request-body overrides (deep-merged; RFC-0017 phase 1). - pub body_overrides: Option, + /// Values of the preset's template parameters (`account_id`, `region`, + /// `project`, ...); only the ones the registry row declares are accepted. + pub params: Option>, } -/// Build a language model for a provider by name. +/// Wire shape of [`ProviderOptions`]: the same keys plus the two removed ones, +/// kept only so their presence can be reported instead of dropped. +#[derive(Deserialize)] +struct ProviderOptionsWire { + base_url: Option, + headers: Option>, + organization: Option, + project: Option, + params: Option>, + max_retries: Option, + body_overrides: Option, +} + +impl TryFrom for ProviderOptions { + type Error = String; + + fn try_from(wire: ProviderOptionsWire) -> Result { + if wire.max_retries.is_some() { + return Err(MAX_RETRIES_REJECTED.to_string()); + } + if wire.body_overrides.is_some() { + return Err(BODY_OVERRIDES_REJECTED.to_string()); + } + Ok(Self { + base_url: wire.base_url, + headers: wire.headers, + organization: wire.organization, + project: wire.project, + params: wire.params, + }) + } +} + +const MAX_RETRIES_REJECTED: &str = "`max_retries` is a call-level option: pass it in the call \ + options (`maxRetries`), not in the provider configuration"; + +const BODY_OVERRIDES_REJECTED: &str = "`body_overrides` is no longer supported: request-body \ + overrides were removed from the provider configuration and the call options"; + +/// Report a provider configuration that still carries `max_retries` or +/// `body_overrides`, for configuration shapes that are not [`ProviderOptions`] +/// JSON (the Node and Python `ProviderConfig`). Pass whether each key was +/// given. /// -/// Lookup order: runtime overlay (RFC-0020 [`register_provider`]) → built-in -/// registry → [`AiMuxError::NoSuchProvider`]. +/// # Errors /// -/// - `api_key = None` reads the provider's env var from the registry entry -/// (or the external entry's `env_var` / `api_key` field). -/// - `options` overrides individual fields of the resolved entry -/// (replaces the retired `with_base_url` etc.). -/// - Unknown names return [`AiMuxError::NoSuchProvider`] naming the requested -/// provider; built-in names are available through [`provider_names`] -/// (overlay-registered names are not). +/// Returns [`AiMuxError::InvalidArgument`] naming the first removed key that +/// was given. +pub fn reject_removed_provider_options( + max_retries: bool, + body_overrides: bool, +) -> Result<(), AiMuxError> { + if max_retries { + return Err(AiMuxError::InvalidArgument(MAX_RETRIES_REJECTED.into())); + } + if body_overrides { + return Err(AiMuxError::InvalidArgument(BODY_OVERRIDES_REJECTED.into())); + } + Ok(()) +} + +/// A language model by provider name: [`provider_handle`] then +/// `language_model(model_id)`. /// /// # Errors /// -/// Returns [`AiMuxError::NoSuchProvider`] for unknown names, -/// `InvalidArgument` for invalid entries/options, and key-resolution errors -/// from the registry env var. +/// As [`provider_handle`], plus whatever the provider returns for the id. pub fn provider( name: impl AsRef, api_key: Option, model_id: &str, options: Option, -) -> Result, AiMuxError> { +) -> Result, AiMuxError> { let p = provider_handle(name, api_key, options)?; p.language_model(model_id) } @@ -168,10 +199,12 @@ pub fn provider( /// [`load_providers_from_json`]. Overrides a same-named built-in entry (whole /// replacement, not deep merge) or adds a new one. /// -/// Only OpenAI-compatible providers can be registered this way — native -/// protocols (anthropic/google/bedrock…) are code implementations and cannot -/// be described by config data. -#[derive(Debug, Clone, Deserialize)] +/// Only OpenAI-compatible providers can be registered this way - native +/// protocols (anthropic/google/bedrock...) are code implementations and cannot +/// be described by config data. The JSON form accepts exactly these fields; +/// `max_retries` and `body_overrides` (removed) and any unknown field are +/// `InvalidArgument`. +#[derive(Clone, Deserialize)] pub struct ExternalProviderEntry { /// Provider name used for `provider("name", ...)` lookup. Required. pub name: String, @@ -187,7 +220,7 @@ pub struct ExternalProviderEntry { /// Protocol kind. Only `"openai_compat"` is accepted; other values error. #[serde(default = "default_openai_compat")] pub protocol: String, - /// Provider capability profile. All fields optional, defaults to full(). + /// Provider capability profile. All fields optional, defaults to full. #[serde(default)] pub profile: ProviderProfile, // --- Fields equivalent to ProviderOptions (provider-level config) --- @@ -197,14 +230,28 @@ pub struct ExternalProviderEntry { pub organization: Option, /// OpenAI project ID (`OpenAI-Project` header). pub project: Option, - /// Retry count override; `Some(0)` disables retries. - pub max_retries: Option, - /// Request-body overrides (deep-merged; RFC-0017 phase 1). - pub body_overrides: Option, /// Free-form note for the user; the library ignores this. pub comment: Option, } +impl std::fmt::Debug for ExternalProviderEntry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ExternalProviderEntry") + .field("name", &self.name) + .field("display", &self.display) + .field("base_url", &self.base_url) + .field("env_var", &self.env_var) + .field("api_key", &self.api_key.is_some()) + .field("protocol", &self.protocol) + .field("profile", &self.profile) + .field("headers", &self.headers.as_ref().map(HashMap::len)) + .field("organization", &self.organization) + .field("project", &self.project) + .field("comment", &self.comment) + .finish() + } +} + fn default_openai_compat() -> String { "openai_compat".to_string() } @@ -214,22 +261,32 @@ struct ProvidersConfig { providers: Vec, } +/// The keys of an external entry's JSON object. +const EXTERNAL_ENTRY_KEYS: [&str; 11] = [ + "name", + "display", + "base_url", + "env_var", + "api_key", + "protocol", + "profile", + "headers", + "organization", + "project", + "comment", +]; + static OVERLAYS: OnceLock>> = OnceLock::new(); fn overlays() -> &'static RwLock> { OVERLAYS.get_or_init(|| RwLock::new(HashMap::new())) } -/// Test-only helper: remove a name from the overlay so tests are hermetic. -#[cfg(test)] -pub(crate) fn clear_overlay(name: &str) { - overlays().write().unwrap().remove(name); -} - /// Register (or replace) an external provider entry. /// -/// Validation failures (empty name, non-`http(s)://` base_url, unsupported -/// protocol) return [`AiMuxError::InvalidArgument`] — they never panic. +/// Validation failures (empty name, non-`http(s)://` base_url, templated +/// base_url, unsupported protocol, unusable profile) return +/// [`AiMuxError::InvalidArgument`] - they never panic. /// /// # Errors /// @@ -255,10 +312,38 @@ pub fn is_external_provider(name: &str) -> bool { /// /// # Errors /// -/// Returns `AiMuxError::JsonParse` for malformed JSON and propagates each +/// Returns `AiMuxError::JsonParse` for malformed JSON or a document of the +/// wrong shape, `InvalidArgument` for an entry carrying a removed +/// (`max_retries`, `body_overrides`) or unknown field, and propagates each /// entry's `register_provider` validation error. pub fn load_providers_from_json(json: &str) -> Result<(), AiMuxError> { - let config: ProvidersConfig = serde_json::from_str(json).map_err(|e| { + let value: Value = serde_json::from_str(json).map_err(|e| { + AiMuxError::JsonParse(format!("failed to parse external providers config: {e}")) + })?; + if let Some(entries) = value.get("providers").and_then(Value::as_array) { + for entry in entries.iter().filter_map(Value::as_object) { + let name = entry + .get("name") + .and_then(Value::as_str) + .unwrap_or(""); + reject_removed_provider_options( + entry.contains_key("max_retries"), + entry.contains_key("body_overrides"), + ) + .map_err(|e| AiMuxError::InvalidArgument(format!("external provider '{name}': {e}")))?; + if let Some(unknown) = entry + .keys() + .find(|key| !EXTERNAL_ENTRY_KEYS.contains(&key.as_str())) + { + return Err(AiMuxError::InvalidArgument(format!( + "external provider '{name}' has an unknown field `{unknown}`; accepted \ + fields: {}", + EXTERNAL_ENTRY_KEYS.join(", ") + ))); + } + } + } + let config: ProvidersConfig = serde_json::from_value(value).map_err(|e| { AiMuxError::JsonParse(format!("failed to parse external providers config: {e}")) })?; for entry in config.providers { @@ -286,193 +371,131 @@ fn validate_external_entry(entry: &ExternalProviderEntry) -> Result<(), AiMuxErr entry.name, entry.base_url ))); } + if entry.base_url.contains(['{', '}', '<', '>']) { + return Err(AiMuxError::InvalidArgument(format!( + "external provider '{}' base_url {:?} has a placeholder; an external entry takes a \ + concrete URL", + entry.name, entry.base_url + ))); + } if entry.protocol != "openai_compat" { return Err(AiMuxError::InvalidArgument(format!( "external provider '{}' has unsupported protocol {:?}; only \"openai_compat\" is supported", entry.name, entry.protocol ))); } + entry.profile.chat_profile().map_err(|e| { + AiMuxError::InvalidArgument(format!("external provider '{}': {e}", entry.name)) + })?; Ok(()) } -// ── Unified lookup (built-in registry + overlay layer) ────────────────────── - -/// A resolved provider entry — the common shape both the built-in registry -/// and the runtime overlay produce, so the downstream construction logic -/// (placeholder check, key resolution, config build) is shared. -struct ResolvedEntry { - name: String, - display: String, - base_url: String, - env_var: String, - profile: ProviderProfile, - /// Entry-level api_key (external entries only). `"env:VAR"` references are - /// resolved at lookup time; literal strings are used as-is. Built-in - /// registry entries always have `None` here. - api_key: Option, - /// Provider-level config carried by external entries (None for built-ins, - /// which rely on ProviderOptions for per-call overrides). - headers: Option>, - organization: Option, - project: Option, - max_retries: Option, - body_overrides: Option, -} - -impl ResolvedEntry { - /// Resolve a built-in [`RegistryEntry`] into the common shape. - fn from_registry(entry: &RegistryEntry) -> Self { - Self { - name: entry.name.clone(), - display: entry.display.clone(), - base_url: entry.base_url.clone(), - env_var: entry.env_var.clone(), - profile: entry.profile.clone(), - api_key: None, - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - } - } +// ── By-name creation for the bindings ──────────────────────────────────────── - /// Resolve an [`ExternalProviderEntry`] into the common shape. - fn from_external(entry: &ExternalProviderEntry) -> Self { - Self { - name: entry.name.clone(), - display: entry.display.clone().unwrap_or_else(|| entry.name.clone()), - base_url: entry.base_url.clone(), - env_var: entry.env_var.clone().unwrap_or_default(), - profile: entry.profile.clone(), - api_key: entry.api_key.clone(), - headers: entry.headers.clone(), - organization: entry.organization.clone(), - project: entry.project.clone(), - max_retries: entry.max_retries, - body_overrides: entry.body_overrides.clone(), - } - } -} - -/// Build a **provider handle** for a built-in or externally-registered provider -/// by name (RFC-0027 + RFC-0020 overlay). -/// -/// Lookup order: runtime overlay (RFC-0020) → native OpenAI / built-in registry -/// → NoSuchProvider. -/// -/// Unlike [`provider`] (which binds to a single `model_id` and returns a -/// `LanguageModel`), this returns the [`Provider`] itself, so callers can call -/// [`Provider::list_models`] for runtime discovery, then -/// [`Provider::language_model`] or [`Provider::decision_model`] on a chosen id. -/// -/// Same key/options semantics as [`provider`]. +/// Create a provider by name for a binding: a runtime-registered external +/// entry first, otherwise [`create_provider`](crate::create_provider) (the +/// vendor packages and the registry rows). /// /// # Errors /// -/// Returns [`AiMuxError::NoSuchProvider`] for unknown names, `InvalidArgument` -/// for an unexpanded templated base URL or invalid options, and key-resolution -/// errors. +/// [`AiMuxError::NoSuchProvider`] for an unknown name, `InvalidArgument` for +/// invalid options. pub fn provider_handle( name: impl AsRef, api_key: Option, options: Option, -) -> Result, AiMuxError> { +) -> Result, AiMuxError> { let name = name.as_ref(); + let external = overlays().read().unwrap().get(name).cloned(); + match external { + Some(entry) => Ok(Arc::new(build_external(&entry, api_key, options)?)), + None => crate::create_provider(name, preset_settings(api_key, options)), + } +} - // 1. Runtime overlay (RFC-0020) — registered entries take precedence. - let resolved = if let Some(ext) = overlays().read().unwrap().get(name) { - ResolvedEntry::from_external(ext) - } else if let Some((default_url, url_env)) = match name { - "ollama" => Some(("http://127.0.0.1:11434/v1", "OLLAMA_BASE_URL")), - "llamacpp" => Some(("http://127.0.0.1:8080/v1", "LLAMACPP_BASE_URL")), - "localai" => Some(("http://127.0.0.1:8080/v1", "LOCALAI_BASE_URL")), - "vllm" => Some(("http://127.0.0.1:8000/v1", "VLLM_BASE_URL")), - "sglang" => Some(("http://127.0.0.1:30000/v1", "SGLANG_BASE_URL")), - "laya" => Some(("http://127.0.0.1:8000/v1", "LAYA_BASE_URL")), - _ => None, - } { - ResolvedEntry::from_registry(&RegistryEntry { - name: name.into(), - display: name.into(), - base_url: std::env::var(url_env) - .ok() - .filter(|s| !s.trim().is_empty()) - .unwrap_or_else(|| default_url.into()), - env_var: String::new(), - profile: ProviderProfile::default(), - }) - } else if name == "openai" { - // Native OpenAI is not in the compatibility registry. Expose the same - // provider handle/config path for its language and decision factories. - ResolvedEntry::from_registry(&RegistryEntry { - name: "openai".into(), - display: "OpenAI".into(), - base_url: OpenAIConfig::new("").base_url, - env_var: "OPENAI_API_KEY".into(), - profile: ProviderProfile::default(), - }) - } else { - // 2. Built-in registry. - let entry = registry().iter().find(|e| e.name == name).ok_or_else(|| { - AiMuxError::NoSuchProvider { - // Display derives from the id alone; valid names are discoverable - // via `provider_names()` — listing 250 names here would - // ride along in every error, across the C ABI. - provider_id: name.to_string(), - } - })?; - ResolvedEntry::from_registry(entry) +/// `PresetSettings` from the by-name arguments. Organization and project +/// become their headers, below the caller's own headers. +fn preset_settings(api_key: Option, options: Option) -> PresetSettings { + let mut settings = PresetSettings { + api_key: api_key.map(Resolvable::Value), + ..PresetSettings::default() }; + if let Some(options) = options { + let mut headers = aimux_provider_utils::HeaderMapOpt::new(); + if let Some(org) = options.organization { + headers.insert("OpenAI-Organization".to_string(), Some(org)); + } + if let Some(project) = options.project { + headers.insert("OpenAI-Project".to_string(), Some(project)); + } + for (name, value) in options.headers.into_iter().flatten() { + headers.insert(name, Some(value)); + } + settings.base_url = options.base_url; + settings.headers = (!headers.is_empty()).then_some(headers); + settings.params = options.params.unwrap_or_default(); + } + settings +} - // Reject base_urls that still carry unexpanded placeholders - // (e.g. cloudflare's `{CLOUDFLARE_ACCOUNT_ID}`, snowflake's - // ``). These entries are intentionally templated — - // the caller MUST supply a concrete base_url via ProviderOptions. - let base_url_overridden = options.as_ref().is_some_and(|o| o.base_url.is_some()); - if !base_url_overridden && base_url_has_placeholder(&resolved.base_url) { +/// An external entry as a compatible provider: its profile, its URL (the +/// per-call override wins), its headers and its key, which is resolved now. +fn build_external( + entry: &ExternalProviderEntry, + api_key: Option, + options: Option, +) -> Result { + let display = entry.display.clone().unwrap_or_else(|| entry.name.clone()); + let options = options.unwrap_or_default(); + if let Some(params) = &options.params + && !params.is_empty() + { return Err(AiMuxError::InvalidArgument(format!( - "provider '{name}' has a templated base_url {:?} with an unexpanded \ - placeholder; pass a concrete `base_url` via ProviderOptions to use it", - resolved.base_url + "external provider '{}' has no template parameters; `params` does not apply", + entry.name ))); } - // Resolve the api key. Priority: explicit parameter > entry-level api_key - // (supports "env:VAR" references) > entry env_var (read from environment). - let (key, source) = if api_key.is_none() - && resolved.env_var.is_empty() - && resolved.api_key.is_none() - && matches!( - name, - "ollama" | "llamacpp" | "localai" | "vllm" | "sglang" | "laya" - ) { - let key_env = format!("{}_API_KEY", name.to_uppercase()); - match std::env::var(&key_env).ok().filter(|key| !key.is_empty()) { - Some(key) => (key, Some(format!("env:{key_env}"))), - None => ("not-needed".into(), Some("none".into())), - } - } else { - resolve_key(&resolved, api_key)? - }; + let key = resolve_external_key(entry, &display, api_key)?; + + let mut fixed_headers = Vec::new(); + if let Some(org) = options.organization.or_else(|| entry.organization.clone()) { + fixed_headers.push(("OpenAI-Organization".to_string(), org)); + } + if let Some(project) = options.project.or_else(|| entry.project.clone()) { + fixed_headers.push(("OpenAI-Project".to_string(), project)); + } + let headers = options + .headers + .or_else(|| entry.headers.clone()) + .map(|headers| headers.into_iter().map(|(k, v)| (k, Some(v))).collect()); + let base_url = validate_base_url(options.base_url.as_deref().unwrap_or(&entry.base_url))?; - let mut config = build_resolved_config(&resolved, key, options); - config = config.with_api_key_source(source.as_deref()); - Ok(Box::new(OpenAIProvider::new(config))) + OpenAICompatibleProvider::assemble(Assembly { + name: entry.name.clone(), + base_url: BaseUrl::Fixed(base_url), + credential: Credential::Explicit(Resolvable::Value(key)), + fixed_headers, + headers, + query_params: None, + fetch: None, + transform_request_body: None, + profile: entry.profile.chat_profile()?, + }) } -/// Resolve the api key for a [`ResolvedEntry`]. Returns `(key, source)` where -/// `source` is the origin tag for RFC-0023 replay reconstruction. +/// Resolve the api key of an external entry. /// /// Priority: explicit `api_key` parameter > entry-level `api_key` field /// (supports `"env:VAR"` references, resolved against the environment) > /// entry `env_var` (read from the environment via [`load_api_key`]). -fn resolve_key( - entry: &ResolvedEntry, +fn resolve_external_key( + entry: &ExternalProviderEntry, + display: &str, api_key: Option, -) -> Result<(String, Option), AiMuxError> { +) -> Result { if let Some(key) = api_key { - return Ok((key, Some("explicit".to_string()))); + return Ok(key); } if let Some(entry_key) = &entry.api_key && !entry_key.is_empty() @@ -491,669 +514,15 @@ fn resolve_key( entry.name )) })?; - return Ok((val, Some(format!("env:{var}")))); + return Ok(val); } - return Ok((entry_key.clone(), Some("explicit".to_string()))); + return Ok(entry_key.clone()); } - if entry.env_var.is_empty() { - return Err(AiMuxError::InvalidArgument(format!( + match entry.env_var.as_deref().filter(|var| !var.is_empty()) { + Some(var) => load_api_key(None, var, display), + None => Err(AiMuxError::InvalidArgument(format!( "provider '{}' has no api_key parameter, entry-level api_key, or env_var to read from", entry.name - ))); - } - let key = aimux_provider_utils::load_api_key(None, &entry.env_var, &entry.display)?; - Ok((key, Some(format!("env:{}", entry.env_var)))) -} - -/// Resolve a [`ResolvedEntry`] + per-call [`ProviderOptions`] into a fully-wired -/// `OpenAIConfig`. Provider-level config from the entry (headers/org/project/ -/// retries/body_overrides) is applied first; per-call options override them. -fn build_resolved_config( - entry: &ResolvedEntry, - key: String, - options: Option, -) -> OpenAIConfig { - let mut config = OpenAIConfig::new(key) - .with_base_url(entry.base_url.clone()) - .with_provider(entry.name.clone()) - .with_profile(profile_from_registry(&entry.profile)); - - // Provider-level config (from the entry, built-in or external). - if let Some(headers) = &entry.headers { - config = config.with_headers(headers.clone()); - } - if let Some(org) = &entry.organization { - config = config.with_org_id(org.clone()); - } - if let Some(project) = &entry.project { - config = config.with_project(project.clone()); - } - if let Some(max_retries) = entry.max_retries { - config.retry_config.max_retries = max_retries; - } - if let Some(overrides) = &entry.body_overrides { - config = config.with_body_overrides(overrides.clone()); - } - - // Per-call ProviderOptions override on top. - if let Some(opts) = options { - if let Some(url) = opts.base_url { - config = config.with_base_url(url); - } - if let Some(headers) = opts.headers { - config = config.with_headers(headers); - } - if let Some(org) = opts.organization { - config = config.with_org_id(org); - } - if let Some(project) = opts.project { - config = config.with_project(project); - } - if let Some(max_retries) = opts.max_retries { - config.retry_config.max_retries = max_retries; - } - if let Some(overrides) = opts.body_overrides { - config = config.with_body_overrides(overrides); - } - } - config -} - -/// Translate a provider profile into the runtime profile. -/// -/// `stream_usage_key` / `max_tokens_key` are `&'static str` in -/// `OpenAICompatProfile`; the registry strings are leaked once at first use -/// (bounded: at most one string per field per entry). -fn profile_from_registry(p: &ProviderProfile) -> OpenAICompatProfile { - OpenAICompatProfile { - supports_top_k: p.supports_top_k, - supports_tools: p.supports_tools, - supports_response_format: p.supports_response_format, - stream_usage_key: p.stream_usage_key.as_deref().map(|s| { - let leaked: &'static str = Box::leak(s.to_string().into_boxed_str()); - leaked - }), - max_tokens_key: p.max_tokens_key.as_deref().map(|s| { - let leaked: &'static str = Box::leak(s.to_string().into_boxed_str()); - leaked - }), - } -} - -/// Convenience: build from the env-var key of the registry entry. -/// -/// # Errors -/// -/// Propagates the errors of [`provider`] (unknown provider, missing env-var -/// key, invalid options). -pub fn provider_from_env( - name: impl AsRef, - model_id: &str, - options: Option, -) -> Result, AiMuxError> { - provider(name, None, model_id, options) -} - -/// Names of all built-in registry providers. -pub fn provider_names() -> impl Iterator { - registry().iter().map(|entry| entry.name.as_str()) -} - -/// The environment variable [`provider`] / [`provider_handle`] read `name`'s -/// API key from when no key is passed: the runtime overlay's `api_key` -/// `env:` reference or its `env_var`, `{NAME}_API_KEY` for the local -/// OpenAI-compatible servers (ollama, llamacpp, …), or the registry entry's -/// `env_var`. `None` for an unknown name or an entry without one. -/// -/// Hosts that let a caller pick a key by variable name (the aimux-web -/// console's `env:VAR`) use it to allow only the provider's own variable. -#[must_use] -pub fn provider_api_key_env(name: &str) -> Option { - if let Some(ext) = overlays().read().unwrap().get(name) { - return ext - .api_key - .as_deref() - .and_then(|key| key.strip_prefix("env:")) - .or(ext.env_var.as_deref()) - .filter(|var| !var.is_empty()) - .map(str::to_string); - } - match name { - "ollama" | "llamacpp" | "localai" | "vllm" | "sglang" | "laya" => { - Some(format!("{}_API_KEY", name.to_uppercase())) - } - "openai" => Some("OPENAI_API_KEY".to_string()), - _ => registry() - .iter() - .find(|entry| entry.name == name) - .map(|entry| entry.env_var.clone()), - } -} - -/// Public lookup of a registered provider's runtime profile — used by tests -/// that assert registry wiring (e.g. `max_tokens_key`) without constructing a -/// model. Returns `None` for unknown provider names. -#[must_use] -pub fn provider_registry_entry(name: &str) -> Option { - registry() - .iter() - .find(|e| e.name == name) - .map(|e| profile_from_registry(&e.profile)) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn provider_api_key_env_names_the_variable_the_provider_reads() { - assert_eq!( - provider_api_key_env("groq").as_deref(), - Some("GROQ_API_KEY") - ); - assert_eq!( - provider_api_key_env("openai").as_deref(), - Some("OPENAI_API_KEY") - ); - assert_eq!( - provider_api_key_env("ollama").as_deref(), - Some("OLLAMA_API_KEY") - ); - assert_eq!(provider_api_key_env("no-such-provider"), None); - - let name = "key-env-overlay-test"; - let entry = |api_key: Option<&str>, env_var: Option<&str>| ExternalProviderEntry { - name: name.into(), - display: None, - base_url: "https://overlay.example/v1".into(), - env_var: env_var.map(str::to_string), - api_key: api_key.map(str::to_string), - protocol: "openai_compat".into(), - profile: ProviderProfile::default(), - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }; - register_provider(entry(None, Some("OVERLAY_KEY"))).unwrap(); - assert_eq!(provider_api_key_env(name).as_deref(), Some("OVERLAY_KEY")); - // An `env:` api_key reference wins, as in key resolution. - register_provider(entry(Some("env:OVERLAY_REF"), Some("OVERLAY_KEY"))).unwrap(); - assert_eq!(provider_api_key_env(name).as_deref(), Some("OVERLAY_REF")); - clear_overlay(name); - } - - #[test] - fn provider_builds_groq_model() { - let model = match provider("groq", Some("sk-test".into()), "llama-3.3-70b", None) { - Ok(m) => m, - Err(e) => panic!("groq should construct: {e}"), - }; - assert_eq!(model.provider(), "groq"); // RFC-0023 C2: registry provider surfaces its real name, not "openai" - assert_eq!(model.model_id(), "llama-3.3-70b"); - } - - #[test] - fn provider_applies_registry_profile() { - // groq entry: supports_top_k=false, stream_usage_key="x_groq", - // max_tokens_key="max_completion_tokens" — verify the config carries them. - let entry = registry().iter().find(|e| e.name == "groq").unwrap(); - let profile = profile_from_registry(&entry.profile); - assert!(!profile.supports_top_k); - assert_eq!(profile.stream_usage_key, Some("x_groq")); - assert_eq!(profile.max_tokens_key, Some("max_completion_tokens")); - - // stepfun entry: max_tokens_key="max_tokens". - let entry = registry().iter().find(|e| e.name == "stepfun").unwrap(); - let profile = profile_from_registry(&entry.profile); - assert_eq!(profile.max_tokens_key, Some("max_tokens")); - } - - #[test] - fn provider_accepts_string_names() { - let model = match provider("groq", Some("sk-test".into()), "llama-3.3-70b", None) { - Ok(m) => m, - Err(e) => panic!("string name should construct: {e}"), - }; - assert_eq!(model.model_id(), "llama-3.3-70b"); - } - - #[test] - fn provider_unknown_name_reports_provider_id() { - let err = match provider("no-such-provider", Some("k".into()), "m", None) { - Ok(_) => panic!("unknown name must fail"), - Err(e) => e, - }; - match err { - AiMuxError::NoSuchProvider { ref provider_id } => { - assert_eq!(provider_id, "no-such-provider"); - // Display derives from the single stored fact. - assert_eq!(err.to_string(), "No such provider: no-such-provider"); - // The 250 names stay out of the error — they ride the C ABI. - let text = err.to_string(); - assert!(!text.contains("groq"), "must not list the registry: {text}"); - } - other => panic!("expected NoSuchProvider, got {other:?}"), - } - } - - #[test] - fn provider_missing_env_key_fails() { - // env var almost certainly unset in CI; the error must be the missing-key - // error, not a registry problem. - let err = match provider("abacus", None, "m", None) { - Ok(_) => panic!("missing env key must fail"), - Err(e) => e, - }; - assert!(err.to_string().to_lowercase().contains("api key")); - } - - #[test] - fn registry_entries_are_valid() { - let entries = registry(); - assert_eq!(entries.len(), 251); - for e in entries { - assert!(e.name.starts_with(|c: char| c.is_ascii_lowercase())); - assert!(!e.base_url.is_empty()); - assert!(!e.env_var.is_empty()); - } - } - - #[test] - fn registry_no_corrupt_base_urls() { - // Issue #90 R2: no registry base_url may carry non-ASCII pollution - // or be a non-URL fragment. Templated placeholders (cloudflare/neon/ - // snowflake/oci) are allowed but rejected at construct time — see - // provider_handle's base_url_has_placeholder check. - for e in registry() { - assert!( - e.base_url.starts_with("https://") || e.base_url.starts_with("http://"), - "registry entry '{}' has a non-URL base_url: {:?}", - e.name, - e.base_url - ); - assert!( - e.base_url.is_ascii(), - "registry entry '{}' base_url has non-ASCII chars: {:?}", - e.name, - e.base_url - ); - } - } - - #[test] - fn registry_fixed_base_urls_are_correct() { - // Issue #90 R2: pin the corrected base_urls for entries that were - // broken (regression guard against reverting to the old bad values). - let by_name: std::collections::HashMap<&str, &str> = registry() - .iter() - .map(|e| (e.name.as_str(), e.base_url.as_str())) - .collect(); - assert_eq!( - by_name.get("xpersona").copied(), - Some("https://www.xpersona.co/v1"), - "xpersona base_url was '/v1' (truncated); must be the full URL" - ); - assert_eq!( - by_name.get("moonshotai_cn").copied(), - Some("https://api.moonshot.cn/anthropic/v1"), - "moonshotai_cn base_url had a leaked '（Anthropic' annotation suffix" - ); - assert_eq!( - by_name.get("zhipuai_coding_plan").copied(), - Some("https://open.bigmodel.cn/api/coding/paas/v4"), - "zhipuai_coding_plan base_url was a docs page, not the API endpoint" - ); - assert_eq!( - by_name.get("the_grid_ai").copied(), - Some("https://api.thegrid.ai/v1"), - "the_grid_ai base_url was a docs page, not the API endpoint" - ); - } - - #[test] - fn provider_rejects_templated_base_url_without_override() { - // Both placeholder forms must be rejected without an override: - // cloudflare uses `{CLOUDFLARE_ACCOUNT_ID}`, snowflake uses - // `` — exercise both so `base_url_has_placeholder` - // can't silently lose one form. - for name in ["cloudflare", "snowflake"] { - let err = match provider(name, Some("dummy".into()), "m", None) { - Ok(_) => panic!("templated base_url for '{name}' without override must fail"), - Err(e) => e, - }; - assert!( - matches!(err, AiMuxError::InvalidArgument(_)), - "expected InvalidArgument for '{name}' templated base_url, got {err:?}" - ); - } - } - - #[test] - fn provider_accepts_templated_base_url_with_override() { - // With a concrete base_url override, the templated entry must construct - // fine (key is dummy; we only assert it gets past validation). - let res = provider( - "cloudflare", - Some("dummy".into()), - "m", - Some(ProviderOptions { - base_url: Some("https://example.com/v1".into()), - ..Default::default() - }), - ); - assert!(res.is_ok(), "override should bypass placeholder check"); - } - - #[test] - fn provider_names_list_the_registry() { - let names: Vec<_> = provider_names().collect(); - assert_eq!(names.len(), registry().len()); - assert!(names.contains(&"groq")); - } - - // ── RFC-0020: external provider overlay ──────────────────────────────── - - #[test] - fn register_and_lookup_external_provider() { - clear_overlay("test-relay-new"); - register_provider(ExternalProviderEntry { - name: "test-relay-new".into(), - display: Some("Test Relay".into()), - base_url: "https://relay.test.example/v1".into(), - env_var: Some("TEST_RELAY_KEY".into()), - api_key: Some("dummy-key".into()), - protocol: "openai_compat".into(), - profile: ProviderProfile::default(), - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }) - .unwrap(); - let model = provider("test-relay-new", None, "test-model", None).unwrap(); - assert_eq!(model.provider(), "test-relay-new"); - clear_overlay("test-relay-new"); - } - - #[test] - fn external_provider_overrides_builtin() { - // Register an overlay for the built-in "groq" name with a different - // base_url; provider() must resolve to the overlay, not the registry. - clear_overlay("groq"); // hermetic start (other tests may use "groq") - register_provider(ExternalProviderEntry { - name: "groq".into(), - display: Some("Groq Override".into()), - base_url: "https://my-groq-relay.example/v1".into(), - env_var: Some("GROQ_API_KEY".into()), - api_key: Some("dummy".into()), - protocol: "openai_compat".into(), - profile: ProviderProfile::default(), - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }) - .unwrap(); - // Actually call provider() (not just read the overlay map) to verify - // the lookup path routes to the overlay. The registry entry for groq - // has base_url "https://api.groq.com/openai/v1"; the overlay's is - // different, so config_snapshot().base_url tells us which path ran. - let model = provider("groq", None, "llama-3.3-70b", None).unwrap(); - let snap = model.config_snapshot(); - assert_eq!( - snap.base_url.as_deref(), - Some("https://my-groq-relay.example/v1"), - "provider() must resolve via the overlay, not the built-in registry" - ); - clear_overlay("groq"); - } - - #[test] - fn load_providers_from_json_parses_and_registers() { - clear_overlay("test-json-1"); - clear_overlay("test-json-2"); - let json = r#"{ - "providers": [ - { - "name": "test-json-1", - "base_url": "https://a.test/v1", - "api_key": "dummy" - }, - { - "name": "test-json-2", - "base_url": "https://b.test/v1", - "env_var": "TEST_JSON_2_KEY" - } - ] - }"#; - load_providers_from_json(json).unwrap(); - assert!(overlays().read().unwrap().contains_key("test-json-1")); - assert!(overlays().read().unwrap().contains_key("test-json-2")); - clear_overlay("test-json-1"); - clear_overlay("test-json-2"); - } - - #[test] - fn register_provider_rejects_bad_base_url() { - clear_overlay("test-bad-url"); - let err = register_provider(ExternalProviderEntry { - name: "test-bad-url".into(), - base_url: "ftp://nope".into(), - env_var: None, - api_key: None, - protocol: "openai_compat".into(), - display: None, - profile: ProviderProfile::default(), - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }) - .unwrap_err(); - assert!( - matches!(err, AiMuxError::InvalidArgument(ref m) if m.contains("http(s)://")), - "expected InvalidArgument about scheme, got {err:?}" - ); - clear_overlay("test-bad-url"); - } - - #[test] - fn register_provider_rejects_bad_protocol() { - clear_overlay("test-bad-proto"); - let err = register_provider(ExternalProviderEntry { - name: "test-bad-proto".into(), - base_url: "https://ok.example/v1".into(), - protocol: "anthropic".into(), - env_var: None, - api_key: None, - display: None, - profile: ProviderProfile::default(), - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }) - .unwrap_err(); - assert!( - matches!(err, AiMuxError::InvalidArgument(ref m) if m.contains("openai_compat")), - "expected InvalidArgument about protocol, got {err:?}" - ); - clear_overlay("test-bad-proto"); - } - - #[test] - fn register_provider_rejects_empty_name() { - let err = register_provider(ExternalProviderEntry { - name: " ".into(), - base_url: "https://ok.example/v1".into(), - protocol: "openai_compat".into(), - env_var: None, - api_key: None, - display: None, - profile: ProviderProfile::default(), - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }) - .unwrap_err(); - assert!( - matches!(err, AiMuxError::InvalidArgument(ref m) if m.contains("name")), - "expected InvalidArgument about name, got {err:?}" - ); - } - - #[test] - fn external_provider_profile_applied() { - clear_overlay("test-profile"); - register_provider(ExternalProviderEntry { - name: "test-profile".into(), - base_url: "https://profile.test/v1".into(), - api_key: Some("dummy".into()), - protocol: "openai_compat".into(), - profile: ProviderProfile { - supports_top_k: false, - supports_tools: false, - supports_response_format: true, - stream_usage_key: Some("x_custom".into()), - max_tokens_key: Some("max_completion_tokens".into()), - }, - env_var: None, - display: None, - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }) - .unwrap(); - let entry = overlays() - .read() - .unwrap() - .get("test-profile") - .unwrap() - .clone(); - let p = profile_from_registry(&entry.profile); - assert!(!p.supports_top_k); - assert!(!p.supports_tools); - assert_eq!(p.stream_usage_key, Some("x_custom")); - assert_eq!(p.max_tokens_key, Some("max_completion_tokens")); - clear_overlay("test-profile"); - } - - #[test] - fn external_provider_profile_defaults_to_full() { - // When `profile` is omitted entirely from the JSON, the entry must - // still default to the OpenAI-compatible baseline (all three - // supports_* = true). Regression guard: the derive(Default) on - // ProviderProfile was previously yielding false for these. - clear_overlay("test-default-profile"); - load_providers_from_json(r#"{ "providers": [ { "name": "test-default-profile", "base_url": "https://x.test/v1", "api_key": "dummy" } ] }"#) - .unwrap(); - let entry = overlays() - .read() - .unwrap() - .get("test-default-profile") - .unwrap() - .clone(); - let p = profile_from_registry(&entry.profile); - assert!( - p.supports_top_k, - "omitted profile → supports_top_k must be true" - ); - assert!( - p.supports_tools, - "omitted profile → supports_tools must be true" - ); - assert!( - p.supports_response_format, - "omitted profile → supports_response_format must be true" - ); - clear_overlay("test-default-profile"); - } - - #[test] - fn external_provider_env_var_api_key_resolves() { - // entry-level api_key = "env:VAR" must read the env var at lookup time. - clear_overlay("test-envkey"); - // SAFETY: test-only; no other thread is reading this var concurrently. - unsafe { std::env::set_var("AIMUX_TEST_OVERLAY_KEY", "secret-from-env") }; - register_provider(ExternalProviderEntry { - name: "test-envkey".into(), - base_url: "https://envkey.test/v1".into(), - api_key: Some("env:AIMUX_TEST_OVERLAY_KEY".into()), - protocol: "openai_compat".into(), - env_var: None, - display: None, - profile: ProviderProfile::default(), - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }) - .unwrap(); - // provider() with api_key=None must resolve via the "env:" reference. - let model = provider("test-envkey", None, "m", None).unwrap(); - assert_eq!(model.provider(), "test-envkey"); - // SAFETY: test-only cleanup. - unsafe { std::env::remove_var("AIMUX_TEST_OVERLAY_KEY") }; - clear_overlay("test-envkey"); - } - - #[test] - fn external_provider_env_var_api_key_missing_env_fails() { - clear_overlay("test-envkey-missing"); - // SAFETY: test-only cleanup. - unsafe { std::env::remove_var("AIMUX_TEST_OVERLAY_MISSING") }; - register_provider(ExternalProviderEntry { - name: "test-envkey-missing".into(), - base_url: "https://missing.test/v1".into(), - api_key: Some("env:AIMUX_TEST_OVERLAY_MISSING".into()), - protocol: "openai_compat".into(), - env_var: None, - display: None, - profile: ProviderProfile::default(), - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }) - .unwrap(); - let err = match provider("test-envkey-missing", None, "m", None) { - Ok(_) => panic!("missing env var must fail"), - Err(e) => e, - }; - assert!( - matches!(err, AiMuxError::InvalidArgument(ref m) if m.contains("AIMUX_TEST_OVERLAY_MISSING")), - "expected InvalidArgument naming the env var, got {err:?}" - ); - clear_overlay("test-envkey-missing"); - } - - #[test] - fn load_providers_from_json_rejects_invalid_json() { - let err = load_providers_from_json("not json at all").unwrap_err(); - assert!( - matches!(err, AiMuxError::JsonParse(ref m) if m.contains("parse")), - "expected JsonParse for malformed input, got {err:?}" - ); + ))), } } diff --git a/aimux-providers/src/provider_registry.json b/aimux-providers/src/provider_registry.json index 210bcdd44..38796a969 100644 --- a/aimux-providers/src/provider_registry.json +++ b/aimux-providers/src/provider_registry.json @@ -237,6 +237,23 @@ "env_var": "BASETEN_API_KEY", "profile": {} }, + { + "name": "bedrock_mantle", + "display": "Bedrock Mantle", + "base_url": "https://bedrock-mantle.{region}.api.aws/v1", + "env_var": "BEDROCK_MANTLE_API_KEY", + "params": [ + { + "name": "region", + "env": [ + "BEDROCK_MANTLE_REGION", + "AWS_REGION" + ], + "default": "us-east-1" + } + ], + "profile": {} + }, { "name": "berget", "display": "Berget.AI", @@ -352,15 +369,31 @@ { "name": "cloudflare", "display": "Cloudflare", - "base_url": "https://api.cloudflare.com/client/v4/accounts/{CLOUDFLARE_ACCOUNT_ID}/ai/v1", + "base_url": "https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1", "env_var": "CLOUDFLARE_API_KEY", + "params": [ + { + "name": "account_id", + "env": [ + "CLOUDFLARE_ACCOUNT_ID" + ] + } + ], "profile": {} }, { "name": "cloudflare_workers_ai", "display": "Cloudflare Workers AI", - "base_url": "https://api.cloudflare.com/client/v4/accounts/{CLOUDFLARE_ACCOUNT_ID}/ai/v1", + "base_url": "https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1", "env_var": "CLOUDFLARE_API_KEY", + "params": [ + { + "name": "account_id", + "env": [ + "CLOUDFLARE_ACCOUNT_ID" + ] + } + ], "profile": {} }, { @@ -433,6 +466,14 @@ "env_var": "CRUSOE_API_KEY", "profile": {} }, + { + "name": "cybertron", + "display": "Cybertron", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "CYBERTRON_BASE_URL", + "profile": {} + }, { "name": "daoxe", "display": "DaoXE", @@ -475,13 +516,6 @@ "env_var": "DEEPINFRA_API_KEY", "profile": {} }, - { - "name": "deepseek", - "display": "DeepSeek", - "base_url": "https://api.deepseek.com/v1", - "env_var": "DEEPSEEK_API_KEY", - "profile": {} - }, { "name": "digitalocean", "display": "DigitalOcean", @@ -496,6 +530,14 @@ "env_var": "DINFERENCE_API_KEY", "profile": {} }, + { + "name": "docker_model_runner", + "display": "Docker Model Runner", + "base_url": "http://model-runner.docker.internal/engines/llama.cpp/v1", + "auth": "none", + "base_url_env": "DOCKER_MODEL_RUNNER_BASE_URL", + "profile": {} + }, { "name": "doubao", "display": "Doubao", @@ -615,6 +657,14 @@ "env_var": "GALADRIEL_API_KEY", "profile": {} }, + { + "name": "gaudi", + "display": "Intel Gaudi", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "GAUDI_BASE_URL", + "profile": {} + }, { "name": "gdc", "display": "GDC", @@ -664,17 +714,6 @@ "env_var": "GRADIENT_API_KEY", "profile": {} }, - { - "name": "groq", - "display": "Groq", - "base_url": "https://api.groq.com/openai/v1", - "env_var": "GROQ_API_KEY", - "profile": { - "supports_top_k": false, - "stream_usage_key": "x_groq", - "max_tokens_key": "max_completion_tokens" - } - }, { "name": "helicone", "display": "Helicone", @@ -782,6 +821,14 @@ "env_var": "JIEKOU_API_KEY", "profile": {} }, + { + "name": "jlama", + "display": "Jlama", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "JLAMA_BASE_URL", + "profile": {} + }, { "name": "kenari", "display": "Kenari", @@ -845,6 +892,15 @@ "env_var": "LAMBDA_API_KEY", "profile": {} }, + { + "name": "laya", + "display": "Laya", + "base_url": "http://127.0.0.1:8000/v1", + "env_var": "LAYA_API_KEY", + "auth": "none", + "base_url_env": "LAYA_BASE_URL", + "profile": {} + }, { "name": "lemonade", "display": "Lemonade", @@ -880,6 +936,14 @@ "env_var": "LINGYIWANWU_API_KEY", "profile": {} }, + { + "name": "litellm_proxy", + "display": "LiteLLM Proxy", + "base_url": "http://127.0.0.1:4000/v1", + "auth": "none", + "base_url_env": "LITELLM_PROXY_BASE_URL", + "profile": {} + }, { "name": "llama", "display": "Llama", @@ -887,6 +951,23 @@ "env_var": "LLAMA_API_KEY", "profile": {} }, + { + "name": "llamacpp", + "display": "llama.cpp", + "base_url": "http://127.0.0.1:8080/v1", + "env_var": "LLAMACPP_API_KEY", + "auth": "none", + "base_url_env": "LLAMACPP_BASE_URL", + "profile": {} + }, + { + "name": "llamafile", + "display": "Llamafile", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "LLAMAFILE_BASE_URL", + "profile": {} + }, { "name": "llamagate", "display": "Llamagate", @@ -908,6 +989,31 @@ "env_var": "LLMTR_API_KEY", "profile": {} }, + { + "name": "lmstudio", + "display": "LM Studio", + "base_url": "http://127.0.0.1:1234/v1", + "auth": "none", + "base_url_env": "LMSTUDIO_BASE_URL", + "profile": {} + }, + { + "name": "local", + "display": "Local LLM", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "LOCAL_LLM_BASE_URL", + "profile": {} + }, + { + "name": "localai", + "display": "LocalAI", + "base_url": "http://127.0.0.1:8080/v1", + "env_var": "LOCALAI_API_KEY", + "auth": "none", + "base_url_env": "LOCALAI_BASE_URL", + "profile": {} + }, { "name": "longcat", "display": "LongCat", @@ -1006,6 +1112,14 @@ "env_var": "MIRA_API_KEY", "profile": {} }, + { + "name": "mistralrs", + "display": "Mistral.rs", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "MISTRALRS_BASE_URL", + "profile": {} + }, { "name": "mixlayer", "display": "Mixlayer", @@ -1013,6 +1127,14 @@ "env_var": "MIXLAYER_API_KEY", "profile": {} }, + { + "name": "mlx", + "display": "MLX (Apple Silicon)", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "MLX_BASE_URL", + "profile": {} + }, { "name": "moark", "display": "Moark", @@ -1093,8 +1215,13 @@ { "name": "neon", "display": "Neon", - "base_url": "https:///v1", + "base_url": "https://{branch_host}/v1", "env_var": "NEON_AI_GATEWAY_TOKEN", + "params": [ + { + "name": "branch_host" + } + ], "profile": {} }, { @@ -1149,8 +1276,13 @@ { "name": "oci", "display": "OCI", - "base_url": "https://inference.generativeai.${region}.oci.oraclecloud.com/openai/v1", + "base_url": "https://inference.generativeai.{region}.oci.oraclecloud.com/openai/v1", "env_var": "OCI_API_KEY", + "params": [ + { + "name": "region" + } + ], "profile": {} }, { @@ -1167,6 +1299,15 @@ "env_var": "OHMYGPT_API_KEY", "profile": {} }, + { + "name": "ollama", + "display": "Ollama", + "base_url": "http://127.0.0.1:11434/v1", + "env_var": "OLLAMA_API_KEY", + "auth": "none", + "base_url_env": "OLLAMA_BASE_URL", + "profile": {} + }, { "name": "ollama_cloud", "display": "Ollama Cloud", @@ -1174,6 +1315,30 @@ "env_var": "OLLAMA_CLOUD_API_KEY", "profile": {} }, + { + "name": "omlx", + "display": "OMLX / MLX LM", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "OMLX_BASE_URL", + "profile": {} + }, + { + "name": "onnx", + "display": "ONNX Runtime", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "ONNX_BASE_URL", + "profile": {} + }, + { + "name": "oobabooba", + "display": "Oobabooga Text Generation WebUI", + "base_url": "http://127.0.0.1:5000/v1", + "auth": "none", + "base_url_env": "OOBABOOBA_BASE_URL", + "profile": {} + }, { "name": "openaimax", "display": "OpenAIMax", @@ -1209,6 +1374,21 @@ "env_var": "OPENCODE_ZEN_API_KEY", "profile": {} }, + { + "name": "openrouter", + "display": "OpenRouter", + "base_url": "https://openrouter.ai/api/v1", + "env_var": "OPENROUTER_API_KEY", + "profile": {} + }, + { + "name": "openvino", + "display": "OpenVINO", + "base_url": "http://127.0.0.1:8080/v1", + "auth": "none", + "base_url_env": "OPENVINO_BASE_URL", + "profile": {} + }, { "name": "orcarouter", "display": "OrcaRouter", @@ -1420,6 +1600,15 @@ "env_var": "SCX_AI_API_KEY", "profile": {} }, + { + "name": "sglang", + "display": "SGLang", + "base_url": "http://127.0.0.1:30000/v1", + "env_var": "SGLANG_API_KEY", + "auth": "none", + "base_url_env": "SGLANG_BASE_URL", + "profile": {} + }, { "name": "siliconflow", "display": "SiliconFlow", @@ -1432,15 +1621,25 @@ { "name": "snowflake", "display": "Snowflake", - "base_url": "https://.snowflakecomputing.com/api/v2/cortex/v1", + "base_url": "https://{account_identifier}.snowflakecomputing.com/api/v2/cortex/v1", "env_var": "SNOWFLAKE_PAT", + "params": [ + { + "name": "account_identifier" + } + ], "profile": {} }, { "name": "snowflake_cortex", "display": "Snowflake Cortex", - "base_url": "https://.snowflakecomputing.com/api/v2/cortex/v1", + "base_url": "https://{account_identifier}.snowflakecomputing.com/api/v2/cortex/v1", "env_var": "SNOWFLAKE_CORTEX_PAT", + "params": [ + { + "name": "account_identifier" + } + ], "profile": {} }, { @@ -1655,6 +1854,346 @@ "env_var": "VERCEL_API_KEY", "profile": {} }, + { + "name": "vertex_ai_ai21_models", + "display": "AI21 models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, + { + "name": "vertex_ai_anthropic_models", + "display": "Anthropic Claude models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, + { + "name": "vertex_ai_deepseek_models", + "display": "DeepSeek models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, + { + "name": "vertex_ai_llama_models", + "display": "Meta Llama models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, + { + "name": "vertex_ai_minimax_models", + "display": "MiniMax models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, + { + "name": "vertex_ai_mistral_models", + "display": "Mistral models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, + { + "name": "vertex_ai_moonshot_models", + "display": "Moonshot AI models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, + { + "name": "vertex_ai_openai_models", + "display": "OpenAI models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, + { + "name": "vertex_ai_qwen_models", + "display": "Qwen models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, + { + "name": "vertex_ai_zai_models", + "display": "Z.AI models on Vertex AI", + "base_url": "https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi", + "env_var": "GOOGLE_VERTEX_ACCESS_TOKEN", + "params": [ + { + "name": "project", + "env": [ + "GOOGLE_VERTEX_PROJECT" + ] + }, + { + "name": "location", + "env": [ + "GOOGLE_VERTEX_LOCATION" + ], + "default": "global" + }, + { + "name": "host", + "derive": { + "from": "location", + "map": { + "global": "aiplatform.googleapis.com", + "eu": "aiplatform.eu.rep.googleapis.com", + "us": "aiplatform.us.rep.googleapis.com" + }, + "otherwise": "{location}-aiplatform.googleapis.com" + } + } + ], + "profile": {} + }, { "name": "vivgrid", "display": "Vivgrid", @@ -1662,6 +2201,15 @@ "env_var": "VIVGRID_API_KEY", "profile": {} }, + { + "name": "vllm", + "display": "vLLM", + "base_url": "http://127.0.0.1:8000/v1", + "env_var": "VLLM_API_KEY", + "auth": "none", + "base_url_env": "VLLM_BASE_URL", + "profile": {} + }, { "name": "volc_engine", "display": "VolcEngine", @@ -1718,6 +2266,14 @@ "env_var": "XIAOMI_API_KEY", "profile": {} }, + { + "name": "xinference", + "display": "Xinference", + "base_url": "http://127.0.0.1:9997/v1", + "auth": "none", + "base_url_env": "XINFERENCE_BASE_URL", + "profile": {} + }, { "name": "xpersona", "display": "Xpersona", diff --git a/aimux-providers/src/recraft.rs b/aimux-providers/src/recraft.rs index 0196615c2..962e35786 100644 --- a/aimux-providers/src/recraft.rs +++ b/aimux-providers/src/recraft.rs @@ -8,8 +8,14 @@ //! - Base URL: `https://external.api.recraft.ai/v1` //! - Response: `{ "data": [{ "url" | "b64_json" }] }` (OpenAI Images shape) //! - No streaming. +//! +//! [`create_recraft`] takes [`RecraftProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`RecraftProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `RECRAFT_API_TOKEN`. +//! [`recraft()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::Duration; use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -18,10 +24,14 @@ use aimux_core::error::AiMuxError; use aimux_core::image_model::{ ImageCallOptions, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; -use aimux_core::provider::Provider; -use aimux_core::retry; -use aimux_core::shared::{SharedProviderOptions, Warning}; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_core::shared::Warning; +use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{ + Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, poll_interval_ms, provider_headers, + retry_download, +}; fn recraft_failed_response_handler() -> aimux_provider_utils::ResponseHandler { aimux_provider_utils::create_json_error_response_handler(|data| { @@ -41,111 +51,133 @@ fn recraft_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, -} +pub(crate) mod options; -impl RecraftConfig { - /// Create from an API key, using the default Recraft base URL. - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: DEFAULT_BASE_URL.to_string(), - headers: None, - } - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://external.api.recraft.ai/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "RECRAFT_API_TOKEN"; +const DEFAULT_NAME: &str = "recraft"; - /// Override the base URL (useful for tests / self-hosted endpoints). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Settings of [`create_recraft`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct RecraftProviderSettings { + /// Base URL for the API calls. Default `https://external.api.recraft.ai/v1`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `RECRAFT_API_TOKEN` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.image"`). + /// Default `"recraft"`. The providerOptions key stays `recraft`. + 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, +} - /// Set additional static HTTP headers sent with every request. - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self +impl std::fmt::Debug for RecraftProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RecraftProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Create from the `RECRAFT_API_TOKEN` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `RECRAFT_API_TOKEN` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, ENV_VAR, "Recraft")?; - Ok(Self::new(api_key)) - } +/// Create a Recraft provider. +/// +/// # 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. +pub fn create_recraft(settings: RecraftProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(RecraftProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Recraft"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) } -/// Recraft provider — creates [`RecraftImageModel`] instances. +/// The default provider: `create_recraft` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn recraft() -> &'static RecraftProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_recraft(RecraftProviderSettings::default()) + .expect("default Recraft settings are always valid") + }) +} + +/// A Recraft provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct RecraftProvider { - config: RecraftConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl RecraftProvider { - #[must_use] - pub fn new(config: RecraftConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create an image generation model instance (e.g. `"recraftv3"`). + /// An image model (e.g. `"recraftv3"`); `provider()` is `"{name}.image"`. #[must_use] pub fn image(&self, model_id: &str) -> RecraftImageModel { - RecraftImageModel::new(model_id.to_string(), self.config.clone()) + RecraftImageModel::from_config(model_id.to_string(), self.model_config("image")) } } -impl Provider for RecraftProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } -} +crate::impl_single_modality_provider!(RecraftProvider, image_model, |p, id| p.image(id)); /// A Recraft image generation model. pub struct RecraftImageModel { model_id: String, - config: RecraftConfig, + config: EndpointConfig, } impl RecraftImageModel { - #[must_use] - pub fn new(model_id: String, config: RecraftConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut h = HashMap::new(); - h.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - h.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - h.insert(k.clone(), v.clone()); - } - } - h - } - - fn generations_endpoint(&self) -> String { - format!("{}/images/generations", self.config.base_url) - } } /// Recraft generation provider options (camelCase → snake_case mapping). @@ -161,8 +193,10 @@ struct RecraftOptions { /// /// Unknown fields are ignored (safe degradation); values are passed through as /// raw JSON so that unrecognised-but-valid enum values do not cause errors. -fn parse_recraft_options(provider_options: &SharedProviderOptions) -> RecraftOptions { - let recraft = provider_options.get("recraft"); +fn parse_recraft_options( + provider_options: &aimux_core::shared::SharedProviderOptions, +) -> RecraftOptions { + let recraft = options::recraft_options(Some(provider_options)); RecraftOptions { style: recraft.and_then(|o| o.get("style")).cloned(), style_id: recraft.and_then(|o| o.get("styleId")).cloned(), @@ -226,9 +260,9 @@ fn build_generation_body( /// - If neither is present, returns an empty [`ImageOutputs::Base64`]. async fn extract_images( response: &Value, - retries: &retry::PreparedRetries, abort_signal: Option, base_url: &str, + retry_interval: Duration, ) -> Result { let items = response.get("data").and_then(|d| d.as_array()); @@ -260,27 +294,21 @@ async fn extract_images( for url in &urls { // data[].url is a generated-image URL from the response body, so it // goes through the SSRF download guard. - let resp = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: vec![], - - abort_signal: abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: true, - trusted_origin: Some(base_url.to_string()), - credentialed_origin: Some(base_url.to_string()), - }, - aimux_provider_utils::create_binary_response_handler(), - recraft_failed_response_handler(), - ) - }) - .await?; + let resp = retry_download(abort_signal.as_ref(), retry_interval, || { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + abort_signal: abort_signal.clone(), + validate_url: true, + trusted_origin: Some(base_url.to_string()), + credentialed_origin: Some(base_url.to_string()), + ..Default::default() + }, + aimux_provider_utils::create_binary_response_handler(), + recraft_failed_response_handler(), + ) + }) + .await?; binaries.push(resp.value.to_vec()); } Ok(ImageOutputs::Binary(binaries)) @@ -289,7 +317,7 @@ async fn extract_images( #[async_trait] impl ImageModel for RecraftImageModel { fn provider(&self) -> &str { - PROVIDER_NAME + &self.config.provider } fn model_id(&self) -> &str { @@ -316,23 +344,10 @@ impl ImageModel for RecraftImageModel { let recraft_opts = parse_recraft_options(&options.provider_options); let body = build_generation_body(&self.model_id, options, &recraft_opts); - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.generations_endpoint(), - headers: header_list, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, + exchange.request(exchange.url("/images/generations"), options), Value::Object(body), aimux_provider_utils::create_json_response_handler(), recraft_failed_response_handler(), @@ -341,17 +356,16 @@ impl ImageModel for RecraftImageModel { let response_headers = resp.response_headers; let value: Value = resp.value; - let retries = retry::prepare_retries( - options.max_retries, - self.retry_config(), - options.abort_signal.clone(), - ); let images = extract_images( &value, - &retries, options.abort_signal.clone(), - &self.config.base_url, + exchange.base_url(), + Duration::from_millis(poll_interval_ms( + options::recraft_options(Some(&options.provider_options)), + POLL_INTERVAL_MS_KEY, + DOWNLOAD_RETRY_INTERVAL_MS, + )), ) .await?; diff --git a/aimux-providers/src/recraft/options.rs b/aimux-providers/src/recraft/options.rs new file mode 100644 index 000000000..7f89fb0de --- /dev/null +++ b/aimux-providers/src/recraft/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Recraft. +//! +//! Recraft reads `providerOptions.recraft`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "recraft"; + +/// The Recraft options in a providerOptions map. +pub(crate) fn recraft_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/replay.rs b/aimux-providers/src/replay.rs index ccfa4e552..2635c880b 100644 --- a/aimux-providers/src/replay.rs +++ b/aimux-providers/src/replay.rs @@ -1,508 +1,118 @@ -//! RFC-0023 请求回放(P4,层 2):按 `ProviderRecord` 自动重建 provider。 +//! RFC-0023 请求回放(P4,层 2):按 `ProviderRecord` 重建 model。 //! //! 拆层(评审 R4,2026-08-06):自动构造需要 `aimux-providers` 的构造能力, //! 放这里避免 core→providers→core 循环依赖。core 侧的 `replay_with_model` -//! 是 provider 无关的输入重建 + 重发;本模块负责"按录制配置重建 model"。 +//! 是 provider 无关的输入重建 + 重发;本模块负责"按录制身份重建 model"。 //! -//! **能力边界(定稿 R8)**:MVP 仅覆盖 **OpenAI 兼容族**——`LanguageModel::provider()` -//! 返回协议实现名,所有 OpenAI 兼容注册表 provider 都报 `"openai"`,所以按 -//! `provider.provider == "openai"` 判定 + 录制的 `base_url` 区分具体厂商。 -//! 其他协议(anthropic/google/azure/bedrock/...)明确 `Unsupported`: -//! 调用方需传 model 实例给 [`replay_with_model`](aimux_core::replay::replay_with_model)。 -//! -//! api_key 来源(`ProviderRecord.api_key_source`): -//! - `explicit` → 必须传 `api_key`(录制不存明文,隐私强制)。 -//! - `env:VAR` → 读该环境变量。 -//! - `none`(本地模型)→ `"not-needed"` 占位。 -//! - `unknown`(provider 尚未覆盖 `config_snapshot` 时的默认)→ 用传入 -//! `api_key`,未传则报错并提示。 +//! 录制只记身份(`provider_id` / `provider` / `model_id`),不记配置 +//! (RFC-0036 §3.3)。重建按 `provider_id` 从调用方的 registry 取 provider, +//! 保留其配置,再按 `model_id` 取语言模型。重建出的 model 若与录制时的 +//! `provider` 不是同一个 +//! (例如录制用的是 `openai.responses`,默认语言模型是 `openai.chat`),拒绝 +//! 重建:调用方需传 model 实例给 +//! [`replay_with_model`](aimux_core::replay::replay_with_model)。 + +use std::sync::Arc; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; +use aimux_core::provider_registry::ProviderRegistry; use aimux_core::recording::ProviderRecord; -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIProvider}; -use crate::provider::ProviderOptions; - -/// Rebuild an official decision model from a unified recording. +/// Rebuild a decision model using the caller's registered provider configuration. /// # Errors -/// Rejects unsupported providers, invalid model IDs and missing credentials. +/// Rejects missing identities, missing providers and a different model provider. pub fn rebuild_decision_provider( p: &ProviderRecord, - api_key: Option<&str>, -) -> Result, AiMuxError> { - if p.provider == "openai" || crate::runtime_decision::supported(&p.provider) { - let mut config = OpenAIConfig::new(resolve_api_key(p, api_key)?) - .with_provider(&p.provider) - .with_api_key_source(Some(&p.api_key_source)); - if let Some(base_url) = &p.base_url { - config = config.with_base_url(base_url); - } - if let Some(options) = &p.provider_options { - let options: ProviderOptions = serde_json::from_value(options.clone()) - .map_err(|e| AiMuxError::InvalidArgument(format!("decision replay config: {e}")))?; - config.org_id = options.organization; - config.project = options.project; - config.headers = options.headers.map(|mut headers| { - headers.retain(|_, value| value != "[REDACTED]"); - headers - }); - config.body_overrides = options.body_overrides; - if let Some(max_retries) = options.max_retries { - config.retry_config.max_retries = max_retries; - } - } - return OpenAIProvider::new(config).decision_model(&p.model_id); - } - if p.provider != "jev" { - return Err(AiMuxError::UnsupportedFunctionality(format!( - "decision replay: unsupported provider {}", - p.provider - ))); - } - let mut config = crate::JevConfig::new(resolve_api_key(p, api_key)?); - if let Some(endpoint) = &p.base_url { - config = config.with_endpoint(endpoint); - } - if let Some(headers) = p.provider_options.as_ref().and_then(|v| v.get("headers")) { - config.headers = serde_json::from_value(headers.clone()) - .map_err(|e| AiMuxError::InvalidArgument(format!("decision replay headers: {e}")))?; - // A recorded redaction is not a usable credential. The freshly supplied - // API key is installed by Jev; other secret headers require configuration. - config.headers.retain(|_, value| value != "[REDACTED]"); + registry: &ProviderRegistry, +) -> Result, AiMuxError> { + let provider_id = p.registry_id()?; + if p.model_id.is_empty() { + return Err(AiMuxError::InvalidArgument( + "decision replay: empty provider/model identity".into(), + )); } - if let Some(source) = p - .profile - .as_ref() - .and_then(|v| v.get("decision_capabilities")) - .and_then(|v| v.get("probability_source")) - { - config.probability_source = serde_json::from_value(source.clone()) - .map_err(|e| AiMuxError::InvalidArgument(e.to_string()))?; + let model = registry + .provider(provider_id)? + .decision_model(&p.model_id)?; + if model.provider() != p.provider { + return Err(AiMuxError::InvalidArgument( + "decision replay: registered model provider differs from recording".into(), + )); } - crate::JevProvider::new(config).decision_model(&p.model_id) + Ok(model) } -/// 按 `ProviderRecord` 重建 provider(OpenAI 兼容族,MVP)。 +/// 按 `ProviderRecord` 的 `provider_id` + `model_id` 重建 model。 /// -/// - `provider.provider != "openai"` → `Unsupported`(R8:不猜测原生协议)。 -/// - `base_url` 用录制值(OpenAI 兼容族区分厂商的关键);`provider_options` -/// 中的 org/project/headers/body_overrides 一并恢复;`profile` 从录制恢复。 -/// - 注意:录制的 headers 已脱敏,真实头无法从此恢复——需要时调用方重建后 -/// 再通过 `ProviderOptions.headers` 补。 +/// 使用调用方 registry 中的 provider 配置。 /// /// # Errors /// -/// Returns `Unsupported` for non-OpenAI-compatible recordings, -/// `InvalidArgument` for a missing `base_url`, and key-resolution errors. +/// `InvalidArgument` for an empty `provider_id` (a recording written by aimux +/// 0.5 or earlier, see [`ProviderRecord::registry_id`]) or `model_id`, or when the +/// provider's default language model is not the one the recording was made +/// with; [`AiMuxError::NoSuchProvider`] when `provider_id` is not registered. pub fn rebuild_provider( p: &ProviderRecord, - api_key: Option<&str>, -) -> Result, AiMuxError> { - if !is_openai_compatible_provider(&p.provider) { - return Err(AiMuxError::UnsupportedFunctionality(format!( - "mock replay: rebuild_provider covers only the OpenAI-compatible family \ - (recorded provider '{}'); pass a model instance to replay_with_model", - p.provider - ))); - } + registry: &ProviderRegistry, +) -> Result, AiMuxError> { + p.registry_id()?; if p.model_id.is_empty() { return Err(AiMuxError::InvalidArgument( "mock replay: provider record has empty model_id".into(), )); } - - let key = resolve_api_key(p, api_key)?; - // Preserve the recorded provider name (registry entry / thin-wrapper name) - // so the rebuilt model keeps its real identity (RFC-0023 C2). Direct OpenAI - // recordings carry `"openai"` and round-trip unchanged. - let mut config = OpenAIConfig::new(key) - .with_provider(&p.provider) - .with_profile(profile_from_value(p.profile.as_ref())); - - if let Some(base_url) = &p.base_url { - config = config.with_base_url(base_url); - } - - if let Some(opts) = &p.provider_options { - // provider_options 形状未知时静默降级为默认配置(不因多余/未知字段失败)。 - if let Ok(opts) = serde_json::from_value::(opts.clone()) { - if let Some(url) = opts.base_url { - config = config.with_base_url(url); - } - if let Some(headers) = opts.headers { - config = config.with_headers(headers); - } - if let Some(org) = opts.organization { - config = config.with_org_id(org); - } - if let Some(project) = opts.project { - config = config.with_project(project); - } - if let Some(max_retries) = opts.max_retries { - config.retry_config.max_retries = max_retries; - } - if let Some(overrides) = opts.body_overrides { - config = config.with_body_overrides(overrides); - } - } - } - - OpenAIProvider::new(config).language_model(&p.model_id) -} - -/// Whether a recorded provider name belongs to the OpenAI-compatible family -/// (rebuildable via `OpenAIProvider`). After RFC-0023 C2, OpenAI-compatible -/// models surface their *real* provider name (registry entry or thin-wrapper -/// name) instead of a hardcoded `"openai"`, so the replay guard must recognise -/// all of them. Native protocols (anthropic / google / azure / bedrock / …) -/// are rejected — their snapshots are not OpenAI-shaped. -/// -/// Covers: -/// - `"openai"` (direct `OpenAIProvider`), -/// - any name in `provider_registry.json` (cloud OpenAI-compatible providers), -/// - the fixed set of local/server thin wrappers that wrap `OpenAIProvider`. -fn is_openai_compatible_provider(name: &str) -> bool { - if name == "openai" { - return true; - } - // Registry-backed cloud OpenAI-compatible providers (deepseek, groq, …). - if crate::provider::provider_registry_entry(name).is_some() { - return true; - } - // Externally-registered providers via the RFC-0020 overlay. - if crate::provider::is_external_provider(name) { - return true; - } - // Local / self-hosted inference servers that wrap `OpenAIProvider`. - // Keep in sync with the thin-wrapper modules (each sets `with_provider`). - const THIN_WRAPPERS: &[&str] = &[ - "ollama", - "lmstudio", - "vllm", - "llamafile", - "mistralrs", - "llamacpp", - "localai", - "local", - "mlx", - "omlx", - "onnx", - "oobabooba", - "openvino", - "sglang", - "xinference", - "cybertron", - "docker_model_runner", - "gaudi", - "jlama", - "litellm_proxy", - "openrouter", - ]; - THIN_WRAPPERS.contains(&name) -} - -/// 从录制 JSON 恢复 `OpenAICompatProfile`;缺失/未知字段回退默认(`full()`)。 -/// `&'static str` 字段从 String leak(bounded:每录制文件每条最多一次,同 -/// registry `profile_from_registry` 模式)。 -fn profile_from_value(v: Option<&serde_json::Value>) -> OpenAICompatProfile { - let Some(v) = v else { - return OpenAICompatProfile::full(); - }; - let s = |k: &str| { - v[k].as_str().map(|s| { - let leaked: &'static str = Box::leak(s.to_string().into_boxed_str()); - leaked - }) - }; - OpenAICompatProfile { - supports_top_k: v["supports_top_k"].as_bool().unwrap_or(true), - supports_tools: v["supports_tools"].as_bool().unwrap_or(true), - supports_response_format: v["supports_response_format"].as_bool().unwrap_or(true), - stream_usage_key: s("stream_usage_key"), - max_tokens_key: s("max_tokens_key"), - } -} - -/// 按 api_key_source 解析 key;explicit/unknown 缺 key 时报清晰错误。 -fn resolve_api_key(p: &ProviderRecord, api_key: Option<&str>) -> Result { - let source = p.api_key_source.trim(); - if let Some(var) = source.strip_prefix("env:") { - return std::env::var(var).map_err(|_| { - AiMuxError::InvalidArgument(format!( - "mock replay: env var '{var}' not set (needed to rebuild provider '{}')", - p.provider - )) - }); - } - match source { - // 录制不存明文(隐私强制),调用方必须补 key。 - "explicit" => api_key.map(str::to_string).ok_or_else(|| { - AiMuxError::InvalidArgument( - "mock replay: recording uses an explicit api key (not stored); \ - pass api_key to rebuild_provider" - .to_string(), - ) - }), - // 本地模型无需 key。 - "none" => Ok("not-needed".to_string()), - // "unknown"(默认快照)或空:用调用方给的 key,否则报错兜底。 - _ => api_key.map(str::to_string).ok_or_else(|| { - AiMuxError::InvalidArgument(format!( - "mock replay: no api key source for provider '{}' (api_key_source={:?}); \ - pass api_key or provide a model instance to replay_with_model", - p.provider, p.api_key_source - )) - }), + let model = registry + .provider(&p.provider_id)? + .language_model(&p.model_id)?; + if model.provider() != p.provider { + return Err(AiMuxError::InvalidArgument(format!( + "mock replay: the recording was made with `{}` but provider '{}' rebuilds `{}`; \ + pass the model to replay_with_model", + p.provider, + p.provider_id, + model.provider() + ))); } + Ok(model) } #[cfg(test)] mod tests { use super::*; - use aimux_core::recording::Recorder; - - fn openai_record(api_key_source: &str, model_id: &str) -> ProviderRecord { - ProviderRecord { - provider: "openai".into(), - model_id: model_id.into(), - base_url: Some("https://api.openai.com/v1".into()), - api_key_source: api_key_source.into(), - profile: None, - provider_options: None, - } - } - - #[test] - fn explicit_key_builds_model() { - let p = openai_record("explicit", "gpt-4o"); - let model = rebuild_provider(&p, Some("sk-test")).unwrap(); - assert_eq!(model.provider(), "openai"); - assert_eq!(model.model_id(), "gpt-4o"); - } - - #[test] - fn explicit_key_missing_errors() { - let p = openai_record("explicit", "gpt-4o"); - let err = rebuild_provider(&p, None).err().unwrap(); - assert!(matches!(err, AiMuxError::InvalidArgument(_)), "{err}"); - assert!(err.to_string().contains("explicit"), "{err}"); - } - - #[test] - fn env_key_reads_var() { - let p = openai_record("env:AIMUX_REPLAY_TEST_KEY", "gpt-4o"); - unsafe { std::env::set_var("AIMUX_REPLAY_TEST_KEY", "sk-env") }; - let model = rebuild_provider(&p, None).unwrap(); - assert_eq!(model.model_id(), "gpt-4o"); - unsafe { std::env::remove_var("AIMUX_REPLAY_TEST_KEY") }; - } - - #[test] - fn env_key_missing_errors() { - unsafe { std::env::remove_var("AIMUX_REPLAY_TEST_MISSING") }; - let p = openai_record("env:AIMUX_REPLAY_TEST_MISSING", "gpt-4o"); - let err = rebuild_provider(&p, None).err().unwrap(); - assert!(matches!(err, AiMuxError::InvalidArgument(_)), "{err}"); - assert!( - err.to_string().contains("AIMUX_REPLAY_TEST_MISSING"), - "{err}" - ); - } - - #[test] - fn none_source_uses_placeholder() { - let p = openai_record("none", "llama-3.3-70b"); - let model = rebuild_provider(&p, None).unwrap(); - assert_eq!(model.model_id(), "llama-3.3-70b"); - } - - #[test] - fn unknown_source_falls_back_to_caller_key() { - // provider 未覆盖 config_snapshot 时 api_key_source 为 "unknown"。 - let p = openai_record("unknown", "gpt-4o"); - let model = rebuild_provider(&p, Some("sk-caller")).unwrap(); - assert_eq!(model.model_id(), "gpt-4o"); - } - - #[test] - fn unknown_source_without_key_errors() { - let p = openai_record("unknown", "gpt-4o"); - let err = rebuild_provider(&p, None).err().unwrap(); - assert!(matches!(err, AiMuxError::InvalidArgument(_)), "{err}"); - } - - #[test] - fn non_openai_protocol_is_unsupported() { - let p = ProviderRecord { - provider: "anthropic".into(), - model_id: "claude-3-5-sonnet".into(), - base_url: None, - api_key_source: "env:ANTHROPIC_API_KEY".into(), - profile: None, - provider_options: None, - }; - let err = rebuild_provider(&p, Some("sk")).err().unwrap(); - assert!( - matches!(err, AiMuxError::UnsupportedFunctionality(_)), - "{err}" - ); - assert!(err.to_string().contains("replay_with_model"), "{err}"); - } - - #[test] - fn external_overlay_provider_is_replayable() { - // RFC-0020 overlay: a provider registered at runtime must be - // recognized by the replay path (is_openai_compatible_provider), - // otherwise rebuild_provider returns Unsupported and replay silently - // fails for newly-registered external providers. - let name = "test-replay-overlay"; - crate::provider::register_provider(crate::provider::ExternalProviderEntry { - name: name.into(), - display: None, - base_url: "https://relay.test.example/v1".into(), - env_var: None, - api_key: Some("dummy".into()), - protocol: "openai_compat".into(), - profile: crate::provider::ProviderProfile::default(), - headers: None, - organization: None, - project: None, - max_retries: None, - body_overrides: None, - comment: None, - }) - .unwrap(); - assert!( - is_openai_compatible_provider(name), - "external overlay provider must be replay-recognizable" - ); - // rebuild_provider should succeed (not Unsupported) for the overlay name. - let p = ProviderRecord { - provider: name.into(), - model_id: "m".into(), - base_url: Some("https://relay.test.example/v1".into()), - api_key_source: "explicit".into(), - profile: None, - provider_options: None, - }; - let result = rebuild_provider(&p, Some("dummy")); - assert!( - result.is_ok(), - "rebuild_provider should accept overlay name" - ); - // Clean up the overlay so other tests are unaffected. - crate::provider::clear_overlay(name); - } #[test] fn empty_model_id_errors() { - let p = openai_record("explicit", ""); - let err = rebuild_provider(&p, Some("sk")).err().unwrap(); + let registry = + aimux_core::create_provider_registry(crate::default_providers(), Default::default()); + let err = rebuild_provider(&ProviderRecord::new("groq", "groq.chat", ""), ®istry) + .err() + .unwrap(); assert!(matches!(err, AiMuxError::InvalidArgument(_)), "{err}"); } + /// A recording written by master (aimux 0.5, schema 2) has no + /// `provider_id`: both rebuilds refuse it at replay time, naming why. #[test] - fn provider_options_restore_org_and_project() { - let mut p = openai_record("explicit", "gpt-4o"); - p.provider_options = Some(serde_json::json!({ - "organization": "org-123", - "project": "proj-456", - })); - let model = rebuild_provider(&p, Some("sk")).unwrap(); - assert_eq!(model.model_id(), "gpt-4o"); - } - - #[test] - fn config_snapshot_roundtrips_through_rebuild() { - // 构造带非默认 profile + base_url 的 OpenAI 兼容 model(groq 风格)。 - let config = crate::openai::OpenAIConfig::new("sk-test") - .with_base_url("https://api.groq.com/openai/v1") - .with_provider("groq") - .with_profile(crate::openai::OpenAICompatProfile::groq()); - let model = crate::openai::OpenAIProvider::new(config).model("llama-3.3-70b"); - let snap = model.config_snapshot(); - assert_eq!(snap.provider, "groq"); // RFC-0023 C2: surfaces the registry name, not "openai" - assert_eq!( - snap.base_url.as_deref(), - Some("https://api.groq.com/openai/v1") - ); - assert_eq!(snap.api_key_source, "explicit"); - assert!(snap.profile.is_some(), "profile must be recorded"); - - // 从 snapshot 重建,再 snapshot → 配置等价(round-trip)。 - let rebuilt = rebuild_provider(&snap, Some("sk-test")).unwrap(); - let snap2 = rebuilt.config_snapshot(); - assert_eq!(snap2.provider, snap.provider); - assert_eq!(snap2.model_id, snap.model_id); - assert_eq!(snap2.base_url, snap.base_url); - assert_eq!(snap2.profile, snap.profile); - assert_eq!(snap2.api_key_source, snap.api_key_source); - assert_eq!(snap2.provider_options, snap.provider_options); - } - - #[test] - fn config_snapshot_marks_env_source() { - let config = crate::openai::OpenAIConfig::from_env().unwrap_or_else(|_| { - crate::openai::OpenAIConfig::new("sk").with_api_key_source(Some("env:OPENAI_API_KEY")) - }); - let model = crate::openai::OpenAIProvider::new(config).model("gpt-4o"); - assert_eq!(model.config_snapshot().api_key_source, "env:OPENAI_API_KEY"); - } - - #[test] - fn local_wrapper_marks_none_source() { - // ollama 等本地 wrapper:占位 key → "none",回放重建无需 key。 - let config = crate::ollama::OllamaConfig::from_env().unwrap(); - let model = crate::ollama::OllamaProvider::new(config).model("llama3.2"); - let snap = model.config_snapshot(); - assert_eq!(snap.api_key_source, "none", "{snap:?}"); - // "none" 来源重建:不传 key 也应成功。 - let rebuilt = rebuild_provider(&snap, None).unwrap(); - assert_eq!(rebuilt.model_id(), "llama3.2"); - } - - #[test] - fn recording_boundary_redacts_sensitive_headers() { - // 脱敏验证(P5):config.headers 带 Authorization 明文,snapshot 内存中 - // 可含(录制边界为门控),但经 RingRecorder 录制导出后必须无明文。 - use std::collections::HashMap; - let mut headers = HashMap::new(); - headers.insert("Authorization".to_string(), "Bearer sk-secret".to_string()); - headers.insert("X-Custom".to_string(), "ok".to_string()); - let config = crate::openai::OpenAIConfig::new("sk-secret").with_headers(headers); - let model = crate::openai::OpenAIProvider::new(config).model("gpt-4o"); - let snap = model.config_snapshot(); - assert!(snap.provider_options.is_some(), "headers must be recorded"); - - let ring = aimux_core::recording::RingRecorder::with_capacity(8); - let options = aimux_core::generate::GenerateTextOptions::default().into_call_options(vec![ - aimux_core::language_model_message::LanguageModelMessage::user_text("ping"), - ]); - ring.record_input("c1", &options, "openai", "gpt-4o"); - ring.record_provider("c1", &snap); - ring.record_outcome( - "c1", - &aimux_core::recording::OutcomeRecord { - status: aimux_core::recording::OutcomeStatus::Success, - finish_reason: Some("stop".into()), - error: None, - error_value: None, - decision_result: None, - usage: None, - }, - ); - ring.record_transport_closed("c1"); - - let mut buf = Vec::new(); - ring.export_jsonl(&mut buf).unwrap(); - let text = String::from_utf8(buf).unwrap(); - assert!(!text.contains("sk-secret"), "plaintext key leaked: {text}"); - assert!(text.contains("[REDACTED]"), "must be redacted: {text}"); - assert!( - text.contains("\"X-Custom\":\"ok\""), - "non-sensitive kept: {text}" - ); + fn a_recording_written_by_aimux_0_5_is_refused_at_rebuild() { + let line = include_str!("../../aimux-core/tests/fixtures/master_language_recordings.jsonl") + .lines() + .next() + .unwrap(); + let recording: aimux_core::recording::Recording = serde_json::from_str(line).unwrap(); + let registry = + aimux_core::create_provider_registry(crate::default_providers(), Default::default()); + + let err = rebuild_provider(&recording.provider, ®istry) + .err() + .unwrap(); + assert!(matches!(err, AiMuxError::InvalidArgument(_)), "{err}"); + assert!(err.to_string().contains("no provider_id"), "{err}"); + assert!(err.to_string().contains("aimux 0.5"), "{err}"); + let err = rebuild_decision_provider(&recording.provider, ®istry) + .err() + .unwrap(); + assert!(err.to_string().contains("no provider_id"), "{err}"); } } diff --git a/aimux-providers/src/replicate.rs b/aimux-providers/src/replicate.rs index b0dca270d..40a5447f4 100644 --- a/aimux-providers/src/replicate.rs +++ b/aimux-providers/src/replicate.rs @@ -2,8 +2,13 @@ //! //! Aligned with Vercel AI SDK `ReplicateImageModel` //! (`reference/ai/packages/replicate/src/replicate-image-model.ts`). +//! +//! [`create_replicate`] takes [`ReplicateProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`ReplicateProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `REPLICATE_API_TOKEN`. +//! [`replicate()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use std::time::Duration; use async_trait::async_trait; @@ -14,9 +19,14 @@ use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; -use aimux_core::retry; use aimux_core::shared::Warning; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; + +use crate::shared::{ + AuthScheme, Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, credential_headers, + is_poll_control_key, poll_interval_ms, retry_download, +}; /// Replicate's `prefer: wait` holds the connection for at most this long when /// no explicit duration is given (their documented default and maximum). @@ -40,84 +50,181 @@ fn replicate_failed_response_handler() -> aimux_provider_utils::ResponseHandler< }) } -/// Configuration for the Replicate provider. -#[derive(Debug, Clone)] -pub struct ReplicateConfig { - pub api_key: String, - pub base_url: String, - pub headers: Option>, -} +/// Milliseconds between two attempts to download a finished prediction's output +/// (`providerOptions.replicate.pollIntervalMs` overrides it for one call). +const DOWNLOAD_RETRY_INTERVAL_MS: u64 = 1_000; -impl ReplicateConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.replicate.com/v1".to_string(), - headers: None, - } - } +pub(crate) mod options; - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.replicate.com/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "REPLICATE_API_TOKEN"; +const DEFAULT_NAME: &str = "replicate"; - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } +/// Settings of [`create_replicate`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct ReplicateProviderSettings { + /// Base URL for the API calls. Default `https://api.replicate.com/v1`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `REPLICATE_API_TOKEN` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.image"`, `"{name}.video"`). + /// Default `"replicate"`. The providerOptions key stays `replicate`. + 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, +} - /// Create from the `REPLICATE_API_TOKEN` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "REPLICATE_API_TOKEN", "Replicate")?; - Ok(Self::new(api_key)) +impl std::fmt::Debug for ReplicateProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ReplicateProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } -/// Replicate provider — creates `ReplicateImageModel` instances. +/// Create a Replicate provider. +/// +/// # 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. +pub fn create_replicate( + settings: ReplicateProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(ReplicateProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + credential: Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Replicate"), + user_headers: settings.headers, + fetch: settings.fetch, + }) +} + +/// The default provider: `create_replicate` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn replicate() -> &'static ReplicateProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_replicate(ReplicateProviderSettings::default()) + .expect("default Replicate settings are always valid") + }) +} + +/// A Replicate provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct ReplicateProvider { - config: ReplicateConfig, + name: String, + base_url: String, + credential: Credential, + user_headers: Option, + fetch: Option, } impl ReplicateProvider { - #[must_use] - pub fn new(config: ReplicateConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + match method { + "image" => credential_headers( + self.credential.clone(), + AuthScheme::Bearer, + Vec::new(), + self.user_headers.clone(), + ), + _ => credential_headers( + self.credential.clone(), + AuthScheme::Scheme("Token"), + Vec::new(), + self.user_headers.clone(), + ), + }, + self.fetch.clone(), + None, + ) } + /// An image model (e.g. `"black-forest-labs/flux-schnell"`); `provider()` is `"{name}.image"`. #[must_use] pub fn image(&self, model_id: &str) -> ReplicateImageModel { - ReplicateImageModel::new(model_id.to_string(), self.config.clone()) + ReplicateImageModel::from_config(model_id.to_string(), self.model_config("image")) } - /// Create a video generation model instance for the given model name - /// (e.g. `"wan-lab/wan-2.1-t2v-14b"`). + /// A video model; `provider()` is `"{name}.video"`. #[must_use] pub fn video(&self, model_id: &str) -> ReplicateVideoModel { - ReplicateVideoModel::new(model_id.to_string(), self.config.clone()) + ReplicateVideoModel::from_config(model_id.to_string(), self.model_config("video")) + } +} + +impl ::aimux_core::Provider for ReplicateProvider { + fn language_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "languageModel")) + } + + fn embedding_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model( + &self, + model_id: &str, + ) -> Result<::std::sync::Arc, AiMuxError> { + Ok(::std::sync::Arc::new(self.image(model_id))) + } + + fn video_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(::std::sync::Arc::new(self.video(model_id)))) } } -/// Flux-2 model pattern for max input images. const FLUX_2_PATTERN: &str = "black-forest-labs/flux-2-"; const MAX_FLUX_2_INPUT_IMAGES: u32 = 8; /// A Replicate image generation model. pub struct ReplicateImageModel { model_id: String, - config: ReplicateConfig, + config: EndpointConfig, } impl ReplicateImageModel { - #[must_use] - pub fn new(model_id: String, config: ReplicateConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } @@ -125,25 +232,6 @@ impl ReplicateImageModel { self.model_id.starts_with(FLUX_2_PATTERN) } - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut h = HashMap::new(); - h.insert( - "Authorization".into(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - h.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - h.insert(k.clone(), v.clone()); - } - } - h - } - fn file_to_data_uri(file: &ImageFile) -> Result { match file { ImageFile::Url { url } => Ok(url.clone()), @@ -159,22 +247,21 @@ impl ReplicateImageModel { } } - fn endpoint(&self) -> String { - if let Some((model_id, _version)) = self.model_id.split_once(':') { - format!("{}/models/{}/predictions", self.config.base_url, model_id) - } else { - format!( - "{}/models/{}/predictions", - self.config.base_url, self.model_id - ) - } + /// The path of the prediction endpoint of this model; a `:version` suffix + /// of the model id is sent in the body, not in the path. + fn endpoint_path(&self) -> String { + let model = self + .model_id + .split_once(':') + .map_or(self.model_id.as_str(), |(model, _version)| model); + format!("/models/{model}/predictions") } } #[async_trait] impl ImageModel for ReplicateImageModel { fn provider(&self) -> &str { - "replicate" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id @@ -194,7 +281,7 @@ impl ImageModel for ReplicateImageModel { .split_once(':') .unwrap_or((&self.model_id, "")); - let replicate_opts = options.provider_options.get("replicate"); + let replicate_opts = options::replicate_options(Some(&options.provider_options)); let max_wait = replicate_opts .and_then(|o| o.get("maxWaitTimeInSeconds")) .and_then(serde_json::Value::as_u64); @@ -265,10 +352,11 @@ impl ImageModel for ReplicateImageModel { input.insert("mask".into(), json!(m)); } - // Forward replicate provider options (excluding maxWaitTimeInSeconds) + // Forward replicate provider options (excluding maxWaitTimeInSeconds and + // the pacing keys, which are read here, never sent) if let Some(ro) = replicate_opts { for (k, v) in ro { - if k != "maxWaitTimeInSeconds" { + if k != "maxWaitTimeInSeconds" && !is_poll_control_key(k) { input.insert(k.clone(), v.clone()); } } @@ -280,37 +368,27 @@ impl ImageModel for ReplicateImageModel { body.insert("version".into(), json!(version)); } - // Build headers with prefer - let mut headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let prefer = if let Some(mw) = max_wait { format!("wait={mw}") } else { "wait".to_string() }; - headers.insert("prefer".into(), prefer); - - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let mut request = exchange.request(exchange.url(&self.endpoint_path()), options); + request + .headers + .retain(|(name, _)| !name.eq_ignore_ascii_case("prefer")); + request.headers.push(("prefer".to_string(), prefer)); + // This exchange legitimately holds the connection for the whole wait + // window; widen the hang guard past it so a slow-but-alive generation + // is not misread as a dead transport (retryable -> re-create -> double + // billing). + request.response_timeout = Some(Duration::from_secs( + max_wait.unwrap_or(DEFAULT_WAIT_SECONDS) + WAIT_TRANSPORT_MARGIN_SECONDS, + )); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers: header_list, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - // This exchange legitimately holds the connection for the - // whole wait window; widen the hang guard past it so a - // slow-but-alive generation is not misread as a dead - // transport (retryable → re-create → double billing). - response_timeout: Some(Duration::from_secs( - max_wait.unwrap_or(DEFAULT_WAIT_SECONDS) + WAIT_TRANSPORT_MARGIN_SECONDS, - )), - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, + request, Value::Object(body), aimux_provider_utils::create_json_response_handler(), replicate_failed_response_handler(), @@ -319,11 +397,6 @@ impl ImageModel for ReplicateImageModel { let rh = resp.response_headers; let rb: Value = resp.value; - let retries = retry::prepare_retries( - options.max_retries, - self.retry_config(), - options.abort_signal.clone(), - ); // Extract output (string or array of strings) let urls: Vec = match &rb["output"] { @@ -336,30 +409,33 @@ impl ImageModel for ReplicateImageModel { }; // Download images; output URLs come from the prediction response - // body, so they go through the SSRF download guard. + // body, so they go through the SSRF download guard. A transient failure + // repeats the download, never the prediction. let mut downloaded: Vec> = Vec::new(); for url in &urls { - let ir = retries - .retry(|| { + let ir = retry_download( + options.abort_signal.as_ref(), + Duration::from_millis(poll_interval_ms( + replicate_opts, + POLL_INTERVAL_MS_KEY, + DOWNLOAD_RETRY_INTERVAL_MS, + )), + || { aimux_provider_utils::get_from_api( HttpRequest { url: url.clone(), - headers: vec![], - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, validate_url: true, - trusted_origin: Some(self.config.base_url.clone()), - credentialed_origin: Some(self.config.base_url.clone()), + trusted_origin: Some(exchange.base_url().to_string()), + credentialed_origin: Some(exchange.base_url().to_string()), + ..Default::default() }, aimux_provider_utils::create_binary_response_handler(), replicate_failed_response_handler(), ) - }) - .await?; + }, + ) + .await?; downloaded.push(ir.value.to_vec()); } @@ -395,33 +471,13 @@ use aimux_core::video_model::{ /// status polling via `do_status`. pub struct ReplicateVideoModel { model_id: String, - config: ReplicateConfig, + config: EndpointConfig, } impl ReplicateVideoModel { - #[must_use] - pub fn new(model_id: String, config: ReplicateConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Token {}", self.config.api_key), - ); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } } fn video_file_to_url(file: &VideoFile) -> Result { @@ -440,7 +496,7 @@ fn video_file_to_url(file: &VideoFile) -> Result { #[async_trait] impl VideoModel for ReplicateVideoModel { fn provider(&self) -> &str { - "replicate" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id @@ -471,24 +527,11 @@ impl VideoModel for ReplicateVideoModel { "input": Value::Object(input), }); - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; // Submit prediction. let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: format!("{}/predictions", self.config.base_url), - headers: header_list, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, + exchange.request(exchange.url("/predictions"), options), body, aimux_provider_utils::create_json_response_handler(), replicate_failed_response_handler(), @@ -531,24 +574,11 @@ impl VideoModel for ReplicateVideoModel { ) })?; - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); - let poll_url = format!("{}/predictions/{}", self.config.base_url, prediction_id); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let poll_url = exchange.url(&format!("/predictions/{prediction_id}")); let resp = aimux_provider_utils::get_from_api( - HttpRequest { - url: poll_url.clone(), - headers: header_list, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - response_timeout: None, - max_json_response_bytes: None, - validate_url: false, - trusted_origin: None, - credentialed_origin: None, - }, + exchange.request(poll_url.clone(), options), aimux_provider_utils::create_json_response_handler::(), replicate_failed_response_handler(), ) diff --git a/aimux-providers/src/replicate/options.rs b/aimux-providers/src/replicate/options.rs new file mode 100644 index 000000000..7d8f68522 --- /dev/null +++ b/aimux-providers/src/replicate/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Replicate. +//! +//! Replicate reads `providerOptions.replicate`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "replicate"; + +/// The Replicate options in a providerOptions map. +pub(crate) fn replicate_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/revai.rs b/aimux-providers/src/revai.rs index 8e2a3c145..109f06b91 100644 --- a/aimux-providers/src/revai.rs +++ b/aimux-providers/src/revai.rs @@ -7,8 +7,14 @@ //! 1. POST `/speechtotext/v1/jobs` with multipart form (media file + config JSON) //! 2. GET `/speechtotext/v1/jobs/{id}` to poll until status is `transcribed` //! 3. GET `/speechtotext/v1/jobs/{id}/transcript` to fetch the final transcript +//! +//! [`create_revai`] takes [`RevaiProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`RevaiProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `REVAI_API_KEY`. +//! [`revai()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::Duration; use async_trait::async_trait; use serde::Deserialize; @@ -16,15 +22,17 @@ use serde_json::{Value, json}; use aimux_core::error::AiMuxError; use aimux_core::error::ApiCallError; -use aimux_core::retry; use aimux_core::shared::Warning; use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::{ - HttpRequest, MultipartForm, load_api_key, media_type_to_extension, sleep_or_abort, - without_trailing_slash, +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{MultipartForm, media_type_to_extension}; + +use crate::shared::{ + Credential, EndpointConfig, POLL_INTERVAL_MS_KEY, PollStep, poll_interval_ms, poll_until, + provider_headers, retry_download, }; fn revai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -49,62 +57,130 @@ fn revai_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, +/// Milliseconds between two polls of a job +/// (`providerOptions.revai.pollIntervalMs` overrides it for one call). +const POLL_INTERVAL_MS: u64 = 100; +/// How many times a job is polled before the call gives up (ten minutes at the +/// default interval). +const MAX_POLL_ATTEMPTS: u32 = 6_000; + +pub(crate) mod options; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.rev.ai"; +pub(crate) const API_KEY_ENV_VAR: &str = "REVAI_API_KEY"; +const DEFAULT_NAME: &str = "revai"; + +/// Settings of [`create_revai`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct RevaiProviderSettings { + /// Base URL for the API calls. Default `https://api.rev.ai`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `REVAI_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.transcription"`). + /// Default `"revai"`. The providerOptions key stays `revai`. + 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, } -impl RevaiConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.rev.ai".to_string(), - headers: None, - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self +impl std::fmt::Debug for RevaiProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RevaiProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } +/// Create a Rev.ai provider. +/// +/// # 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. +pub fn create_revai(settings: RevaiProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(RevaiProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Rev.ai"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `REVAI_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "REVAI_API_KEY", "Rev.ai")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_revai` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn revai() -> &'static RevaiProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_revai(RevaiProviderSettings::default()) + .expect("default Rev.ai settings are always valid") + }) } +/// A Rev.ai provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct RevaiProvider { - config: RevaiConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl RevaiProvider { - #[must_use] - pub fn new(config: RevaiConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// A transcription (STT) model (e.g. `"machine"`); `provider()` is `"{name}.transcription"`. #[must_use] pub fn transcription(&self, model_id: &str) -> RevaiTranscriptionModel { - RevaiTranscriptionModel::new(model_id.to_string(), self.config.clone()) + RevaiTranscriptionModel::from_config( + model_id.to_string(), + self.model_config("transcription"), + ) } } +crate::impl_single_modality_provider!(RevaiProvider, transcription_model, |p, id| p + .transcription(id)); + // ── Response schema ───────────────────────────────────────────────────────── #[derive(Debug, Deserialize)] @@ -148,59 +224,25 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid base64: {e}"))) } + AudioInput::Url(_) => Err(AudioInput::url_error()), } } pub struct RevaiTranscriptionModel { model_id: String, - config: RevaiConfig, + config: EndpointConfig, } impl RevaiTranscriptionModel { - #[must_use] - pub fn new(model_id: String, config: RevaiConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn jobs_url(&self) -> String { - format!("{}/speechtotext/v1/jobs", self.config.base_url) - } - - fn job_status_url(&self, job_id: &str) -> String { - format!("{}/speechtotext/v1/jobs/{}", self.config.base_url, job_id) - } - - fn transcript_url(&self, job_id: &str) -> String { - format!( - "{}/speechtotext/v1/jobs/{}/transcript", - self.config.base_url, job_id - ) - } } #[async_trait] impl TranscriptionModel for RevaiTranscriptionModel { fn provider(&self) -> &str { - "revai" + &self.config.provider } fn model_id(&self) -> &str { @@ -224,19 +266,13 @@ impl TranscriptionModel for RevaiTranscriptionModel { let mut form = MultipartForm::new(); form.file("media", &filename, &options.media_type, &audio_bytes)?; form.text("config", &config)?; - let headers = self.build_headers(options.headers.as_ref()); - let submit_url = self.jobs_url(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let submit_url = exchange.url("/speechtotext/v1/jobs"); - // Submit job. + // Submit job. This creates the job and is sent exactly once; only the + // status poll and the transcript fetch below repeat. let resp = aimux_provider_utils::post_form_data_to_api( - HttpRequest::new( - submit_url.clone(), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), + exchange.request(submit_url.clone(), options), form, aimux_provider_utils::create_json_response_handler::(), revai_failed_response_handler(), @@ -265,74 +301,56 @@ impl TranscriptionModel for RevaiTranscriptionModel { ) })?; let submission_language = submit_response.language; - let retries = retry::prepare_retries( - options.max_retries, - self.retry_config(), - options.abort_signal.clone(), - ); + let poll_interval = Duration::from_millis(poll_interval_ms( + options::revai_options(options.provider_options.as_ref()), + POLL_INTERVAL_MS_KEY, + POLL_INTERVAL_MS, + )); // Poll for completion. - let job_status: RevaiJobResponse; - loop { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - - let poll_url = self.job_status_url(&job_id); - let resp = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest::new( - poll_url.clone(), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), - aimux_provider_utils::create_json_response_handler::(), - revai_failed_response_handler(), - ) - }) - .await?; - - let response_body = resp.raw_value.as_ref().map(ToString::to_string); - let poll: RevaiJobResponse = resp.value; - - if poll.status.as_deref() == Some("transcribed") { - job_status = poll; - break; - } - if poll.status.as_deref() == Some("failed") { - return Err(AiMuxError::ApiCall(Box::new(ApiCallError { - status_code: Some(200), - provider_code: Some("failed".to_string()), - response_body, - ..ApiCallError::new("Transcription job failed", poll_url, serde_json::json!({})) - }))); - } - } - let _ = job_status; - - // Fetch transcript. - let resp = retries - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest::new( - self.transcript_url(&job_id), - headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - options, - ), - aimux_provider_utils::create_json_response_handler::(), + let poll_url = exchange.url(&format!("/speechtotext/v1/jobs/{job_id}")); + poll_until( + &format!("revai job {job_id}"), + options.abort_signal.as_ref(), + poll_interval, + MAX_POLL_ATTEMPTS, + || async { + let resp = aimux_provider_utils::get_from_api( + exchange.request(poll_url.clone(), options), + aimux_provider_utils::create_json_response_handler::(), revai_failed_response_handler(), ) - }) - .await?; + .await?; + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + match resp.value.status.as_deref() { + Some("transcribed") => Ok(PollStep::Ready(())), + Some("failed") => Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some("failed".to_string()), + response_body, + ..ApiCallError::new( + "Transcription job failed", + poll_url.clone(), + serde_json::json!({}), + ) + }))), + _ => Ok(PollStep::Pending), + } + }, + ) + .await?; + + // Fetch transcript. A transient failure repeats this request, never the + // submit. + let transcript_url = exchange.url(&format!("/speechtotext/v1/jobs/{job_id}/transcript")); + let resp = retry_download(options.abort_signal.as_ref(), poll_interval, || { + aimux_provider_utils::get_from_api( + exchange.request(transcript_url.clone(), options), + aimux_provider_utils::create_json_response_handler::(), + revai_failed_response_handler(), + ) + }) + .await?; let response_headers = resp.response_headers; let raw_body = resp.raw_value.unwrap_or(Value::Null); diff --git a/aimux-providers/src/revai/options.rs b/aimux-providers/src/revai/options.rs new file mode 100644 index 000000000..b6aa9042b --- /dev/null +++ b/aimux-providers/src/revai/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Rev.ai. +//! +//! Rev.ai reads `providerOptions.revai`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "revai"; + +/// The Rev.ai options in a providerOptions map. +pub(crate) fn revai_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/runtime_decision.rs b/aimux-providers/src/runtime_decision.rs index dc0441a25..2c4edd904 100644 --- a/aimux-providers/src/runtime_decision.rs +++ b/aimux-providers/src/runtime_decision.rs @@ -3,13 +3,12 @@ use std::collections::{BTreeMap, HashMap}; use aimux_core::AiMuxError; use aimux_core::decision_model::*; -use aimux_core::recording::ProviderRecord; use aimux_core::types::{InputTokenUsage, OutputTokenUsage, Usage}; use async_trait::async_trait; use serde_json::{Map, Value, json}; -use crate::decision_support::{self as support, invalid}; -use crate::openai::OpenAIConfig; +use crate::decision_support::{self as support, DecisionConfig, invalid}; +use crate::openai_compatible::config::CompatModelConfig; pub(crate) fn supported(provider: &str) -> bool { matches!( @@ -25,18 +24,24 @@ pub(crate) fn supported(provider: &str) -> bool { ) } +/// Decisions of an official runtime, made from its OpenAI-compatible preset. pub(crate) struct RuntimeDecisionModel { - pub(crate) config: OpenAIConfig, + /// The preset name; it selects the runtime's protocol and limits. + pub(crate) provider: String, + pub(crate) config: DecisionConfig, pub(crate) model_id: String, } impl RuntimeDecisionModel { - pub(crate) fn new(model_id: &str, config: OpenAIConfig) -> Result { - if !supported(&config.provider) || config.body_overrides.is_some() { - return Err(AiMuxError::UnsupportedFunctionality( - "native decisions require an official adapter without generation body_overrides" - .into(), - )); + pub(crate) fn new( + provider: &str, + model_id: &str, + config: CompatModelConfig, + ) -> Result { + if !supported(provider) { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "provider '{provider}' does not provide decision models" + ))); } if model_id.trim().is_empty() { return Err(AiMuxError::InvalidArgument( @@ -44,7 +49,8 @@ impl RuntimeDecisionModel { )); } Ok(Self { - config, + provider: provider.into(), + config: DecisionConfig::Compatible(config), model_id: model_id.into(), }) } @@ -153,7 +159,8 @@ impl RuntimeDecisionModel { } copy_options(&mut body, options, allowed)?; self.add_images(request, &mut body)?; - let base = self.config.base_url.trim_end_matches('/'); + let base = self.config.base_url()?; + let base = base.as_str(); let url = if matches!(provider, "cloudflare" | "cloudflare_workers_ai") { let model = self .model_id @@ -616,14 +623,11 @@ fn sglang_response( #[async_trait] impl DecisionModel for RuntimeDecisionModel { fn provider(&self) -> &str { - &self.config.provider + &self.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } fn capabilities(&self) -> DecisionCapabilities { let provider = self.provider(); DecisionCapabilities { @@ -661,24 +665,13 @@ impl DecisionModel for RuntimeDecisionModel { }, } } - fn config_snapshot(&self) -> ProviderRecord { - let mut record = crate::openai::config_snapshot_from_config( - self.provider(), - self.model_id(), - &self.config, - ); - record.profile = Some( - json!({"decision_default_protocol":"systemone", "decision_capabilities":self.capabilities()}), - ); - record - } async fn do_decide(&self, request: &DecisionCallOptions) -> Result { let mut options = support::options(request, self.provider())?; let protocol = options.remove("protocol").unwrap_or(json!("systemone")); match (self.provider(), protocol.as_str()) { ("sglang", Some("decisions")) => { let body = self.sglang_request(request, options)?; - let url = format!("{}/decisions", self.config.base_url.trim_end_matches('/')); + let url = format!("{}/decisions", self.config.base_url()?); support::post(&self.config, request, &url, body, |raw, headers| { sglang_response(request, raw, headers) }) diff --git a/aimux-providers/src/runwayml.rs b/aimux-providers/src/runwayml.rs index 9413245f6..ae39f2154 100644 --- a/aimux-providers/src/runwayml.rs +++ b/aimux-providers/src/runwayml.rs @@ -8,24 +8,28 @@ //! //! Authentication is a `Bearer` token and every request must carry the //! `X-Runway-Version: 2024-11-06` header. +//! +//! [`create_runwayml`] takes [`RunwaymlProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`RunwaymlProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `RUNWAYML_API_SECRET`. +//! [`runwayml()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Map, Value, json}; use aimux_core::error::{AiMuxError, ApiCallError}; -use aimux_core::provider::Provider; use aimux_core::video_model::{ VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoFrameType, VideoModel, VideoOperationStart, VideoOperationStatus, VideoResponse, VideoResult, }; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; const PROVIDER_NAME: &str = "runwayml"; -const DEFAULT_BASE_URL: &str = "https://api.dev.runwayml.com"; -const ENV_VAR: &str = "RUNWAYML_API_SECRET"; const RUNWAY_VERSION: &str = "2024-11-06"; /// RunwayML returns errors as a flat `{"error": ""}` object. @@ -44,71 +48,116 @@ fn runwayml_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.dev.runwayml.com"; +pub(crate) const API_KEY_ENV_VAR: &str = "RUNWAYML_API_SECRET"; +const DEFAULT_NAME: &str = "runwayml"; + +/// Settings of [`create_runwayml`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct RunwaymlProviderSettings { + /// Base URL for the API calls. Default `https://api.dev.runwayml.com`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `RUNWAYML_API_SECRET` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.video"`). + /// Default `"runwayml"`. The providerOptions key stays `runwayml`. + 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, } -impl RunwaymlConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: DEFAULT_BASE_URL.to_string(), - headers: None, - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self +impl std::fmt::Debug for RunwaymlProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RunwaymlProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Create from the `RUNWAYML_API_SECRET` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, ENV_VAR, "RunwayML")?; - Ok(Self::new(api_key)) - } +/// Create a RunwayML provider. +/// +/// # 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. +pub fn create_runwayml(settings: RunwaymlProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(RunwaymlProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "RunwayML"), + vec![("X-Runway-Version".to_string(), RUNWAY_VERSION.to_string())], + settings.headers, + ), + fetch: settings.fetch, + }) } -// ── Provider ──────────────────────────────────────────────────────────────── +/// The default provider: `create_runwayml` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn runwayml() -> &'static RunwaymlProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_runwayml(RunwaymlProviderSettings::default()) + .expect("default RunwayML settings are always valid") + }) +} +/// A RunwayML provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct RunwaymlProvider { - config: RunwaymlConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl RunwaymlProvider { - #[must_use] - pub fn new(config: RunwaymlConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a video generation model instance for the given model ID. + /// A video model (e.g. `"gen4_turbo"`); `provider()` is `"{name}.video"`. #[must_use] pub fn video(&self, model_id: &str) -> RunwaymlVideoModel { - RunwaymlVideoModel::new(model_id.to_string(), self.config.clone()) + RunwaymlVideoModel::from_config(model_id.to_string(), self.model_config("video")) } } -impl Provider for RunwaymlProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } -} +crate::impl_single_modality_provider!(RunwaymlProvider, video_model, |p, id| p.video(id)); // ── Helpers ───────────────────────────────────────────────────────────────── @@ -146,40 +195,19 @@ struct RunwaymlTaskDetailsResponse { /// A RunwayML video generation model — implements [`VideoModel`]. pub struct RunwaymlVideoModel { model_id: String, - config: RunwaymlConfig, + config: EndpointConfig, } impl RunwaymlVideoModel { - #[must_use] - pub fn new(model_id: String, config: RunwaymlConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - headers.insert("X-Runway-Version".to_string(), RUNWAY_VERSION.to_string()); - if let Some(ref config_headers) = self.config.headers { - for (k, v) in config_headers { - headers.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } } #[async_trait] impl VideoModel for RunwaymlVideoModel { fn provider(&self) -> &str { - PROVIDER_NAME + &self.config.provider } fn model_id(&self) -> &str { @@ -228,19 +256,18 @@ impl VideoModel for RunwaymlVideoModel { body.insert("ratio".to_string(), json!(ar.to_string())); } - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let submit_path = if is_image_to_video { "/v1/image_to_video" } else { "/v1/text_to_video" }; - let submit_url = format!("{}{submit_path}", self.config.base_url); + let submit_url = exchange.url(submit_path); // Submit the task. let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(submit_url, header_list, options), + exchange.request(submit_url, options), Value::Object(body), aimux_provider_utils::create_json_response_handler(), runwayml_failed_response_handler(), @@ -276,12 +303,11 @@ impl VideoModel for RunwaymlVideoModel { )) })?; - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); - let poll_url = format!("{}/v1/tasks/{}", self.config.base_url, task_id); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let poll_url = exchange.url(&format!("/v1/tasks/{task_id}")); let resp = aimux_provider_utils::get_from_api( - HttpRequest::new(poll_url.clone(), header_list, options), + exchange.request(poll_url.clone(), options), aimux_provider_utils::create_json_response_handler::(), runwayml_failed_response_handler(), ) diff --git a/aimux-providers/src/searxng.rs b/aimux-providers/src/searxng.rs index f4e58db88..9e27cf271 100644 --- a/aimux-providers/src/searxng.rs +++ b/aimux-providers/src/searxng.rs @@ -3,92 +3,162 @@ //! Implements the `SearchModel` trait against a self-hosted SearXNG instance //! (`GET {SEARXNG_URL}/search?q=...&format=json`). //! -//! SearXNG is a modality-specific, unauthenticated provider: there is no API -//! key. The instance URL is supplied via the `SEARXNG_URL` environment -//! variable (required — there is no default). A `403` response typically -//! indicates that the `json` output format is not enabled on the instance. +//! SearXNG is a self-hosted, normally unauthenticated provider: there is no +//! default instance. The instance URL is the `base_url` setting or the +//! `SEARXNG_URL` environment variable, read when a request is made (a missing +//! URL fails that request with `AiMuxError::LoadSetting`). A `403` response +//! typically indicates that the `json` output format is not enabled on the +//! instance. +//! +//! [`create_searxng`] takes [`SearxngProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`SearxngProvider`]. The instance URL is not read +//! there: the instance URL is resolved for every request. +//! [`searxng()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::{Arc, OnceLock}; use async_trait::async_trait; use serde::Deserialize; use serde_json::Value; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::HttpRequest; -use aimux_provider_utils::without_trailing_slash; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, load_setting, validate_base_url, +}; + +use crate::shared::{Credential, Endpoint, EndpointConfig, provider_headers}; /// Fixed model ID for the SearXNG search model. const MODEL_ID: &str = "searxng-search"; -/// Configuration for the SearXNG provider. +const BASE_URL_ENV_VAR: &str = "SEARXNG_URL"; +const DEFAULT_NAME: &str = "searxng"; + +/// Settings of [`create_searxng`]. /// -/// SearXNG is unauthenticated, so there is no API key — only the instance -/// base URL. -#[derive(Debug, Clone)] -pub struct SearxngConfig { - pub base_url: String, +/// Every field is optional. The instance URL is validated when the provider is +/// created if it is given; `api_key` and `headers` are evaluated on every +/// request. +#[derive(Clone, Default)] +pub struct SearxngProviderSettings { + /// The instance URL; a trailing slash is removed. `None` reads + /// `SEARXNG_URL` when a request is made and fails that request with + /// `AiMuxError::LoadSetting` if it is unset: there is no default instance. + pub base_url: Option, + /// An optional bearer token for an instance behind an authenticating + /// proxy. `None` sends no `Authorization` header. A [`Resolvable::Future`] + /// is awaited once, an [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header. + /// Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` string + /// (`"{name}.search"`). Default `"searxng"`. + 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, } -impl SearxngConfig { - /// Create from an explicit instance base URL. - pub fn new(base_url: impl Into) -> Self { - Self { - base_url: without_trailing_slash(&base_url.into()), - } - } - - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - - /// Create from the `SEARXNG_URL` environment variable (required). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `SEARXNG_URL` is not set. - pub fn from_env() -> Result { - let base_url = std::env::var("SEARXNG_URL").map_err(|_| { - AiMuxError::InvalidArgument( - "SEARXNG_URL environment variable is required for SearXNG".to_string(), +impl std::fmt::Debug for SearxngProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SearxngProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), ) - })?; - Ok(Self::new(base_url)) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } -/// SearXNG provider — creates `SearxngSearchModel` instances. +/// Create a SearXNG provider. /// -/// SearXNG is a search-only provider; it does not support language models. +/// # Errors +/// +/// Returns `AiMuxError::InvalidArgument` when `base_url` is given and is not +/// an `http(s)` URL with a host. That is the only way this fails: a missing +/// instance URL is reported by the first request. +pub fn create_searxng(settings: SearxngProviderSettings) -> Result { + let base_url = settings + .base_url + .as_deref() + .map(validate_base_url) + .transpose()?; + let credential = match settings.api_key { + Some(key) => Credential::Explicit(key), + None => Credential::None, + }; + Ok(SearxngProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers(credential, Vec::new(), settings.headers), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_searxng` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing instance URL surfaces from the first request instead. +pub fn searxng() -> &'static SearxngProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_searxng(SearxngProviderSettings::default()) + .expect("default SearXNG settings are always valid") + }) +} + +/// A SearXNG provider. Search only; it holds no HTTP client. pub struct SearxngProvider { - config: SearxngConfig, + name: String, + base_url: Option, + headers: HeadersFn, + fetch: Option, } impl SearxngProvider { - #[must_use] - pub fn new(config: SearxngConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + let base_url = self.base_url.clone(); + let headers = self.headers.clone(); + EndpointConfig::dynamic( + format!("{}.{method}", self.name), + Arc::new(move || { + let base_url = base_url.clone(); + let headers = headers.clone(); + Box::pin(async move { + let base_url = match base_url { + Some(url) => url, + None => validate_base_url(&load_setting( + None, + BASE_URL_ENV_VAR, + "SearXNG instance URL", + )?)?, + }; + Ok(Endpoint { + base_url, + headers: headers.resolve().await?, + }) + }) + }), + self.fetch.clone(), + ) } - /// Create a search model instance. + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> SearxngSearchModel { - SearxngSearchModel::new(self.config.clone()) + SearxngSearchModel::from_config(self.model_config("search")) } } -impl Provider for SearxngProvider { - fn name(&self) -> &str { - "searxng" - } -} +crate::impl_single_modality_provider!(SearxngProvider, search_model, |p, _id| p.search_model()); /// A single SearXNG result entry. All fields are optional so unknown-but-legal /// values degrade safely. @@ -128,24 +198,19 @@ fn map_results(entries: Vec) -> Vec { /// A SearXNG search model. pub struct SearxngSearchModel { - config: SearxngConfig, + config: EndpointConfig, } impl SearxngSearchModel { - #[must_use] - pub fn new(config: SearxngConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn endpoint(&self) -> String { - format!("{}/search", self.config.base_url) - } } #[async_trait] impl SearchModel for SearxngSearchModel { fn provider(&self) -> &str { - "searxng" + &self.config.provider } fn model_id(&self) -> &str { @@ -153,31 +218,16 @@ impl SearchModel for SearxngSearchModel { } async fn do_search(&self, options: &SearchCallOptions) -> Result { - // SearXNG is unauthenticated; only forward user-supplied extra headers. - let headers: Vec<(String, String)> = options - .headers - .as_ref() - .map(|extra: &HashMap| { - extra.iter().map(|(k, v)| (k.clone(), v.clone())).collect() - }) - .unwrap_or_default(); - - let mut url = url::Url::parse(&self.endpoint()) + let exchange = self.config.exchange(options.headers.as_ref()).await?; + + let mut url = url::Url::parse(&exchange.url("/search")) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid searxng endpoint: {e}")))?; url.query_pairs_mut() .append_pair("q", &options.query) .append_pair("format", "json"); let resp = aimux_provider_utils::get_from_api( - HttpRequest { - url: url.to_string(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(url.to_string(), options), aimux_provider_utils::create_json_response_handler::(), aimux_provider_utils::create_status_code_error_response_handler(), ) diff --git a/aimux-providers/src/serper.rs b/aimux-providers/src/serper.rs index 5c87a9126..caa15e7de 100644 --- a/aimux-providers/src/serper.rs +++ b/aimux-providers/src/serper.rs @@ -3,77 +3,161 @@ //! Implements the `SearchModel` trait against the Serper search API //! (`POST https://google.serper.dev/search`). Uses `X-API-KEY` header auth //! via `SERPER_API_KEY`. +//! +//! [`create_serper`] takes [`SerperProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`SerperProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `SERPER_API_KEY`. +//! [`serper()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; + +/// Failed-response handler of the Serper endpoint: `message`, and `type` +/// else `code` as the provider code, from the body's `error` object, or from +/// the top level when there is none. This is the mapping the endpoint used +/// with the removed shared handler; it is not an AI SDK schema. +fn serper_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(serde_json::Value::as_str) + .map(str::to_owned), + } + }) +} const MODEL_ID: &str = "serper-search"; -/// Configuration for the Serper provider. -#[derive(Debug, Clone)] -pub struct SerperConfig { - pub api_key: String, - pub base_url: String, +pub(crate) const DEFAULT_BASE_URL: &str = "https://google.serper.dev"; +pub(crate) const API_KEY_ENV_VAR: &str = "SERPER_API_KEY"; +const DEFAULT_NAME: &str = "serper"; + +/// Settings of [`create_serper`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct SerperProviderSettings { + /// Base URL for the API calls. Default `https://google.serper.dev`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `SERPER_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.search"`). + /// Default `"serper"`. The providerOptions key stays `serper`. + 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, } -impl SerperConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://google.serper.dev".to_string(), - } +impl std::fmt::Debug for SerperProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SerperProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Create a Serper provider. +/// +/// # 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. +pub fn create_serper(settings: SerperProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(SerperProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Serper"), + AuthScheme::Header("X-API-KEY"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `SERPER_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "SERPER_API_KEY", "Serper")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_serper` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn serper() -> &'static SerperProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_serper(SerperProviderSettings::default()) + .expect("default Serper settings are always valid") + }) } -/// Serper provider — search-only. +/// A Serper provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct SerperProvider { - config: SerperConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl SerperProvider { - #[must_use] - pub fn new(config: SerperConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> SerperSearchModel { - SerperSearchModel::new(self.config.clone()) + SerperSearchModel::from_config(self.model_config("search")) } } -impl Provider for SerperProvider { - fn name(&self) -> &str { - "serper" - } -} +crate::impl_single_modality_provider!(SerperProvider, search_model, |p, _id| p.search_model()); fn build_request_body(options: &SearchCallOptions) -> Value { let mut body = json!({ @@ -115,39 +199,21 @@ fn map_results(entries: Vec) -> Vec { .collect() } -/// Serper search model — implements `SearchModel`. +/// A Serper search model. pub struct SerperSearchModel { - config: SerperConfig, + config: EndpointConfig, } impl SerperSearchModel { - #[must_use] - pub fn new(config: SerperConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn endpoint(&self) -> String { - format!("{}/search", self.config.base_url) - } - - fn build_headers(&self, extra: Option<&HashMap>) -> Vec<(String, String)> { - let mut headers = vec![ - ("Content-Type".to_string(), "application/json".to_string()), - ("X-API-KEY".to_string(), self.config.api_key.clone()), - ]; - if let Some(extra) = extra { - for (k, v) in extra { - headers.push((k.clone(), v.clone())); - } - } - headers - } } #[async_trait] impl SearchModel for SerperSearchModel { fn provider(&self) -> &str { - "serper" + &self.config.provider } fn model_id(&self) -> &str { @@ -156,21 +222,13 @@ impl SearchModel for SerperSearchModel { async fn do_search(&self, options: &SearchCallOptions) -> Result { let body = build_request_body(options); - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(exchange.url("/search"), options), body, aimux_provider_utils::create_json_response_handler(), - aimux_provider_utils::create_standard_json_error_response_handler(), + serper_failed_response_handler(), ) .await?; let response_headers = resp.response_headers; diff --git a/aimux-providers/src/sglang.rs b/aimux-providers/src/sglang.rs deleted file mode 100644 index d8ce8d5dd..000000000 --- a/aimux-providers/src/sglang.rs +++ /dev/null @@ -1,86 +0,0 @@ -//! SGLang provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:30000/v1`. The `SGLANG_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:30000/v1"; -const ENV_VAR: &str = "SGLANG_BASE_URL"; -const PROVIDER_NAME: &str = "sglang"; -const PLACEHOLDER_API_KEY: &str = "sglang"; - -pub struct SglangConfig(OpenAIConfig); - -impl SglangConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()), - ) - } - - /// Create from the `SGLANG_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:30000/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let mut config = Self::new(PLACEHOLDER_API_KEY); - config.0 = config.0.with_api_key_source(Some("none")); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct SglangProvider(OpenAIProvider); - -impl SglangProvider { - #[must_use] - pub fn new(config: SglangConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for SglangProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - fn decision_model( - &self, - model_id: &str, - ) -> Result, AiMuxError> { - self.0.decision_model(model_id) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/shared/discovery.rs b/aimux-providers/src/shared/discovery.rs new file mode 100644 index 000000000..541bb4f93 --- /dev/null +++ b/aimux-providers/src/shared/discovery.rs @@ -0,0 +1,59 @@ +//! `GET {base_url}/models` for the providers whose catalogue is the +//! OpenAI-shaped `{ "data": [{ "id", "owned_by", "created" }] }`. + +use aimux_core::AiMuxError; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_provider_utils::{HttpRequest, ResponseHandler}; + +use super::EndpointConfig; + +#[derive(serde::Deserialize)] +struct ModelsList { + #[serde(default)] + data: Vec, +} + +#[derive(serde::Deserialize)] +struct ModelEntry { + id: String, + #[serde(default)] + owned_by: Option, + #[serde(default)] + created: Option, +} + +/// One `GET {base_url}/models` exchange: no retry, no recording. Discovery is +/// not a Core operation and the AI SDK has no equivalent, so a failure is +/// reported to the caller as it happened. +/// +/// # Errors +/// +/// Returns the header-resolution error (a missing key is `LoadApiKey`), +/// `ApiCall` for HTTP/transport failures and `JsonParse` when the body does +/// not deserialize into the models list. +pub(crate) async fn list_data_models( + config: &EndpointConfig, + failed_response_handler: ResponseHandler, +) -> Result, AiMuxError> { + let exchange = config.exchange(None).await?; + let resp = aimux_provider_utils::get_from_api( + exchange.with_transport(HttpRequest { + url: exchange.url("/models"), + headers: exchange.headers(), + ..Default::default() + }), + aimux_provider_utils::create_json_response_handler(), + failed_response_handler, + ) + .await?; + let parsed: ModelsList = resp.value; + Ok(parsed + .data + .into_iter() + .map(|entry| RuntimeModel { + id: entry.id, + owned_by: entry.owned_by, + created: entry.created, + }) + .collect()) +} diff --git a/aimux-providers/src/shared/exchange.rs b/aimux-providers/src/shared/exchange.rs new file mode 100644 index 000000000..d6a941890 --- /dev/null +++ b/aimux-providers/src/shared/exchange.rs @@ -0,0 +1,183 @@ +//! The private per-model configuration every Gemini-API model reads. +//! +//! This is the Rust shape of the object `createGoogle` hands each model class +//! (`{ provider, baseURL, headers, fetch, supportedUrls }`), and the Vertex +//! package hands its models the same one. The model asks it for an +//! [`Exchange`] before every request: the base URL and the request +//! headers are evaluated then, so a key, a project or a token is loaded when a +//! call is made and never when the provider is created. It has no getters for +//! the credential or the settings that produced it. + +use std::collections::HashMap; +use std::sync::Arc; + +use futures::future::BoxFuture; +use serde_json::Value; + +use aimux_core::AiMuxError; +use aimux_core::language_model::SupportedUrls; +use aimux_provider_utils::{ + ExchangeContext, FetchFunction, HeaderMapOpt, HeadersFn, HttpRequest, combine_headers, + normalize_headers, +}; + +use super::TransformRequestBody; + +/// What one request is addressed to: the base URL and the provider headers, +/// both resolved for this request. +pub(crate) struct Endpoint { + pub(crate) base_url: String, + pub(crate) headers: HeaderMapOpt, +} + +/// Resolves the endpoint of one request. +pub(crate) type EndpointFn = + Arc BoxFuture<'static, Result> + Send + Sync>; + +/// The URL patterns a model fetches itself, for a model id. +pub(crate) type SupportedUrlsFn = Arc SupportedUrls + Send + Sync>; + +/// What a model needs to talk to the API. Built by the provider that owns the +/// model. +#[derive(Clone)] +pub(crate) struct EndpointConfig { + /// The identity the model reports from `provider()`. + pub(crate) provider: String, + /// Base URL and provider headers, resolved on every request. + pub(crate) endpoint: EndpointFn, + /// Transport; `None` uses the process default, resolved per request. + pub(crate) fetch: Option, + /// URLs the model fetches itself, by model id. Empty: the caller + /// downloads them. + pub(crate) supported_urls: SupportedUrlsFn, + /// Provider-level request-body rewrite. + pub(crate) transform_request_body: Option, +} + +impl EndpointConfig { + /// A config for a provider whose base URL is fixed when it is created and + /// whose headers are evaluated on every request (the common case). + pub(crate) fn fixed( + provider: String, + base_url: String, + headers: HeadersFn, + fetch: Option, + transform_request_body: Option, + ) -> Self { + Self { + provider, + endpoint: Arc::new(move || { + let base_url = base_url.clone(); + let headers = headers.clone(); + Box::pin(async move { + Ok(Endpoint { + base_url, + headers: headers.resolve().await?, + }) + }) + }), + fetch, + supported_urls: Arc::new(|_| SupportedUrls::default()), + transform_request_body, + } + } + + /// A config whose base URL is itself resolved on every request (a + /// self-hosted instance named by an environment variable). + pub(crate) fn dynamic( + provider: String, + endpoint: EndpointFn, + fetch: Option, + ) -> Self { + Self { + provider, + endpoint, + fetch, + supported_urls: Arc::new(|_| SupportedUrls::default()), + transform_request_body: None, + } + } + + /// The same config with the URLs the model fetches itself. + #[must_use] + pub(crate) fn with_supported_urls(mut self, urls: SupportedUrlsFn) -> Self { + self.supported_urls = urls; + self + } + + /// Resolve the endpoint and layer the per-call headers over the provider + /// headers (case-insensitive, later wins, `None` removes). + /// + /// # Errors + /// + /// Returns the error of the credential, header or setting producer, for + /// instance `AiMuxError::LoadApiKey` when no key was given and the + /// environment variable is unset. + pub(crate) async fn exchange( + &self, + call_headers: Option<&HashMap>, + ) -> Result { + let endpoint = (self.endpoint)().await?; + let call: HeaderMapOpt = call_headers + .map(|headers| { + headers + .iter() + .map(|(name, value)| (name.clone(), Some(value.clone()))) + .collect() + }) + .unwrap_or_default(); + Ok(Exchange { + base_url: endpoint.base_url, + headers: normalize_headers(combine_headers(&[&endpoint.headers, &call])), + fetch: self.fetch.clone(), + transform_request_body: self.transform_request_body.clone(), + }) + } +} + +/// One request's resolved address, headers and transport. +pub(crate) struct Exchange { + base_url: String, + headers: Vec<(String, String)>, + fetch: Option, + transform_request_body: Option, +} + +impl Exchange { + /// The base URL of this request, no trailing slash. + pub(crate) fn base_url(&self) -> &str { + &self.base_url + } + + /// `base_url` followed by `path` (which starts with `/`). + pub(crate) fn url(&self, path: &str) -> String { + format!("{}{path}", self.base_url) + } + + /// The request headers that go on the wire. + pub(crate) fn headers(&self) -> Vec<(String, String)> { + self.headers.clone() + } + + /// An exchange that inherits the operation's cancellation and recording + /// context and carries this request's headers, transport and credential + /// origin. + pub(crate) fn request(&self, url: String, options: &impl ExchangeContext) -> HttpRequest { + self.with_transport(HttpRequest::new(url, self.headers(), options)) + } + + /// Attach the transport and the credential origin to a request. + pub(crate) fn with_transport(&self, mut request: HttpRequest) -> HttpRequest { + request.fetch = self.fetch.clone(); + request.credentialed_origin = Some(self.base_url.clone()); + request + } + + /// Run the provider-level body rewrite, when there is one. + pub(crate) fn transform_body(&self, body: Value) -> Value { + match &self.transform_request_body { + Some(transform) => transform(body), + None => body, + } + } +} diff --git a/aimux-providers/src/shared/mod.rs b/aimux-providers/src/shared/mod.rs new file mode 100644 index 000000000..9072bb487 --- /dev/null +++ b/aimux-providers/src/shared/mod.rs @@ -0,0 +1,183 @@ +//! Helpers the native packages and the presets share. +//! +//! Only transport plumbing lives here: how the provider-level request headers +//! are produced. Nothing in this module knows a vendor. + +mod discovery; +mod exchange; +mod poll; + +pub(crate) use discovery::list_data_models; +pub(crate) use exchange::{Endpoint, EndpointConfig}; +pub(crate) use poll::{ + POLL_INTERVAL_MILLIS_KEY, POLL_INTERVAL_MS_KEY, PollStep, is_poll_control_key, + poll_interval_ms, poll_until, retry_download, +}; + +use std::sync::Arc; + +use serde_json::Value; + +use aimux_core::AiMuxError; +use aimux_provider_utils::{HeaderMapOpt, HeadersFn, Resolvable, combine_headers, load_api_key}; + +/// A host part (a Vertex location, an Azure resource name) is one DNS label: letters, digits and hyphens, not starting or +/// ending with a hyphen (`isValidHostnamePart`). +pub(crate) fn is_valid_hostname_part(part: &str) -> bool { + let bytes = part.as_bytes(); + !bytes.is_empty() + && bytes.len() <= 63 + && bytes + .iter() + .all(|b| b.is_ascii_alphanumeric() || *b == b'-') + && bytes[0] != b'-' + && bytes[bytes.len() - 1] != b'-' +} + +/// The base URL of a native package: the explicit setting, else the +/// environment variable `var` read when the provider is created (the AI SDK's +/// `loadOptionalSetting`), else `default`. +/// +/// An invalid explicit value fails the factory (`Err` from the outer +/// result). An invalid environment value does not: the AI SDK does not +/// validate it when the provider is created and the first request fails, so +/// the inner `Err` is kept and reported by every request of that provider. +/// Creating a default provider therefore never fails on the environment. +pub(crate) fn base_url_setting( + explicit: Option<&str>, + var: &str, + default: &str, + normalize: impl Fn(&str) -> Result, +) -> Result, AiMuxError> { + if let Some(url) = explicit { + return normalize(url).map(Ok); + } + Ok( + match aimux_provider_utils::load_optional_setting(None, var) { + Some(url) => normalize(&url) + .map_err(|error| AiMuxError::InvalidArgument(format!("{var}: {error}"))), + None => Ok(default.to_string()), + }, + ) +} + +/// A provider-level rewrite of every JSON request body, called once after the +/// body is serialized and before it is sent. +pub type TransformRequestBody = Arc Value + Send + Sync>; + +/// Where the credential of a provider comes from. +#[derive(Clone, Debug)] +pub(crate) enum Credential { + /// No credential: no `Authorization` header is produced and no key is + /// resolved (a local server, or `auth: none` in a preset). + None, + /// The caller's key. It is used as given, `""` included, and never falls + /// 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 }, + /// Read the environment variable on every request; unset or empty sends + /// no credential (a local runtime started with or without a key). + OptionalEnv(String), +} + +impl Credential { + /// `Explicit` when the caller gave a key, otherwise `Env`. + pub(crate) fn explicit_or_env( + api_key: Option>, + var: &str, + description: &str, + ) -> Self { + match api_key { + Some(key) => Self::Explicit(key), + None => Self::Env { + var: var.to_string(), + description: description.to_string(), + }, + } + } + + /// The secret of this request, `None` for [`Credential::None`]. A + /// package that puts the key somewhere other than a header (a query + /// parameter, a JWT claim) asks for it here, per request. + pub(crate) async fn secret(&self) -> Result, AiMuxError> { + 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())), + } + } +} + +/// How the credential is put on the wire. +#[derive(Clone, Copy, Debug)] +pub(crate) enum AuthScheme { + /// `Authorization: Bearer `. + Bearer, + /// `: ` (`x-api-key`). + Header(&'static str), + /// `Authorization: ` (`Token `, `Key `). + Scheme(&'static str), +} + +/// Provider headers with a bearer credential: [`credential_headers`] with +/// [`AuthScheme::Bearer`]. +pub(crate) fn provider_headers( + credential: Credential, + fixed: Vec<(String, String)>, + user: Option, +) -> HeadersFn { + credential_headers(credential, AuthScheme::Bearer, fixed, user) +} + +/// Provider headers, evaluated on every request: the credential first, then +/// the fixed headers (organization, project, version, ...), then the user's +/// headers, which may override or remove any of them (case-insensitively, +/// `None` removes). +pub(crate) fn credential_headers( + credential: Credential, + scheme: AuthScheme, + fixed: Vec<(String, String)>, + user: Option, +) -> HeadersFn { + Resolvable::from_async_fn(move || { + let credential = credential.clone(); + let fixed = fixed.clone(); + let user = user.clone(); + async move { + let mut layer = HeaderMapOpt::new(); + if let Some(secret) = credential.secret().await? { + match scheme { + AuthScheme::Bearer => { + layer.insert( + "Authorization".to_string(), + Some(format!("Bearer {secret}")), + ); + } + AuthScheme::Header(name) => { + layer.insert(name.to_string(), Some(secret)); + } + AuthScheme::Scheme(scheme) => { + layer.insert( + "Authorization".to_string(), + Some(format!("{scheme} {secret}")), + ); + } + } + } + for (name, value) in fixed { + layer.insert(name, Some(value)); + } + Ok(match &user { + Some(user) => combine_headers(&[&layer, user]), + None => layer, + }) + } + }) +} diff --git a/aimux-providers/src/shared/poll.rs b/aimux-providers/src/shared/poll.rs new file mode 100644 index 000000000..26243c318 --- /dev/null +++ b/aimux-providers/src/shared/poll.rs @@ -0,0 +1,373 @@ +//! The bounded stages of an asynchronous job: poll and download. +//! +//! A vendor whose API answers a submit with a job id polls that job until it +//! finishes, then downloads the result. Both stages are safe to repeat; the +//! submit is not, so nothing here (or in the callers) ever sends it a second +//! time. The budgets are fixed constants of each package: an attempt count +//! nobody can configure, and a poll interval that `providerOptions` may +//! shorten or lengthen (`pollIntervalMs`, the same key the Google Files +//! package reads). There is no call-level retry count in this module and no +//! use of Core's operation-level retry: a poll that fails with a transient error +//! (408/409/429/5xx, a connection failure) simply spends one attempt, and the +//! loop gives up when the attempts are gone. + +use std::future::Future; +use std::time::Duration; + +use serde_json::Value; + +use aimux_core::shared::JsonObject; +use aimux_core::{AbortSignal, AiMuxError}; +use aimux_provider_utils::sleep_or_abort; + +/// How many times a download of a finished job's result is attempted. +pub(crate) const DOWNLOAD_ATTEMPTS: u32 = 3; + +/// The longest a `Retry-After` hint can stretch the wait between two attempts. +const MAX_RETRY_AFTER: Duration = Duration::from_secs(30); + +/// What one poll saw. +pub(crate) enum PollStep { + /// The job finished; the loop returns this value. + Ready(T), + /// The job is still running; wait and ask again. + Pending, +} + +/// The providerOptions key that sets the poll interval of a package with no +/// upstream counterpart. +pub(crate) const POLL_INTERVAL_MS_KEY: &str = "pollIntervalMs"; + +/// The key `@ai-sdk/luma` and `@ai-sdk/black-forest-labs` read for the same +/// thing; those two packages read it instead of [`POLL_INTERVAL_MS_KEY`]. +pub(crate) const POLL_INTERVAL_MILLIS_KEY: &str = "pollIntervalMillis"; + +/// Whether a providerOptions field is pacing control rather than request data. +/// +/// These four keys are never forwarded to a vendor, in any package, whether or +/// not the package reads them: only the interval key of the package is read, +/// the rest (`pollTimeoutMillis`, `maxPollAttempts`) are ignored because the +/// attempt budget is a package constant. +pub(crate) fn is_poll_control_key(key: &str) -> bool { + matches!( + key, + "pollIntervalMs" | "pollIntervalMillis" | "pollTimeoutMillis" | "maxPollAttempts" + ) +} + +/// The poll interval a package's providerOptions namespace sets under `key`, +/// or `default_ms`. +pub(crate) fn poll_interval_ms(namespace: Option<&JsonObject>, key: &str, default_ms: u64) -> u64 { + namespace + .and_then(|options| options.get(key)) + .and_then(Value::as_u64) + .unwrap_or(default_ms) +} + +fn check_aborted(abort: Option<&AbortSignal>) -> Result<(), AiMuxError> { + match abort { + Some(signal) if signal.is_aborted() => Err(AiMuxError::from_abort_signal(signal)), + _ => Ok(()), + } +} + +/// Wait before the next attempt. A zero wait (tests, or a caller that wants +/// back-to-back polls) only yields: a zero-length timer would still cost a +/// full timer tick. +async fn wait(duration: Duration, abort: Option<&AbortSignal>) -> Result<(), AiMuxError> { + if duration.is_zero() { + tokio::task::yield_now().await; + return check_aborted(abort); + } + sleep_or_abort(duration, abort).await +} + +/// The wait before the next attempt: the interval, stretched to a +/// `Retry-After` hint when the failed response carried one. +fn wait_after(interval: Duration, error: Option<&AiMuxError>) -> Duration { + let hint = error + .and_then(AiMuxError::retry_after_hint) + .and_then(|ms| u64::try_from(ms).ok()) + .map(Duration::from_millis) + .map(|hint| hint.min(MAX_RETRY_AFTER)); + interval.max(hint.unwrap_or_default()) +} + +/// Poll until `step` reports [`PollStep::Ready`]. +/// +/// At most `max_attempts` polls are made. Cancellation is checked before every +/// poll and interrupts every wait. A transient error ([`AiMuxError::is_retryable`]) +/// spends the attempt like a pending answer would; any other error ends the +/// loop. When the attempts are gone the error of the last poll is returned if +/// that poll failed, otherwise `AiMuxError::Timeout`. +/// +/// # Errors +/// +/// Returns the first non-transient error of `step`, the cancellation, the +/// last transient error once the attempts are spent, or `Timeout` when the job +/// was still pending after the last attempt. +pub(crate) async fn poll_until( + what: &str, + abort: Option<&AbortSignal>, + interval: Duration, + max_attempts: u32, + mut step: F, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future, AiMuxError>>, +{ + let mut last_error: Option = None; + for attempt in 0..max_attempts { + check_aborted(abort)?; + match step().await { + Ok(PollStep::Ready(value)) => return Ok(value), + Ok(PollStep::Pending) => last_error = None, + Err(error) if error.is_retryable() => last_error = Some(error), + Err(error) => return Err(error), + } + if attempt + 1 < max_attempts { + wait(wait_after(interval, last_error.as_ref()), abort).await?; + } + } + match last_error { + Some(error) => Err(error), + None => Err(AiMuxError::Timeout(format!( + "{what} polling timed out after {max_attempts} attempts ({}ms apart)", + interval.as_millis() + ))), + } +} + +/// Download a finished job's result: [`DOWNLOAD_ATTEMPTS`] tries, a transient +/// error waits `interval` and tries again, any other error ends it. +/// +/// # Errors +/// +/// Returns the first non-transient error, the cancellation, or the error of +/// the last attempt. +pub(crate) async fn retry_download( + abort: Option<&AbortSignal>, + interval: Duration, + mut attempt: F, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut tries = 0; + loop { + check_aborted(abort)?; + tries += 1; + match attempt().await { + Ok(value) => return Ok(value), + Err(error) if error.is_retryable() && tries < DOWNLOAD_ATTEMPTS => { + wait(wait_after(interval, Some(&error)), abort).await?; + } + Err(error) => return Err(error), + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicU32, Ordering}; + + use aimux_core::ApiCallError; + + use super::*; + + fn transient() -> AiMuxError { + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(503), + is_retryable: true, + ..ApiCallError::new("unavailable", "http://x/job", serde_json::json!({})) + })) + } + + fn fatal() -> AiMuxError { + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(400), + ..ApiCallError::new("bad", "http://x/job", serde_json::json!({})) + })) + } + + #[test] + fn interval_override_reads_the_namespace_only() { + let options = serde_json::json!({ "pollIntervalMs": 7, "maxPollAttempts": 1 }); + assert_eq!( + poll_interval_ms(options.as_object(), POLL_INTERVAL_MS_KEY, 500), + 7 + ); + // Only the key asked for is read. + assert_eq!( + poll_interval_ms(options.as_object(), POLL_INTERVAL_MILLIS_KEY, 500), + 500 + ); + assert_eq!( + poll_interval_ms(serde_json::json!({}).as_object(), POLL_INTERVAL_MS_KEY, 500), + 500 + ); + assert_eq!(poll_interval_ms(None, POLL_INTERVAL_MS_KEY, 500), 500); + // A string is not a number: the default stays. + assert_eq!( + poll_interval_ms( + serde_json::json!({ "pollIntervalMs": "7" }).as_object(), + POLL_INTERVAL_MS_KEY, + 500 + ), + 500 + ); + } + + #[test] + fn the_four_pacing_keys_are_control_keys_and_nothing_else_is() { + for key in [ + "pollIntervalMs", + "pollIntervalMillis", + "pollTimeoutMillis", + "maxPollAttempts", + ] { + assert!(is_poll_control_key(key), "{key}"); + } + for key in ["prompt", "maxWaitTimeInSeconds", "poll", "pollInterval"] { + assert!(!is_poll_control_key(key), "{key}"); + } + } + + #[tokio::test] + async fn returns_the_first_ready_value() { + let polls = AtomicU32::new(0); + let value = poll_until("job", None, Duration::ZERO, 10, || async { + Ok(if polls.fetch_add(1, Ordering::SeqCst) == 2 { + PollStep::Ready("done") + } else { + PollStep::Pending + }) + }) + .await + .unwrap(); + assert_eq!(value, "done"); + assert_eq!(polls.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn a_transient_error_spends_an_attempt_and_the_next_poll_can_succeed() { + let polls = AtomicU32::new(0); + let value = poll_until("job", None, Duration::ZERO, 10, || async { + match polls.fetch_add(1, Ordering::SeqCst) { + 0 => Err(transient()), + _ => Ok(PollStep::Ready(1)), + } + }) + .await + .unwrap(); + assert_eq!(value, 1); + assert_eq!(polls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn a_fatal_error_ends_the_loop_at_once() { + let polls = AtomicU32::new(0); + let error = poll_until::<(), _, _>("job", None, Duration::ZERO, 10, || async { + polls.fetch_add(1, Ordering::SeqCst); + Err(fatal()) + }) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::ApiCall(ref d) if d.status_code == Some(400))); + assert_eq!(polls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn exhausted_attempts_return_the_last_transient_error() { + let polls = AtomicU32::new(0); + let error = poll_until::<(), _, _>("job", None, Duration::ZERO, 4, || async { + polls.fetch_add(1, Ordering::SeqCst); + Err(transient()) + }) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::ApiCall(ref d) if d.status_code == Some(503))); + assert_eq!(polls.load(Ordering::SeqCst), 4); + } + + #[tokio::test] + async fn exhausted_attempts_while_pending_are_a_timeout() { + let polls = AtomicU32::new(0); + let error = poll_until::<(), _, _>("job", None, Duration::ZERO, 3, || async { + polls.fetch_add(1, Ordering::SeqCst); + Ok(PollStep::Pending) + }) + .await + .unwrap_err(); + assert!( + matches!(error, AiMuxError::Timeout(ref m) if m.contains("job") && m.contains('3')) + ); + assert_eq!(polls.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn an_aborted_signal_stops_before_the_next_poll() { + let signal = AbortSignal::new(); + let polls = AtomicU32::new(0); + let error = poll_until::<(), _, _>("job", Some(&signal), Duration::ZERO, 10, || async { + polls.fetch_add(1, Ordering::SeqCst); + signal.abort(); + Ok(PollStep::Pending) + }) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::Aborted(_))); + assert_eq!(polls.load(Ordering::SeqCst), 1); + + // Aborted before the first poll: no poll at all. + let polls = AtomicU32::new(0); + let error = poll_until::<(), _, _>("job", Some(&signal), Duration::ZERO, 10, || async { + polls.fetch_add(1, Ordering::SeqCst); + Ok(PollStep::Pending) + }) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::Aborted(_))); + assert_eq!(polls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn download_tries_a_bounded_number_of_times() { + let tries = AtomicU32::new(0); + let error = retry_download::<(), _, _>(None, Duration::ZERO, || async { + tries.fetch_add(1, Ordering::SeqCst); + Err(transient()) + }) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::ApiCall(_))); + assert_eq!(tries.load(Ordering::SeqCst), DOWNLOAD_ATTEMPTS); + } + + #[tokio::test] + async fn download_recovers_from_a_transient_error_and_stops_on_a_fatal_one() { + let tries = AtomicU32::new(0); + let value = retry_download(None, Duration::ZERO, || async { + if tries.fetch_add(1, Ordering::SeqCst) == 0 { + Err(transient()) + } else { + Ok(5) + } + }) + .await + .unwrap(); + assert_eq!((value, tries.load(Ordering::SeqCst)), (5, 2)); + + let tries = AtomicU32::new(0); + let error = retry_download::<(), _, _>(None, Duration::ZERO, || async { + tries.fetch_add(1, Ordering::SeqCst); + Err(fatal()) + }) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::ApiCall(ref d) if d.status_code == Some(400))); + assert_eq!(tries.load(Ordering::SeqCst), 1); + } +} diff --git a/aimux-providers/src/stability.rs b/aimux-providers/src/stability.rs index a5aa918f1..544867edd 100644 --- a/aimux-providers/src/stability.rs +++ b/aimux-providers/src/stability.rs @@ -8,22 +8,27 @@ //! `stable-image-core` → `/generate/core`, `sd3` → `/generate/sd3`. Any other //! model ID is used as-is for the sub-path. Responses are image binary //! (`Accept: image/*`) or base64-encoded JSON. +//! +//! [`create_stability`] takes [`StabilityProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`StabilityProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `STABILITY_API_KEY`. +//! [`stability()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use bytes::Bytes; use serde_json::Value; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::image_model::{ ImageCallOptions, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; use aimux_core::shared::Warning; -use aimux_provider_utils::{ - HttpBody, HttpRequest, MultipartForm, load_api_key, without_trailing_slash, -}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; +use aimux_provider_utils::{HttpBody, MultipartForm}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; /// Stability error response structure: `{ "id": "...", "name": "...", "errors": ["..."] }`. /// @@ -109,106 +114,131 @@ fn model_id_to_subpath(model_id: &str) -> &str { } } -/// Configuration for the Stability provider. -#[derive(Debug, Clone)] -pub struct StabilityConfig { - pub api_key: String, - pub base_url: String, - pub headers: Option>, +pub(crate) mod options; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.stability.ai"; +pub(crate) const API_KEY_ENV_VAR: &str = "STABILITY_API_KEY"; +const DEFAULT_NAME: &str = "stability"; + +/// Settings of [`create_stability`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct StabilityProviderSettings { + /// Base URL for the API calls. Default `https://api.stability.ai`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `STABILITY_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.image"`). + /// Default `"stability"`. The providerOptions key stays `stability`. + 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, } -impl StabilityConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.stability.ai".to_string(), - headers: None, - } - } - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self - } - /// Create from the `STABILITY_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "STABILITY_API_KEY", "Stability")?; - Ok(Self::new(api_key)) +impl std::fmt::Debug for StabilityProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("StabilityProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } } +/// Create a Stability provider. +/// +/// # 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. +pub fn create_stability( + settings: StabilityProviderSettings, +) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(StabilityProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Stability"), + vec![("Accept".to_string(), "image/*".to_string())], + settings.headers, + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_stability` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn stability() -> &'static StabilityProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_stability(StabilityProviderSettings::default()) + .expect("default Stability settings are always valid") + }) +} + +/// A Stability provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct StabilityProvider { - config: StabilityConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl StabilityProvider { - #[must_use] - pub fn new(config: StabilityConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + + /// An image model (e.g. `"stable-image-core"`); `provider()` is `"{name}.image"`. #[must_use] pub fn image(&self, model_id: &str) -> StabilityImageModel { - StabilityImageModel::new(model_id.to_string(), self.config.clone()) + StabilityImageModel::from_config(model_id.to_string(), self.model_config("image")) } } -impl Provider for StabilityProvider { - fn name(&self) -> &str { - "stability" - } -} +crate::impl_single_modality_provider!(StabilityProvider, image_model, |p, id| p.image(id)); /// A Stability image generation model. pub struct StabilityImageModel { model_id: String, - config: StabilityConfig, + config: EndpointConfig, } impl StabilityImageModel { - #[must_use] - pub fn new(model_id: String, config: StabilityConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut h = HashMap::new(); - h.insert( - "Authorization".into(), - format!("Bearer {}", self.config.api_key), - ); - // Request the image binary directly; a caller may override this with - // `application/json` via config/extra headers to receive base64 JSON. - h.insert("Accept".into(), "image/*".into()); - if let Some(ref ch) = self.config.headers { - for (k, v) in ch { - h.insert(k.clone(), v.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - h.insert(k.clone(), v.clone()); - } - } - h - } - - fn endpoint(&self) -> String { - format!( - "{}/v2beta/stable-image/generate/{}", - self.config.base_url, - model_id_to_subpath(&self.model_id) - ) - } } /// Greatest common divisor of two `u32` values. @@ -229,7 +259,7 @@ fn convert_size_to_aspect_ratio(size: &aimux_core::shared::Size) -> Option &str { - "stability" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id @@ -241,7 +271,7 @@ impl ImageModel for StabilityImageModel { async fn do_generate(&self, options: &ImageCallOptions) -> Result { let mut warnings: Vec = Vec::new(); - let stability_opts = options.provider_options.get("stability"); + let stability_opts = options::stability_options(Some(&options.provider_options)); // Resolve aspect ratio: explicit option > provider option > derived from size. let aspect_ratio: Option = if let Some(ar) = options.aspect_ratio { @@ -323,11 +353,16 @@ impl ImageModel for StabilityImageModel { let (body_bytes, content_type) = form.finish(); - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_to_api( - HttpRequest::new(self.endpoint(), header_list, options), + exchange.request( + exchange.url(&format!( + "/v2beta/stable-image/generate/{}", + model_id_to_subpath(&self.model_id) + )), + options, + ), HttpBody::Bytes(body_bytes, content_type), stability_successful_response_handler(), stability_failed_response_handler(), diff --git a/aimux-providers/src/stability/options.rs b/aimux-providers/src/stability/options.rs new file mode 100644 index 000000000..1c561ba51 --- /dev/null +++ b/aimux-providers/src/stability/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Stability. +//! +//! Stability reads `providerOptions.stability`, whatever the provider is named. This +//! module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "stability"; + +/// The Stability options in a providerOptions map. +pub(crate) fn stability_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/tavily.rs b/aimux-providers/src/tavily.rs index a01ab7ec3..9bc22c5dd 100644 --- a/aimux-providers/src/tavily.rs +++ b/aimux-providers/src/tavily.rs @@ -2,19 +2,25 @@ //! //! Implements the `SearchModel` trait against the Tavily search API //! (`POST https://api.tavily.com/search`). Bearer auth via `TAVILY_API_KEY`. +//! +//! [`create_tavily`] takes [`TavilyProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`TavilyProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `TAVILY_API_KEY`. +//! [`tavily()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; const MODEL_ID: &str = "tavily-search"; @@ -33,61 +39,116 @@ fn tavily_failed_response_handler() -> aimux_provider_utils::ResponseHandler, + /// The API key. `None` loads `TAVILY_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.search"`). + /// Default `"tavily"`. The providerOptions key stays `tavily`. + 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, } -impl TavilyConfig { - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.tavily.com".to_string(), - } +impl std::fmt::Debug for TavilyProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TavilyProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Create a Tavily provider. +/// +/// # 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. +pub fn create_tavily(settings: TavilyProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(TavilyProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Tavily"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} - /// Create from the `TAVILY_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when the environment variable is not - /// set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "TAVILY_API_KEY", "Tavily")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_tavily` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn tavily() -> &'static TavilyProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_tavily(TavilyProviderSettings::default()) + .expect("default Tavily settings are always valid") + }) } -/// Tavily provider — search-only. +/// A Tavily provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct TavilyProvider { - config: TavilyConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl TavilyProvider { - #[must_use] - pub fn new(config: TavilyConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> TavilySearchModel { - TavilySearchModel::new(self.config.clone()) + TavilySearchModel::from_config(self.model_config("search")) } } -impl Provider for TavilyProvider { - fn name(&self) -> &str { - "tavily" - } -} +crate::impl_single_modality_provider!(TavilyProvider, search_model, |p, _id| p.search_model()); fn build_request_body(options: &SearchCallOptions) -> Value { let mut body = json!({ @@ -145,42 +206,21 @@ fn map_results(entries: Vec) -> Vec { .collect() } -/// Tavily search model — implements `SearchModel`. +/// A Tavily search model. pub struct TavilySearchModel { - config: TavilyConfig, + config: EndpointConfig, } impl TavilySearchModel { - #[must_use] - pub fn new(config: TavilyConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn endpoint(&self) -> String { - format!("{}/search", self.config.base_url) - } - - fn build_headers(&self, extra: Option<&HashMap>) -> Vec<(String, String)> { - let mut headers = vec![ - ("Content-Type".to_string(), "application/json".to_string()), - ( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ), - ]; - if let Some(extra) = extra { - for (k, v) in extra { - headers.push((k.clone(), v.clone())); - } - } - headers - } } #[async_trait] impl SearchModel for TavilySearchModel { fn provider(&self) -> &str { - "tavily" + &self.config.provider } fn model_id(&self) -> &str { @@ -189,18 +229,10 @@ impl SearchModel for TavilySearchModel { async fn do_search(&self, options: &SearchCallOptions) -> Result { let body = build_request_body(options); - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url: self.endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(exchange.url("/search"), options), body, aimux_provider_utils::create_json_response_handler(), tavily_failed_response_handler(), diff --git a/aimux-providers/src/tinyfish.rs b/aimux-providers/src/tinyfish.rs index 06bb368c2..6647e256b 100644 --- a/aimux-providers/src/tinyfish.rs +++ b/aimux-providers/src/tinyfish.rs @@ -6,26 +6,49 @@ //! Authentication uses the `X-API-Key` header (env `TINYFISH_API_KEY`). //! TinyFish is a search-only provider: it exposes a web search protocol //! (query → results) and does not support language models. +//! +//! [`create_tinyfish`] takes [`TinyfishProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`TinyfishProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `TINYFISH_API_KEY`. +//! [`tinyfish()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::Value; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_core::shared::SharedHeaders; - -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; -/// Provider canonical name. -const PROVIDER_NAME: &str = "tinyfish"; +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; /// Fixed model id for the TinyFish search model. +/// Failed-response handler of the TinyFish endpoint: `message`, and `type` +/// else `code` as the provider code, from the body's `error` object, or from +/// the top level when there is none. This is the mapping the endpoint used +/// with the removed shared handler; it is not an AI SDK schema. +fn tinyfish_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(serde_json::Value::as_str) + .map(str::to_owned), + } + }) +} + const MODEL_ID: &str = "tinyfish-search"; /// Default number of results when `max_results` is unset. @@ -33,79 +56,117 @@ const DEFAULT_COUNT: u32 = 10; // ── Config ─────────────────────────────────────────────────────────────────── -/// Configuration for the TinyFish provider. -/// -/// The `Debug` implementation redacts the API key so credentials never appear -/// in logs or error messages. -#[derive(Clone)] -pub struct TinyfishConfig { - api_key: String, - base_url: String, -} +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.search.tinyfish.ai"; +pub(crate) const API_KEY_ENV_VAR: &str = "TINYFISH_API_KEY"; +const DEFAULT_NAME: &str = "tinyfish"; -impl TinyfishConfig { - /// Create from an API key (uses the default TinyFish base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.search.tinyfish.ai".to_string(), - } - } - - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - - /// Create from the `TINYFISH_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `TINYFISH_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "TINYFISH_API_KEY", "TinyFish")?; - Ok(Self::new(api_key)) - } +/// Settings of [`create_tinyfish`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct TinyfishProviderSettings { + /// Base URL for the API calls. Default `https://api.search.tinyfish.ai`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `TINYFISH_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.search"`). + /// Default `"tinyfish"`. The providerOptions key stays `tinyfish`. + 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, } -impl std::fmt::Debug for TinyfishConfig { +impl std::fmt::Debug for TinyfishProviderSettings { + /// Never prints the key or header values. fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TinyfishConfig") - .field("api_key", &"") + f.debug_struct("TinyfishProviderSettings") .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) .finish() } } -// ── Provider ───────────────────────────────────────────────────────────────── - -/// TinyFish provider — creates [`TinyfishSearchModel`] instances. +/// Create a TinyFish provider. /// -/// TinyFish is a search-only provider; it does not support language models. +/// # 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. +pub fn create_tinyfish(settings: TinyfishProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(TinyfishProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "TinyFish"), + AuthScheme::Header("X-API-Key"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_tinyfish` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn tinyfish() -> &'static TinyfishProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_tinyfish(TinyfishProviderSettings::default()) + .expect("default TinyFish settings are always valid") + }) +} + +/// A TinyFish provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct TinyfishProvider { - config: TinyfishConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl TinyfishProvider { - #[must_use] - pub fn new(config: TinyfishConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a search model instance. + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> TinyfishSearchModel { - TinyfishSearchModel::new(self.config.clone()) + TinyfishSearchModel::from_config(self.model_config("search")) } } -impl Provider for TinyfishProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } -} +crate::impl_single_modality_provider!(TinyfishProvider, search_model, |p, _id| p.search_model()); // ── Request builder ────────────────────────────────────────────────────────── @@ -158,36 +219,19 @@ fn map_result(r: TinyfishResult) -> SearchResultItem { /// A TinyFish search model. pub struct TinyfishSearchModel { - config: TinyfishConfig, + config: EndpointConfig, } impl TinyfishSearchModel { - #[must_use] - pub fn new(config: TinyfishConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn build_headers(&self, extra: Option<&SharedHeaders>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("X-API-Key".to_string(), self.config.api_key.clone()); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - /// The search endpoint (root path; query parameters are appended at call time). - fn endpoint(&self) -> String { - self.config.base_url.clone() - } } #[async_trait] impl SearchModel for TinyfishSearchModel { fn provider(&self) -> &str { - PROVIDER_NAME + &self.config.provider } fn model_id(&self) -> &str { @@ -196,29 +240,18 @@ impl SearchModel for TinyfishSearchModel { async fn do_search(&self, options: &SearchCallOptions) -> Result { let count = resolve_count(options.max_results); - let headers: Vec<(String, String)> = self - .build_headers(options.headers.as_ref()) - .into_iter() - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; - let mut url = url::Url::parse(&self.endpoint()) + let mut url = url::Url::parse(exchange.base_url()) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid tinyfish endpoint: {e}")))?; url.query_pairs_mut() .append_pair("query", &options.query) .append_pair("count", &count.to_string()); let resp = aimux_provider_utils::get_from_api( - HttpRequest { - url: url.to_string(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(url.to_string(), options), aimux_provider_utils::create_json_response_handler::(), - aimux_provider_utils::create_standard_json_error_response_handler(), + tinyfish_failed_response_handler(), ) .await?; diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index 3a1d8e57f..136ba0c0e 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -1,624 +1,127 @@ -//! Anthropic partner models on Vertex AI — implements `LanguageModel`. +//! Anthropic partner models on Vertex AI. //! -//! Serves Anthropic Claude models (e.g. `claude-sonnet-4-20250514`) through the -//! Vertex AI `rawPredict` / `streamRawPredict` endpoints. These endpoints proxy -//! the standard Anthropic Messages API: the request body is an Anthropic -//! Messages request wrapped in an `anthropic_version` envelope — with the model -//! identity carried by the URL path (`publishers/anthropic/models/{model}`), -//! not the body — and the response is a verbatim Anthropic Messages response. +//! Claude on Vertex is the Anthropic Messages API behind the +//! `publishers/anthropic/models/{model}:rawPredict` / `:streamRawPredict` +//! endpoints, so its model is the shared +//! [`AnthropicMessagesModel`](crate::anthropic::AnthropicMessagesModel), the +//! Rust form of `createVertexAnthropic`'s `AnthropicLanguageModel`. Only what +//! Vertex does differently is configured here: //! -//! This reuses the shared [`crate::anthropic::convert`] message conversion -//! logic and [`crate::anthropic::types`] response types. Only the endpoint -//! construction and authentication differ from the standard Anthropic provider, -//! and both come from the parent Vertex provider ([`super::VertexAuth`]). +//! - the URL carries the model and the mode, +//! - 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), +//! - errors are Google-shaped, +//! - no URL sources, no structured-output beta and no `strict` tool +//! definitions. //! //! Reference: -use aimux_core::tool::RawToolCall; -use std::collections::HashMap; +use std::sync::Arc; -use async_trait::async_trait; -use futures::StreamExt; +use futures::future::BoxFuture; use serde_json::{Value, json}; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateResult, StreamResult}; -use aimux_core::shared::provider_namespace; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::{ - FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, -}; +use aimux_core::AiMuxError; +use aimux_core::language_model::SupportedUrls; +use aimux_provider_utils::{FetchFunction, Resolvable}; -use aimux_provider_utils::{HttpRequest, RetryConfig}; - -use crate::anthropic::convert::{build_request_body_with_warnings, parse_stop_reason}; -use crate::anthropic::stream::{ - finalize_streamed_tool_input, initial_tool_input, server_tool_provider_name, - stream_parts_for_result_block, tool_call_caller_metadata, -}; -use crate::anthropic::tool_name_mapping::ToolNameMapping; -use crate::anthropic::types::{AnthropicResponse, ContentBlock, StreamEvent}; - -use super::VertexAuth; +use crate::anthropic::AnthropicMessagesModel; +use crate::anthropic::config::{AnthropicModelConfig, AnthropicModelHooks, TransformRequestBody}; +use crate::anthropic::options::CANONICAL; +use crate::shared::Endpoint; /// `anthropic_version` envelope value required by the Vertex AI `rawPredict` / /// `streamRawPredict` endpoints. const ANTHROPIC_VERTEX_VERSION: &str = "vertex-2023-10-16"; -/// Google/Vertex error structure: `{ "error": { "message": "...", "status": "..." } }`. -/// -/// Configuration for an Anthropic-on-Vertex model instance. -/// -/// `base_url` is the Vertex AI base URL *without* a `/publishers/{publisher}` -/// suffix (e.g. `.../projects/{project}/locations/{location}`); the publisher -/// (`anthropic`) is appended by the endpoint helpers. The parent -/// [`super::VertexProvider::anthropic_model`] strips any existing -/// `/publishers/google` suffix from its configured base URL before constructing -/// this config. -#[derive(Debug, Clone)] -pub struct VertexAnthropicConfig { - pub base_url: String, - pub auth: VertexAuth, - /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 - pub api_key_source: Option, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, -} - -/// An Anthropic Claude language model served via Vertex AI `rawPredict`. -/// -/// 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 VertexAnthropicModel { - model_id: String, - config: VertexAnthropicConfig, -} - -impl VertexAnthropicModel { - #[must_use] - pub fn new(model_id: String, config: VertexAnthropicConfig) -> Self { - Self { model_id, config } - } - - /// `…/publishers/anthropic/models/{model}:rawPredict` - fn generate_endpoint(&self) -> String { - format!( - "{}/publishers/anthropic/models/{}:rawPredict", - self.config.base_url, self.model_id - ) - } - - /// `…/publishers/anthropic/models/{model}:streamRawPredict` - fn stream_endpoint(&self) -> String { - format!( - "{}/publishers/anthropic/models/{}:streamRawPredict", - self.config.base_url, self.model_id - ) - } - - /// Build the request headers, reusing the parent Vertex auth (Bearer token - /// or API key). Mirrors [`super::VertexModel::build_headers`]. - fn build_headers(&self, extra: Option<&HashMap>) -> Vec<(String, String)> { - let mut headers = vec![("Content-Type".to_string(), "application/json".to_string())]; - match &self.config.auth { - VertexAuth::BearerToken(token) => { - headers.push(("Authorization".to_string(), format!("Bearer {token}"))); - } - VertexAuth::ApiKey(key) => { - headers.push(("x-goog-api-key".to_string(), key.clone())); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.push((k.clone(), v.clone())); - } - } - headers - } - - /// Wrap a standard Anthropic Messages request body in the Vertex AI - /// `rawPredict` envelope: prepend `anthropic_version` and drop the `model` - /// field (the model identity is carried by the URL path, not the body). - fn wrap_raw_predict_body(&self, anthropic_body: Value) -> Value { - let mut envelope = serde_json::Map::new(); - envelope.insert( - "anthropic_version".to_string(), - json!(ANTHROPIC_VERTEX_VERSION), - ); - if let Value::Object(map) = anthropic_body { - for (k, v) in map { - if k == "model" { - continue; - } - envelope.insert(k, v); - } - } - Value::Object(envelope) +/// The provider string of every Anthropic model on Vertex. +const PROVIDER: &str = "googleVertex.anthropic.messages"; + +/// An Anthropic Claude language model served via Vertex AI. +pub type VertexAnthropicModel = AnthropicMessagesModel; + +/// Resolves the endpoint of a request: the `.../publishers/anthropic/models` +/// base URL and the provider headers. +pub(super) type EndpointFn = + Arc BoxFuture<'static, Result> + Send + Sync>; + +/// Wrap a standard Messages request body in the `rawPredict` envelope: drop +/// `model` (the URL carries it) and add `anthropic_version`. +fn raw_predict_envelope(body: Value) -> Value { + let mut envelope = serde_json::Map::new(); + envelope.insert( + "anthropic_version".to_string(), + json!(ANTHROPIC_VERTEX_VERSION), + ); + if let Value::Object(map) = body { + envelope.extend(map.into_iter().filter(|(key, _)| key != "model")); } + Value::Object(envelope) } -#[async_trait] -impl LanguageModel for VertexAnthropicModel { - fn provider(&self) -> &str { - "google.vertex" - } - - fn model_id(&self) -> &str { - &self.model_id - } - - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - // M2b: record identity + credential source + auth kind. Never serialize - // the bearer token or API key plaintext. - let auth_kind = match &self.config.auth { - VertexAuth::BearerToken(_) => "bearer_token", - VertexAuth::ApiKey(_) => "api_key", - }; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.base_url.clone()), - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: Some(serde_json::json!({ - "auth_kind": auth_kind, - "anthropic_version": ANTHROPIC_VERTEX_VERSION, +/// The configuration of one request against `endpoint`. +fn request_config( + endpoint: Endpoint, + fetch: Option, + transform_request_body: Option, +) -> AnthropicModelConfig { + let Endpoint { base_url, headers } = endpoint; + let models = base_url.clone(); + let url_base = base_url.clone(); + AnthropicModelConfig { + provider: PROVIDER.to_string(), + url: Arc::new(move |path| format!("{url_base}{path}")), + headers: Resolvable::Value(headers), + fetch, + supported_urls: SupportedUrls::default(), + transform_request_body, + base_url, + provider_options_name: CANONICAL.to_string(), + hooks: AnthropicModelHooks { + request_url: Some(Arc::new(move |model_id, stream| { + let method = if stream { + "streamRawPredict" + } else { + "rawPredict" + }; + format!("{models}/{model_id}:{method}") })), - } - } - - async fn do_generate(&self, options: &CallOptions) -> Result { - let req = build_request_body_with_warnings(&self.model_id, options, false)?; - let body = self.wrap_raw_predict_body(req.body); - let url = self.generate_endpoint(); - let headers = self.build_headers(options.headers.as_ref()); - let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(url.clone(), headers, options), - body.clone(), - aimux_provider_utils::create_json_response_handler(), - crate::google::google_failed_response_handler(), - ) - .await?; - - let response_body = resp.raw_value; - let response_headers = resp.response_headers; - let data: AnthropicResponse = resp.value; - - let content = crate::anthropic::stream::parse_anthropic_content( - &data.content, - &ToolNameMapping::new(options.tools.as_deref()), - &crate::anthropic::stream::generate_id, - ); - - let finish_reason = data - .stop_reason - .as_deref() - .map(parse_stop_reason) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }); - - // RFC-0015 P0-2: fill cache fields + raw (same shape as Anthropic). - let usage = crate::anthropic::usage::usage_from_anthropic(&data.usage); - - Ok(GenerateResult { - content, - finish_reason, - usage, - warnings: req.warnings, - provider_metadata: None, - response: Some(aimux_core::shared::ResponseInfo { - id: Some(data.id), - timestamp: None, - model_id: Some(data.model), - headers: Some(response_headers), - body: response_body, - }), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - }) - } - - async fn do_stream(&self, options: &CallOptions) -> Result { - let req = build_request_body_with_warnings(&self.model_id, options, true)?; - let body = self.wrap_raw_predict_body(req.body); - let warnings = req.warnings; - let tool_names = ToolNameMapping::new(options.tools.as_deref()); - let url = self.stream_endpoint(); - let headers = self.build_headers(options.headers.as_ref()); - let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(url.clone(), headers, options), - body.clone(), - aimux_provider_utils::create_event_source_response_handler::(), - crate::google::google_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, - }; - if let Some(Ok(StreamEvent::Error { error })) = first_event.as_ref() { - return Err(crate::anthropic::stream::anthropic_stream_error( - error, - &url, - body.clone(), - response_headers, - )); - } - let stream_error_url = url; - let stream_request_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - - let mut sse = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - let mut blocks: HashMap = HashMap::new(); - let mut final_usage = Usage::default(); - let mut final_finish_reason: Option = None; - let mut response_meta_emitted = false; - let mut mcp_tool_calls: HashMap = HashMap::new(); - let mut server_tool_calls: HashMap = HashMap::new(); - - while let Some(event) = sse.next().await { - match event { - Ok(stream_event) => { - match stream_event { - StreamEvent::MessageStart { message } => { - if let Some(usage) = &message.usage { - // RFC-0015 P0-2: full input side incl. - // cache fields + raw (same as Anthropic). - final_usage = - crate::anthropic::usage::usage_from_anthropic(usage); - } - if !response_meta_emitted { - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: Some(message.id.clone()), - timestamp: None, - model_id: Some(message.model.clone()), - })); - response_meta_emitted = true; - } - } - StreamEvent::ContentBlockStart { index, content_block } => { - match content_block { - ContentBlock::Text { .. } => { - blocks.insert(index, BlockState::Text { started: false }); - } - ContentBlock::Thinking { .. } => { - blocks.insert(index, BlockState::Thinking { started: false }); - } - ContentBlock::ToolUse { - id, - name, - input, - caller, - } => { - let custom_name = tool_names - .to_custom_tool_name(&name) - .to_string(); - let initial_input = initial_tool_input(&input); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: custom_name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - blocks.insert(index, BlockState::ToolUse { - id, - name: custom_name, - first_delta: initial_input.is_empty(), - accumulated_json: initial_input, - provider_executed: None, - dynamic: None, - provider_tool_name: None, - provider_tool_input_type: None, - provider_metadata: tool_call_caller_metadata(caller.as_ref()), - }); - } - ContentBlock::ServerToolUse { id, name, input } => { - if matches!( - name.as_str(), - "tool_search_tool_regex" | "tool_search_tool_bm25" - ) { - server_tool_calls.insert(id.clone(), name.clone()); - } - let provider_name = server_tool_provider_name(&name); - let custom_name = tool_names - .to_custom_tool_name(provider_name) - .to_string(); - let dynamic = (provider_name == "code_execution" - && tool_names.mark_code_execution_dynamic()) - .then_some(true); - let initial_input = initial_tool_input(&input); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: custom_name.clone(), - provider_executed: Some(true), - dynamic, - title: None, - provider_metadata: None, - }); - blocks.insert(index, BlockState::ToolUse { - id, - name: custom_name, - first_delta: initial_input.is_empty(), - accumulated_json: initial_input, - provider_executed: Some(true), - dynamic, - provider_tool_name: Some(provider_name.to_string()), - provider_tool_input_type: match name.as_str() { - "text_editor_code_execution" | "bash_code_execution" => { - Some(name) - } - "code_execution" => { - Some("programmatic-tool-call".to_string()) - } - _ => None, - }, - provider_metadata: None, - }); - } - ContentBlock::McpToolUse { id, name, input, server_name } => { - mcp_tool_calls - .insert(id.clone(), (name.clone(), server_name.clone())); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id.clone(), - tool_name: name.clone(), - input: input.to_string(), - provider_executed: Some(true), - dynamic: Some(true), - provider_metadata: Some(provider_namespace("anthropic", json!({ - "type": "mcp-tool-use", - "serverName": server_name, - })).expect("provider metadata must be an object")), - })); - } - ContentBlock::RedactedThinking { data } => { - yield Ok(StreamPart::ReasoningStart { - id: index.to_string(), - provider_metadata: Some(provider_namespace("anthropic", json!({ - "redactedData": data - })).expect("provider metadata must be an object")), - }); - blocks.insert(index, BlockState::Thinking { started: true }); - } - other => { - for part in stream_parts_for_result_block( - &other, - &tool_names, - &mcp_tool_calls, - &server_tool_calls, - &crate::anthropic::stream::generate_id, - ) { - yield Ok(part); - } - } - } - } - StreamEvent::ContentBlockDelta { index, delta } => { - if let Some(text) = delta.text { - 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}); - } - yield Ok(StreamPart::TextDelta { - id: index.to_string(), - delta: text, - provider_metadata: None, - }); - } - if let Some(partial) = delta.partial_json { - let delta_event: Option<(String, String)> = match blocks.get_mut(&index) { - Some(BlockState::ToolUse { - id, - accumulated_json, - provider_tool_input_type, - first_delta, - .. - }) if !partial.is_empty() => { - let emitted_delta = if *first_delta { - if let Some(input_type) = provider_tool_input_type { - format!( - "{{\"type\": \"{input_type}\",{}", - partial.strip_prefix('{').unwrap_or(&partial) - ) - } else { - partial - } - } else { - partial - }; - accumulated_json.push_str(&emitted_delta); - *first_delta = false; - Some((id.clone(), emitted_delta)) - } - _ => None, - }; - if let Some((id, delta)) = delta_event { - yield Ok(StreamPart::ToolInputDelta { - id, - delta, - provider_metadata: None, - }); - } - } - if let Some(thinking) = delta.thinking { - 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, - }); - } - } - StreamEvent::ContentBlockStop { index } => { - if let Some(state) = blocks.remove(&index) { - match state { - BlockState::Text { started: true } => { - yield Ok(StreamPart::TextEnd { - id: index.to_string(), - provider_metadata: None, - }); - } - BlockState::Text { started: false } => {} - BlockState::Thinking { started: true } => { - yield Ok(StreamPart::ReasoningEnd { - id: index.to_string(), - provider_metadata: None, - }); - } - BlockState::Thinking { started: false } => {} - BlockState::ToolUse { - id, - name, - accumulated_json, - provider_executed, - dynamic, - provider_tool_name, - provider_tool_input_type, - provider_metadata, - .. - } => { - yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None}); - let input = finalize_streamed_tool_input( - accumulated_json, - provider_tool_name.as_deref(), - provider_tool_input_type.as_deref(), - ); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: name, - input, - provider_executed, - dynamic, - provider_metadata, - })); - } - } - } - } - StreamEvent::MessageDelta { delta, usage } => { - if let Some(reason) = delta.stop_reason { - final_finish_reason = Some(parse_stop_reason(&reason)); - } - if let Some(u) = usage { - // Reuse the shared conversion so the output - // reasoning/text split (from - // output_tokens_details) is preserved, matching - // anthropic/stream.rs. Only the output side is - // applied — MessageDelta carries output tokens - // only, so input-side usage from MessageStart is - // kept intact. - final_usage.output_tokens = - crate::anthropic::usage::usage_from_anthropic(&u) - .output_tokens; - } - } - StreamEvent::MessageStop => break, - StreamEvent::Error { error } => { - yield Ok(StreamPart::Error { - error: crate::anthropic::stream::anthropic_stream_error( - &error, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); - return; - } - _ => {} - } - } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; - } - } - } - } - - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }), - usage: final_usage, - provider_metadata: None, - }); - }; - - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) + prepare_body: Some(Arc::new(raw_predict_envelope)), + failed_response_handler: crate::google::google_failed_response_handler, + supports_native_structured_output: false, + supports_strict_tools: false, + }, + resolve: None, } } -/// Per-content-block state during streaming (mirrors the Anthropic provider). -enum BlockState { - Text { - started: bool, - }, - ToolUse { - id: String, - name: String, - accumulated_json: String, - provider_executed: Option, - dynamic: Option, - provider_tool_name: Option, - provider_tool_input_type: Option, - provider_metadata: Option, - first_delta: bool, - }, - Thinking { - started: bool, - }, +/// The model `model_id` on the Vertex endpoint `endpoint` resolves to for each +/// request. +pub(super) fn model( + model_id: &str, + endpoint: EndpointFn, + fetch: Option, + transform_request_body: Option, +) -> VertexAnthropicModel { + let resolve_fetch = fetch.clone(); + let resolve_transform = transform_request_body.clone(); + let mut config = request_config( + Endpoint { + base_url: String::new(), + headers: Default::default(), + }, + fetch, + transform_request_body, + ); + // The model's own fields only describe its identity; every request is + // built from the endpoint resolved for it. + config.resolve = Some(Arc::new(move || { + let endpoint = endpoint.clone(); + let fetch = resolve_fetch.clone(); + let transform = resolve_transform.clone(); + Box::pin(async move { Ok(request_config(endpoint().await?, fetch, transform)) }) + })); + AnthropicMessagesModel::with_config(model_id.to_string(), config) } diff --git a/aimux-providers/src/vertex/embedding.rs b/aimux-providers/src/vertex/embedding.rs index 5ebd732bb..0807adb06 100644 --- a/aimux-providers/src/vertex/embedding.rs +++ b/aimux-providers/src/vertex/embedding.rs @@ -8,12 +8,12 @@ //! `POST {base_url}/models/{model}:embedContent` (single value only) //! - Others: `POST {base_url}/models/{model}:predict` (batch) -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; use crate::google::google_failed_response_handler; +use crate::google::options::Namespace; +use crate::shared::EndpointConfig; use aimux_core::embedding_model::{ EmbeddingCallOptions, EmbeddingModel, EmbeddingResponse, EmbeddingResult, EmbeddingUsage, @@ -21,43 +21,19 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::HttpRequest; - -use super::VertexAuth; -use super::model::VertexConfig; - /// A Google Vertex AI embedding model (e.g. `"textembedding-gecko@001"`). /// /// 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 VertexEmbeddingModel { model_id: String, - config: VertexConfig, + config: EndpointConfig, } impl VertexEmbeddingModel { - #[must_use] - pub fn new(model_id: String, config: VertexConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> Vec<(String, String)> { - let mut headers = vec![("Content-Type".to_string(), "application/json".to_string())]; - match &self.config.auth { - VertexAuth::BearerToken(token) => { - headers.push(("Authorization".to_string(), format!("Bearer {token}"))); - } - VertexAuth::ApiKey(key) => { - headers.push(("x-goog-api-key".to_string(), key.clone())); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.push((k.clone(), v.clone())); - } - } - headers - } } /// Returns `true` for models that only support the `:embedContent` endpoint @@ -69,17 +45,13 @@ fn uses_embed_content_endpoint(model_id: &str) -> bool { #[async_trait] impl EmbeddingModel for VertexEmbeddingModel { fn provider(&self) -> &str { - "google.vertex" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - fn max_embeddings_per_call(&self) -> Option { if uses_embed_content_endpoint(&self.model_id) { Some(1) @@ -96,10 +68,10 @@ impl EmbeddingModel for VertexEmbeddingModel { &self, options: &EmbeddingCallOptions, ) -> Result { - // Parse provider options: try "googleVertex", then "vertex", then "google". + // Parse provider options: `googleVertex`, then `google`. let vertex_options = parse_vertex_provider_options(options.provider_options.as_ref()); - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; if uses_embed_content_endpoint(&self.model_id) { // gemini-embedding-2: use :embedContent endpoint (single value). @@ -136,22 +108,11 @@ impl EmbeddingModel for VertexEmbeddingModel { Value::Object(embed_config), ); - let url = format!( - "{}/models/{}:embedContent", - self.config.base_url, self.model_id - ); + let url = exchange.url(&format!("/models/{}:embedContent", self.model_id)); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url, - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, - Value::Object(body), + exchange.request(url, options), + exchange.transform_body(Value::Object(body)), aimux_provider_utils::create_json_response_handler(), google_failed_response_handler(), ) @@ -221,19 +182,11 @@ impl EmbeddingModel for VertexEmbeddingModel { body.insert("instances".to_string(), Value::Array(instances)); body.insert("parameters".to_string(), Value::Object(parameters)); - let url = format!("{}/models/{}:predict", self.config.base_url, self.model_id); + let url = exchange.url(&format!("/models/{}:predict", self.model_id)); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url, - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, - Value::Object(body), + exchange.request(url, options), + exchange.transform_body(Value::Object(body)), aimux_provider_utils::create_json_response_handler(), crate::google::google_failed_response_handler(), ) @@ -302,16 +255,12 @@ struct VertexEmbeddingProviderOptions { /// Parse Vertex embedding provider options. /// -/// Tries the `"googleVertex"` key first, then `"vertex"`, then `"google"` -/// (matching the TS fallback chain). +/// Tries the `googleVertex` key first, then `google` (the shared Gemini +/// model's key). fn parse_vertex_provider_options( options: Option<&SharedProviderOptions>, ) -> VertexEmbeddingProviderOptions { - let opts = options; - let provider_opts = opts - .and_then(|o| o.get("googleVertex")) - .or_else(|| opts.and_then(|o| o.get("vertex"))) - .or_else(|| opts.and_then(|o| o.get("google"))); + let provider_opts = Namespace::Vertex.read_in(options); VertexEmbeddingProviderOptions { output_dimensionality: provider_opts diff --git a/aimux-providers/src/vertex/image.rs b/aimux-providers/src/vertex/image.rs index f4dc4a1b6..eb23c746e 100644 --- a/aimux-providers/src/vertex/image.rs +++ b/aimux-providers/src/vertex/image.rs @@ -7,23 +7,19 @@ //! - **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}; use crate::google::google_failed_response_handler; +use crate::google::options::{vertex_metadata_map, vertex_options}; +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::{Warning, provider_namespace}; - -use aimux_provider_utils::HttpRequest; - -use super::{VertexAuth, VertexConfig}; +use aimux_core::shared::Warning; fn is_gemini_model(model_id: &str) -> bool { model_id.starts_with("gemini-") @@ -35,43 +31,25 @@ fn is_gemini_model(model_id: &str) -> bool { /// `Client` internally (RFC-0009 §4.1). pub struct VertexImageModel { model_id: String, - config: VertexConfig, + config: EndpointConfig, } impl VertexImageModel { - #[must_use] - pub fn new(model_id: String, config: VertexConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - fn build_headers(&self, extra: Option<&HashMap>) -> Vec<(String, String)> { - let mut h = vec![("Content-Type".into(), "application/json".into())]; - match &self.config.auth { - VertexAuth::BearerToken(token) => { - h.push(("Authorization".into(), format!("Bearer {token}"))); - } - VertexAuth::ApiKey(key) => { - h.push(("x-goog-api-key".into(), key.clone())); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - h.push((k.clone(), v.clone())); - } - } - h + fn predict_path(&self) -> String { + format!("/models/{}:predict", self.model_id) } - fn predict_endpoint(&self) -> String { - format!("{}/models/{}:predict", self.config.base_url, self.model_id) - } - fn generate_content_endpoint(&self) -> String { + fn generate_content_path(&self) -> String { let mp = if self.model_id.contains('/') { self.model_id.clone() } else { format!("models/{}", self.model_id) }; - format!("{}/{}:generateContent", self.config.base_url, mp) + format!("/{mp}:generateContent") } fn get_base64_data(file: &ImageFile) -> Result { @@ -104,11 +82,8 @@ impl VertexImageModel { }); } - // Parse provider options (googleVertex or vertex) - let gv_opts = options - .provider_options - .get("googleVertex") - .or_else(|| options.provider_options.get("vertex")); + // 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")) @@ -177,11 +152,11 @@ impl VertexImageModel { json!({ "instances": [{ "prompt": options.prompt }], "parameters": parameters }) }; - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.predict_endpoint(), headers, options), - body, + 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(), ) @@ -221,11 +196,7 @@ impl VertexImageModel { .unwrap_or_default(); let payload = json!({ "images": image_metas }); - let mut metadata = provider_namespace("googleVertex", payload.clone()) - .expect("provider metadata must be an object"); - metadata.extend( - provider_namespace("vertex", payload).expect("provider metadata must be an object"), - ); + let metadata = vertex_metadata_map(&payload); Ok(ImageResult { images: ImageOutputs::Base64(images), @@ -302,10 +273,7 @@ impl VertexImageModel { } // Only declared generation-config options survive the upstream language schema. - let gv_opts = options - .provider_options - .get("googleVertex") - .or_else(|| options.provider_options.get("vertex")); + 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) { @@ -318,11 +286,11 @@ impl VertexImageModel { } let body = json!({ "contents": [{ "role": "user", "parts": parts }], "generationConfig": Value::Object(gc) }); - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.generate_content_endpoint(), headers, options), - body, + exchange.request(exchange.url(&self.generate_content_path()), options), + exchange.transform_body(body), aimux_provider_utils::create_json_response_handler(), google_failed_response_handler(), ) @@ -372,11 +340,7 @@ impl VertexImageModel { }); let payload = json!({ "images": images.iter().map(|_| json!({})).collect::>() }); - let mut metadata = provider_namespace("googleVertex", payload.clone()) - .expect("provider metadata must be an object"); - metadata.extend( - provider_namespace("vertex", payload).expect("provider metadata must be an object"), - ); + let metadata = vertex_metadata_map(&payload); Ok(ImageResult { images: ImageOutputs::Base64(images), @@ -395,14 +359,11 @@ impl VertexImageModel { #[async_trait] impl ImageModel for VertexImageModel { fn provider(&self) -> &str { - "google.vertex" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } fn max_images_per_call(&self) -> Option { if is_gemini_model(&self.model_id) { Some(10) diff --git a/aimux-providers/src/vertex/mod.rs b/aimux-providers/src/vertex/mod.rs index 9c828aa70..4da5b5a58 100644 --- a/aimux-providers/src/vertex/mod.rs +++ b/aimux-providers/src/vertex/mod.rs @@ -1,23 +1,53 @@ //! Google Vertex AI provider. //! -//! Implements the `LanguageModel` trait against the Vertex AI API -//! (`{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}/publishers/google/models/{model}:generateContent`). +//! [`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. //! -//! Vertex AI uses the same request/response format as the Google Gemini API, -//! so this provider reuses the shared `google::convert` message conversion -//! logic. The differences are: -//! - **Endpoint**: Vertex uses a project/location-scoped URL instead of the -//! public Gemini API URL. -//! - **Authentication**: Vertex uses a Google Cloud OAuth2 Bearer token -//! (obtained from a service account or `gcloud auth print-access-token`) -//! instead of an API key. Express mode (API key) is also supported. +//! 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: +use std::sync::Arc; + +use futures::future::BoxFuture; + +use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::{RetryConfig, without_trailing_slash}; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::{LanguageModel, SupportedUrls}; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_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, +}; +use std::sync::OnceLock; + +use crate::shared::{Endpoint, EndpointConfig, TransformRequestBody, is_valid_hostname_part}; mod anthropic_model; mod embedding; @@ -26,401 +56,471 @@ mod model; mod transcription; mod video; -pub use anthropic_model::{VertexAnthropicConfig, VertexAnthropicModel}; +pub use anthropic_model::VertexAnthropicModel; pub use embedding::VertexEmbeddingModel; pub use image::VertexImageModel; -pub use model::{VertexConfig, VertexModel}; +pub use model::VertexModel; pub use transcription::VertexTranscriptionModel; pub use video::VertexVideoModel; -/// Authentication method for the Vertex AI provider. -#[derive(Debug, Clone)] -pub enum VertexAuth { - /// Google Cloud OAuth2 Bearer token (standard Vertex AI auth). - BearerToken(String), - /// API key for Vertex AI Express mode. - ApiKey(String), -} +/// The Express-mode base URL. +const EXPRESS_MODE_BASE_URL: &str = "https://aiplatform.googleapis.com/v1/publishers/google"; +pub(crate) const API_KEY_ENV_VAR: &str = "GOOGLE_VERTEX_API_KEY"; +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"; -/// Configuration for the Google Vertex AI provider. -#[derive(Debug, Clone)] -pub struct VertexProviderConfig { - /// The base URL for Vertex AI API calls. - pub base_url: String, - /// Google Cloud project ID (used for Speech-to-Text transcription). - pub project: Option, - /// Default location/region (used for Speech-to-Text transcription). +/// 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). +/// +/// 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. +#[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. + pub api_key: Option>, + /// The location (region) of standard mode: a single DNS label 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 + /// custom endpoint. pub location: Option, - /// Authentication method. - pub auth: VertexAuth, - /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 - pub api_key_source: Option, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, + /// The Google Cloud project of standard mode. `None` loads + /// `GOOGLE_VERTEX_PROJECT` when a request is made and fails that request + /// with `AiMuxError::LoadSetting` if it is unset. + pub project: Option, + /// Base URL for the API calls, replacing the host/project/location URL in + /// both modes. A trailing slash is removed. + pub base_url: Option, + /// Extra headers on every request, as a value or a producer called on + /// every request. A `None` value removes the header. In standard mode they + /// 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. + pub access_token: Option>, + /// The transport: a mock, a signing decorator, a proxy-aware client. + /// `None` uses the process default, resolved per request. + pub fetch: Option, + /// Rewrites every JSON request body once, after it is serialized and + /// before it is sent. + pub transform_request_body: Option, } -impl VertexProviderConfig { - /// Create a config using a Google Cloud OAuth2 Bearer token. - /// - /// # Arguments - /// - `access_token` - A Google Cloud access token (e.g. from - /// `gcloud auth print-access-token` or a service account JWT exchange). - /// - `project` - The Google Cloud project ID. - /// - `location` - The Vertex AI location (e.g. `"us-central1"`). - pub fn new( - access_token: impl Into, - project: impl Into, - location: impl Into, - ) -> Self { - let location = location.into(); - let project = project.into(); - let base_url = build_base_url(&project, &location, false); - Self { - base_url, - project: Some(project), - location: Some(location), - auth: VertexAuth::BearerToken(access_token.into()), - api_key_source: None, - retry_config: RetryConfig::default(), - } +impl std::fmt::Debug for VertexProviderSettings { + /// Never prints keys, tokens or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VertexProviderSettings") + .field("api_key", &self.api_key) + .field("location", &self.location) + .field("project", &self.project) + .field("base_url", &self.base_url) + .field("headers", &self.headers.is_some()) + .field("access_token", &self.access_token) + .field("fetch", &self.fetch.is_some()) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ) + .finish() } +} - /// Create a config using an API key (Vertex AI Express mode). - pub fn with_api_key(api_key: impl Into) -> Self { - Self { - base_url: "https://aiplatform.googleapis.com/v1/publishers/google".to_string(), - project: None, - location: None, - auth: VertexAuth::ApiKey(api_key.into()), - api_key_source: None, - retry_config: RetryConfig::default(), - } +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(), + )) } +} - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self +/// 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 { + match location { + "global" => "aiplatform.googleapis.com".to_string(), + "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), + _ => format!("{location}-aiplatform.googleapis.com"), } +} - /// 标注凭证来源(RFC-0023 回放重建用)。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self +/// Which publisher's models a base URL addresses. +#[derive(Clone, Copy)] +enum Publisher { + Google, + Anthropic, +} + +/// A project and a location, resolved for one request. +pub(crate) struct ProjectLocation { + pub(crate) project: String, + pub(crate) location: String, +} + +/// Resolves the project and location of a request (Speech-to-Text and the +/// operations that name them), failing for Express mode. +pub(crate) type ProjectLocationFn = + Arc BoxFuture<'static, Result> + Send + Sync>; + +/// The settings that decide, per request, which mode applies, where requests +/// go and what authenticates them. +struct Resolver { + api_key: Option>, + location: Option, + project: Option, + base_url: Option, + headers: Option>, + access_token: Option>, +} + +impl Resolver { + /// The Express-mode key, when there is a non-empty one. + 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())) } - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. - #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self + fn location(&self) -> Result { + let location = load_setting(self.location.as_deref(), LOCATION_ENV_VAR, "location")?; + validate_location(&location)?; + Ok(location) } - /// Create from environment variables. - /// - /// Checks for `GOOGLE_VERTEX_API_KEY` (Express mode) first, then falls - /// back to `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION`. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when neither `GOOGLE_VERTEX_API_KEY` - /// nor the `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` pair is - /// set. - pub fn from_env() -> Result { - // Express mode with API key. - if let Ok(key) = std::env::var("GOOGLE_VERTEX_API_KEY") - && !key.trim().is_empty() - { - return Ok( - Self::with_api_key(key).with_api_key_source(Some("env:GOOGLE_VERTEX_API_KEY")) - ); + fn project(&self) -> Result { + load_setting(self.project.as_deref(), PROJECT_ENV_VAR, "project") + } + + /// The base URL of `publisher`'s models. + fn base_url(&self, express: bool, publisher: Publisher) -> Result { + match publisher { + Publisher::Google => { + if let Some(url) = &self.base_url { + return Ok(url.clone()); + } + if express { + return Ok(EXPRESS_MODE_BASE_URL.to_string()); + } + let location = self.location()?; + let project = self.project()?; + Ok(format!( + "https://{}/v1beta1/projects/{project}/locations/{location}/publishers/google", + location_host(&location) + )) + } + 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) + ) + }; + Ok(format!("{root}/publishers/anthropic/models")) + } } + } - // Standard auth with access token. - let access_token = std::env::var("GOOGLE_VERTEX_ACCESS_TOKEN").map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN (or GOOGLE_VERTEX_API_KEY) environment variable is required for Vertex AI" - .to_string(), - ) - })?; - let project = std::env::var("GOOGLE_VERTEX_PROJECT").map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI".to_string(), - ) - })?; - let location = - std::env::var("GOOGLE_VERTEX_LOCATION").unwrap_or_else(|_| "us-central1".to_string()); + /// 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 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")?, + }; + let mut auth = HeaderMapOpt::new(); + auth.insert("Authorization".to_string(), Some(format!("Bearer {token}"))); + Ok(combine_headers(&[&auth, &user])) + } - Ok(Self::new(access_token, project, location) - .with_api_key_source(Some("env:GOOGLE_VERTEX_ACCESS_TOKEN"))) + /// The endpoint of `publisher`'s models for one request. + async fn endpoint(&self, publisher: Publisher) -> Result { + let express_key = self.express_key().await?; + let base_url = self.base_url(express_key.is_some(), publisher)?; + let headers = self.headers(express_key.as_deref()).await?; + let headers = aimux_provider_utils::headers::with_user_agent_suffix_fn( + Resolvable::Value(headers), + "google-vertex", + "5.0.98", + ) + .resolve() + .await?; + Ok(Endpoint { base_url, headers }) } } -/// Build the Vertex AI base URL for a project/location. +/// Create a Google Vertex AI provider. /// -/// - `global` → `https://aiplatform.googleapis.com/v1beta1/projects/{project}/locations/global/publishers/google` -/// - `eu`/`us` → `https://aiplatform.{region}.rep.googleapis.com/v1beta1/projects/{project}/locations/{region}/publishers/google` -/// - other → `https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}/publishers/google` -fn build_base_url(project: &str, location: &str, _endpoint: bool) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1beta1/projects/{project}/locations/{location}/publishers/google") +/// # 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. +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)?; + } + Ok(VertexProvider { + resolver: Arc::new(Resolver { + api_key: settings.api_key, + location: settings.location, + project: settings.project, + base_url, + headers: settings.headers, + access_token: settings.access_token, + }), + fetch: settings.fetch, + transform_request_body: settings.transform_request_body, + }) } -/// Google Vertex AI provider — creates [`VertexModel`] instances. -/// -/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the -/// process-wide shared `Client` internally (RFC-0009 §4.1). +/// 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 +/// request instead. +pub fn google_vertex() -> &'static VertexProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_google_vertex(VertexProviderSettings::default()) + .expect("default Vertex settings are always valid") + }) +} + +/// A Google Vertex AI provider (the AI SDK's `GoogleVertexProvider`). Cheap to +/// clone the models out of; it holds no HTTP client. pub struct VertexProvider { - config: VertexProviderConfig, + resolver: Arc, + fetch: Option, + transform_request_body: Option, +} + +/// The URL patterns the Vertex models fetch themselves +/// (`supportedUrls` in `createGoogleVertex`). +fn supported_urls() -> SupportedUrls { + let pattern = |source: &str| regex::Regex::new(source).expect("static pattern"); + SupportedUrls( + [( + "*".to_string(), + vec![pattern(r"^https?://.*$"), pattern(r"^gs://.*$")], + )] + .into_iter() + .collect(), + ) } impl VertexProvider { - #[must_use] - pub fn new(config: VertexProviderConfig) -> Self { - Self { config } + /// The configuration of a Gemini-publisher model reporting `provider`. + /// With `tuned`, the request fails in Express mode (tuned models are + /// served from a deployed endpoint, which an API key cannot reach). + fn model_config(&self, provider: &str, tuned: bool) -> EndpointConfig { + let resolver = self.resolver.clone(); + EndpointConfig { + provider: provider.to_string(), + endpoint: Arc::new(move || { + let resolver = resolver.clone(); + Box::pin(async move { + if tuned && resolver.express_key().await?.is_some() { + return Err(AiMuxError::InvalidArgument( + "Google Vertex tuned models do not support Express Mode API keys. \ + Use standard Google Cloud credentials instead." + .to_string(), + )); + } + resolver.endpoint(Publisher::Google).await + }) + }), + fetch: self.fetch.clone(), + supported_urls: Arc::new(|_| supported_urls()), + transform_request_body: self.transform_request_body.clone(), + } } - /// Create a model instance for the given Vertex AI model name - /// (e.g. `"gemini-2.0-flash"`). - /// - /// Returns an error when Express Mode (API key auth) is used with a tuned - /// model (`endpoints/…`), which the Vertex AI API does not support. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when Express Mode (API key auth) is - /// combined with a tuned model (`endpoints/…`), which the Vertex AI API does - /// not support. - pub fn model(&self, model_id: &str) -> Result { - if model_id.starts_with("endpoints/") && matches!(self.config.auth, VertexAuth::ApiKey(_)) { - return Err(AiMuxError::InvalidArgument( - "Google Vertex tuned models do not support Express Mode API keys. Use standard Google Cloud credentials instead.".to_string(), - )); - } - Ok(VertexModel::new( - model_id.to_string(), - VertexConfig { - base_url: self.config.base_url.clone(), - auth: self.config.auth.clone(), - api_key_source: self.config.api_key_source.clone(), - retry_config: self.config.retry_config, - }, - )) + /// The project and location of a request, for the Speech-to-Text models. + fn project_location(&self, what: &'static str) -> ProjectLocationFn { + let resolver = self.resolver.clone(); + Arc::new(move || { + let resolver = resolver.clone(); + Box::pin(async move { + if resolver.express_key().await?.is_some() { + return Err(AiMuxError::InvalidArgument(format!( + "Google Vertex {what} models do not support Express Mode API keys. \ + Use standard Google Cloud credentials instead." + ))); + } + Ok(ProjectLocation { + location: resolver.location()?, + project: resolver.project()?, + }) + }) + }) } - /// Create an Anthropic partner model instance served via Vertex AI - /// `rawPredict` (e.g. `"claude-sonnet-4-20250514"`). - /// - /// The model identity is carried by the request URL - /// (`publishers/anthropic/models/{model}`); the configured base URL's - /// `/publishers/google` suffix (if present) is stripped so the - /// `/publishers/anthropic` suffix can be appended by the model. Auth is - /// reused from this provider (Bearer token or API key). - /// - /// # Errors + /// A Gemini model (e.g. `"gemini-2.0-flash"`, or a tuned + /// `"endpoints/{id}"`); `provider()` is `"google.vertex"`. /// - /// Propagates configuration errors of the underlying `VertexAnthropicModel` - /// construction (none in the current implementation). - pub fn anthropic_model(&self, model_id: &str) -> Result { - let base_url = self - .config - .base_url - .strip_suffix("/publishers/google") - .map(std::string::ToString::to_string) - .unwrap_or_else(|| self.config.base_url.clone()); - Ok(VertexAnthropicModel::new( + /// A tuned model fails its requests with `AiMuxError::InvalidArgument` in + /// Express mode. + #[must_use] + pub fn chat(&self, model_id: &str) -> VertexModel { + VertexModel::from_config( model_id.to_string(), - VertexAnthropicConfig { - base_url, - auth: self.config.auth.clone(), - api_key_source: self.config.api_key_source.clone(), - retry_config: self.config.retry_config, - }, - )) + self.model_config(PROVIDER, model_id.starts_with(ENDPOINT_MODEL_PREFIX)), + ) } - /// Create an embedding model instance for the given Vertex AI model name - /// (e.g. `"textembedding-gecko@001"`). + /// An embedding model (e.g. `"textembedding-gecko@001"`); `provider()` is + /// `"google.vertex"`. #[must_use] - pub fn embedding_model(&self, model_id: &str) -> VertexEmbeddingModel { - VertexEmbeddingModel::new( - model_id.to_string(), - VertexConfig { - base_url: self.config.base_url.clone(), - auth: self.config.auth.clone(), - api_key_source: self.config.api_key_source.clone(), - retry_config: self.config.retry_config, - }, - ) + pub fn embedding(&self, model_id: &str) -> VertexEmbeddingModel { + VertexEmbeddingModel::from_config(model_id.to_string(), self.model_config(PROVIDER, false)) } - /// Create an image generation model instance for the given Vertex AI model - /// name (e.g. `"imagen-4.0-generate-001"` or `"gemini-2.5-flash-image"`). + /// An image generation model (e.g. `"imagen-4.0-generate-001"` or + /// `"gemini-2.5-flash-image"`); `provider()` is `"google.vertex"`. #[must_use] pub fn image(&self, model_id: &str) -> VertexImageModel { - VertexImageModel::new( + VertexImageModel::from_config(model_id.to_string(), self.model_config(PROVIDER, false)) + } + + /// A video generation model (e.g. `"veo-3.0-generate-001"`); + /// `provider()` is `"google.vertex.video"`. + #[must_use] + pub fn video(&self, model_id: &str) -> VertexVideoModel { + VertexVideoModel::from_config( model_id.to_string(), - VertexConfig { - base_url: self.config.base_url.clone(), - auth: self.config.auth.clone(), - api_key_source: self.config.api_key_source.clone(), - retry_config: self.config.retry_config, - }, + self.model_config("google.vertex.video", false), ) } - /// Create a transcription (STT) model instance for the given Speech-to-Text - /// model name (e.g. `"chirp_3"`). Requires project and location to be set - /// in the config (use [`VertexProviderConfig::new`]). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `project` is not set in the - /// config. - pub fn transcription(&self, model_id: &str) -> Result { - let project = self.config.project.clone().ok_or_else(|| { - AiMuxError::InvalidArgument( - "Google Vertex transcription requires project and location. Use VertexProviderConfig::new() instead of with_api_key()." - .to_string(), - ) - })?; - let location = self - .config - .location - .clone() - .unwrap_or_else(|| "us-central1".to_string()); - Ok(VertexTranscriptionModel::new( + /// A Speech-to-Text transcription model (e.g. `"chirp_3"`); `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( model_id.to_string(), - project, - location, - self.config.auth.clone(), - self.config.base_url.clone(), + self.model_config("google.vertex.transcription", false), + self.project_location("transcription"), ) - .with_retry_config(self.config.retry_config)) } - /// Create a video generation model instance for the given Vertex AI model - /// name (e.g. `"veo-3.0-generate-001"`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `project` is not set in the - /// config. - pub fn video(&self, model_id: &str) -> Result { - let project = self.config.project.clone().ok_or_else(|| { - AiMuxError::InvalidArgument( - "Google Vertex video requires project and location. Use VertexProviderConfig::new() instead of with_api_key()." - .to_string(), - ) - })?; - let location = self - .config - .location - .clone() - .unwrap_or_else(|| "us-central1".to_string()); - Ok(VertexVideoModel::new( - model_id.to_string(), - project, - location, - self.config.auth.clone(), - self.config.base_url.clone(), + /// 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(), ) - .with_retry_config(self.config.retry_config)) + } + + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; the same model as [`chat`](Self::chat). + #[must_use] + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.chat(model_id)) } } impl Provider for VertexProvider { - fn name(&self) -> &str { - "google.vertex" + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) + } + + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id)?)) + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.image(model_id))) } - /// List models via `GET {base_url}/models` (Gemini native, RFC-0027). - fn list_models( + fn transcription_model( &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let base_url = self.config.base_url.clone(); - let auth = self.config.auth.clone(); - let retry_config = self.config.retry_config; - Box::pin(async move { - let base = base_url.trim_end_matches('/'); - let url = format!("{base}/models"); - let mut headers = Vec::new(); - match &auth { - VertexAuth::BearerToken(token) => { - headers.push(("Authorization".to_string(), format!("Bearer {token}"))); - } - VertexAuth::ApiKey(key) => { - headers.push(("x-goog-api-key".to_string(), key.clone())); - } - } - headers.push(("Content-Type".to_string(), "application/json".to_string())); - - use aimux_provider_utils::HttpRequest; - // Retry rationale: see `openai::model::execute_list_models`. - let resp = aimux_core::retry::prepare_retries(None, retry_config, None) - .retry(|| { - aimux_provider_utils::get_from_api( - HttpRequest { - url: url.clone(), - headers: headers.clone(), - abort_signal: None, - call_id: None, - recording_context: None, - ..Default::default() - }, - aimux_provider_utils::create_json_response_handler(), - crate::google::google_failed_response_handler(), - ) - }) - .await?; + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.transcription(model_id)))) + } - #[derive(serde::Deserialize)] - struct Resp { - #[serde(default)] - models: Vec, - } - #[derive(serde::Deserialize)] - struct Entry { - name: String, - #[serde(default)] - display_name: Option, - } - let parsed: Resp = resp.value; - let runtime: Vec = parsed - .models - .into_iter() - .map(|e| { - let id = e - .name - .strip_prefix("models/") - .unwrap_or(&e.name) - .to_string(); - aimux_core::model_catalogue::RuntimeModel { - id, - owned_by: e.display_name, - created: None, - } - }) - .collect(); - Ok(runtime) - }) + fn video_model(&self, model_id: &str) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.video(model_id)))) + } +} + +impl ProviderDiscovery for VertexProvider { + /// `GET {base_url}/models` (Gemini native): one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config("google.vertex.models", false); + Box::pin(async move { crate::google::list_models_once(&config).await }) } } diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs index 0877f7a26..55214928d 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -6,8 +6,6 @@ use aimux_core::tool::RawToolCall; use aimux_core::tool::ToolResult; -use std::collections::HashMap; - use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; @@ -18,32 +16,21 @@ use aimux_core::options::CallOptions; use aimux_core::result::{ GenerateContent, GenerateResult, GeneratedFile, ReasoningOutput, Source, StreamResult, }; -use aimux_core::shared::{FileBytes, GeneratedFileData, provider_namespace}; +use aimux_core::shared::{FileBytes, GeneratedFileData}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, }; -use aimux_provider_utils::{HttpRequest, RetryConfig}; - use crate::google::convert::{ build_vertex_request_body, code_execution_tool_name, convert_usage, extract_sources, parse_finish_reason, }; use crate::google::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; -use super::VertexAuth; - -/// Configuration for a Vertex model instance (cloned from the provider). -#[derive(Debug, Clone)] -pub struct VertexConfig { - pub base_url: String, - pub auth: VertexAuth, - /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 - pub api_key_source: Option, - /// Retry settings used by Core model operations. - pub retry_config: RetryConfig, -} +use crate::google::options::Namespace; +use crate::shared::EndpointConfig; +use aimux_core::language_model::SupportedUrls; /// A Google Vertex AI language model. /// @@ -51,72 +38,50 @@ pub struct VertexConfig { /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct VertexModel { model_id: String, - config: VertexConfig, + config: EndpointConfig, } impl VertexModel { - #[must_use] - pub fn new(model_id: String, config: VertexConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - fn build_headers(&self, extra: Option<&HashMap>) -> Vec<(String, String)> { - let mut headers = vec![("Content-Type".to_string(), "application/json".to_string())]; - match &self.config.auth { - VertexAuth::BearerToken(token) => { - headers.push(("Authorization".to_string(), format!("Bearer {token}"))); - } - VertexAuth::ApiKey(key) => { - headers.push(("x-goog-api-key".to_string(), key.clone())); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - headers.push((k.clone(), v.clone())); - } - } - headers - } - /// 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 - /// configured base URL. Mirrors the TS `loadBaseURL({ endpoint: true })`. - fn effective_base_url(&self) -> &str { - if self.model_id.starts_with("endpoints/") { - // Strip the trailing `/publishers/google` (if present). - if let Some(stripped) = self.config.base_url.strip_suffix("/publishers/google") { - return stripped; - } + /// 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; } - &self.config.base_url + base_url } - /// `…/models/{model}:generateContent` - fn generate_endpoint(&self) -> String { - let model_path = if self.model_id.contains('/') { + fn model_path(&self) -> String { + if self.model_id.contains('/') { self.model_id.clone() } else { format!("models/{}", self.model_id) - }; + } + } + + /// `…/models/{model}:generateContent` + fn generate_endpoint(&self, base_url: &str) -> String { format!( "{}/{}:generateContent", - self.effective_base_url(), - model_path + self.effective_base_url(base_url), + self.model_path() ) } /// `…/models/{model}:streamGenerateContent?alt=sse` - fn stream_endpoint(&self) -> String { - let model_path = if self.model_id.contains('/') { - self.model_id.clone() - } else { - format!("models/{}", self.model_id) - }; + fn stream_endpoint(&self, base_url: &str) -> String { format!( "{}/{}:streamGenerateContent?alt=sse", - self.effective_base_url(), - model_path + self.effective_base_url(base_url), + self.model_path() ) } } @@ -124,39 +89,15 @@ impl VertexModel { #[async_trait] impl LanguageModel for VertexModel { fn provider(&self) -> &str { - "google.vertex" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - use aimux_core::recording::ProviderRecord; - // M2b: record identity + credential source + auth kind. Never serialize - // the bearer token or API key plaintext. - let auth_kind = match &self.config.auth { - VertexAuth::BearerToken(_) => "bearer_token", - VertexAuth::ApiKey(_) => "api_key", - }; - ProviderRecord { - provider: self.provider().to_string(), - model_id: self.model_id.clone(), - base_url: Some(self.config.base_url.clone()), - api_key_source: self - .config - .api_key_source - .clone() - .unwrap_or_else(|| "explicit".to_string()), - profile: None, - provider_options: Some(serde_json::json!({ - "auth_kind": auth_kind, - })), - } + async fn supported_urls(&self) -> SupportedUrls { + (self.config.supported_urls)(&self.model_id) } async fn do_generate(&self, options: &CallOptions) -> Result { @@ -228,9 +169,10 @@ impl LanguageModel for VertexModel { crate::google::convert::validate_tool_result_files(&options.prompt)?; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); let body = build_vertex_request_body(&self.model_id, options)?; - let headers = self.build_headers(options.headers.as_ref()); + 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( - HttpRequest::new(self.generate_endpoint(), headers, options), + exchange.request(self.generate_endpoint(exchange.base_url()), options), body.clone(), aimux_provider_utils::create_json_response_handler(), crate::google::google_failed_response_handler(), @@ -361,10 +303,11 @@ impl LanguageModel for VertexModel { crate::google::convert::validate_tool_result_files(&options.prompt)?; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); let body = build_vertex_request_body(&self.model_id, options)?; - let headers = self.build_headers(options.headers.as_ref()); - let endpoint = self.stream_endpoint(); + 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( - HttpRequest::new(endpoint.clone(), headers, options), + exchange.request(endpoint.clone(), options), body.clone(), aimux_provider_utils::create_event_source_response_handler::(), crate::google::google_failed_response_handler(), @@ -829,12 +772,7 @@ impl LanguageModel for VertexModel { // ── Helpers ────────────────────────────────────────────────────────────────── fn vertex_provider_metadata(payload: Value) -> ProviderMetadata { - let mut metadata = provider_namespace("googleVertex", payload.clone()) - .expect("provider metadata must be an object"); - metadata.extend( - provider_namespace("vertex", payload).expect("provider metadata must be an object"), - ); - metadata + Namespace::Vertex.metadata(payload) } fn vertex_server_tool_metadata( diff --git a/aimux-providers/src/vertex/transcription.rs b/aimux-providers/src/vertex/transcription.rs index 14f8e8a30..5c978683d 100644 --- a/aimux-providers/src/vertex/transcription.rs +++ b/aimux-providers/src/vertex/transcription.rs @@ -10,21 +10,20 @@ //! features). It returns `results[]` with `alternatives[]` containing //! `transcript` and `words[]` with timing offsets. -use std::collections::HashMap; - use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; +use crate::google::options::Namespace; +use crate::shared::EndpointConfig; use aimux_core::error::AiMuxError; use aimux_core::shared::Warning; use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::{HttpRequest, RetryConfig}; -use super::VertexAuth; +use super::ProjectLocationFn; // ── Response schema ───────────────────────────────────────────────────────── @@ -94,6 +93,7 @@ fn audio_input_to_base64(audio: &AudioInput) -> Result { &base64::engine::general_purpose::STANDARD, bytes, )), + AudioInput::Url(_) => Err(AudioInput::url_error()), } } @@ -101,69 +101,30 @@ fn audio_input_to_base64(audio: &AudioInput) -> Result { pub struct VertexTranscriptionModel { model_id: String, - project: String, - location: String, - auth: VertexAuth, - base_url: String, - retry_config: RetryConfig, + config: EndpointConfig, + project_location: ProjectLocationFn, } impl VertexTranscriptionModel { - #[must_use] - pub fn new( + pub(crate) fn from_config( model_id: String, - project: String, - location: String, - auth: VertexAuth, - base_url: String, + config: EndpointConfig, + project_location: ProjectLocationFn, ) -> Self { Self { model_id, - project, - location, - auth, - base_url, - retry_config: RetryConfig::default(), - } - } - - pub(crate) fn with_retry_config(mut self, retry_config: RetryConfig) -> Self { - self.retry_config = retry_config; - self - } - - fn auth_header(&self) -> HashMap { - let mut headers = HashMap::new(); - match &self.auth { - VertexAuth::BearerToken(token) => { - headers.insert("Authorization".to_string(), format!("Bearer {token}")); - } - VertexAuth::ApiKey(key) => { - headers.insert("x-goog-api-key".to_string(), key.clone()); - } + config, + project_location, } - headers } - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = self.auth_header(); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self, region: &str) -> String { - // When base_url is a test server (localhost / 127.0.0.1), use it - // directly so wiremock can intercept the request. - if self.base_url.starts_with("http://127.0.0.1") - || self.base_url.starts_with("http://localhost") - { + /// The Speech-to-Text v2 recognizer endpoint. A local `base_url` (a test + /// server or a proxy on the loopback interface) is used directly; the + /// real service has its own regional host. + fn endpoint(&self, base_url: &str, project: &str, region: &str) -> String { + if base_url.starts_with("http://127.0.0.1") || base_url.starts_with("http://localhost") { return format!( - "{}/v2/projects/{}/locations/{}/recognizers/_:recognize", - self.base_url, self.project, region + "{base_url}/v2/projects/{project}/locations/{region}/recognizers/_:recognize" ); } let host = if region == "global" { @@ -171,41 +132,35 @@ impl VertexTranscriptionModel { } else { format!("{region}-speech.googleapis.com") }; - format!( - "https://{host}/v2/projects/{}/locations/{}/recognizers/_:recognize", - self.project, region - ) + format!("https://{host}/v2/projects/{project}/locations/{region}/recognizers/_:recognize") } } #[async_trait] impl TranscriptionModel for VertexTranscriptionModel { fn provider(&self) -> &str { - "google.vertex" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.retry_config - } - async fn do_generate( &self, options: &TranscriptionCallOptions, ) -> Result { let warnings: Vec = Vec::new(); - // Parse provider options (may be under googleVertex, vertex, or google). - let mut region = self.location.clone(); + // Parse provider options (`googleVertex`, then `google`). + let target = (self.project_location)().await?; + let mut region = target.location.clone(); let mut language_codes: Vec = vec!["auto".to_string()]; let mut enable_word_time_offsets = true; let mut enable_automatic_punctuation = true; if let Some(ref po) = options.provider_options { - for key in &["googleVertex", "vertex", "google"] { + for key in Namespace::Vertex.read_keys() { if let Some(gv) = po.get(*key) { if let Some(r) = gv.get("region").and_then(|v| v.as_str()) { region = r.to_string(); @@ -248,25 +203,12 @@ impl TranscriptionModel for VertexTranscriptionModel { "content": content, }); - let headers = self.build_headers(options.headers.as_ref()); - let url = self.endpoint(®ion); - - let header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let url = self.endpoint(exchange.base_url(), &target.project, ®ion); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest { - url, - headers: header_list, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, - request_body.clone(), + exchange.request(url, options), + exchange.transform_body(request_body.clone()), aimux_provider_utils::create_json_response_handler::(), crate::google::google_failed_response_handler(), ) diff --git a/aimux-providers/src/vertex/video.rs b/aimux-providers/src/vertex/video.rs index a5139f4cf..63a182032 100644 --- a/aimux-providers/src/vertex/video.rs +++ b/aimux-providers/src/vertex/video.rs @@ -5,11 +5,9 @@ //! //! Uses the Long Running Operations API: //! 1. POST `{base_url}/models/{model}:predictLongRunning` → returns operation name -//! 2. GET operation — polled by Core via `do_status` until `done: true` +//! 2. POST `:fetchPredictOperation` — polled by Core until `done: true` //! 3. Return video URL(s) -use std::collections::HashMap; - use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -19,9 +17,7 @@ use aimux_core::video_model::{ VideoResponse, VideoResult, }; -use aimux_provider_utils::{HttpRequest, RetryConfig}; - -use super::VertexAuth; +use crate::shared::EndpointConfig; /// A Google Vertex AI video generation model. /// @@ -29,96 +25,25 @@ use super::VertexAuth; /// `Client` internally (RFC-0009 §4.1). pub struct VertexVideoModel { model_id: String, - project: String, - location: String, - auth: VertexAuth, - base_url: String, - retry_config: RetryConfig, + config: EndpointConfig, } impl VertexVideoModel { - #[must_use] - pub fn new( - model_id: String, - project: String, - location: String, - auth: VertexAuth, - base_url: String, - ) -> Self { - Self { - model_id, - project, - location, - auth, - base_url, - retry_config: RetryConfig::default(), - } - } - - pub(crate) fn with_retry_config(mut self, retry_config: RetryConfig) -> Self { - self.retry_config = retry_config; - self - } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut h = HashMap::new(); - match &self.auth { - VertexAuth::BearerToken(token) => { - h.insert("Authorization".to_string(), format!("Bearer {token}")); - } - VertexAuth::ApiKey(key) => { - h.insert("x-goog-api-key".to_string(), key.clone()); - } - } - if let Some(extra) = extra { - for (k, v) in extra { - h.insert(k.clone(), v.clone()); - } - } - h - } - - fn predict_url(&self) -> String { - if self.base_url.starts_with("http://127.0.0.1") - || self.base_url.starts_with("http://localhost") - { - return format!( - "{}/models/{}:predictLongRunning", - self.base_url, self.model_id - ); - } - format!( - "https://{}-aiplatform.googleapis.com/v1beta1/projects/{}/locations/{}/publishers/google/models/{}:predictLongRunning", - self.location, self.project, self.location, self.model_id - ) - } - - fn operation_url(&self, name: &str) -> String { - if self.base_url.starts_with("http://127.0.0.1") - || self.base_url.starts_with("http://localhost") - { - return format!("{}/{}", self.base_url, name); - } - format!( - "https://{}-aiplatform.googleapis.com/v1beta1/{}", - self.location, name - ) + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { + Self { model_id, config } } } #[async_trait] impl VideoModel for VertexVideoModel { fn provider(&self) -> &str { - "google.vertex" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.retry_config - } fn max_videos_per_call(&self) -> Option { - Some(1) + Some(4) } async fn do_start( @@ -133,6 +58,7 @@ impl VideoModel for VertexVideoModel { } 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())); } @@ -151,15 +77,12 @@ impl VideoModel for VertexVideoModel { "parameters": Value::Object(parameters), }); - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; + let url = exchange.url(&format!("/models/{}:predictLongRunning", self.model_id)); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.predict_url(), header_list.clone(), options), - body, + exchange.request(url, options), + exchange.transform_body(body), aimux_provider_utils::create_json_response_handler(), crate::google::google_failed_response_handler(), ) @@ -171,9 +94,7 @@ impl VideoModel for VertexVideoModel { .get("name") .and_then(|v| v.as_str()) .ok_or_else(|| { - AiMuxError::InvalidResponseData( - "Vertex video prediction missing operation name".to_string(), - ) + AiMuxError::InvalidResponseData("No operation name returned from API".to_string()) })? .to_string(); @@ -203,15 +124,12 @@ impl VideoModel for VertexVideoModel { ) })?; - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - - let poll_url = self.operation_url(operation_name); - let resp = aimux_provider_utils::get_from_api( - HttpRequest::new(poll_url.clone(), header_list, options), + 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), aimux_provider_utils::create_json_response_handler::(), crate::google::google_failed_response_handler(), ) @@ -220,8 +138,9 @@ impl VideoModel for VertexVideoModel { let response_headers = resp.response_headers; let response_body = resp.raw_value.as_ref().map(ToString::to_string); let raw_body: Value = resp.value; - // Check the in-band error first: a terminal response may carry both - // done:true and an error object (provider-declared failure). + 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") { let msg = err .get("message") @@ -237,39 +156,84 @@ impl VideoModel for VertexVideoModel { ..ApiCallError::new(msg, poll_url, serde_json::json!({})) }))); } - if raw_body.get("done").and_then(serde_json::Value::as_bool) != Some(true) { - return Ok(VideoOperationStatus::Pending); - } - let videos: Vec = raw_body + 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(|v| v.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|v| { - v.get("gcsUri") - .or_else(|| v.get("url")) - .and_then(|u| u.as_str()) - .map(|url| VideoData::Url { - url: url.to_string(), - media_type: "video/mp4".to_string(), - }) - }) - .collect() - }) - .unwrap_or_default(); + .and_then(Value::as_array) + .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); + } + video_metadata.push(Value::Object(metadata)); + } + } if videos.is_empty() { return Err(AiMuxError::InvalidResponseData( - "Vertex video operation completed without any video output".to_string(), + "No valid videos in response".to_string(), )); } Ok(VideoOperationStatus::Completed(VideoResult { videos, warnings: Vec::new(), - provider_metadata: None, + 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(), + ), response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), diff --git a/aimux-providers/src/vertex_ai_ai21_models.rs b/aimux-providers/src/vertex_ai_ai21_models.rs deleted file mode 100644 index a6de8ca78..000000000 --- a/aimux-providers/src/vertex_ai_ai21_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! AI21 models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"ai21/jamba-1.5-large"`, `"ai21/jamba-1.5-mini"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_ai21_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiAi21ModelsConfig::with_base_url`]; prefer -/// [`VertexAiAi21ModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the AI21 Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiAi21ModelsConfig(OpenAIConfig); - -impl VertexAiAi21ModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// AI21 Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiAi21ModelsProvider(OpenAIProvider); - -impl VertexAiAi21ModelsProvider { - #[must_use] - pub fn new(config: VertexAiAi21ModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"ai21/jamba-1.5-large"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiAi21ModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vertex_ai_anthropic_models.rs b/aimux-providers/src/vertex_ai_anthropic_models.rs deleted file mode 100644 index a1717b8a2..000000000 --- a/aimux-providers/src/vertex_ai_anthropic_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! Anthropic Claude models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"anthropic/claude-sonnet-4"`, `"anthropic/claude-opus-4"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_anthropic_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiAnthropicModelsConfig::with_base_url`]; prefer -/// [`VertexAiAnthropicModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the Anthropic Claude Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiAnthropicModelsConfig(OpenAIConfig); - -impl VertexAiAnthropicModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// Anthropic Claude Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiAnthropicModelsProvider(OpenAIProvider); - -impl VertexAiAnthropicModelsProvider { - #[must_use] - pub fn new(config: VertexAiAnthropicModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"anthropic/claude-sonnet-4"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiAnthropicModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vertex_ai_deepseek_models.rs b/aimux-providers/src/vertex_ai_deepseek_models.rs deleted file mode 100644 index 1d2c5b015..000000000 --- a/aimux-providers/src/vertex_ai_deepseek_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! DeepSeek models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"deepseek-ai/deepseek-v3.1-maas"`, `"deepseek-ai/deepseek-r1-0528-maas"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_deepseek_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiDeepseekModelsConfig::with_base_url`]; prefer -/// [`VertexAiDeepseekModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the DeepSeek Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiDeepseekModelsConfig(OpenAIConfig); - -impl VertexAiDeepseekModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// DeepSeek Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiDeepseekModelsProvider(OpenAIProvider); - -impl VertexAiDeepseekModelsProvider { - #[must_use] - pub fn new(config: VertexAiDeepseekModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"deepseek-ai/deepseek-v3.1-maas"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiDeepseekModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vertex_ai_llama_models.rs b/aimux-providers/src/vertex_ai_llama_models.rs deleted file mode 100644 index 76edcb6b6..000000000 --- a/aimux-providers/src/vertex_ai_llama_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! Meta Llama models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"meta/llama-4-scout-17b-16e-instruct-maas"`, `"meta/llama-4-maverick-17b-128e-instruct-maas"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_llama_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiLlamaModelsConfig::with_base_url`]; prefer -/// [`VertexAiLlamaModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the Meta Llama Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiLlamaModelsConfig(OpenAIConfig); - -impl VertexAiLlamaModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// Meta Llama Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiLlamaModelsProvider(OpenAIProvider); - -impl VertexAiLlamaModelsProvider { - #[must_use] - pub fn new(config: VertexAiLlamaModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"meta/llama-4-scout-17b-16e-instruct-maas"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiLlamaModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vertex_ai_minimax_models.rs b/aimux-providers/src/vertex_ai_minimax_models.rs deleted file mode 100644 index 4e2e17ffb..000000000 --- a/aimux-providers/src/vertex_ai_minimax_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! MiniMax models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"minimax/minimax-m2-maas"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_minimax_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiMinimaxModelsConfig::with_base_url`]; prefer -/// [`VertexAiMinimaxModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the MiniMax Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiMinimaxModelsConfig(OpenAIConfig); - -impl VertexAiMinimaxModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// MiniMax Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiMinimaxModelsProvider(OpenAIProvider); - -impl VertexAiMinimaxModelsProvider { - #[must_use] - pub fn new(config: VertexAiMinimaxModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"minimax/minimax-m2-maas"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiMinimaxModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vertex_ai_mistral_models.rs b/aimux-providers/src/vertex_ai_mistral_models.rs deleted file mode 100644 index 0ba7d47b2..000000000 --- a/aimux-providers/src/vertex_ai_mistral_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! Mistral models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"mistralai/mistral-large-2411"`, `"mistralai/codestral-2501"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_mistral_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiMistralModelsConfig::with_base_url`]; prefer -/// [`VertexAiMistralModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the Mistral Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiMistralModelsConfig(OpenAIConfig); - -impl VertexAiMistralModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// Mistral Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiMistralModelsProvider(OpenAIProvider); - -impl VertexAiMistralModelsProvider { - #[must_use] - pub fn new(config: VertexAiMistralModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"mistralai/mistral-large-2411"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiMistralModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vertex_ai_moonshot_models.rs b/aimux-providers/src/vertex_ai_moonshot_models.rs deleted file mode 100644 index 63a66529f..000000000 --- a/aimux-providers/src/vertex_ai_moonshot_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! Moonshot AI models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"moonshotai/kimi-k2-thinking-maas"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_moonshot_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiMoonshotModelsConfig::with_base_url`]; prefer -/// [`VertexAiMoonshotModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the Moonshot AI Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiMoonshotModelsConfig(OpenAIConfig); - -impl VertexAiMoonshotModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// Moonshot AI Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiMoonshotModelsProvider(OpenAIProvider); - -impl VertexAiMoonshotModelsProvider { - #[must_use] - pub fn new(config: VertexAiMoonshotModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"moonshotai/kimi-k2-thinking-maas"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiMoonshotModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vertex_ai_openai_models.rs b/aimux-providers/src/vertex_ai_openai_models.rs deleted file mode 100644 index 94ac6e330..000000000 --- a/aimux-providers/src/vertex_ai_openai_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! OpenAI models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"openai/gpt-oss-120b-maas"`, `"openai/gpt-oss-20b-maas"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_openai_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiOpenaiModelsConfig::with_base_url`]; prefer -/// [`VertexAiOpenaiModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the OpenAI Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiOpenaiModelsConfig(OpenAIConfig); - -impl VertexAiOpenaiModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// OpenAI Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiOpenaiModelsProvider(OpenAIProvider); - -impl VertexAiOpenaiModelsProvider { - #[must_use] - pub fn new(config: VertexAiOpenaiModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"openai/gpt-oss-120b-maas"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiOpenaiModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vertex_ai_qwen_models.rs b/aimux-providers/src/vertex_ai_qwen_models.rs deleted file mode 100644 index a8ec610dd..000000000 --- a/aimux-providers/src/vertex_ai_qwen_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! Qwen models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"qwen/qwen3-coder-480b-a35b-instruct-maas"`, `"qwen/qwen3-next-80b-a3b-instruct-maas"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_qwen_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiQwenModelsConfig::with_base_url`]; prefer -/// [`VertexAiQwenModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the Qwen Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiQwenModelsConfig(OpenAIConfig); - -impl VertexAiQwenModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// Qwen Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiQwenModelsProvider(OpenAIProvider); - -impl VertexAiQwenModelsProvider { - #[must_use] - pub fn new(config: VertexAiQwenModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"qwen/qwen3-coder-480b-a35b-instruct-maas"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiQwenModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vertex_ai_zai_models.rs b/aimux-providers/src/vertex_ai_zai_models.rs deleted file mode 100644 index 8470bb990..000000000 --- a/aimux-providers/src/vertex_ai_zai_models.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! Z.AI models on Vertex AI MaaS — a thin OpenAI-compatible wrapper. -//! -//! Vertex AI serves partner and open models (Anthropic, AI21, DeepSeek, Llama, -//! MiniMax, Mistral, Moonshot, OpenAI, Qwen, Z.AI) through an OpenAI-compatible -//! Chat Completions endpoint — the "Model as a Service" (MaaS) OpenAPI surface -//! — rather than the native `rawPredict` path: -//! -//! `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` -//! -//! The host is derived from the location: `global` uses -//! `aiplatform.googleapis.com`, `eu`/`us` use `aiplatform.{loc}.rep.googleapis.com`, -//! and any other location uses `{loc}-aiplatform.googleapis.com`. Authentication -//! uses a Google Cloud OAuth2 Bearer token (the same `GOOGLE_VERTEX_ACCESS_TOKEN` -//! used by the native Vertex provider), sent as `Authorization: Bearer `. -//! -//! Because the endpoint is OpenAI-compatible, this provider is a thin wrapper -//! over [`OpenAIProvider`]: only the base URL, -//! the Bearer-token env var, and the provider name differ. The shared -//! `OpenAIProvider` appends `/chat/completions` to the configured base URL. -//! Sample model ids: `"zai-org/glm-4.7-maas"`, `"zai-org/glm-5-maas"`. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const PROVIDER_NAME: &str = "vertex_ai_zai_models"; -const TOKEN_ENV_VAR: &str = "GOOGLE_VERTEX_ACCESS_TOKEN"; -const PROJECT_ENV_VAR: &str = "GOOGLE_VERTEX_PROJECT"; -const LOCATION_ENV_VAR: &str = "GOOGLE_VERTEX_LOCATION"; -const DEFAULT_LOCATION: &str = "global"; -/// Fallback project when `GOOGLE_VERTEX_PROJECT` is unset and no base URL is -/// supplied via [`VertexAiZaiModelsConfig::with_base_url`]; prefer -/// [`VertexAiZaiModelsConfig::from_env`] or set the project explicitly for real usage. -const DEFAULT_PROJECT: &str = "your-project"; - -/// Build the Vertex AI MaaS OpenAI-compatible base URL for a project/location. -/// -/// - `global` → `https://aiplatform.googleapis.com/v1/projects/{p}/locations/global/endpoints/openapi` -/// - `eu`/`us` → `https://aiplatform.{loc}.rep.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -/// - other → `https://{loc}-aiplatform.googleapis.com/v1/projects/{p}/locations/{loc}/endpoints/openapi` -fn build_maas_base_url(project: &str, location: &str) -> String { - let host = match location { - "global" => "aiplatform.googleapis.com".to_string(), - "eu" | "us" => format!("aiplatform.{location}.rep.googleapis.com"), - _ => format!("{location}-aiplatform.googleapis.com"), - }; - format!("https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi") -} - -/// Assemble the shared [`OpenAIConfig`] for the given token + project/location. -fn build_config(api_key: String, project: &str, location: &str) -> OpenAIConfig { - OpenAIConfig::new(api_key) - .with_base_url(build_maas_base_url(project, location)) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) -} - -/// Configuration for the Z.AI Vertex AI MaaS provider (wraps [`OpenAIConfig`]). -pub struct VertexAiZaiModelsConfig(OpenAIConfig); - -impl VertexAiZaiModelsConfig { - /// Create from a Google Cloud Bearer access token, constructing the base - /// URL from `GOOGLE_VERTEX_PROJECT` / `GOOGLE_VERTEX_LOCATION` (with - /// `global` / `your-project` fallbacks). Override the URL with - /// [`Self::with_base_url`] for tests or proxies. - pub fn new(api_key: impl Into) -> Self { - let project = - std::env::var(PROJECT_ENV_VAR).unwrap_or_else(|_| DEFAULT_PROJECT.to_string()); - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Self(build_config(api_key.into(), &project, &location)) - } - - /// Create from `GOOGLE_VERTEX_ACCESS_TOKEN` + `GOOGLE_VERTEX_PROJECT` + - /// `GOOGLE_VERTEX_LOCATION` (location defaults to `global`). - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `GOOGLE_VERTEX_ACCESS_TOKEN` or - /// `GOOGLE_VERTEX_PROJECT` is not set. - pub fn from_env() -> Result { - let token = std::env::var(TOKEN_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_ACCESS_TOKEN environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let project = std::env::var(PROJECT_ENV_VAR).map_err(|_| { - AiMuxError::InvalidArgument( - "GOOGLE_VERTEX_PROJECT environment variable is required for Vertex AI MaaS" - .to_string(), - ) - })?; - let location = - std::env::var(LOCATION_ENV_VAR).unwrap_or_else(|_| DEFAULT_LOCATION.to_string()); - Ok(Self(build_config(token, &project, &location))) - } - - /// Override the base URL (useful for tests / proxies). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -/// Z.AI Vertex AI MaaS provider — creates [`OpenAIModel`] instances pointed -/// at the Vertex AI MaaS OpenAPI endpoint. -pub struct VertexAiZaiModelsProvider(OpenAIProvider); - -impl VertexAiZaiModelsProvider { - #[must_use] - pub fn new(config: VertexAiZaiModelsConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - /// Create a model instance for the given Vertex AI MaaS model id - /// (e.g. `"zai-org/glm-4.7-maas"`). - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VertexAiZaiModelsProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/vllm.rs b/aimux-providers/src/vllm.rs deleted file mode 100644 index 4cf018fa6..000000000 --- a/aimux-providers/src/vllm.rs +++ /dev/null @@ -1,86 +0,0 @@ -//! vLLM provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:8000/v1`. The `VLLM_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:8000/v1"; -const ENV_VAR: &str = "VLLM_BASE_URL"; -const PROVIDER_NAME: &str = "vllm"; -const PLACEHOLDER_API_KEY: &str = "vllm"; - -pub struct VllmConfig(OpenAIConfig); - -impl VllmConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()), - ) - } - - /// Create from the `VLLM_BASE_URL` environment variable, falling back to - /// `http://127.0.0.1:8000/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let mut config = Self::new(PLACEHOLDER_API_KEY); - config.0 = config.0.with_api_key_source(Some("none")); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct VllmProvider(OpenAIProvider); - -impl VllmProvider { - #[must_use] - pub fn new(config: VllmConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for VllmProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - fn decision_model( - &self, - model_id: &str, - ) -> Result, AiMuxError> { - self.0.decision_model(model_id) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/voyage/embedding.rs b/aimux-providers/src/voyage/embedding.rs index 1e3666939..b52d2aa4f 100644 --- a/aimux-providers/src/voyage/embedding.rs +++ b/aimux-providers/src/voyage/embedding.rs @@ -19,45 +19,26 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::HttpRequest; +use crate::shared::EndpointConfig; -use super::{VoyageConfig, voyage_failed_response_handler}; +use super::voyage_failed_response_handler; /// A Voyage embedding model (e.g. `"voyage-3.5"`). pub struct VoyageEmbeddingModel { model_id: String, - config: VoyageConfig, + config: EndpointConfig, } impl VoyageEmbeddingModel { - #[must_use] - pub fn new(model_id: String, config: VoyageConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/embeddings", self.config.base_url) - } } #[async_trait] impl EmbeddingModel for VoyageEmbeddingModel { fn provider(&self) -> &str { - "voyage" + &self.config.provider } fn model_id(&self) -> &str { @@ -94,19 +75,13 @@ impl EmbeddingModel for VoyageEmbeddingModel { body.insert("output_dtype".to_string(), json!(output_dtype)); } - let headers = self.build_headers(options.headers.as_ref()); - let header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; // `send()` returns Ok only for 2xx; non-2xx responses are mapped to an - // error internally using the shared error structure. `HttpBody::Json` - // sets `Content-Type: application/json`, so it is intentionally not - // added to the header list above. + // error internally using the shared error structure. let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), header_list, options), - Value::Object(body), + exchange.request(exchange.url("/embeddings"), options), + exchange.transform_body(Value::Object(body)), aimux_provider_utils::create_json_response_handler(), voyage_failed_response_handler(), ) @@ -177,7 +152,7 @@ struct VoyageEmbeddingProviderOptions { fn parse_voyage_provider_options( options: Option<&SharedProviderOptions>, ) -> VoyageEmbeddingProviderOptions { - let provider_opts = options.and_then(|opts| opts.get("voyage")); + let provider_opts = super::options::voyage_options(options); VoyageEmbeddingProviderOptions { input_type: provider_opts .and_then(|o| o.get("inputType")) diff --git a/aimux-providers/src/voyage/mod.rs b/aimux-providers/src/voyage/mod.rs index 7ac94c982..eed86624b 100644 --- a/aimux-providers/src/voyage/mod.rs +++ b/aimux-providers/src/voyage/mod.rs @@ -1,17 +1,34 @@ //! Voyage AI provider. //! +//! [`create_voyage`] takes [`VoyageProviderSettings`], validates the base URL, +//! fixes the provider name and returns a [`VoyageProvider`]. The API key is not +//! read there; it is loaded in the request headers of every call, from the +//! setting or from `VOYAGE_API_KEY`. [`voyage()`] is the default instance. The +//! AI SDK has no Voyage package; the factory follows the shape of its +//! embedding and reranking vendors. +//! //! Implements the `EmbeddingModel` trait against the Voyage AI API //! (`api.voyageai.com/v1/embeddings`). pub mod embedding; +pub(crate) mod options; pub mod reranking; +pub use crate::shared::TransformRequestBody; pub use embedding::VoyageEmbeddingModel; pub use reranking::VoyageRerankingModel; +use std::sync::{Arc, OnceLock}; + +use aimux_core::embedding_model::EmbeddingModel; use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; -use aimux_provider_utils::{load_api_key, without_trailing_slash}; +use aimux_core::reranking_model::RerankingModel; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; pub(crate) fn voyage_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -27,68 +44,150 @@ pub(crate) fn voyage_failed_response_handler() -> aimux_provider_utils::Response }) } -/// Configuration for the Voyage AI provider. -#[derive(Debug, Clone)] -pub struct VoyageConfig { - pub api_key: String, - pub base_url: String, +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.voyageai.com/v1"; +pub(crate) const API_KEY_ENV_VAR: &str = "VOYAGE_API_KEY"; +const DEFAULT_NAME: &str = "voyage"; + +/// Settings of [`create_voyage`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct VoyageProviderSettings { + /// Base URL for the API calls. Default `https://api.voyageai.com/v1`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `VOYAGE_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including `Authorization`. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings + /// (`"{name}.embedding"`, `"{name}.reranking"`). Default `"voyage"`. The + /// providerOptions key stays `voyage`. + 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, + /// Rewrites every JSON request body once, after it is serialized and + /// before it is sent. + pub transform_request_body: Option, } -impl VoyageConfig { - /// Create from an API key (uses default Voyage AI base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://api.voyageai.com/v1".to_string(), - } +impl std::fmt::Debug for VoyageProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VoyageProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) + .field( + "transform_request_body", + &self.transform_request_body.is_some(), + ) + .finish() } +} - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } +/// Create a Voyage AI provider. +/// +/// # 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. +pub fn create_voyage(settings: VoyageProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(VoyageProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: provider_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "Voyage"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + transform_request_body: settings.transform_request_body, + }) +} - /// Create from the `VOYAGE_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `VOYAGE_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "VOYAGE_API_KEY", "Voyage")?; - Ok(Self::new(api_key)) - } +/// The default provider: `create_voyage` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn voyage() -> &'static VoyageProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_voyage(VoyageProviderSettings::default()) + .expect("default Voyage settings are always valid") + }) } -/// Voyage AI provider — creates `VoyageEmbeddingModel` instances. +/// A Voyage AI provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct VoyageProvider { - config: VoyageConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, + transform_request_body: Option, } impl VoyageProvider { - #[must_use] - pub fn new(config: VoyageConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + self.transform_request_body.clone(), + ) } - /// Create an embedding model instance for the given model name (e.g. - /// `"voyage-3.5"`). + /// An embedding model (e.g. `"voyage-3.5"`); `provider()` is + /// `"{name}.embedding"`. #[must_use] - pub fn embedding_model(&self, model_id: &str) -> VoyageEmbeddingModel { - VoyageEmbeddingModel::new(model_id.to_string(), self.config.clone()) + pub fn embedding(&self, model_id: &str) -> VoyageEmbeddingModel { + VoyageEmbeddingModel::from_config(model_id.to_string(), self.model_config("embedding")) } - /// Create a reranking model instance for the given model name (e.g. - /// `"rerank-2.5"`). + /// A reranking model (e.g. `"rerank-2.5"`); `provider()` is + /// `"{name}.reranking"`. #[must_use] - pub fn reranking_model(&self, model_id: &str) -> reranking::VoyageRerankingModel { - reranking::VoyageRerankingModel::new(model_id.to_string(), self.config.clone()) + pub fn reranking(&self, model_id: &str) -> VoyageRerankingModel { + VoyageRerankingModel::from_config(model_id.to_string(), self.model_config("reranking")) } } impl Provider for VoyageProvider { - fn name(&self) -> &str { - "voyage" + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "languageModel")) + } + + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } + + fn reranking_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.reranking(model_id)))) } } diff --git a/aimux-providers/src/voyage/options.rs b/aimux-providers/src/voyage/options.rs new file mode 100644 index 000000000..7276472f6 --- /dev/null +++ b/aimux-providers/src/voyage/options.rs @@ -0,0 +1,16 @@ +//! The providerOptions namespace of Voyage AI. +//! +//! Voyage reads `providerOptions.voyage`, whatever the provider is named. +//! This module is the one place that spells the key. + +use aimux_core::shared::{JsonObject, SharedProviderOptions}; + +/// The providerOptions key. +pub(crate) const NAMESPACE: &str = "voyage"; + +/// The Voyage options in a providerOptions map. +pub(crate) fn voyage_options( + provider_options: Option<&SharedProviderOptions>, +) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} diff --git a/aimux-providers/src/voyage/reranking.rs b/aimux-providers/src/voyage/reranking.rs index 8fffe0b86..834dfa59c 100644 --- a/aimux-providers/src/voyage/reranking.rs +++ b/aimux-providers/src/voyage/reranking.rs @@ -3,8 +3,6 @@ //! Aligned with Vercel AI SDK `VoyageRerankingModel` //! (`reference/ai/packages/voyage/src/reranking/voyage-reranking-model.ts`). -use std::collections::HashMap; - use async_trait::async_trait; use serde::Deserialize; use serde_json::{Value, json}; @@ -17,9 +15,9 @@ use aimux_core::reranking_model::{ use aimux_core::shared::SharedProviderOptions; use aimux_core::types::Warning; -use aimux_provider_utils::HttpRequest; +use crate::shared::EndpointConfig; -use super::{VoyageConfig, voyage_failed_response_handler}; +use super::voyage_failed_response_handler; /// Voyage provider-specific reranking options. #[derive(Debug, Clone, Default)] @@ -32,9 +30,7 @@ fn parse_voyage_reranking_options( provider_options: Option<&SharedProviderOptions>, ) -> VoyageRerankingOptions { let mut opts = VoyageRerankingOptions::default(); - if let Some(po) = provider_options - && let Some(voyage) = po.get("voyage") - { + if let Some(voyage) = super::options::voyage_options(provider_options) { if let Some(v) = voyage .get("returnDocuments") .and_then(serde_json::Value::as_bool) @@ -66,38 +62,19 @@ struct VoyageRerankingResult { /// A Voyage reranking model. pub struct VoyageRerankingModel { model_id: String, - config: VoyageConfig, + config: EndpointConfig, } impl VoyageRerankingModel { - #[must_use] - pub fn new(model_id: String, config: VoyageConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/rerank", self.config.base_url) - } } #[async_trait] impl RerankingModel for VoyageRerankingModel { fn provider(&self) -> &str { - "voyage" + &self.config.provider } fn model_id(&self) -> &str { @@ -141,15 +118,11 @@ impl RerankingModel for VoyageRerankingModel { body["truncation"] = json!(truncation); } - let headers = self.build_headers(options.headers.as_ref()); - let mut header_list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - header_list.push(("Content-Type".to_string(), "application/json".to_string())); + 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( - HttpRequest::new(self.endpoint(), header_list, options), + exchange.request(exchange.url("/rerank"), options), body.clone(), aimux_provider_utils::create_json_response_handler::(), voyage_failed_response_handler(), diff --git a/aimux-providers/src/xai/convert.rs b/aimux-providers/src/xai/convert.rs index 7f4c53977..c28f643d6 100644 --- a/aimux-providers/src/xai/convert.rs +++ b/aimux-providers/src/xai/convert.rs @@ -1,47 +1,13 @@ -//! Conversion between `LanguageModelPrompt` and xAI API format. +//! Helpers shared by the xAI Responses converter. //! -//! Mirrors the TS xai package's `convert-to-xai-chat-messages.ts`, -//! `convert-xai-chat-usage.ts`, `xai-prepare-tools.ts`, -//! `supports-reasoning-effort.ts`, `map-xai-finish-reason.ts`, and -//! `remove-additional-properties.ts`. +//! 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. -use aimux_core::error::AiMuxError; -use aimux_core::language_model_message::{ - AssistantPart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, ToolPart, - ToolResultPart, UserPart, -}; -use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat}; -use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; -use aimux_core::tool::Tool; -use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Warning}; +use aimux_core::shared::SharedProviderReference; +use serde_json::Value; -use serde_json::{Value, json}; - -use super::types::XaiUsageResponse; - -// ── Finish reason ──────────────────────────────────────────────────────────── - -/// Map an xAI finish reason string to the unified enum. -#[must_use] -pub fn map_xai_finish_reason(s: &str) -> FinishReasonUnified { - match s { - "stop" => FinishReasonUnified::Stop, - "length" => FinishReasonUnified::Length, - "tool_calls" | "function_call" => FinishReasonUnified::ToolCalls, - "content_filter" => FinishReasonUnified::ContentFilter, - _ => FinishReasonUnified::Other, - } -} - -#[must_use] -pub fn parse_finish_reason(s: &str) -> FinishReason { - FinishReason { - unified: map_xai_finish_reason(s), - raw: Some(s.to_string()), - } -} - -// ── Reasoning effort support ───────────────────────────────────────────────── +// ── Reasoning effort ───────────────────────────────────────────────────────── /// Models that reject the `reasoning_effort` parameter. /// Matches `^grok-4\.20(-\d{4})?-(non-)?reasoning$`. @@ -67,126 +33,7 @@ pub fn supports_reasoning_effort(model_id: &str) -> bool { !is_model_without_reasoning_effort(model_id) } -// ── Usage conversion ───────────────────────────────────────────────────────── - -#[must_use] -pub fn convert_xai_usage(usage: &XaiUsageResponse) -> aimux_core::types::Usage { - let prompt_tokens = usage.prompt_tokens.unwrap_or(0); - let completion_tokens = usage.completion_tokens.unwrap_or(0); - let cache_read = usage - .prompt_tokens_details - .as_ref() - .and_then(|d| d.cached_tokens) - .unwrap_or(0); - let reasoning_tokens = usage - .completion_tokens_details - .as_ref() - .and_then(|d| d.reasoning_tokens) - .unwrap_or(0); - let prompt_includes_cached = cache_read <= prompt_tokens; - let input_total = if prompt_includes_cached { - prompt_tokens - } else { - prompt_tokens + cache_read - }; - let input_no_cache = if prompt_includes_cached { - prompt_tokens - cache_read - } else { - prompt_tokens - }; - let output_total = completion_tokens + reasoning_tokens; - - aimux_core::types::Usage { - input_tokens: aimux_core::types::InputTokenUsage { - total: Some(input_total), - no_cache: Some(input_no_cache), - cache_read: Some(cache_read), - cache_write: None, - }, - output_tokens: aimux_core::types::OutputTokenUsage { - total: Some(output_total), - text: Some(completion_tokens), - reasoning: Some(reasoning_tokens), - }, - // RFC-0015 P0-3: keep the raw provider usage payload. - raw: serde_json::to_value(usage) - .ok() - .and_then(|value| value.as_object().cloned()), - } -} - -// ── Tool preparation ───────────────────────────────────────────────────────── - -#[derive(Debug, Clone)] -pub struct PreparedTools { - pub tools: Option>, - pub tool_choice: Option, - pub tool_warnings: Vec, -} - -#[must_use] -pub fn prepare_tools( - tools: &Option>, - tool_choice: Option<&LanguageModelV4ToolChoice>, -) -> PreparedTools { - 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 xai_tools: Vec = tools - .iter() - .filter_map(|t| match t { - Tool::Function(ft) => { - let mut func = json!({ - "name": ft.name, - "parameters": remove_additional_properties_false(&ft.input_schema), - }); - if let Some(ref desc) = ft.description { - func["description"] = json!(desc); - } - if let Some(strict) = ft.strict { - func["strict"] = json!(strict); - } - Some(json!({ "type": "function", "function": func })) - } - Tool::Provider(pt) => { - tool_warnings.push(Warning::Unsupported { - feature: format!("provider-defined tool {}", pt.name), - details: None, - }); - None - } - }) - .collect(); - if xai_tools.is_empty() { - None - } else { - Some(xai_tools) - } - } - }; - - let tool_choice_opt = match (&tools_opt, tool_choice) { - (None, _) => None, - (Some(_), None) => None, - (Some(_), Some(tc)) => match tc { - LanguageModelV4ToolChoice::Auto => Some(json!("auto")), - LanguageModelV4ToolChoice::None => Some(json!("none")), - LanguageModelV4ToolChoice::Required => Some(json!("required")), - LanguageModelV4ToolChoice::Tool { tool_name } => { - Some(json!({ "type": "function", "function": { "name": tool_name } })) - } - }, - }; - - PreparedTools { - tools: tools_opt, - tool_choice: tool_choice_opt, - tool_warnings, - } -} +// ── JSON schema ────────────────────────────────────────────────────────────── pub fn remove_additional_properties_false(value: &Value) -> Value { match value { @@ -207,54 +54,6 @@ pub fn remove_additional_properties_false(value: &Value) -> Value { } } -// ── Provider options helpers ───────────────────────────────────────────────── - -fn xai_option(options: &Option, key: &str) -> Option { - options - .as_ref() - .and_then(|m| m.get("xai")) - .and_then(|o| o.get(key)) - .cloned() -} - -fn get_image_detail(provider_options: &Option) -> Option { - provider_options - .as_ref() - .and_then(|po| po.get("xai")) - .and_then(|o| o.get("imageDetail")) - .cloned() -} - -// ── Media type helpers ─────────────────────────────────────────────────────── - -fn get_top_level_media_type(media_type: &str) -> &str { - media_type.split('/').next().unwrap_or("") -} - -#[must_use] -pub fn resolve_full_media_type(media_type: &str, b64_data: &str) -> String { - let top_level = get_top_level_media_type(media_type); - if top_level == "image" && media_type != "image" && !media_type.ends_with("/*") { - return media_type.to_string(); - } - if top_level == "image" { - if b64_data.starts_with("iVBORw0KGgo") { - return "image/png".to_string(); - } - if b64_data.starts_with("/9j/") { - return "image/jpeg".to_string(); - } - if b64_data.starts_with("R0lGOD") { - return "image/gif".to_string(); - } - if b64_data.starts_with("UklGR") { - return "image/webp".to_string(); - } - return "image/png".to_string(); - } - media_type.to_string() -} - /// Resolve the provider-specific reference string from a file part's /// `provider` map. /// @@ -263,392 +62,17 @@ pub fn resolve_full_media_type(media_type: &str, b64_data: &str) -> String { /// Returns a `String` listing the available providers when the requested key /// is absent. pub fn resolve_provider_reference( - reference: &aimux_core::shared::SharedProviderReference, + reference: &SharedProviderReference, provider: &str, ) -> Result { if let Some(value) = reference.get(provider) { return Ok(value.clone()); } - let available: Vec = reference.keys().cloned().collect(); + let mut available: Vec<&str> = reference.keys().map(String::as_str).collect(); + available.sort_unstable(); Err(format!( "No provider reference found for provider '{}'. Available providers: {}", provider, available.join(", ") )) } - -// ── Message conversion ─────────────────────────────────────────────────────── - -/// Convert a [`LanguageModelPrompt`] into xAI `messages`, returning warnings. -/// -/// # Errors -/// -/// Returns `AiMuxError::InvalidArgument` when a user part cannot be converted -/// (e.g. invalid base64 data or an unsupported file part). -pub fn convert_to_xai_messages( - prompt: &LanguageModelPrompt, -) -> Result<(Vec, Vec), AiMuxError> { - let mut messages = Vec::new(); - let warnings: Vec = Vec::new(); - - for msg in prompt { - match msg { - LanguageModelMessage::System { content, .. } => { - messages.push(json!({ "role": "system", "content": content })); - } - LanguageModelMessage::User { content, .. } => { - if content.len() == 1 - && let UserPart::Text(TextPart { text, .. }) = &content[0] - { - messages.push(json!({ "role": "user", "content": text })); - continue; - } - let mut user_content = Vec::new(); - for (i, part) in content.iter().enumerate() { - user_content.push(convert_user_part(part, i)?); - } - messages.push(json!({ "role": "user", "content": user_content })); - } - LanguageModelMessage::Assistant { content, .. } => { - let mut text = String::new(); - let mut tool_calls: Vec = Vec::new(); - for part in content { - match part { - AssistantPart::Text(TextPart { text: t, .. }) => text.push_str(t), - AssistantPart::ToolCall(ToolCallPart { - tool_call_id, - tool_name, - input, - .. - }) => { - let arguments = input.to_string(); - tool_calls.push(json!({ "id": tool_call_id, "type": "function", "function": { "name": tool_name, "arguments": arguments } })); - } - _ => {} - } - } - let mut msg_obj = json!({ "role": "assistant", "content": text }); - if !tool_calls.is_empty() { - msg_obj["tool_calls"] = json!(tool_calls); - } - messages.push(msg_obj); - } - LanguageModelMessage::Tool { content, .. } => { - for part in content { - let ToolPart::ToolResult(ToolResultPart { - tool_call_id, - output, - .. - }) = part - else { - continue; - }; - let content = crate::openai::convert::tool_result_to_content(output); - messages.push( - json!({ "role": "tool", "tool_call_id": tool_call_id, "content": content }), - ); - } - } - } - } - Ok((messages, warnings)) -} - -fn convert_user_part(part: &UserPart, _index: usize) -> Result { - match part { - UserPart::Text(TextPart { text, .. }) => Ok(json!({ "type": "text", "text": text })), - UserPart::File(file) => match &file.data { - FileData::Data { - data: FileBytes::Binary(bytes), - } => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(bytes); - convert_image_part(&file.media_type, Some(&b64), None, &file.provider_options) - } - FileData::Data { - data: FileBytes::Base64(data), - } => convert_image_part(&file.media_type, Some(data), None, &file.provider_options), - FileData::Url { url, .. } => { - convert_image_part(&file.media_type, None, Some(url), &file.provider_options) - } - FileData::Reference { reference } => { - let file_id = resolve_provider_reference(reference, "xai") - .map_err(AiMuxError::InvalidArgument)?; - Ok(json!({ "type": "file", "file": { "file_id": file_id } })) - } - FileData::Text { .. } => Err(AiMuxError::UnsupportedFunctionality( - "text file parts".into(), - )), - }, - } -} - -fn convert_image_part( - media_type: &str, - b64_data: Option<&str>, - url: Option<&str>, - provider_options: &Option, -) -> Result { - let top_level = get_top_level_media_type(media_type); - if top_level != "image" { - return Err(AiMuxError::UnsupportedFunctionality(format!( - "file part media type {media_type}" - ))); - } - let image_url = if let Some(url_str) = url { - json!({ "url": url_str }) - } else if let Some(b64) = b64_data { - let full_mt = resolve_full_media_type(media_type, b64); - json!({ "url": format!("data:{};base64,{}", full_mt, b64) }) - } else { - json!({ "url": "" }) - }; - let mut image_url_obj = image_url; - if let Some(detail) = get_image_detail(provider_options) { - image_url_obj["detail"] = detail; - } - Ok(json!({ "type": "image_url", "image_url": image_url_obj })) -} - -// ── Request body ───────────────────────────────────────────────────────────── - -#[derive(Debug, Clone)] -pub struct RequestBodyResult { - pub body: Value, - pub warnings: Vec, -} - -/// Convert `CallOptions` to an xAI request body, returning warnings. -/// -/// # Errors -/// -/// Propagates prompt-conversion errors from `convert_to_xai_messages`. -pub fn build_request_body_with_warnings( - model_id: &str, - options: &CallOptions, - stream: bool, -) -> Result { - let mut warnings: Vec = Vec::new(); - let xai_opts = &options.provider_options; - - if options.top_k.is_some() { - warnings.push(Warning::Unsupported { - feature: "topK".to_string(), - details: None, - }); - } - if options.frequency_penalty.is_some() { - warnings.push(Warning::Unsupported { - feature: "frequencyPenalty".to_string(), - details: None, - }); - } - if options.presence_penalty.is_some() { - warnings.push(Warning::Unsupported { - feature: "presencePenalty".to_string(), - details: None, - }); - } - if options.stop_sequences.is_some() { - warnings.push(Warning::Unsupported { - feature: "stopSequences".to_string(), - details: None, - }); - } - - let (messages, message_warnings) = convert_to_xai_messages(&options.prompt)?; - warnings.extend(message_warnings); - - let prepared = prepare_tools(&options.tools, options.tool_choice.as_ref()); - for tw in &prepared.tool_warnings { - warnings.push(tw.clone()); - } - - let mut reasoning_effort: Option = xai_option(xai_opts, "reasoningEffort").map(|v| { - v.as_str() - .map(std::string::ToString::to_string) - .unwrap_or(v.to_string()) - }); - - if reasoning_effort.is_none() && options.reasoning.is_some_and(ReasoningEffort::is_custom) { - let reasoning = options.reasoning.unwrap(); - if !supports_reasoning_effort(model_id) { - warnings.push(Warning::Unsupported { - feature: "reasoning".to_string(), - details: Some(format!( - "reasoning \"{reasoning}\" is not supported by this model." - )), - }); - } else if reasoning == ReasoningEffort::None { - reasoning_effort = Some("none".to_string()); - } else { - reasoning_effort = Some(match reasoning { - ReasoningEffort::Minimal => "low".to_string(), - ReasoningEffort::Low => "low".to_string(), - ReasoningEffort::Medium => "medium".to_string(), - ReasoningEffort::High => "high".to_string(), - ReasoningEffort::Xhigh => "high".to_string(), - _ => reasoning.to_string(), - }); - } - } - - let mut body = json!({ "model": model_id, "messages": messages }); - if stream { - body["stream"] = json!(true); - body["stream_options"] = json!({ "include_usage": true }); - } - if let Some(max_tokens) = options.max_output_tokens { - body["max_completion_tokens"] = json!(max_tokens); - } - if let Some(temp) = options.temperature { - body["temperature"] = json!(temp); - } - if let Some(tp) = options.top_p { - body["top_p"] = json!(tp); - } - if let Some(seed) = options.seed { - body["seed"] = json!(seed); - } - - let logprobs = xai_option(xai_opts, "logprobs") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - let top_logprobs = xai_option(xai_opts, "topLogprobs"); - if logprobs || top_logprobs.is_some() { - body["logprobs"] = json!(true); - } - if let Some(tlp) = top_logprobs { - body["top_logprobs"] = tlp; - } - - if let Some(effort) = &reasoning_effort { - body["reasoning_effort"] = json!(effort); - } - if let Some(val) = xai_option(xai_opts, "parallel_function_calling") { - body["parallel_function_calling"] = val; - } - - if let Some(ref rf) = options.response_format { - match rf { - ResponseFormat::Text => {} - ResponseFormat::Json { - schema, - name, - description: _, - } => { - if schema.is_some() { - let mut schema_obj = json!({}); - if let Some(s) = schema { - schema_obj["schema"] = s.clone(); - } - schema_obj["name"] = - json!(name.clone().unwrap_or_else(|| "response".to_string())); - schema_obj["strict"] = json!(true); - body["response_format"] = - json!({ "type": "json_schema", "json_schema": schema_obj }); - } else { - body["response_format"] = json!({ "type": "json_object" }); - } - } - } - } - - if let Some(sp) = xai_option(xai_opts, "searchParameters") { - body["search_parameters"] = convert_search_parameters(&sp); - } - if let Some(tools) = &prepared.tools { - body["tools"] = json!(tools); - } - if let Some(tc) = prepared.tool_choice { - body["tool_choice"] = tc; - } - - Ok(RequestBodyResult { body, warnings }) -} - -fn convert_search_parameters(sp: &Value) -> Value { - let mut result = json!({}); - if let Some(v) = sp.get("mode") { - result["mode"] = v.clone(); - } - if let Some(v) = sp.get("returnCitations") { - result["return_citations"] = v.clone(); - } - if let Some(v) = sp.get("fromDate") { - result["from_date"] = v.clone(); - } - if let Some(v) = sp.get("toDate") { - result["to_date"] = v.clone(); - } - if let Some(v) = sp.get("maxSearchResults") { - result["max_search_results"] = v.clone(); - } - if let Some(sources) = sp.get("sources").and_then(|s| s.as_array()) { - result["sources"] = json!( - sources - .iter() - .map(convert_search_source) - .collect::>() - ); - } - result -} - -fn convert_search_source(source: &Value) -> Value { - let mut result = json!({}); - if let Some(t) = source.get("type") { - result["type"] = t.clone(); - } - let st = source.get("type").and_then(|v| v.as_str()).unwrap_or(""); - match st { - "web" => { - if let Some(v) = source.get("country") { - result["country"] = v.clone(); - } - if let Some(v) = source.get("excludedWebsites") { - result["excluded_websites"] = v.clone(); - } - if let Some(v) = source.get("allowedWebsites") { - result["allowed_websites"] = v.clone(); - } - if let Some(v) = source.get("safeSearch") { - result["safe_search"] = v.clone(); - } - } - "x" => { - if let Some(v) = source.get("excludedXHandles") { - result["excluded_x_handles"] = v.clone(); - } - if let Some(v) = source.get("includedXHandles") { - result["included_x_handles"] = v.clone(); - } else if let Some(v) = source.get("xHandles") { - result["included_x_handles"] = v.clone(); - } - if let Some(v) = source.get("postFavoriteCount") { - result["post_favorite_count"] = v.clone(); - } - if let Some(v) = source.get("postViewCount") { - result["post_view_count"] = v.clone(); - } - } - "news" => { - if let Some(v) = source.get("country") { - result["country"] = v.clone(); - } - if let Some(v) = source.get("excludedWebsites") { - result["excluded_websites"] = v.clone(); - } - if let Some(v) = source.get("safeSearch") { - result["safe_search"] = v.clone(); - } - } - "rss" => { - if let Some(v) = source.get("links") { - result["links"] = v.clone(); - } - } - _ => {} - } - result -} diff --git a/aimux-providers/src/xai/mod.rs b/aimux-providers/src/xai/mod.rs index 104ea0396..b595f9213 100644 --- a/aimux-providers/src/xai/mod.rs +++ b/aimux-providers/src/xai/mod.rs @@ -1,32 +1,42 @@ -//! xAI (Grok) provider — a thin wrapper with xAI-specific behaviour. +//! xAI (Grok) provider. //! -//! xAI exposes an OpenAI-compatible Chat Completions API at -//! `https://api.x.ai/v1`. While the wire format is OpenAI-compatible, xAI has -//! enough provider-specific behaviour (reasoning content extraction, citations, -//! search parameters, xai-keyed provider options, non-inclusive cached tokens, -//! reasoning-effort model gating, 200-status errors) to warrant its own model -//! implementation ([`XaiModel`]) rather than reusing `OpenAIModel`. +//! [`create_xai`] is the Rust form of the AI SDK's `createXai`: it takes +//! [`XAIProviderSettings`], validates the base URL, 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. +//! +//! The AI SDK's xAI package serves the Responses API only +//! ([`XaiResponsesModel`], `provider()` = `xai.responses`), and +//! [`language_model`](Provider::language_model) returns it. There is no Chat +//! Completions model in this package; xAI's OpenAI-compatible endpoint is +//! reachable through `create_openai_compatible` with `https://api.x.ai/v1`. pub mod convert; -mod model; +pub(crate) mod options; pub mod responses; -mod types; -pub use model::XaiModel; pub use responses::XaiResponsesModel; -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; -use aimux_provider_utils::load_api_key; +use std::sync::{Arc, OnceLock}; + +use futures::future::BoxFuture; use serde::de::DeserializeOwned; use serde_json::Value; -use crate::openai::OpenAIConfig; +use aimux_core::embedding_model::EmbeddingModel; +use aimux_core::error::AiMuxError; +use aimux_core::image_model::ImageModel; +use aimux_core::language_model::LanguageModel; +use aimux_core::model_catalogue::RuntimeModel; +use aimux_core::provider::{Provider, ProviderDiscovery}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; -const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1"; -const ENV_VAR: &str = "XAI_API_KEY"; -const PROVIDER_NAME: &str = "xai"; +use crate::shared::{Credential, EndpointConfig, provider_headers}; + +pub(crate) const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1"; +pub(crate) 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| { @@ -213,122 +223,150 @@ where .streaming() } -/// Configuration for the xAI provider (wraps [`OpenAIConfig`]). -#[derive(Clone)] -pub struct XAIConfig(OpenAIConfig); - -impl XAIConfig { - /// Create from an API key, using the default xAI base URL. - pub fn new(api_key: impl Into) -> Self { - Self(OpenAIConfig::new(api_key).with_base_url(DEFAULT_BASE_URL)) - } - - /// Create from the `XAI_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `XAI_API_KEY` is not set. - pub fn from_env() -> Result { - let key = load_api_key(None, ENV_VAR, "xAI")?; - Ok(Self::new(key).with_api_key_source(Some("env:XAI_API_KEY"))) - } - - /// 标注 api_key 来源(RFC-0023 回放重建用)。透传到内部 `OpenAIConfig`。 - #[must_use] - pub fn with_api_key_source(mut self, source: Option<&str>) -> Self { - self.0 = self.0.with_api_key_source(source); - self - } - - /// 内部 `OpenAIConfig` 引用(config_snapshot 复用 OpenAI helper 用,M2b)。 - pub(crate) fn openai_config(&self) -> &OpenAIConfig { - &self.0 - } +/// Settings of [`create_xai`] (the AI SDK's `XaiProviderSettings`). +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct XAIProviderSettings { + /// Base URL for the API calls. Default `https://api.x.ai/v1`; a trailing + /// slash is removed. + pub base_url: Option, + /// The API key. `None` loads `XAI_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. + pub api_key: Option, + /// Extra headers on every request. A `None` value removes the header, + /// including `Authorization`. Per-call headers win over these. + pub headers: Option, + /// The transport: a mock, a signing decorator, a proxy-aware client. + /// `None` uses the process default, resolved per request. + pub fetch: Option, +} - /// Override the base URL (useful for tests / self-hosted endpoints). - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self +impl std::fmt::Debug for XAIProviderSettings { + /// Never prints the key or header values. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("XAIProviderSettings") + .field("base_url", &self.base_url) + .field("api_key", &self.api_key.is_some()) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("fetch", &self.fetch.is_some()) + .finish() } +} - /// Get the API key. - pub(crate) fn api_key(&self) -> &str { - &self.0.api_key - } +/// Create an xAI provider. +/// +/// # 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. +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(), + }; + Ok(XAIProvider { + 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, + "xAI API key", + ), + Vec::new(), + settings.headers, + ), + options::NAMESPACE, + "5.0.12", + ), + fetch: settings.fetch, + }) +} - /// Get the base URL. - pub(crate) fn base_url(&self) -> &str { - &self.0.base_url - } +/// The default provider: `create_xai` with default settings, created on first +/// use. Creating it reads nothing from the environment and cannot fail; a +/// missing key surfaces from the first request instead. +pub fn xai() -> &'static XAIProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_xai(XAIProviderSettings::default()).expect("default xAI settings are always valid") + }) } -/// xAI provider — creates [`XaiModel`] instances pointed at xAI. -/// -/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the -/// process-wide shared `Client` internally (RFC-0009 §4.1). +/// An xAI provider (the AI SDK's `XaiProvider`). Cheap to clone the models out +/// of; it holds no HTTP client. pub struct XAIProvider { - config: XAIConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl XAIProvider { - #[must_use] - pub fn new(config: XAIConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a model instance for the given xAI model id (e.g. `"grok-2"`). + /// A Responses model (e.g. `"grok-4"`); `provider()` is + /// `"{name}.responses"`. /// - /// Clones the provider config so the model inherits the same - /// `api_key_source` / `max_retries` (M2b: previously reconstructed with - /// `XAIConfig::new`, which dropped the credential source). + /// Uses the xAI `/responses` endpoint with the Responses API wire format + /// (input items, reasoning objects, provider-executed tools, etc.). #[must_use] - pub fn model(&self, model_id: &str) -> XaiModel { - XaiModel::new(model_id.to_string(), self.config.clone()) + pub fn responses(&self, model_id: &str) -> XaiResponsesModel { + XaiResponsesModel::from_config(model_id.to_string(), self.model_config("responses")) } - /// Create a Responses API model instance for the given xAI model id. - /// - /// Uses the xAI `/responses` endpoint with the Responses API wire format - /// (input items, reasoning objects, provider-executed tools, etc.). + /// The provider as a function: the default language model for an id. The + /// AI SDK's callable provider; the same model as + /// [`responses`](Self::responses) and + /// [`language_model`](Provider::language_model). #[must_use] - pub fn responses_model(&self, model_id: &str) -> XaiResponsesModel { - XaiResponsesModel::new(model_id.to_string(), self.config.clone()) + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.responses(model_id)) } } impl Provider for XAIProvider { - fn name(&self) -> &str { - PROVIDER_NAME + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) } - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - /// List models via `GET {base_url}/models` (OpenAI-compatible, RFC-0027). - fn list_models( - &self, - ) -> std::pin::Pin< - Box< - dyn std::future::Future< - Output = Result, AiMuxError>, - > + Send - + '_, - >, - > { - let config = OpenAIConfig::new(self.config.api_key()) - .with_base_url(self.config.base_url()) - .with_provider(PROVIDER_NAME); - // The freshly built OpenAIConfig carries the default retry settings — - // use the user's configured ones from the wrapped config instead. - let retry_config = self.config.0.retry_config; + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "embeddingModel")) + } + + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Err(AiMuxError::no_such_model(model_id, "imageModel")) + } +} + +impl ProviderDiscovery for XAIProvider { + /// `GET {base_url}/models`: one exchange, no retry. + fn list_models(&self) -> BoxFuture<'_, Result, AiMuxError>> { + let config = self.model_config("models"); Box::pin(async move { - let headers = crate::openai::model::build_auth_headers(&config); - let runtime = - crate::openai::model::execute_list_models(&config.base_url, &headers, retry_config) - .await?; - Ok(runtime) + crate::shared::list_data_models(&config, xai_failed_response_handler()).await }) } } diff --git a/aimux-providers/src/xai/model.rs b/aimux-providers/src/xai/model.rs deleted file mode 100644 index 96e37840d..000000000 --- a/aimux-providers/src/xai/model.rs +++ /dev/null @@ -1,598 +0,0 @@ -//! xAI language model — implements `LanguageModel` trait. -//! -//! Mirrors the TS `XaiChatLanguageModel`. Unlike the thin OpenAI-compatible -//! wrappers, xAI has enough provider-specific behaviour (reasoning content, -//! citations, search parameters, xai-keyed provider options, non-inclusive -//! cached tokens, reasoning-effort model gating, 200-status errors) to warrant -//! its own model implementation rather than reusing `OpenAIModel`. - -use aimux_core::tool::RawToolCall; -use std::collections::HashMap; -use std::sync::atomic::{AtomicU64, Ordering}; - -use async_trait::async_trait; -use futures::StreamExt; -use serde_json::Value; - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; - -use aimux_provider_utils::HttpRequest; - -use super::convert::{build_request_body_with_warnings, convert_xai_usage, parse_finish_reason}; -use super::types::{XaiChatResponse, XaiStreamChunk}; - -/// Global counter for generating unique source IDs. -static SOURCE_ID_COUNTER: AtomicU64 = AtomicU64::new(0); - -/// Generate a unique source ID. -fn generate_source_id() -> String { - format!( - "xai-source-{}", - SOURCE_ID_COUNTER.fetch_add(1, Ordering::SeqCst) - ) -} - -/// An xAI language model (Grok). -/// -/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the -/// process-wide shared `Client` internally (RFC-0009 §4.1). -pub struct XaiModel { - model_id: String, - config: super::XAIConfig, -} - -impl XaiModel { - #[must_use] - pub fn new(model_id: String, config: super::XAIConfig) -> Self { - Self { model_id, config } - } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key()), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/chat/completions", self.config.base_url()) - } -} - -/// Build the header list for a JSON POST: auth/extra headers + `Content-Type`. -/// -/// Returns a `Vec<(String, String)>` for `HttpRequest` — no reqwest types. -fn build_header_list(headers: &HashMap) -> Vec<(String, String)> { - let mut list: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - list.push(("Content-Type".to_string(), "application/json".to_string())); - list -} - -#[async_trait] -impl LanguageModel for XaiModel { - fn provider(&self) -> &str { - "xai.chat" - } - - fn model_id(&self) -> &str { - &self.model_id - } - - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.openai_config().retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - // M2b: xAI wraps OpenAIConfig — reuse the OpenAI snapshot helper with - // xAI's own provider name. api_key_source/profile come from the inner - // config. - crate::openai::config_snapshot_from_config( - self.provider(), - &self.model_id, - self.config.openai_config(), - ) - } - - async fn do_generate(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()); - let request_result = build_request_body_with_warnings(&self.model_id, options, false)?; - let body = request_result.body; - - let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), build_header_list(&headers), options), - body.clone(), - super::xai_successful_response_handler::(), - super::xai_failed_response_handler(), - ) - .await?; - - // Capture response headers. - let response_headers = resp.response_headers; - - let response_value = resp.raw_value.unwrap_or(Value::Null); - let data = resp.value; - - let choice = data - .choices - .and_then(|mut c| c.drain(..).next()) - .ok_or_else(|| AiMuxError::InvalidResponseData("no choices in response".to_string()))?; - - // Build content array. - let mut content = Vec::new(); - - // Extract text content - if let Some(text) = choice.message.content - && !text.is_empty() - { - let mut text = text; - // Skip if this content duplicates the last assistant message - if let Some(last_msg) = body - .get("messages") - .and_then(|m| m.as_array()) - .and_then(|arr| arr.last()) - && last_msg.get("role").and_then(|v| v.as_str()) == Some("assistant") - && last_msg.get("content").and_then(|v| v.as_str()) == Some(&text) - { - text = String::new(); - } - if !text.is_empty() { - content.push(GenerateContent::Text { - text, - provider_metadata: None, - }); - } - } - - // Extract reasoning content - if let Some(reasoning) = choice.message.reasoning_content - && !reasoning.is_empty() - { - content.push(GenerateContent::Reasoning(ReasoningOutput { - text: reasoning, - provider_metadata: None, - })); - } - - // Extract tool calls - if let Some(tool_calls) = choice.message.tool_calls { - for tc in tool_calls { - let input = tc.function.arguments; - content.push(GenerateContent::ToolCall(RawToolCall { - tool_call_id: tc.id, - tool_name: tc.function.name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - } - - // Extract citations - if let Some(citations) = &data.citations { - for url in citations { - content.push(GenerateContent::Source(Source::Url { - id: generate_source_id(), - url: url.clone(), - title: None, - provider_metadata: None, - })); - } - } - - let finish_reason = choice - .finish_reason - .as_deref() - .map(parse_finish_reason) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }); - - let usage = data.usage.as_ref().map(convert_xai_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, - }); - - let timestamp = data - .created - .and_then(|c| chrono::DateTime::from_timestamp(c as i64, 0).map(|dt| dt.to_rfc3339())); - - Ok(GenerateResult { - content, - finish_reason, - usage, - warnings: request_result.warnings, - provider_metadata: None, - response: Some(aimux_core::shared::ResponseInfo { - id: data.id, - timestamp, - model_id: data.model, - headers: Some(response_headers), - body: Some(response_value), - }), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - }) - } - - async fn do_stream(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()); - let request_result = build_request_body_with_warnings(&self.model_id, options, true)?; - let body = request_result.body; - let warnings = request_result.warnings; - let endpoint = self.endpoint(); - - let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.clone(), build_header_list(&headers), options), - body.clone(), - super::xai_event_source_response_handler::(), - super::xai_failed_response_handler(), - ) - .await?; - - let response_headers = resp.response_headers; - - let mut sse_stream = resp.value; - - // 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 val)) = first_event - && val.get("error").is_some_and(|value| !value.is_null()) - { - return Err(super::xai_stream_error( - val, - &endpoint, - body.clone(), - response_headers.clone(), - )); - } - - let messages_for_dedup = body - .get("messages") - .and_then(|m| m.as_array()) - .cloned() - .unwrap_or_default(); - let stream_error_url = endpoint.clone(); - let stream_request_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); - - let mut final_usage: Option = None; - let mut final_finish_reason: Option = None; - let mut response_metadata_emitted = false; - - // Content block tracking - let mut content_blocks: HashMap = HashMap::new(); // block_id -> ended - let mut last_reasoning_deltas: HashMap = HashMap::new(); - let mut active_reasoning_block_id: Option = None; - - // Tool-call accumulators keyed by index (for incremental streaming). - let mut tool_calls: HashMap = HashMap::new(); // (id, name, arguments) - let mut tool_call_order: Vec = Vec::new(); - - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); - - while let Some(event) = event_iter.next().await { - match event { - Ok(parsed) => { - // Check for mid-stream error. - if parsed.get("error").is_some() { - yield Ok(StreamPart::Error { - error: super::xai_stream_error( - &parsed, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); - break; - } - - let chunk: XaiStreamChunk = match serde_json::from_value(parsed) { - Ok(c) => c, - Err(e) => { - yield Err(AiMuxError::from(e)); - continue; - } - }; - - // Emit response metadata on first chunk. - if !response_metadata_emitted - && (chunk.id.is_some() || chunk.model.is_some()) - { - response_metadata_emitted = true; - let timestamp = chunk.created.and_then(|c| { - chrono::DateTime::from_timestamp(c as i64, 0) - .map(|dt| dt.to_rfc3339()) - }); - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: chunk.id.clone(), - timestamp, - model_id: chunk.model.clone(), - })); - } - - // Emit citations as sources. - if let Some(citations) = &chunk.citations { - for url in citations { - yield Ok(StreamPart::Source(Source::Url { - id: generate_source_id(), - url: url.clone(), - title: None, - provider_metadata: None, - })); - } - } - - // Update usage. - if let Some(usage) = &chunk.usage { - final_usage = Some(convert_xai_usage(usage)); - } - - // Process choices. - for choice in &chunk.choices { - let choice_index = choice.index.unwrap_or(0); - let chunk_id = chunk.id.clone().unwrap_or_else(|| choice_index.to_string()); - - let delta = &choice.delta; - - // Process text content. - if let Some(text_content) = &delta.content - && !text_content.is_empty() - { - // End active reasoning block. - if let Some(ref active_id) = active_reasoning_block_id { - if !content_blocks.get(active_id).copied().unwrap_or(false) { - yield Ok(StreamPart::ReasoningEnd { - id: active_id.clone(), - provider_metadata: None, - }); - content_blocks.insert(active_id.clone(), true); - } - active_reasoning_block_id = None; - } - - // Skip if duplicates last assistant message. - let last_msg = messages_for_dedup.last(); - let is_dup = last_msg - .and_then(|m| m.get("role").and_then(|v| v.as_str())) - == Some("assistant") - && last_msg - .and_then(|m| m.get("content").and_then(|v| v.as_str())) - == Some(text_content.as_str()); - - if !is_dup { - let block_id = format!("text-{chunk_id}"); - if !content_blocks.contains_key(&block_id) { - content_blocks.insert(block_id.clone(), false); - yield Ok(StreamPart::TextStart { - id: block_id.clone(), - provider_metadata: None, - }); - } - yield Ok(StreamPart::TextDelta { - id: block_id, - delta: text_content.clone(), - provider_metadata: None, - }); - } - } - - // Process reasoning content. - if let Some(reasoning_content) = &delta.reasoning_content - && !reasoning_content.is_empty() - { - let block_id = format!("reasoning-{chunk_id}"); - - // Skip if duplicates last delta. - if last_reasoning_deltas.get(&block_id).map(std::string::String::as_str) - != Some(reasoning_content.as_str()) - { - last_reasoning_deltas - .insert(block_id.clone(), reasoning_content.clone()); - - if !content_blocks.contains_key(&block_id) { - content_blocks.insert(block_id.clone(), false); - active_reasoning_block_id = Some(block_id.clone()); - yield Ok(StreamPart::ReasoningStart { - id: block_id.clone(), - provider_metadata: None, - }); - } - yield Ok(StreamPart::ReasoningDelta { - id: block_id, - delta: reasoning_content.clone(), - provider_metadata: None, - }); - } - } - - // Process tool calls. - // xAI typically sends tool calls in one piece, but - // we also support incremental (OpenAI-style) streaming. - if let Some(tool_call_deltas) = &delta.tool_calls { - // End active reasoning block. - if let Some(ref active_id) = active_reasoning_block_id { - if !content_blocks.get(active_id).copied().unwrap_or(false) { - yield Ok(StreamPart::ReasoningEnd { - id: active_id.clone(), - provider_metadata: None, - }); - content_blocks.insert(active_id.clone(), true); - } - active_reasoning_block_id = None; - } - - for dtc in tool_call_deltas { - let idx = dtc.index.unwrap_or(0); - let func = dtc.function.clone().unwrap_or_default(); - - let is_new = !tool_calls.contains_key(&idx); - if is_new { - let id = dtc.id.clone().unwrap_or_default(); - let name = func.name.unwrap_or_default(); - tool_calls.insert(idx, (id.clone(), name.clone(), String::new())); - tool_call_order.push(idx); - yield Ok(StreamPart::ToolInputStart { - id, - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - - if let Some(args) = func.arguments - && (!is_new || !args.is_empty()) - && let Some(&mut (ref id, _, ref mut acc)) = tool_calls.get_mut(&idx) - { - acc.push_str(&args); - yield Ok(StreamPart::ToolInputDelta { - id: id.clone(), - delta: args, - provider_metadata: None, - }); - } - } - } - - // Finish reason: close any open tool calls. - if let Some(reason) = &choice.finish_reason { - final_finish_reason = Some(parse_finish_reason(reason)); - - // Close any open tool calls. - for &idx in &tool_call_order { - if let Some((id, name, args)) = tool_calls.get(&idx) { - yield Ok(StreamPart::ToolInputEnd { - id: id.clone(), - provider_metadata: None, - }); - let input = args.clone(); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id.clone(), - tool_name: name.clone(), - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - } - tool_calls.clear(); - tool_call_order.clear(); - } - } - } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; - } - } - } - } - - // End any remaining open blocks. - for (block_id, ended) in &content_blocks { - if !ended { - // Determine if it's a reasoning or text block by prefix. - if block_id.starts_with("reasoning-") { - yield Ok(StreamPart::ReasoningEnd { - id: block_id.clone(), - provider_metadata: None, - }); - } else { - yield Ok(StreamPart::TextEnd { - id: block_id.clone(), - provider_metadata: None, - }); - } - } - } - - // Close any remaining open tool calls (no finish_reason was received). - for &idx in &tool_call_order { - if let Some((id, name, args)) = tool_calls.get(&idx) { - yield Ok(StreamPart::ToolInputEnd { - id: id.clone(), - provider_metadata: None, - }); - let input = args.clone(); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id.clone(), - tool_name: name.clone(), - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - } - - // Final part: Finish. - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }), - 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, - }), - provider_metadata: None, - }); - }; - - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) - } -} diff --git a/aimux-providers/src/xai/options.rs b/aimux-providers/src/xai/options.rs new file mode 100644 index 000000000..1078e3228 --- /dev/null +++ b/aimux-providers/src/xai/options.rs @@ -0,0 +1,24 @@ +//! The providerOptions and providerMetadata namespace of xAI. +//! +//! `@ai-sdk/xai` reads `providerOptions.xai` and writes response metadata +//! under `xai`, 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; + +/// The providerOptions / providerMetadata key (and the provider-reference key +/// of uploaded files). +pub(crate) const NAMESPACE: &str = "xai"; + +/// The xAI options in a providerOptions container (a call's map, a part's +/// object). +pub(crate) fn xai_options(provider_options: Option<&SharedProviderOptions>) -> Option<&JsonObject> { + provider_options?.get(NAMESPACE) +} + +/// `{ "xai": payload }`. +pub(crate) fn xai_metadata(payload: Value) -> ProviderMetadata { + provider_namespace(NAMESPACE, payload).expect("provider metadata must be an object") +} diff --git a/aimux-providers/src/xai/responses/convert.rs b/aimux-providers/src/xai/responses/convert.rs index a8c27f216..17f7c83ab 100644 --- a/aimux-providers/src/xai/responses/convert.rs +++ b/aimux-providers/src/xai/responses/convert.rs @@ -8,7 +8,7 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model_message::{ - AssistantPart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, ToolCallPart, ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, }; use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat}; @@ -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 serde_json::{Value, json}; use super::types::XaiResponsesUsage; use crate::xai::convert::{ - remove_additional_properties_false, resolve_full_media_type, resolve_provider_reference, - supports_reasoning_effort, + remove_additional_properties_false, resolve_provider_reference, supports_reasoning_effort, }; // ── Finish reason ──────────────────────────────────────────────────────────── @@ -276,7 +276,7 @@ fn prepare_provider_tool( fn xai_option(options: &Option, key: &str) -> Option { options .as_ref() - .and_then(|m| m.get("xai")) + .and_then(|m| crate::xai::options::xai_options(Some(m))) .and_then(|o| o.get(key)) .cloned() } @@ -323,37 +323,25 @@ pub fn convert_to_xai_responses_input( use base64::Engine; let b64 = base64::engine::general_purpose::STANDARD.encode(bytes); - convert_image_part( - &file.media_type, - Some(&b64), - None, - &file.provider_options, - ) + convert_image_part(file, Some(&b64), None)? } FileData::Data { data: FileBytes::Base64(data), - } => convert_image_part( - &file.media_type, - Some(data), - None, - &file.provider_options, - ), + } => convert_image_part(file, Some(data), None)?, FileData::Url { url, .. } - if file.media_type.split('/').next() == Some("image") => + if get_top_level_media_type(&file.media_type) == "image" => { - convert_image_part( - &file.media_type, - None, - Some(url), - &file.provider_options, - ) + convert_image_part(file, None, Some(url))? } FileData::Url { url, .. } => { json!({ "type": "input_file", "file_url": url }) } FileData::Reference { reference } => { - let file_id = resolve_provider_reference(reference, "xai") - .map_err(AiMuxError::InvalidArgument)?; + let file_id = resolve_provider_reference( + reference, + crate::xai::options::NAMESPACE, + ) + .map_err(AiMuxError::InvalidArgument)?; json!({ "type": "input_file", "file_id": file_id }) } FileData::Text { .. } => { @@ -373,11 +361,10 @@ pub fn convert_to_xai_responses_input( AssistantPart::Text(TextPart { text, provider_options, - .. }) => { let id = provider_options .as_ref() - .and_then(|po| po.get("xai")) + .and_then(|po| crate::xai::options::xai_options(Some(po))) .and_then(|x| x.get("itemId")) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string); @@ -401,7 +388,7 @@ pub fn convert_to_xai_responses_input( let is_provider_executed = provider_executed.unwrap_or_else(|| { provider_options .as_ref() - .and_then(|po| po.get("xai")) + .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) @@ -411,7 +398,7 @@ pub fn convert_to_xai_responses_input( } let id = provider_options .as_ref() - .and_then(|po| po.get("xai")) + .and_then(|po| crate::xai::options::xai_options(Some(po))) .and_then(|x| x.get("itemId")) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string); @@ -434,13 +421,13 @@ pub fn convert_to_xai_responses_input( }) => { let item_id = provider_options .as_ref() - .and_then(|po| po.get("xai")) + .and_then(|po| crate::xai::options::xai_options(Some(po))) .and_then(|x| x.get("itemId")) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string); let encrypted_content = provider_options .as_ref() - .and_then(|po| po.get("xai")) + .and_then(|po| crate::xai::options::xai_options(Some(po))) .and_then(|x| x.get("reasoningEncryptedContent")) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string); @@ -466,9 +453,17 @@ pub fn convert_to_xai_responses_input( }); } } - _ => { + AssistantPart::File(_) + | AssistantPart::Custom(_) + | AssistantPart::ReasoningFile(_) => { + let kind = match part { + AssistantPart::File(_) => "file", + AssistantPart::Custom(_) => "custom", + AssistantPart::ReasoningFile(_) => "reasoning-file", + _ => unreachable!(), + }; warnings.push(Warning::Other { - message: "xAI Responses API does not support this content type in assistant messages".to_string(), + message: format!("xAI Responses API does not support {kind} in assistant messages"), }); } } @@ -484,7 +479,7 @@ pub fn convert_to_xai_responses_input( else { continue; }; - let output_value = convert_tool_result_output(output); + let output_value = convert_tool_result_output(output)?; input.push(json!({ "type": "function_call_output", "call_id": tool_call_id, @@ -498,8 +493,8 @@ pub fn convert_to_xai_responses_input( Ok((input, warnings)) } -fn convert_tool_result_output(output: &ToolResultOutput) -> Value { - match output { +fn convert_tool_result_output(output: &ToolResultOutput) -> Result { + Ok(match output { ToolResultOutput::ExecutionDenied { reason, .. } => { json!(reason.as_deref().unwrap_or("tool execution denied")) } @@ -523,11 +518,7 @@ fn convert_tool_result_output(output: &ToolResultOutput) -> Value { } FileBytes::Base64(data) => data.clone(), }; - format!( - "data:{};base64,{}", - resolve_full_media_type(&part.media_type, &b64), - b64 - ) + format!("data:{};base64,{}", resolve_full_media_type(part)?, b64) } _ => continue, }; @@ -539,19 +530,24 @@ fn convert_tool_result_output(output: &ToolResultOutput) -> Value { json!(parts) } _ => crate::openai::convert::tool_result_to_content(output), - } + }) } fn convert_image_part( - media_type: &str, + file_part: &FilePart, b64_data: Option<&str>, url: Option<&str>, - provider_options: &Option, -) -> Value { +) -> Result { + let media_type = &file_part.media_type; + if get_top_level_media_type(media_type) != "image" { + return Err(AiMuxError::UnsupportedFunctionality(format!( + "file part media type {media_type} as inline data (xAI Responses requires a URL or a Files API reference for non-image files)" + ))); + } let image_url = if let Some(url_str) = url { url_str.to_string() } else if let Some(b64) = b64_data { - let full_mt = resolve_full_media_type(media_type, b64); + let full_mt = resolve_full_media_type(file_part)?; format!("data:{full_mt};base64,{b64}") } else { String::new() @@ -563,15 +559,16 @@ fn convert_image_part( }); // Image detail provider option. - if let Some(detail) = provider_options + if let Some(detail) = file_part + .provider_options .as_ref() - .and_then(|po| po.get("xai")) + .and_then(|po| crate::xai::options::xai_options(Some(po))) .and_then(|x| x.get("imageDetail")) { part["detail"] = detail.clone(); } - part + Ok(part) } // ── Request body builder ───────────────────────────────────────────────────── @@ -600,9 +597,11 @@ pub fn build_responses_request_body( .as_ref() .map(|namespaces| { let mut namespaces = namespaces.clone(); - if let Some(raw) = namespaces.get("xai") { - let parsed = - crate::openai::convert::parse_option_fields(raw, "xai", |key, value| { + 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") @@ -643,14 +642,29 @@ pub fn build_responses_request_body( _ => return None, }; Some(if valid { Ok(value.clone()) } else { Err(()) }) - })?; - namespaces.insert("xai".to_string(), parsed); + }, + )?; + namespaces.insert(crate::xai::options::NAMESPACE.to_string(), parsed); } Ok::<_, AiMuxError>(namespaces) }) .transpose()?; let xai_opts = &parsed_options; + if options.frequency_penalty.is_some() { + warnings.push(Warning::Unsupported { + feature: "frequencyPenalty".to_string(), + details: None, + }); + } + + if options.presence_penalty.is_some() { + warnings.push(Warning::Unsupported { + feature: "presencePenalty".to_string(), + details: None, + }); + } + if options.stop_sequences.is_some() { warnings.push(Warning::Unsupported { feature: "stopSequences".to_string(), @@ -744,6 +758,9 @@ pub fn build_responses_request_body( if let Some(tp) = options.top_p { body["top_p"] = json!(tp); } + if let Some(tk) = options.top_k { + body["top_k"] = json!(tk); + } if let Some(seed) = options.seed { body["seed"] = json!(seed); } diff --git a/aimux-providers/src/xai/responses/mod.rs b/aimux-providers/src/xai/responses/mod.rs index fce9e7198..290108c21 100644 --- a/aimux-providers/src/xai/responses/mod.rs +++ b/aimux-providers/src/xai/responses/mod.rs @@ -28,14 +28,10 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; -use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::HttpRequest; - -use super::super::XAIConfig; -use crate::openai::responses::responses_convert::build_header_list; +use crate::shared::EndpointConfig; use convert::{ build_responses_request_body, convert_xai_responses_usage, get_tool_input, map_xai_responses_finish_reason, resolve_tool_name, @@ -54,65 +50,33 @@ fn generate_source_id() -> String { /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct XaiResponsesModel { model_id: String, - config: XAIConfig, + config: EndpointConfig, } impl XaiResponsesModel { - #[must_use] - pub fn new(model_id: String, config: XAIConfig) -> Self { + pub(crate) fn from_config(model_id: String, config: EndpointConfig) -> Self { Self { model_id, config } } - - fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert( - "Authorization".to_string(), - format!("Bearer {}", self.config.api_key()), - ); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/responses", self.config.base_url()) - } } #[async_trait] impl LanguageModel for XaiResponsesModel { fn provider(&self) -> &str { - "xai.responses" + &self.config.provider } fn model_id(&self) -> &str { &self.model_id } - fn retry_config(&self) -> aimux_core::retry::RetryConfig { - self.config.openai_config().retry_config - } - - fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { - // M2b: reuse the OpenAI snapshot helper with xAI's provider name. - crate::openai::config_snapshot_from_config( - self.provider(), - &self.model_id, - self.config.openai_config(), - ) - } - async fn do_generate(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let request_result = build_responses_request_body(&self.model_id, options, false)?; - let body = request_result.body; + let body = exchange.transform_body(request_result.body); let provider_tool_names = request_result.provider_tool_names; let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(self.endpoint(), build_header_list(&headers), options), + exchange.request(exchange.url("/responses"), options), body.clone(), super::xai_successful_response_handler::(), super::xai_failed_response_handler(), @@ -302,10 +266,7 @@ impl LanguageModel for XaiResponsesModel { } content.push(GenerateContent::Reasoning(ReasoningOutput { text: reasoning_text, - provider_metadata: Some( - provider_namespace("xai", meta) - .expect("provider metadata must be an object"), - ), + provider_metadata: Some(crate::xai::options::xai_metadata(json!(meta))), })); } } @@ -327,10 +288,9 @@ impl LanguageModel for XaiResponsesModel { let (usage, provider_metadata) = if let Some(u) = &data.usage { let meta = if u.cost_in_usd_ticks.is_some() { - Some( - provider_namespace("xai", json!({ "costInUsdTicks": u.cost_in_usd_ticks })) - .expect("provider metadata must be an object"), - ) + Some(crate::xai::options::xai_metadata( + json!({ "costInUsdTicks": u.cost_in_usd_ticks }), + )) } else { None }; @@ -361,15 +321,15 @@ impl LanguageModel for XaiResponsesModel { } async fn do_stream(&self, options: &CallOptions) -> Result { - let headers = self.build_headers(options.headers.as_ref()); + let exchange = self.config.exchange(options.headers.as_ref()).await?; let request_result = build_responses_request_body(&self.model_id, options, true)?; - let body = request_result.body; + let body = exchange.transform_body(request_result.body); let warnings = request_result.warnings; let provider_tool_names = request_result.provider_tool_names; - let endpoint = self.endpoint(); + let endpoint = exchange.url("/responses"); let resp = aimux_provider_utils::post_json_to_api( - HttpRequest::new(endpoint.clone(), build_header_list(&headers), options), + exchange.request(endpoint.clone(), options), body.clone(), super::xai_event_source_response_handler::(), super::xai_failed_response_handler(), @@ -448,7 +408,7 @@ impl LanguageModel for XaiResponsesModel { active_reasoning.insert(item_id.to_string(), ()); yield Ok(StreamPart::ReasoningStart { id: block_id, - provider_metadata: Some(provider_namespace("xai", json!({ "itemId": item_id })).expect("provider metadata must be an object")), + provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), }); } continue; @@ -462,7 +422,7 @@ impl LanguageModel for XaiResponsesModel { yield Ok(StreamPart::ReasoningDelta { id: block_id, delta: delta.to_string(), - provider_metadata: Some(provider_namespace("xai", json!({ "itemId": item_id })).expect("provider metadata must be an object")), + provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), }); continue; } @@ -481,13 +441,13 @@ impl LanguageModel for XaiResponsesModel { active_reasoning.insert(item_id.to_string(), ()); yield Ok(StreamPart::ReasoningStart { id: block_id.clone(), - provider_metadata: Some(provider_namespace("xai", json!({ "itemId": item_id })).expect("provider metadata must be an object")), + provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), }); } yield Ok(StreamPart::ReasoningDelta { id: block_id, delta: delta.to_string(), - provider_metadata: Some(provider_namespace("xai", json!({ "itemId": item_id })).expect("provider metadata must be an object")), + provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), }); continue; } @@ -670,7 +630,7 @@ impl LanguageModel for XaiResponsesModel { active_reasoning.insert(part_id.to_string(), ()); yield Ok(StreamPart::ReasoningStart { id: block_id.clone(), - provider_metadata: Some(provider_namespace("xai", json!({ "itemId": part_id })).expect("provider metadata must be an object")), + provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": part_id }))), }); } @@ -680,7 +640,7 @@ impl LanguageModel for XaiResponsesModel { } yield Ok(StreamPart::ReasoningEnd { id: block_id, - provider_metadata: Some(provider_namespace("xai", meta).expect("provider metadata must be an object")), + provider_metadata: Some(crate::xai::options::xai_metadata(json!(meta))), }); active_reasoning.remove(part_id); } @@ -903,7 +863,7 @@ impl LanguageModel for XaiResponsesModel { } // Final part: Finish. - let provider_meta = cost_in_usd_ticks.map(|cost| provider_namespace("xai", json!({ "costInUsdTicks": cost })).expect("provider metadata must be an object")); + let provider_meta = cost_in_usd_ticks.map(|cost| crate::xai::options::xai_metadata(json!({ "costInUsdTicks": cost }))); yield Ok(StreamPart::Finish { finish_reason: final_finish_reason, diff --git a/aimux-providers/src/xai/types.rs b/aimux-providers/src/xai/types.rs deleted file mode 100644 index ec5b285c5..000000000 --- a/aimux-providers/src/xai/types.rs +++ /dev/null @@ -1,162 +0,0 @@ -//! xAI API response types. - -#![allow(dead_code)] - -use serde::{Deserialize, Serialize}; - -// ── Non-streaming response ── - -#[derive(Debug, Deserialize)] -pub struct XaiChatResponse { - #[serde(default)] - pub id: Option, - #[serde(default)] - pub model: Option, - #[serde(default)] - pub created: Option, - #[serde(default)] - pub choices: Option>, - #[serde(default)] - pub usage: Option, - /// Citations (URLs) returned when search is enabled. - #[serde(default)] - pub citations: Option>, - /// Error code — present when xAI returns an error with HTTP 200. - #[serde(default)] - pub code: Option, - /// Error message — present when xAI returns an error with HTTP 200. - #[serde(default)] - pub error: Option, -} - -#[derive(Debug, Deserialize)] -pub struct XaiChoice { - pub message: XaiMessageResponse, - #[serde(default)] - pub finish_reason: Option, -} - -#[derive(Debug, Deserialize)] -pub struct XaiMessageResponse { - pub role: String, - #[serde(default)] - pub content: Option, - /// Reasoning / thinking content produced by the model. - #[serde(default)] - pub reasoning_content: Option, - #[serde(default)] - pub tool_calls: Option>, -} - -#[derive(Debug, Deserialize)] -pub struct XaiToolCallResponse { - pub id: String, - pub function: XaiFunctionCallResponse, -} - -#[derive(Debug, Deserialize)] -pub struct XaiFunctionCallResponse { - pub name: String, - pub arguments: String, -} - -// ── Usage ── - -#[derive(Debug, Deserialize, Clone, Default, Serialize)] -pub struct XaiUsageResponse { - #[serde(default)] - pub prompt_tokens: Option, - #[serde(default)] - pub completion_tokens: Option, - #[serde(default)] - pub total_tokens: Option, - #[serde(default)] - pub prompt_tokens_details: Option, - #[serde(default)] - pub completion_tokens_details: Option, -} - -#[derive(Debug, Deserialize, Clone, Default, Serialize)] -pub struct XaiPromptTokensDetails { - #[serde(default)] - pub text_tokens: Option, - #[serde(default)] - pub audio_tokens: Option, - #[serde(default)] - pub image_tokens: Option, - #[serde(default)] - pub cached_tokens: Option, -} - -#[derive(Debug, Deserialize, Clone, Default, Serialize)] -pub struct XaiCompletionTokensDetails { - #[serde(default)] - pub reasoning_tokens: Option, - #[serde(default)] - pub audio_tokens: Option, - #[serde(default)] - pub accepted_prediction_tokens: Option, - #[serde(default)] - pub rejected_prediction_tokens: Option, -} - -// ── Streaming response ── - -#[derive(Debug, Deserialize)] -pub struct XaiStreamChunk { - #[serde(default)] - pub id: Option, - #[serde(default)] - pub model: Option, - #[serde(default)] - pub created: Option, - #[serde(default)] - pub choices: Vec, - #[serde(default)] - pub usage: Option, - #[serde(default)] - pub citations: Option>, -} - -#[derive(Debug, Deserialize, Default)] -pub struct XaiStreamChoice { - #[serde(default)] - pub delta: XaiDelta, - #[serde(default)] - pub finish_reason: Option, - #[serde(default)] - pub index: Option, -} - -#[derive(Debug, Deserialize, Default)] -pub struct XaiDelta { - #[serde(default)] - pub role: Option, - #[serde(default)] - pub content: Option, - #[serde(default)] - pub reasoning_content: Option, - #[serde(default)] - pub tool_calls: Option>, -} - -#[derive(Debug, Deserialize)] -pub struct XaiDeltaToolCall { - #[serde(default)] - pub index: Option, - #[serde(default)] - pub id: Option, - #[serde(default)] - pub function: Option, -} - -#[derive(Debug, Deserialize, Clone, Default)] -pub struct XaiDeltaFunction { - #[serde(default)] - pub name: Option, - #[serde(default)] - pub arguments: Option, -} - -// Suppress unused warning for Value import. -// (No longer needed — Value import removed.) diff --git a/aimux-providers/src/xinference.rs b/aimux-providers/src/xinference.rs deleted file mode 100644 index 01607cb2e..000000000 --- a/aimux-providers/src/xinference.rs +++ /dev/null @@ -1,79 +0,0 @@ -//! Xinference provider — a thin OpenAI-compatible wrapper. -//! -//! See for API documentation. Exposes an OpenAI-compatible -//! Chat Completions API at `http://127.0.0.1:9997/v1`. The `XINFERENCE_BASE_URL` environment -//! variable holds a *base URL* (not an API key); when unset, the default -//! endpoint is used. A placeholder API key is sent in the `Authorization` -//! header — the shared `OpenAIProvider` requires a non-empty key string. - -use aimux_core::error::AiMuxError; -use aimux_core::language_model::LanguageModel; -use aimux_core::provider::Provider; - -use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIModel, OpenAIProvider}; - -const DEFAULT_BASE_URL: &str = "http://127.0.0.1:9997/v1"; -const ENV_VAR: &str = "XINFERENCE_BASE_URL"; -const PROVIDER_NAME: &str = "xinference"; -const PLACEHOLDER_API_KEY: &str = "xinference"; - -pub struct XinferenceConfig(OpenAIConfig); - -impl XinferenceConfig { - pub fn new(api_key: impl Into) -> Self { - Self( - OpenAIConfig::new(api_key) - .with_base_url(DEFAULT_BASE_URL) - .with_provider(PROVIDER_NAME) - .with_profile(OpenAICompatProfile::full()) - .with_api_key_source(Some("none")), - ) - } - - /// Create from the `XINFERENCE_BASE_URL` environment variable, falling back - /// to `http://127.0.0.1:9997/v1`. - /// - /// # Errors - /// - /// Never returns an error; an unset or empty variable falls back to the - /// default local endpoint. - pub fn from_env() -> Result { - let config = Self::new(PLACEHOLDER_API_KEY); - match std::env::var(ENV_VAR) { - Ok(url) if !url.trim().is_empty() => Ok(config.with_base_url(url)), - _ => Ok(config), - } - } - - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.0 = self.0.with_base_url(url); - self - } -} - -pub struct XinferenceProvider(OpenAIProvider); - -impl XinferenceProvider { - #[must_use] - pub fn new(config: XinferenceConfig) -> Self { - Self(OpenAIProvider::new(config.0)) - } - - #[must_use] - pub fn model(&self, model_id: &str) -> OpenAIModel { - self.0.model(model_id) - } -} - -impl Provider for XinferenceProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(model_id))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/you_com.rs b/aimux-providers/src/you_com.rs index dcab66c32..3a421776a 100644 --- a/aimux-providers/src/you_com.rs +++ b/aimux-providers/src/you_com.rs @@ -7,26 +7,49 @@ //! base URL is `ydc-index.io`, not `you.com`. You.com is a search-only //! provider: it exposes a web search protocol (query → results) and does not //! support language models. +//! +//! [`create_you_com`] takes [`YouComProviderSettings`], validates the base URL, fixes +//! the provider name and returns a [`YouComProvider`]. The credential is not read +//! there: it is loaded for every request, from the setting or from `YDC_API_KEY`. +//! [`you_com()`] is the default instance; it reads nothing and cannot fail. -use std::collections::HashMap; +use std::sync::OnceLock; use async_trait::async_trait; use serde::Deserialize; use serde_json::Value; use aimux_core::error::AiMuxError; -use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_core::shared::SharedHeaders; - -use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt, HeadersFn, Resolvable, validate_base_url}; -/// Provider canonical name. -const PROVIDER_NAME: &str = "you_com"; +use crate::shared::{AuthScheme, Credential, EndpointConfig, credential_headers}; /// Fixed model id for the You.com search model. +/// Failed-response handler of the You.com endpoint: `message`, and `type` +/// else `code` as the provider code, from the body's `error` object, or from +/// the top level when there is none. This is the mapping the endpoint used +/// with the removed shared handler; it is not an AI SDK schema. +fn you_com_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(serde_json::Value::as_str) + .map(str::to_owned), + } + }) +} + const MODEL_ID: &str = "youcom-search"; /// Default number of results when `max_results` is unset. @@ -34,79 +57,117 @@ const DEFAULT_COUNT: u32 = 10; // ── Config ─────────────────────────────────────────────────────────────────── -/// Configuration for the You.com provider. -/// -/// The `Debug` implementation redacts the API key so credentials never appear -/// in logs or error messages. -#[derive(Clone)] -pub struct YouComConfig { - api_key: String, - base_url: String, -} +pub(crate) const DEFAULT_BASE_URL: &str = "https://ydc-index.io"; +pub(crate) const API_KEY_ENV_VAR: &str = "YDC_API_KEY"; +const DEFAULT_NAME: &str = "you_com"; -impl YouComConfig { - /// Create from an API key (uses the default You.com / YDC base URL). - pub fn new(api_key: impl Into) -> Self { - Self { - api_key: api_key.into(), - base_url: "https://ydc-index.io".to_string(), - } - } - - /// Use a custom base URL. - #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self - } - - /// Create from the `YDC_API_KEY` environment variable. - /// - /// # Errors - /// - /// Returns `AiMuxError::InvalidArgument` when `YDC_API_KEY` is not set. - pub fn from_env() -> Result { - let api_key = load_api_key(None, "YDC_API_KEY", "You.com")?; - Ok(Self::new(api_key)) - } +/// Settings of [`create_you_com`]. +/// +/// Every field is optional. Nothing here is evaluated when the provider is +/// created except `base_url` and `name`; `api_key` and `headers` are +/// evaluated on every request. +#[derive(Clone, Default)] +pub struct YouComProviderSettings { + /// Base URL for the API calls. Default `https://ydc-index.io`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `YDC_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// explicit value is used as given, `""` included: it never falls back to + /// the environment. A [`Resolvable::Future`] is awaited once, an + /// [`Resolvable::AsyncFn`] on every request. + pub api_key: Option>, + /// Extra headers on every request. A `None` value removes the header, + /// including the credential. Per-call headers win over these. + pub headers: Option, + /// The provider name, the prefix of the `provider()` strings (`"{name}.search"`). + /// Default `"you_com"`. The providerOptions key stays `you_com`. + 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, } -impl std::fmt::Debug for YouComConfig { +impl std::fmt::Debug for YouComProviderSettings { + /// Never prints the key or header values. fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("YouComConfig") - .field("api_key", &"") + f.debug_struct("YouComProviderSettings") .field("base_url", &self.base_url) + .field("api_key", &self.api_key) + .field( + "headers", + &self.headers.as_ref().map(std::collections::HashMap::len), + ) + .field("name", &self.name) + .field("fetch", &self.fetch.is_some()) .finish() } } -// ── Provider ───────────────────────────────────────────────────────────────── - -/// You.com provider — creates [`YouComSearchModel`] instances. +/// Create a You.com provider. /// -/// You.com is a search-only provider; it does not support language models. +/// # 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. +pub fn create_you_com(settings: YouComProviderSettings) -> Result { + let base_url = match settings.base_url.as_deref() { + Some(url) => validate_base_url(url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + Ok(YouComProvider { + name: settings.name.unwrap_or_else(|| DEFAULT_NAME.to_string()), + base_url, + headers: credential_headers( + Credential::explicit_or_env(settings.api_key, API_KEY_ENV_VAR, "You.com"), + AuthScheme::Header("X-API-Key"), + Vec::new(), + settings.headers, + ), + fetch: settings.fetch, + }) +} + +/// The default provider: `create_you_com` with default settings, created on +/// first use. Creating it reads nothing from the environment and cannot fail; +/// a missing key surfaces from the first request instead. +pub fn you_com() -> &'static YouComProvider { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| { + create_you_com(YouComProviderSettings::default()) + .expect("default You.com settings are always valid") + }) +} + +/// A You.com provider. Cheap to clone the models out of; it holds no HTTP +/// client. pub struct YouComProvider { - config: YouComConfig, + name: String, + base_url: String, + headers: HeadersFn, + fetch: Option, } impl YouComProvider { - #[must_use] - pub fn new(config: YouComConfig) -> Self { - Self { config } + fn model_config(&self, method: &str) -> EndpointConfig { + EndpointConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Create a search model instance. + /// The search model; `provider()` is `"{name}.search"`. #[must_use] pub fn search_model(&self) -> YouComSearchModel { - YouComSearchModel::new(self.config.clone()) + YouComSearchModel::from_config(self.model_config("search")) } } -impl Provider for YouComProvider { - fn name(&self) -> &str { - PROVIDER_NAME - } -} +crate::impl_single_modality_provider!(YouComProvider, search_model, |p, _id| p.search_model()); // ── Request builder ────────────────────────────────────────────────────────── @@ -161,35 +222,19 @@ fn map_result(r: YoucomResult) -> SearchResultItem { /// A You.com search model. pub struct YouComSearchModel { - config: YouComConfig, + config: EndpointConfig, } impl YouComSearchModel { - #[must_use] - pub fn new(config: YouComConfig) -> Self { + pub(crate) fn from_config(config: EndpointConfig) -> Self { Self { config } } - - fn build_headers(&self, extra: Option<&SharedHeaders>) -> HashMap { - let mut headers = HashMap::new(); - headers.insert("X-API-Key".to_string(), self.config.api_key.clone()); - if let Some(extra) = extra { - for (k, v) in extra { - headers.insert(k.clone(), v.clone()); - } - } - headers - } - - fn endpoint(&self) -> String { - format!("{}/v1/search", self.config.base_url) - } } #[async_trait] impl SearchModel for YouComSearchModel { fn provider(&self) -> &str { - PROVIDER_NAME + &self.config.provider } fn model_id(&self) -> &str { @@ -198,29 +243,18 @@ impl SearchModel for YouComSearchModel { async fn do_search(&self, options: &SearchCallOptions) -> Result { let count = resolve_count(options.max_results); - let headers: Vec<(String, String)> = self - .build_headers(options.headers.as_ref()) - .into_iter() - .collect(); + let exchange = self.config.exchange(options.headers.as_ref()).await?; - let mut url = url::Url::parse(&self.endpoint()) + let mut url = url::Url::parse(&exchange.url("/v1/search")) .map_err(|e| AiMuxError::InvalidArgument(format!("invalid you_com endpoint: {e}")))?; url.query_pairs_mut() .append_pair("query", &options.query) .append_pair("count", &count.to_string()); let resp = aimux_provider_utils::get_from_api( - HttpRequest { - url: url.to_string(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - ..Default::default() - }, + exchange.request(url.to_string(), options), aimux_provider_utils::create_json_response_handler::(), - aimux_provider_utils::create_standard_json_error_response_handler(), + you_com_failed_response_handler(), ) .await?; diff --git a/aimux-providers/tests/alibaba_test.rs b/aimux-providers/tests/alibaba_test.rs index 6e8f6daf8..a1d2dcf0f 100644 --- a/aimux-providers/tests/alibaba_test.rs +++ b/aimux-providers/tests/alibaba_test.rs @@ -5,7 +5,7 @@ //! - `packages/alibaba/src/alibaba-chat-language-model.test.ts` //! - `packages/alibaba/src/convert-alibaba-usage.test.ts` //! -//! Alibaba is a thin OpenAI-compatible wrapper over [`OpenAIProvider`]. The +//! Alibaba is a registry preset of the OpenAI-compatible package. The //! behaviours verified here are the ones the wrapper is responsible for: //! //! - Provider configuration: name, `ALIBABA_API_KEY` env var, custom API key, @@ -28,6 +28,7 @@ use aimux_core::tool::RawToolCall; use futures::StreamExt; use serde_json::{Value, json}; use serial_test::serial; +use std::sync::Arc; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -39,7 +40,7 @@ use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; use aimux_core::types::FinishReasonUnified; -use aimux_providers::{ProviderOptions, provider, provider_from_env}; +use aimux_providers::{PresetSettings, create_provider}; // ── shared helpers ─────────────────────────────────────────────────────────── @@ -93,16 +94,17 @@ async fn collect_stream(result: aimux_core::result::StreamResult) -> Vec Box { - provider( +fn make_provider(server: &MockServer) -> Arc { + create_provider( "alibaba", - Some("test-api-key".to_string()), - "qwen-plus", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-api-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("qwen-plus") .expect("alibaba provider should build") } @@ -113,8 +115,16 @@ fn make_provider(server: &MockServer) -> Box { /// TS: `createAlibaba()` produces a provider whose name is "alibaba". #[test] fn provider_builds_alibaba_model() { - let model = provider("alibaba", Some("test-key".to_string()), "qwen-plus", None) - .expect("alibaba provider should build"); + let model = create_provider( + "alibaba", + PresetSettings { + api_key: Some("test-key".to_string().into()), + ..Default::default() + }, + ) + .expect("provider should build") + .language_model("qwen-plus") + .expect("alibaba provider should build"); assert_eq!(model.model_id(), "qwen-plus"); } @@ -130,15 +140,16 @@ async fn custom_api_key_used_in_auth_header() { .mount(&server) .await; - let model = provider( + let model = create_provider( "alibaba", - Some("my-custom-key".to_string()), - "qwen-plus", - Some(ProviderOptions { + PresetSettings { + api_key: Some("my-custom-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("qwen-plus") .expect("alibaba provider should build"); model @@ -159,15 +170,16 @@ async fn custom_headers_forwarded() { .mount(&server) .await; - let model = provider( + let model = create_provider( "alibaba", - Some("test-key".to_string()), - "qwen-plus", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("qwen-plus") .expect("alibaba provider should build"); let mut options = default_options(test_prompt()); @@ -194,15 +206,16 @@ async fn language_model_via_trait() { .mount(&server) .await; - let model = provider( + let model = create_provider( "alibaba", - Some("test-key".to_string()), - "qwen-plus", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("qwen-plus") .expect("provider should build a model"); model @@ -220,7 +233,9 @@ fn from_env_loads_alibaba_api_key() { std::env::set_var("ALIBABA_API_KEY", "env-test-key"); } - let model = provider_from_env("alibaba", "qwen-plus", None); + let model = create_provider("alibaba", PresetSettings::default()) + .expect("provider should build") + .language_model("qwen-plus"); assert!(model.is_ok(), "from_env should succeed with env var set"); unsafe { @@ -233,15 +248,19 @@ fn from_env_loads_alibaba_api_key() { /// TS: without the env var, `createAlibaba()` fails. #[serial] -#[test] -fn from_env_fails_without_env_var() { +#[tokio::test] +async fn from_env_fails_without_env_var() { let saved = std::env::var("ALIBABA_API_KEY").ok(); unsafe { std::env::remove_var("ALIBABA_API_KEY"); } - let model = provider_from_env("alibaba", "qwen-plus", None); - assert!(model.is_err(), "from_env should fail without env var"); + let model = create_provider("alibaba", PresetSettings::default()) + .expect("provider should build") + .language_model("qwen-plus") + .expect("provider should build a model"); + let result = model.do_generate(&default_options(test_prompt())).await; + assert!(result.is_err(), "from_env should fail without env var"); unsafe { if let Some(v) = saved { @@ -371,6 +390,10 @@ async fn extracts_usage_with_reasoning_tokens() { assert_eq!(result.usage.output_tokens.text, Some(315)); } +// ════════════════════════════════════════════════════════════════════════════ +// Reasoning (top-level reasoning �?reasoning_effort, shared OpenAI behaviour) +// ════════════════════════════════════════════════════════════════════════════ + // ════════════════════════════════════════════════════════════════════════════ // doGenerate / doStream / errors // ════════════════════════════════════════════════════════════════════════════ @@ -445,7 +468,7 @@ async fn do_generate_extracts_tool_call() { }) => { assert_eq!(tool_call_id, "call_abc"); assert_eq!(tool_name, "get-weather"); - assert_eq!(input, r#"{"city":"SF"}"#); + assert_eq!(input, &Value::String(r#"{"city":"SF"}"#.into())); } other => panic!("expected ToolCall, got {other:?}"), } @@ -554,3 +577,32 @@ async fn status_429_maps_to_rate_limited() { "expected RateLimited, got {result:?}" ); } + +/// TS: response headers are exposed on the generate result. +#[tokio::test] +async fn exposes_response_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("test-header", "test-value") + .set_body_json(text_completion_body()), + ) + .mount(&server) + .await; + + let model = make_provider(&server); + + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("should succeed"); + + let headers = result + .response + .as_ref() + .and_then(|response| response.headers.as_ref()) + .expect("response_headers should be Some"); + assert_eq!(headers.get("test-header"), Some(&"test-value".to_string())); +} diff --git a/aimux-providers/tests/amazon_bedrock_image_test.rs b/aimux-providers/tests/amazon_bedrock_image_test.rs index de55218a4..8a8d30ade 100644 --- a/aimux-providers/tests/amazon_bedrock_image_test.rs +++ b/aimux-providers/tests/amazon_bedrock_image_test.rs @@ -1,14 +1,13 @@ -//! Rust translation of the Amazon Bedrock image model tests. +//! Rust translation of the Amazon Bedrock image model tests. //! Source: `reference/ai/packages/amazon-bedrock/src/amazon-bedrock-image-model.test.ts` use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::image_model::{ - ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, -}; -use aimux_providers::{BedrockProvider, BedrockProviderConfig}; +use aimux_core::image_model::{ImageCallOptions, ImageFile, ImageFileData, ImageOutputs}; +use aimux_core::provider::Provider; +use aimux_providers::{AmazonBedrockProviderSettings, create_amazon_bedrock}; const PROMPT: &str = "A cute baby sea otter"; @@ -18,14 +17,19 @@ fn bedrock_response() -> Value { async fn mock_bedrock(server: &MockServer, body: Value) { Mock::given(method("POST")) - .and(path("/model/amazon.titan-image-generator-v1/invoke")) + .and(path("/model/amazon.nova-canvas-v1:0/invoke")) .respond_with(ResponseTemplate::new(200).set_body_json(body)) .mount(server) .await; } -fn config(server: &MockServer) -> BedrockProviderConfig { - BedrockProviderConfig::with_bearer_token("test-token", "us-east-1").with_base_url(server.uri()) +fn config(server: &MockServer) -> AmazonBedrockProviderSettings { + AmazonBedrockProviderSettings { + api_key: Some("test-token".to_string().into()), + region: Some("us-east-1".to_string()), + base_url: Some(server.uri()), + ..Default::default() + } } fn options(prompt: &str) -> ImageCallOptions { @@ -36,8 +40,8 @@ fn options(prompt: &str) -> ImageCallOptions { async fn should_extract_generated_images() { let server = MockServer::start().await; mock_bedrock(&server, bedrock_response()).await; - let provider = BedrockProvider::new(config(&server)); - let model = provider.image("amazon.titan-image-generator-v1"); + let provider = create_amazon_bedrock(config(&server)).unwrap(); + let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); let result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { ImageOutputs::Base64(imgs) => { @@ -52,8 +56,8 @@ async fn should_extract_generated_images() { async fn should_pass_prompt() { let server = MockServer::start().await; mock_bedrock(&server, bedrock_response()).await; - let provider = BedrockProvider::new(config(&server)); - let model = provider.image("amazon.titan-image-generator-v1"); + let provider = create_amazon_bedrock(config(&server)).unwrap(); + let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&reqs[0].body).unwrap(); @@ -65,8 +69,8 @@ async fn should_pass_prompt() { async fn should_pass_size_and_n() { let server = MockServer::start().await; mock_bedrock(&server, bedrock_response()).await; - let provider = BedrockProvider::new(config(&server)); - let model = provider.image("amazon.titan-image-generator-v1"); + let provider = create_amazon_bedrock(config(&server)).unwrap(); + let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); let mut opts = options(PROMPT); opts.n = 2; opts.size = Some(aimux_core::shared::Size::new(1024, 1024)); @@ -82,8 +86,8 @@ async fn should_pass_size_and_n() { async fn should_pass_seed() { let server = MockServer::start().await; mock_bedrock(&server, bedrock_response()).await; - let provider = BedrockProvider::new(config(&server)); - let model = provider.image("amazon.titan-image-generator-v1"); + let provider = create_amazon_bedrock(config(&server)).unwrap(); + let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); let mut opts = options(PROMPT); opts.n = 1; opts.seed = Some(42); @@ -97,8 +101,8 @@ async fn should_pass_seed() { async fn should_warn_for_aspect_ratio() { let server = MockServer::start().await; mock_bedrock(&server, bedrock_response()).await; - let provider = BedrockProvider::new(config(&server)); - let model = provider.image("amazon.titan-image-generator-v1"); + let provider = create_amazon_bedrock(config(&server)).unwrap(); + let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); let mut opts = options(PROMPT); opts.n = 1; opts.aspect_ratio = Some(aimux_core::shared::AspectRatio::new(16, 9)); @@ -106,31 +110,12 @@ async fn should_warn_for_aspect_ratio() { assert!(result.warnings.iter().any(|w| matches!(w, aimux_core::types::Warning::Unsupported { feature, .. } if feature == "aspectRatio"))); } -#[tokio::test] -async fn should_pass_bearer_auth() { - let server = MockServer::start().await; - mock_bedrock(&server, bedrock_response()).await; - let provider = BedrockProvider::new(config(&server)); - let model = provider.image("amazon.titan-image-generator-v1"); - model.do_generate(&options(PROMPT)).await.unwrap(); - let reqs = server.received_requests().await.unwrap(); - assert_eq!( - reqs[0] - .headers - .get("authorization") - .unwrap() - .to_str() - .unwrap(), - "Bearer test-token" - ); -} - #[tokio::test] async fn should_support_image_variation() { let server = MockServer::start().await; mock_bedrock(&server, bedrock_response()).await; - let provider = BedrockProvider::new(config(&server)); - let model = provider.image("amazon.titan-image-generator-v1"); + let provider = create_amazon_bedrock(config(&server)).unwrap(); + let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); let mut opts = options("Variation"); opts.n = 1; opts.files = Some(vec![ImageFile::File { @@ -148,8 +133,8 @@ async fn should_support_image_variation() { async fn should_support_inpainting() { let server = MockServer::start().await; mock_bedrock(&server, bedrock_response()).await; - let provider = BedrockProvider::new(config(&server)); - let model = provider.image("amazon.titan-image-generator-v1"); + let provider = create_amazon_bedrock(config(&server)).unwrap(); + let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); let mut opts = options("Inpaint"); opts.n = 1; opts.files = Some(vec![ImageFile::File { @@ -171,17 +156,21 @@ async fn should_support_inpainting() { async fn should_handle_moderated_response() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/model/amazon.titan-image-generator-v1/invoke")) + .and(path("/model/amazon.nova-canvas-v1:0/invoke")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "fe7256d1-50d9-4663-8592-85eaf002e80c", "status": "Request Moderated", - "details": { "Moderation Reasons": ["violence", "adult"] } + "result": null, + "progress": null, + "details": { "Moderation Reasons": ["Derivative Works Filter"] }, + "preview": null }))) .mount(&server) .await; - let provider = BedrockProvider::new(config(&server)); - let model = provider.image("amazon.titan-image-generator-v1"); + let provider = create_amazon_bedrock(config(&server)).unwrap(); + let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); let result = model.do_generate(&options(PROMPT)).await; assert!(result.is_err()); let err = result.unwrap_err().to_string(); - assert!(err.contains("moderated")); + assert!(err.contains("Amazon Bedrock request was moderated: Derivative Works Filter")); } diff --git a/aimux-providers/tests/anthropic_factory_test.rs b/aimux-providers/tests/anthropic_factory_test.rs new file mode 100644 index 000000000..e3fe240ce --- /dev/null +++ b/aimux-providers/tests/anthropic_factory_test.rs @@ -0,0 +1,632 @@ +//! `create_anthropic` against the recorded AI SDK behavior. +//! +//! `fixtures/aisdk/anthropic/*.json` hold what `@ai-sdk/anthropic`'s +//! `createAnthropic` sent over a mocked fetch for a fixed input: URL, headers, +//! body and the model's `provider` string. Each fixture is replayed here +//! through an injected [`Fetch`] (the factory's `fetch` setting), so the test +//! also proves the transport is the one the settings name. The expected +//! request comes from the fixture; the inputs are rebuilt from the fixture's +//! recorded `sdk.input`. +//! +//! The recorded `x-api-key` is redacted, so the key the case used is +//! substituted back. The recorded SDK and runtime user-agent identifiers are +//! replaced with the pinned provider package's user-agent suffix. + +use aimux_core::tool::RawToolCall; +use std::collections::{BTreeMap, VecDeque}; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::StreamExt; +use reqwest::StatusCode; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use serde_json::{Value, json}; + +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; +use aimux_core::options::CallOptions; +use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput}; +use aimux_core::shared::provider_namespace; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{FunctionTool, LanguageModelV4ToolChoice, Tool}; +use aimux_core::types::FinishReasonUnified; +use aimux_provider_utils::{Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse}; +use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; + +// ── injected transport ─────────────────────────────────────────────────────── + +const KEY: &str = "sk-test-fixture"; + +#[derive(Clone)] +struct Canned { + status: u16, + headers: Vec<(String, String)>, + body: Vec, +} + +/// One request as the transport saw it. +#[derive(Clone, Debug)] +struct Seen { + method: String, + url: String, + /// Lower-cased names (the `http` crate normalizes them). + headers: BTreeMap, + body: Vec, +} + +impl Seen { + fn json_body(&self) -> Value { + serde_json::from_slice(&self.body).expect("request body is JSON") + } +} + +#[derive(Default)] +struct MockFetch { + canned: Mutex>, + seen: Mutex>, +} + +impl MockFetch { + fn new(canned: Vec) -> Arc { + Arc::new(Self { + canned: Mutex::new(canned.into()), + seen: Mutex::default(), + }) + } + + fn seen(&self) -> Vec { + self.seen.lock().unwrap().clone() + } + + fn transport(self: &Arc) -> FetchFunction { + self.clone() + } +} + +#[async_trait] +impl Fetch for MockFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + self.seen.lock().unwrap().push(Seen { + method: request.method.to_string(), + url: request.url.to_string(), + headers: request + .headers + .iter() + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().expect("ASCII header").to_string(), + ) + }) + .collect(), + body: request.body.to_vec(), + }); + let canned = self + .canned + .lock() + .unwrap() + .pop_front() + .ok_or_else(|| FetchError::Other("no canned response left".into()))?; + let mut headers = HeaderMap::new(); + for (name, value) in &canned.headers { + headers.insert( + HeaderName::try_from(name.as_str()).unwrap(), + HeaderValue::try_from(value.as_str()).unwrap(), + ); + } + Ok(FetchResponse::from_bytes( + StatusCode::from_u16(canned.status).unwrap(), + headers, + request.url, + Bytes::from(canned.body), + )) + } +} + +fn settings(mock: &Arc) -> AnthropicProviderSettings { + AnthropicProviderSettings { + api_key: Some(KEY.to_string()), + fetch: Some(mock.transport()), + ..Default::default() + } +} + +// ── fixtures ───────────────────────────────────────────────────────────────── + +/// Every fixture under `fixtures/aisdk/anthropic/`; the directory listing is +/// compared to this so a new fixture cannot go unreplayed. +const FIXTURES: [&str; 5] = [ + "messages-basic", + "messages-custom-name", + "messages-provider-options", + "messages-stream-basic", + "messages-system-tools", +]; + +fn fixtures_dir() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../fixtures/aisdk/anthropic") +} + +struct Fixture { + name: &'static str, + json: Value, +} + +impl Fixture { + fn load(name: &'static str) -> Self { + assert!(FIXTURES.contains(&name), "unlisted fixture {name}"); + let path = fixtures_dir().join(format!("{name}.json")); + let text = std::fs::read_to_string(&path) + .unwrap_or_else(|e| panic!("cannot read {}: {e}", path.display())); + Self { + name, + json: serde_json::from_str(&text).expect("fixture is JSON"), + } + } + + fn provider_string(&self) -> &str { + self.json["model"]["provider"].as_str().unwrap() + } + + fn model_id(&self) -> &str { + self.json["model"]["modelId"].as_str().unwrap() + } + + fn sdk_input(&self) -> &Value { + &self.json["sdk"]["input"] + } + + /// The transport answering with the recorded response. + fn mock(&self) -> Arc { + let response = &self.json["response"]; + let body = match &response["body"] { + Value::String(text) => text.clone().into_bytes(), + other => serde_json::to_vec(other).unwrap(), + }; + MockFetch::new(vec![Canned { + status: response["status"].as_u64().unwrap() as u16, + headers: response["headers"] + .as_object() + .unwrap() + .iter() + .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string())) + .collect(), + body, + }]) + } + + /// The recorded request must equal what went out: method, URL, headers + /// (names case-insensitively, with the pinned user-agent) and body. + fn assert_request(&self, seen: &Seen) { + let recorded = &self.json["request"]; + assert_eq!( + seen.method, + recorded["method"].as_str().unwrap(), + "{}", + self.name + ); + assert_eq!(seen.url, recorded["url"].as_str().unwrap(), "{}", self.name); + + let mut expected: BTreeMap = recorded["headers"] + .as_object() + .unwrap() + .iter() + .map(|(k, v)| (k.to_ascii_lowercase(), v.as_str().unwrap().to_string())) + .collect(); + expected.insert( + "user-agent".to_string(), + "ai-sdk-anthropic/4.0.68 ai-sdk-provider-utils/5.0.51".to_string(), + ); + assert_eq!( + expected.get("x-api-key").map(String::as_str), + Some(""), + "{}: the recording redacts the key", + self.name + ); + expected.insert("x-api-key".to_string(), KEY.to_string()); + assert_eq!(seen.headers, expected, "{}: headers", self.name); + + assert_eq!(seen.json_body(), recorded["body"], "{}: body", self.name); + } + + /// The call options the recording was made with. + fn call_options(&self) -> CallOptions { + call_options_from(self.sdk_input()) + } +} + +/// Rebuild aimux call options from a recorded `sdk.input` (the subset the +/// Anthropic fixtures use). +fn call_options_from(input: &Value) -> CallOptions { + let mut prompt: LanguageModelPrompt = Vec::new(); + if let Some(system) = input.get("system").and_then(Value::as_str) { + prompt.push(LanguageModelMessage::System { + content: system.to_string(), + provider_options: None, + }); + } + if let Some(text) = input.get("prompt").and_then(Value::as_str) { + prompt.push(LanguageModelMessage::user_text(text)); + } + let mut options = CallOptions::new(prompt); + options.max_output_tokens = input + .get("maxOutputTokens") + .and_then(Value::as_u64) + .map(|v| v as u32); + if let Some(tools) = input.get("tools").and_then(Value::as_object) { + options.tools = Some( + tools + .iter() + .map(|(name, tool)| { + let mut function = FunctionTool::new(name.clone(), tool["inputSchema"].clone()); + if let Some(description) = tool["description"].as_str() { + function = function.with_description(description); + } + Tool::Function(function) + }) + .collect(), + ); + } + options.tool_choice = match input.get("toolChoice").and_then(Value::as_str) { + Some("required") => Some(LanguageModelV4ToolChoice::Required), + Some("none") => Some(LanguageModelV4ToolChoice::None), + Some("auto") => Some(LanguageModelV4ToolChoice::Auto), + None => options + .tools + .as_ref() + .filter(|tools| !tools.is_empty()) + .map(|_| LanguageModelV4ToolChoice::Auto), + Some(other) => panic!("unmapped toolChoice {other}"), + }; + options.provider_options = input + .get("providerOptions") + .map(|options| serde_json::from_value(options.clone()).unwrap()); + options +} + +fn generated_text(result: &GenerateResult) -> String { + result + .content + .iter() + .filter_map(|part| match part { + GenerateContent::Text { text, .. } => Some(text.as_str()), + _ => None, + }) + .collect() +} + +/// Result fields every recorded `generateText` case pins. +fn assert_generate_result(fixture: &Fixture, result: &GenerateResult) { + let recorded = &fixture.json["result"]; + assert_eq!( + generated_text(result), + recorded["text"].as_str().unwrap(), + "{}: text", + fixture.name + ); + assert_eq!( + result.response.as_ref().and_then(|r| r.id.as_deref()), + recorded["response"]["id"].as_str(), + "{}: 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 _), + "{}: input tokens", + fixture.name + ); + assert_eq!( + result.usage.output_tokens.total, + recorded["usage"]["outputTokens"].as_u64().map(|v| v as _), + "{}: output tokens", + fixture.name + ); + let unified = match recorded["finishReason"].as_str().unwrap() { + "stop" => FinishReasonUnified::Stop, + "tool-calls" => FinishReasonUnified::ToolCalls, + other => panic!("unmapped finish reason {other}"), + }; + assert_eq!( + result.finish_reason.unified, unified, + "{}: finish", + fixture.name + ); + assert_eq!( + result.finish_reason.raw.as_deref(), + recorded["rawFinishReason"].as_str(), + "{}: raw finish", + fixture.name + ); + // The result-level providerMetadata: the raw usage, stop sequence, + // iterations, container and context management, under the canonical key + // and the provider's own. + assert_eq!( + serde_json::to_value(&result.provider_metadata).unwrap(), + recorded["providerMetadata"], + "{}: providerMetadata", + fixture.name + ); +} + +// ── the fixtures ───────────────────────────────────────────────────────────── + +#[test] +fn every_fixture_is_replayed() { + let mut on_disk: Vec = std::fs::read_dir(fixtures_dir()) + .expect("fixtures/aisdk/anthropic exists") + .map(|entry| { + let name = entry.unwrap().file_name().into_string().unwrap(); + name.strip_suffix(".json") + .unwrap_or_else(|| panic!("unexpected file {name}")) + .to_string() + }) + .collect(); + on_disk.sort(); + assert_eq!( + on_disk, FIXTURES, + "add a test for the new fixture and list it" + ); +} + +/// Replay a fixture whose request is a plain `generateText`. +async fn replay_generate(name: &'static str) -> (Fixture, Arc, GenerateResult) { + let fixture = Fixture::load(name); + let mock = fixture.mock(); + let provider = create_anthropic(settings(&mock)).unwrap(); + let model = provider.messages(fixture.model_id()); + assert_eq!( + model.provider(), + fixture.provider_string(), + "{name}: provider()" + ); + assert_eq!(model.model_id(), fixture.model_id()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1, "{name}: exactly one request"); + fixture.assert_request(&seen[0]); + assert_generate_result(&fixture, &result); + (fixture, mock, result) +} + +#[tokio::test] +async fn messages_basic() { + replay_generate("messages-basic").await; +} + +#[tokio::test] +async fn messages_provider_options_use_the_canonical_namespace() { + let (_, _, result) = replay_generate("messages-provider-options").await; + // The reasoning block carries its signature under the canonical key. + let reasoning = result + .content + .iter() + .find_map(|part| match part { + GenerateContent::Reasoning(ReasoningOutput { + text, + provider_metadata, + }) => Some((text.clone(), provider_metadata.clone())), + _ => None, + }) + .expect("a reasoning part"); + assert_eq!(reasoning.0, "17 * 23 = 391."); + assert_eq!( + reasoning.1, + Some(provider_namespace("anthropic", json!({ "signature": "sig_fixture" })).unwrap()) + ); +} + +#[tokio::test] +async fn messages_system_and_tools() { + let (fixture, _, result) = replay_generate("messages-system-tools").await; + let recorded = &fixture.json["result"]["content"][0]; + let call = result + .content + .iter() + .find_map(|part| match part { + GenerateContent::ToolCall(RawToolCall { + tool_call_id, + tool_name, + input, + .. + }) => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), + _ => None, + }) + .expect("a tool call"); + assert_eq!(call.0, recorded["toolCallId"].as_str().unwrap()); + assert_eq!(call.1, recorded["toolName"].as_str().unwrap()); + assert_eq!( + serde_json::from_str::(&call.2).unwrap(), + recorded["input"] + ); +} + +#[tokio::test] +async fn messages_custom_name_reads_both_namespaces_and_the_custom_one_wins() { + let fixture = Fixture::load("messages-custom-name"); + let mock = fixture.mock(); + let provider = create_anthropic(AnthropicProviderSettings { + name: Some("myproxy".to_string()), + ..settings(&mock) + }) + .unwrap(); + let model = provider.messages(fixture.model_id()); + // The name is the provider string verbatim, with no `.messages` added. + assert_eq!(model.provider(), "myproxy"); + assert_eq!(model.provider(), fixture.provider_string()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + // The name changes identity and the options namespace, never the URL; the + // recorded body carries the `myproxy` user id, not the `anthropic` one. + fixture.assert_request(&mock.seen()[0]); + assert_generate_result(&fixture, &result); + assert_eq!( + fixture.json["observations"]["bodyMetadataUserId"], + json!("from-myproxy-key") + ); + assert_eq!( + fixture.json["observations"]["modelProvider"], + json!("myproxy") + ); +} + +/// `includeRawChunks`: every SSE chunk is forwarded as a `Raw` part before +/// the parts made from it, a chunk that is not JSON included (it has no raw +/// value), as anthropic@4.0.71 `anthropic-language-model.ts:1956-1963` does; +/// that chunk is then reported as an error and the stream goes on. Without +/// the option no `Raw` part is sent. +#[tokio::test] +async fn messages_stream_emits_raw_chunks_when_asked() { + let fixture = Fixture::load("messages-stream-basic"); + let mock = fixture.mock(); + let body = { + let mut canned = mock.canned.lock().unwrap(); + let body = String::from_utf8(canned[0].body.clone()).unwrap(); + let (first, rest) = body.split_once("\n\n").unwrap(); + let body = format!("{first}\n\ndata: {{not json\n\n{rest}"); + canned[0].body = body.clone().into_bytes(); + body + }; + let expected: Vec = body + .lines() + .filter_map(|line| line.strip_prefix("data: ")) + .map(|data| serde_json::from_str(data).unwrap_or(Value::Null)) + .collect(); + let model = create_anthropic(settings(&mock)) + .unwrap() + .messages(fixture.model_id()); + let options = CallOptions { + include_raw_chunks: Some(true), + ..fixture.call_options() + }; + let items: Vec<_> = model + .do_stream(&options) + .await + .unwrap() + .stream + .collect() + .await; + + let raw: Vec = items + .iter() + .filter_map(|item| match item { + Ok(StreamPart::Raw { raw_value }) => Some(raw_value.clone()), + _ => None, + }) + .collect(); + assert_eq!(raw, expected); + let not_json = items + .iter() + .position(|item| { + matches!( + item, + Ok(StreamPart::Raw { + raw_value: Value::Null + }) + ) + }) + .unwrap(); + assert!( + matches!(&items[not_json + 1], Err(error) if error.is_recoverable_stream_error()), + "{:?}", + items[not_json + 1] + ); + let first_delta = items + .iter() + .position(|item| matches!(item, Ok(StreamPart::TextDelta { .. }))) + .unwrap(); + // The chunk a delta was made from is the last raw part before it. + let source = items[..first_delta] + .iter() + .rev() + .find_map(|item| match item { + Ok(StreamPart::Raw { raw_value }) => Some(raw_value), + _ => None, + }); + assert_eq!(source.unwrap()["delta"]["text"], "Hello"); + let text: String = items + .iter() + .filter_map(|item| match item { + Ok(StreamPart::TextDelta { delta, .. }) => Some(delta.as_str()), + _ => None, + }) + .collect(); + assert_eq!(text, "Hello world"); + + let model = create_anthropic(settings(&fixture.mock())) + .unwrap() + .messages(fixture.model_id()); + let result = model.do_stream(&fixture.call_options()).await.unwrap(); + let parts: Vec<_> = result.stream.collect().await; + assert!( + !parts + .iter() + .any(|item| matches!(item, Ok(StreamPart::Raw { .. }))) + ); +} + +#[tokio::test] +async fn messages_stream_basic() { + let fixture = Fixture::load("messages-stream-basic"); + let mock = fixture.mock(); + let model = create_anthropic(settings(&mock)) + .unwrap() + .messages(fixture.model_id()); + assert_eq!(model.provider(), fixture.provider_string()); + + let result = model.do_stream(&fixture.call_options()).await.unwrap(); + let parts: Vec = result + .stream + .map(|part| part.expect("no stream error")) + .collect() + .await; + + // The recording's user-agent names the provider package; the stream case + // is checked like the others. + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + fixture.assert_request(&seen[0]); + + let text: String = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextDelta { delta, .. } => Some(delta.as_str()), + _ => None, + }) + .collect(); + assert_eq!(text, "Hello world"); + let finish = parts.iter().find_map(|part| match part { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => Some((finish_reason, usage, provider_metadata)), + _ => None, + }); + let (reason, usage, metadata) = finish.expect("a finish part"); + // `message_start`'s usage updated by `message_delta`'s. + let finish_step = fixture.json["result"]["parts"] + .as_array() + .and_then(|parts| parts.iter().find(|p| p["type"] == "finish-step")) + .expect("a recorded finish-step"); + assert_eq!( + serde_json::to_value(metadata).unwrap(), + finish_step["providerMetadata"] + ); + assert_eq!(usage.raw.as_ref(), finish_step["usage"]["raw"].as_object()); + assert_eq!(reason.unified, FinishReasonUnified::Stop); + assert_eq!(reason.raw.as_deref(), Some("end_turn")); + assert_eq!(usage.input_tokens.total, Some(12)); + assert_eq!(usage.output_tokens.total, Some(2)); +} diff --git a/aimux-providers/tests/anthropic_files_test.rs b/aimux-providers/tests/anthropic_files_test.rs index 0dbc8f7e4..3b984e53e 100644 --- a/aimux-providers/tests/anthropic_files_test.rs +++ b/aimux-providers/tests/anthropic_files_test.rs @@ -7,9 +7,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::files_model::{Files, UploadFileCallOptions, UploadFileData}; +use aimux_core::files_model::{UploadFileCallOptions, UploadFileData}; +use aimux_core::provider::Provider; use aimux_core::shared::FileBytes; -use aimux_providers::{AnthropicConfig, AnthropicProvider}; +use aimux_providers::{AnthropicProvider, AnthropicProviderSettings, create_anthropic}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -28,8 +29,12 @@ fn file_response_body() -> Value { /// Build an `AnthropicProvider` pointing at the mock server. fn provider(server: &MockServer) -> AnthropicProvider { - let config = AnthropicConfig::new("test-api-key").with_base_url(server.uri()); - AnthropicProvider::new(config) + create_anthropic(AnthropicProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap() } /// Build `UploadFileCallOptions` with binary data. @@ -57,7 +62,7 @@ async fn sends_post_to_v1_files_with_beta_header() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); files.upload_file(&upload_options()).await.unwrap(); @@ -83,7 +88,7 @@ async fn sends_x_api_key_header() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); files.upload_file(&upload_options()).await.unwrap(); @@ -107,7 +112,7 @@ async fn sends_multipart_form_data_with_file() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); files.upload_file(&upload_options()).await.unwrap(); @@ -129,7 +134,7 @@ async fn uses_default_filename_blob_when_not_specified() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); files.upload_file(&upload_options()).await.unwrap(); @@ -151,7 +156,7 @@ async fn uses_custom_filename_from_options() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let mut opts = upload_options(); opts.filename = Some("custom-name.pdf".to_string()); @@ -175,7 +180,7 @@ async fn uses_media_type_from_options() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let mut opts = upload_options(); opts.media_type = "application/pdf".to_string(); @@ -199,7 +204,7 @@ async fn returns_provider_reference_with_anthropic_key() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await.unwrap(); @@ -219,7 +224,7 @@ async fn returns_provider_metadata_with_response_data() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await.unwrap(); @@ -251,7 +256,7 @@ async fn omits_downloadable_when_null() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await.unwrap(); @@ -273,7 +278,7 @@ async fn handles_base64_string_data() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let opts = UploadFileCallOptions { data: UploadFileData::Data { @@ -295,16 +300,25 @@ async fn handles_base64_string_data() { #[tokio::test] async fn has_specification_version_v4() { - let provider = AnthropicProvider::new(AnthropicConfig::new("test-api-key")); - let files = provider.files(); + let provider = create_anthropic(AnthropicProviderSettings { + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap(); + let files = Provider::files(&provider).unwrap(); assert_eq!(files.specification_version(), "v4"); } #[tokio::test] async fn has_correct_provider_name() { - let provider = AnthropicProvider::new(AnthropicConfig::new("test-api-key")); - let files = provider.files(); + let provider = create_anthropic(AnthropicProviderSettings { + api_key: Some("test-api-key".to_string()), + name: Some("anthropic.files".to_string()), + ..Default::default() + }) + .unwrap(); + let files = Provider::files(&provider).unwrap(); assert_eq!(files.provider(), "anthropic.files"); } diff --git a/aimux-providers/tests/assemblyai_transcription_test.rs b/aimux-providers/tests/assemblyai_transcription_test.rs index afd5a7966..438226ae2 100644 --- a/aimux-providers/tests/assemblyai_transcription_test.rs +++ b/aimux-providers/tests/assemblyai_transcription_test.rs @@ -7,8 +7,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +use aimux_core::AiMuxError; use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; -use aimux_providers::{AssemblyAIConfig, AssemblyAIProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{AssemblyAIProviderSettings, create_assemblyai}; fn mock_audio() -> Vec { vec![1u8, 2, 3] @@ -66,8 +68,12 @@ async fn should_extract_text() { let server = MockServer::start().await; mock_all(&server, &fixture_transcript(), &[]).await; - let config = AssemblyAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = AssemblyAIProvider::new(config); + let config = AssemblyAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_assemblyai(config).unwrap(); let model = provider.transcription("universal-2"); let result = model @@ -89,8 +95,12 @@ async fn should_pass_model_as_speech_models() { let server = MockServer::start().await; mock_all(&server, &fixture_transcript(), &[]).await; - let config = AssemblyAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = AssemblyAIProvider::new(config); + let config = AssemblyAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_assemblyai(config).unwrap(); let model = provider.transcription("universal-2"); model @@ -109,8 +119,12 @@ async fn should_use_speech_model_for_best() { let server = MockServer::start().await; mock_all(&server, &fixture_transcript(), &[]).await; - let config = AssemblyAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = AssemblyAIProvider::new(config); + let config = AssemblyAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_assemblyai(config).unwrap(); let model = provider.transcription("best"); let result = model @@ -134,10 +148,13 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = AssemblyAIConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = AssemblyAIProvider::new(config); + let config = AssemblyAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_assemblyai(config).unwrap(); let model = provider.transcription("universal-2"); let mut opts = options(mock_audio(), "audio/wav"); @@ -173,8 +190,12 @@ async fn should_include_response_data() { ) .await; - let config = AssemblyAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = AssemblyAIProvider::new(config); + let config = AssemblyAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_assemblyai(config).unwrap(); let model = provider.transcription("universal-2"); let result = model @@ -191,8 +212,12 @@ async fn should_use_real_date() { let server = MockServer::start().await; mock_all(&server, &fixture_transcript(), &[]).await; - let config = AssemblyAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = AssemblyAIProvider::new(config); + let config = AssemblyAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_assemblyai(config).unwrap(); let model = provider.transcription("universal-2"); let result = model @@ -217,8 +242,12 @@ async fn should_include_provider_metadata_with_utterances() { }); mock_all(&server, &transcript, &[]).await; - let config = AssemblyAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = AssemblyAIProvider::new(config); + let config = AssemblyAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_assemblyai(config).unwrap(); let model = provider.transcription("universal-2"); let result = model @@ -232,14 +261,11 @@ async fn should_include_provider_metadata_with_utterances() { assert!(aai.get("utterances").is_some()); } -/// A transient 503 on the transcript-submit stage, followed by 200, must -/// succeed without re-uploading the audio: the upload exchange is retried -/// independently of the submit exchange, so Core's outer per-`do_generate` -/// retry never needs to (and must not) replay the upload. +/// A 503 on the transcript-submit stage is reported as it happened: the submit +/// is not retried inside the provider, and the upload before it is not replayed. #[tokio::test] -async fn transient_submit_failure_is_retried_without_re_uploading() { +async fn a_failed_submit_is_not_retried_and_the_upload_is_not_replayed() { let server = MockServer::start().await; - let upload_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); let upload_observed = std::sync::Arc::clone(&upload_attempts); Mock::given(method("POST")) @@ -257,38 +283,29 @@ async fn transient_submit_failure_is_retried_without_re_uploading() { Mock::given(method("POST")) .and(path("/v2/transcript")) .respond_with(move |_: &wiremock::Request| { - if submit_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { - ResponseTemplate::new(503) - .insert_header("retry-after-ms", "0") - .set_body_json(json!({"error": "try again"})) - } else { - ResponseTemplate::new(200) - .set_body_json(json!({"id": "test-id", "status": "queued"})) - } + submit_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + ResponseTemplate::new(503) + .insert_header("retry-after-ms", "0") + .set_body_json(json!({"error": "try again"})) }) .mount(&server) .await; - Mock::given(method("GET")) - .and(path("/v2/transcript/test-id")) - .respond_with(ResponseTemplate::new(200).set_body_json(fixture_transcript())) - .mount(&server) - .await; - - let config = AssemblyAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = AssemblyAIProvider::new(config); - let model = provider.transcription("universal-2"); - - let result = model + let provider = create_assemblyai(AssemblyAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }) + .unwrap(); + let error = provider + .transcription("universal-2") .do_generate(&options(mock_audio(), "audio/wav")) .await - .unwrap(); + .unwrap_err(); - assert_eq!(result.text, "Hello from AssemblyAI."); - assert_eq!( - upload_attempts.load(std::sync::atomic::Ordering::SeqCst), - 1, - "the upload stage must not be replayed by a submit-stage retry" - ); - assert_eq!(submit_attempts.load(std::sync::atomic::Ordering::SeqCst), 2); + // Creating the transcript is the stage that bills: it is sent once and the + // failure is reported as it happened; the upload before it is not replayed. + assert!(matches!(error, AiMuxError::ApiCall(ref d) if d.status_code == Some(503))); + assert_eq!(upload_attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!(submit_attempts.load(std::sync::atomic::Ordering::SeqCst), 1); } diff --git a/aimux-providers/tests/aws_polly_test.rs b/aimux-providers/tests/aws_polly_test.rs index 5541ea346..6d4f33982 100644 --- a/aimux-providers/tests/aws_polly_test.rs +++ b/aimux-providers/tests/aws_polly_test.rs @@ -12,7 +12,7 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::speech_model::{AudioData, SpeechCallOptions, SpeechModel}; -use aimux_providers::{AwsPollyConfig, AwsPollyProvider}; +use aimux_providers::{AwsPollyProviderSettings, create_aws_polly}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -29,8 +29,14 @@ fn mock_audio_bytes() -> Vec { bytes } -fn test_config(server: &MockServer) -> AwsPollyConfig { - AwsPollyConfig::new(TEST_ACCESS_KEY, TEST_SECRET_KEY, TEST_REGION).with_base_url(server.uri()) +fn test_config(server: &MockServer) -> AwsPollyProviderSettings { + AwsPollyProviderSettings { + access_key_id: Some(TEST_ACCESS_KEY.to_string()), + secret_access_key: Some(TEST_SECRET_KEY.to_string()), + region: Some(TEST_REGION.to_string()), + base_url: Some(server.uri().to_string()), + ..Default::default() + } } fn speech_options(text: &str) -> SpeechCallOptions { @@ -59,7 +65,7 @@ async fn do_generate_returns_binary_audio() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); let result = model @@ -79,7 +85,7 @@ async fn requests_correct_url() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); model @@ -99,7 +105,7 @@ async fn request_body_carries_engine_text_voice_and_format() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); let mut options = speech_options("Hello from the AI SDK!"); @@ -129,7 +135,7 @@ async fn model_id_maps_to_engine() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech(model_id); model.do_generate(&speech_options("Hello")).await.unwrap(); @@ -149,7 +155,7 @@ async fn request_carries_sigv4_authorization_header() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); model @@ -181,7 +187,7 @@ async fn request_carries_sigv4_supporting_headers() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); model @@ -213,10 +219,15 @@ async fn session_token_is_forwarded_and_signed() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = AwsPollyConfig::new(TEST_ACCESS_KEY, TEST_SECRET_KEY, TEST_REGION) - .with_base_url(server.uri()) - .with_session_token("test-session-token"); - let provider = AwsPollyProvider::new(config); + let config = AwsPollyProviderSettings { + access_key_id: Some(TEST_ACCESS_KEY.to_string()), + secret_access_key: Some(TEST_SECRET_KEY.to_string()), + region: Some(TEST_REGION.to_string()), + base_url: Some(server.uri().to_string()), + session_token: Some("test-session-token".to_string()), + ..Default::default() + }; + let provider = create_aws_polly(config).unwrap(); let model = provider.speech("aws_polly/neural"); model @@ -248,7 +259,7 @@ async fn language_maps_to_language_code() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); let mut options = speech_options("Hola mundo"); @@ -275,7 +286,7 @@ async fn response_carries_metadata() { .mount(&server) .await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); let result = model @@ -314,7 +325,7 @@ async fn error_401_maps_to_auth_error() { .mount(&server) .await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); let err = model @@ -346,7 +357,7 @@ async fn error_403_keeps_the_observed_status() { .mount(&server) .await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); let err = model @@ -373,7 +384,7 @@ async fn error_404_maps_to_model_not_found() { .mount(&server) .await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); let err = model @@ -393,7 +404,7 @@ async fn unsupported_output_format_emits_warning() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let provider = AwsPollyProvider::new(test_config(&server)); + let provider = create_aws_polly(test_config(&server)).unwrap(); let model = provider.speech("aws_polly/neural"); let mut options = speech_options("Hello"); diff --git a/aimux-providers/tests/baseten_test.rs b/aimux-providers/tests/baseten_test.rs index 7e902bf24..df7db3cf6 100644 --- a/aimux-providers/tests/baseten_test.rs +++ b/aimux-providers/tests/baseten_test.rs @@ -2,7 +2,7 @@ //! //! Translated from `packages/baseten/src/baseten-provider.unit.test.ts`. //! -//! Baseten is a thin OpenAI-compatible wrapper over [`OpenAIProvider`]. The TS +//! Baseten is a registry preset of the OpenAI-compatible package. The TS //! suite is almost entirely provider-configuration unit tests (mocked //! constructor calls asserting the base URL, env var, headers, and model-URL //! routing). The Rust wrapper exposes the same configuration surface, so these @@ -23,6 +23,7 @@ use serde_json::{Value, json}; use serial_test::serial; +use std::sync::Arc; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -30,7 +31,7 @@ use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; -use aimux_providers::{ProviderOptions, provider, provider_from_env}; +use aimux_providers::{PresetSettings, create_provider}; fn test_prompt() -> LanguageModelPrompt { vec![LanguageModelMessage::user_text("Hello")] @@ -55,28 +56,32 @@ fn text_completion_body() -> Value { }) } -fn make_provider(server: &MockServer) -> Box { - provider( +fn make_provider(server: &MockServer) -> Arc { + create_provider( "baseten", - Some("test-api-key".to_string()), - "deepseek-ai/DeepSeek-V3-0324", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-api-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("deepseek-ai/DeepSeek-V3-0324") .expect("baseten provider should build") } /// TS: `createBaseten()` produces a provider whose name is "baseten". #[test] fn provider_builds_baseten_model() { - let model = provider( + let model = create_provider( "baseten", - Some("test-key".to_string()), - "deepseek-ai/DeepSeek-V3-0324", - None, + PresetSettings { + api_key: Some("test-key".to_string().into()), + ..Default::default() + }, ) + .expect("provider should build") + .language_model("deepseek-ai/DeepSeek-V3-0324") .expect("baseten provider should build"); assert_eq!(model.model_id(), "deepseek-ai/DeepSeek-V3-0324"); } @@ -115,15 +120,16 @@ async fn custom_api_key_used_in_auth_header() { .mount(&server) .await; - let model = provider( + let model = create_provider( "baseten", - Some("my-custom-key".to_string()), - "deepseek-ai/DeepSeek-V3-0324", - Some(ProviderOptions { + PresetSettings { + api_key: Some("my-custom-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("deepseek-ai/DeepSeek-V3-0324") .expect("baseten provider should build"); model @@ -144,15 +150,16 @@ async fn custom_headers_forwarded() { .mount(&server) .await; - let model = provider( + let model = create_provider( "baseten", - Some("test-key".to_string()), - "deepseek-ai/DeepSeek-V3-0324", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("deepseek-ai/DeepSeek-V3-0324") .expect("baseten provider should build"); let mut options = default_options(test_prompt()); @@ -179,15 +186,16 @@ async fn language_model_via_trait() { .mount(&server) .await; - let model = provider( + let model = create_provider( "baseten", - Some("test-key".to_string()), - "deepseek-ai/DeepSeek-V3-0324", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("deepseek-ai/DeepSeek-V3-0324") .expect("provider should build a model"); model @@ -205,7 +213,9 @@ fn from_env_loads_baseten_api_key() { std::env::set_var("BASETEN_API_KEY", "env-test-key"); } - let model = provider_from_env("baseten", "deepseek-ai/DeepSeek-V3-0324", None); + let model = create_provider("baseten", PresetSettings::default()) + .expect("provider should build") + .language_model("deepseek-ai/DeepSeek-V3-0324"); assert!(model.is_ok(), "from_env should succeed with env var set"); unsafe { @@ -218,15 +228,19 @@ fn from_env_loads_baseten_api_key() { /// TS: without the env var, `createBaseten()` fails. #[serial] -#[test] -fn from_env_fails_without_env_var() { +#[tokio::test] +async fn from_env_fails_without_env_var() { let saved = std::env::var("BASETEN_API_KEY").ok(); unsafe { std::env::remove_var("BASETEN_API_KEY"); } - let model = provider_from_env("baseten", "deepseek-ai/DeepSeek-V3-0324", None); - assert!(model.is_err(), "from_env should fail without env var"); + let model = create_provider("baseten", PresetSettings::default()) + .expect("provider should build") + .language_model("deepseek-ai/DeepSeek-V3-0324") + .expect("provider should build a model"); + let result = model.do_generate(&default_options(test_prompt())).await; + assert!(result.is_err(), "from_env should fail without env var"); unsafe { if let Some(v) = saved { diff --git a/aimux-providers/tests/bedrock_convert_test.rs b/aimux-providers/tests/bedrock_convert_test.rs index fd3ded59e..e4c1dd133 100644 --- a/aimux-providers/tests/bedrock_convert_test.rs +++ b/aimux-providers/tests/bedrock_convert_test.rs @@ -7,7 +7,7 @@ //! //! Cases the Rust data model cannot express are skipped with an inline //! comment. The main categories of skips: -//! - **System-after-non-system throw**: the converter currently lifts +//! - **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. diff --git a/aimux-providers/tests/bedrock_embedding_test.rs b/aimux-providers/tests/bedrock_embedding_test.rs index ba1a1f5bb..e9727dad0 100644 --- a/aimux-providers/tests/bedrock_embedding_test.rs +++ b/aimux-providers/tests/bedrock_embedding_test.rs @@ -10,8 +10,11 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; -use aimux_providers::{BedrockAuth, BedrockProvider, BedrockProviderConfig}; +use aimux_core::embedding_model::EmbeddingCallOptions; +use aimux_core::provider::Provider; +use aimux_providers::{ + AmazonBedrockProvider, AmazonBedrockProviderSettings, create_amazon_bedrock, +}; const TEST_VALUES: &[&str] = &["sunny day at the beach", "rainy day in the city"]; @@ -22,15 +25,14 @@ fn mock_embeddings() -> Vec> { ] } -fn test_provider(base_url: String) -> BedrockProvider { - let config = BedrockProviderConfig { - base_url, - auth: BedrockAuth::BearerToken("test-auth".to_string()), - region: "us-east-1".to_string(), - retry_config: aimux_provider_utils::RetryConfig::default(), - api_key_source: None, - }; - BedrockProvider::new(config) +fn test_provider(base_url: String) -> AmazonBedrockProvider { + create_amazon_bedrock(AmazonBedrockProviderSettings { + base_url: Some(base_url), + api_key: Some("test-auth".to_string().into()), + region: Some("us-east-1".to_string()), + ..Default::default() + }) + .unwrap() } fn default_options(values: Vec) -> EmbeddingCallOptions { @@ -49,16 +51,20 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { fn should_expose_max_embeddings_per_call() { let provider = test_provider("https://example.com".to_string()); - let titan_model = provider.embedding_model("amazon.titan-embed-text-v2:0"); + let titan_model = provider + .embedding_model("amazon.titan-embed-text-v2:0") + .unwrap(); assert_eq!(titan_model.max_embeddings_per_call(), Some(1)); - let cohere_model = provider.embedding_model("cohere.embed-english-v3"); + let cohere_model = provider.embedding_model("cohere.embed-english-v3").unwrap(); assert_eq!(cohere_model.max_embeddings_per_call(), Some(96)); - let cohere_v4_us_model = provider.embedding_model("us.cohere.embed-v4:0"); + let cohere_v4_us_model = provider.embedding_model("us.cohere.embed-v4:0").unwrap(); assert_eq!(cohere_v4_us_model.max_embeddings_per_call(), Some(96)); - let nova_model = provider.embedding_model("amazon.nova-2-multimodal-embeddings-v1:0"); + let nova_model = provider + .embedding_model("amazon.nova-2-multimodal-embeddings-v1:0") + .unwrap(); assert_eq!(nova_model.max_embeddings_per_call(), Some(1)); } @@ -80,7 +86,9 @@ async fn should_handle_titan_single_input() { .await; let provider = test_provider(server.uri()); - let model = provider.embedding_model("amazon.titan-embed-text-v2:0"); + let model = provider + .embedding_model("amazon.titan-embed-text-v2:0") + .unwrap(); let result = model .do_embed(&default_options(vec![TEST_VALUES[0].to_string()])) @@ -109,7 +117,9 @@ async fn should_extract_titan_usage() { .await; let provider = test_provider(server.uri()); - let model = provider.embedding_model("amazon.titan-embed-text-v2:0"); + let model = provider + .embedding_model("amazon.titan-embed-text-v2:0") + .unwrap(); let result = model .do_embed(&default_options(vec![TEST_VALUES[0].to_string()])) @@ -137,7 +147,7 @@ async fn should_support_cohere_v3() { .await; let provider = test_provider(server.uri()); - let model = provider.embedding_model("cohere.embed-english-v3"); + let model = provider.embedding_model("cohere.embed-english-v3").unwrap(); let result = model .do_embed(&default_options(vec![TEST_VALUES[0].to_string()])) @@ -171,7 +181,7 @@ async fn should_support_cohere_v4() { .await; let provider = test_provider(server.uri()); - let model = provider.embedding_model("cohere.embed-v4:0"); + let model = provider.embedding_model("cohere.embed-v4:0").unwrap(); let result = model .do_embed(&default_options(vec![TEST_VALUES[0].to_string()])) @@ -200,7 +210,7 @@ async fn should_send_multiple_values_cohere_v4() { .await; let provider = test_provider(server.uri()); - let model = provider.embedding_model("cohere.embed-v4:0"); + let model = provider.embedding_model("cohere.embed-v4:0").unwrap(); let result = model .do_embed(&default_options( @@ -241,7 +251,7 @@ async fn should_support_cross_region_cohere() { .await; let provider = test_provider(server.uri()); - let model = provider.embedding_model("us.cohere.embed-v4:0"); + let model = provider.embedding_model("us.cohere.embed-v4:0").unwrap(); let result = model .do_embed(&default_options(vec![TEST_VALUES[0].to_string()])) @@ -269,7 +279,9 @@ async fn should_support_nova_embeddings() { .await; let provider = test_provider(server.uri()); - let model = provider.embedding_model("amazon.nova-2-multimodal-embeddings-v1:0"); + let model = provider + .embedding_model("amazon.nova-2-multimodal-embeddings-v1:0") + .unwrap(); let result = model .do_embed(&default_options(vec![TEST_VALUES[0].to_string()])) diff --git a/aimux-providers/tests/bedrock_stream_replay_test.rs b/aimux-providers/tests/bedrock_stream_replay_test.rs index 5fe4eafa6..5c04e245d 100644 --- a/aimux-providers/tests/bedrock_stream_replay_test.rs +++ b/aimux-providers/tests/bedrock_stream_replay_test.rs @@ -28,7 +28,8 @@ use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::stream_part::StreamPart; -use aimux_providers::{BedrockProvider, BedrockProviderConfig}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{AmazonBedrockProviderSettings, create_amazon_bedrock}; /// The one cassette recorded for this model id, so the `(method, path)` group /// holds exactly one entry and body-scoring cannot pick a different exchange: @@ -60,9 +61,14 @@ async fn body_base64_cassette_decodes_and_parses_into_stream_parts() { "no bedrock cassettes mounted — the replay directory moved or is empty" ); - let config = BedrockProviderConfig::with_bearer_token("test-token", "us-east-1") - .with_base_url(server.uri()); - let model = BedrockProvider::new(config).model(MODEL_ID); + let model = create_amazon_bedrock(AmazonBedrockProviderSettings { + api_key: Some(Resolvable::Value("test-token".to_string())), + region: Some("us-east-1".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .expect("valid settings") + .chat(MODEL_ID); // No error tolerance on purpose. A failed base64 decode, or frames the // eventstream parser rejects, surfaces here — and must fail the test rather diff --git a/aimux-providers/tests/black_forest_labs_image_test.rs b/aimux-providers/tests/black_forest_labs_image_test.rs index cd34d4113..8079b8a93 100644 --- a/aimux-providers/tests/black_forest_labs_image_test.rs +++ b/aimux-providers/tests/black_forest_labs_image_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Black Forest Labs image model tests. +//! Rust translation of the Black Forest Labs image model tests. //! Source: `reference/ai/packages/black-forest-labs/src/black-forest-labs-image-model.test.ts` use serde_json::{Value, json}; @@ -6,7 +6,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; -use aimux_providers::{BlackForestLabsConfig, BlackForestLabsProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{BlackForestLabsProviderSettings, create_black_forest_labs}; const PROMPT: &str = "A cute baby sea otter"; @@ -39,8 +40,14 @@ fn options(prompt: &str) -> ImageCallOptions { async fn should_extract_generated_image() { let server = MockServer::start().await; mock_bfl(&server).await; - let config = BlackForestLabsConfig::new("test-key").with_base_url(server.uri()); - let model = BlackForestLabsProvider::new(config).image("flux-pro-1.1"); + let config = BlackForestLabsProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_black_forest_labs(config) + .unwrap() + .image("flux-pro-1.1"); let result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { ImageOutputs::Binary(imgs) => { @@ -55,8 +62,14 @@ async fn should_extract_generated_image() { async fn should_pass_prompt() { let server = MockServer::start().await; mock_bfl(&server).await; - let config = BlackForestLabsConfig::new("test-key").with_base_url(server.uri()); - let model = BlackForestLabsProvider::new(config).image("flux-pro-1.1"); + let config = BlackForestLabsProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_black_forest_labs(config) + .unwrap() + .image("flux-pro-1.1"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&reqs[0].body).unwrap(); @@ -67,8 +80,14 @@ async fn should_pass_prompt() { async fn should_pass_aspect_ratio() { let server = MockServer::start().await; mock_bfl(&server).await; - let config = BlackForestLabsConfig::new("test-key").with_base_url(server.uri()); - let model = BlackForestLabsProvider::new(config).image("flux-pro-1.1"); + let config = BlackForestLabsProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_black_forest_labs(config) + .unwrap() + .image("flux-pro-1.1"); let mut opts = options(PROMPT); opts.aspect_ratio = Some(aimux_core::shared::AspectRatio::new(16, 9)); model.do_generate(&opts).await.unwrap(); @@ -81,8 +100,14 @@ async fn should_pass_aspect_ratio() { async fn should_pass_seed() { let server = MockServer::start().await; mock_bfl(&server).await; - let config = BlackForestLabsConfig::new("test-key").with_base_url(server.uri()); - let model = BlackForestLabsProvider::new(config).image("flux-pro-1.1"); + let config = BlackForestLabsProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_black_forest_labs(config) + .unwrap() + .image("flux-pro-1.1"); let mut opts = options(PROMPT); opts.seed = Some(42); model.do_generate(&opts).await.unwrap(); @@ -95,8 +120,14 @@ async fn should_pass_seed() { async fn should_pass_auth_headers() { let server = MockServer::start().await; mock_bfl(&server).await; - let config = BlackForestLabsConfig::new("test-key").with_base_url(server.uri()); - let model = BlackForestLabsProvider::new(config).image("flux-pro-1.1"); + let config = BlackForestLabsProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_black_forest_labs(config) + .unwrap() + .image("flux-pro-1.1"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); assert_eq!( @@ -114,8 +145,14 @@ async fn should_pass_auth_headers() { async fn should_return_provider_metadata() { let server = MockServer::start().await; mock_bfl(&server).await; - let config = BlackForestLabsConfig::new("test-key").with_base_url(server.uri()); - let model = BlackForestLabsProvider::new(config).image("flux-pro-1.1"); + let config = BlackForestLabsProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_black_forest_labs(config) + .unwrap() + .image("flux-pro-1.1"); let result = model.do_generate(&options(PROMPT)).await.unwrap(); let meta = result.provider_metadata.unwrap(); let bfl = meta.get("blackForestLabs").unwrap(); diff --git a/aimux-providers/tests/cartesia_speech_test.rs b/aimux-providers/tests/cartesia_speech_test.rs index 0c731bdd8..839821216 100644 --- a/aimux-providers/tests/cartesia_speech_test.rs +++ b/aimux-providers/tests/cartesia_speech_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Cartesia speech (TTS) model tests. +//! Rust translation of the Cartesia speech (TTS) model tests. //! //! Source: `reference/ai/packages/cartesia/src/cartesia-speech-model.test.ts` //! @@ -14,7 +14,8 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::shared::Warning; use aimux_core::speech_model::{AudioData, SpeechCallOptions, SpeechModel}; -use aimux_providers::{CartesiaConfig, CartesiaProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{CartesiaProviderSettings, create_cartesia}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -67,8 +68,12 @@ async fn should_generate_speech_with_required_parameters() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -100,8 +105,12 @@ async fn should_throw_when_no_voice_is_provided() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let result = model.do_generate(&speech_options("Hello, world!")).await; @@ -120,8 +129,12 @@ async fn should_map_wav_output_format() { let server = MockServer::start().await; mock_audio_response(&server, "wav").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -148,8 +161,12 @@ async fn should_map_pcm_output_format_with_sample_rate_suffix() { let server = MockServer::start().await; mock_audio_response(&server, "pcm").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -176,8 +193,12 @@ async fn should_handle_language_parameter() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hola, mundo!"); @@ -198,8 +219,12 @@ async fn should_handle_speed_parameter() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -219,8 +244,12 @@ async fn should_warn_and_ignore_an_out_of_range_generic_speed() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -251,8 +280,12 @@ async fn should_warn_about_unsupported_instructions_parameter() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -282,8 +315,12 @@ async fn should_pass_provider_specific_options() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -323,8 +360,12 @@ async fn should_ignore_encoding_for_mp3_output() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -364,8 +405,12 @@ async fn should_warn_about_an_unsupported_sample_rate_suffix() { let server = MockServer::start().await; mock_audio_response(&server, "wav").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -399,10 +444,18 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = CartesiaConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(provider_headers); - let provider = CartesiaProvider::new(config); + let config = CartesiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some( + (provider_headers) + .into_iter() + .map(|(k, v)| (k, Some(v))) + .collect(), + ), + ..Default::default() + }; + let provider = create_cartesia(config).unwrap(); let model = provider.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -437,8 +490,12 @@ async fn should_return_audio_data() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); @@ -455,8 +512,12 @@ async fn should_include_response_data_with_timestamp_modelid_and_headers() { let server = MockServer::start().await; mock_audio_response_with_headers(&server, "mp3", &[("x-request-id", "test-request-id")]).await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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.speech("sonic-3.5"); let mut options = speech_options("Hello, world!"); diff --git a/aimux-providers/tests/cartesia_transcription_test.rs b/aimux-providers/tests/cartesia_transcription_test.rs index e13cf79a1..6be74b644 100644 --- a/aimux-providers/tests/cartesia_transcription_test.rs +++ b/aimux-providers/tests/cartesia_transcription_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Cartesia transcription model tests. +//! Rust translation of the Cartesia transcription model tests. //! Source: `reference/ai/packages/cartesia/src/cartesia-transcription-model.test.ts` use std::collections::HashMap; @@ -9,7 +9,8 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::error::AiMuxError; use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; -use aimux_providers::{CartesiaConfig, CartesiaProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{CartesiaProviderSettings, create_cartesia}; fn mock_audio() -> Vec { vec![1u8, 2, 3] @@ -76,8 +77,12 @@ async fn should_pass_model() { let server = MockServer::start().await; mock_response(&server, &fixture_response(), &[]).await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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"); model @@ -101,10 +106,13 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = CartesiaConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = CartesiaProvider::new(config); + let config = CartesiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_cartesia(config).unwrap(); let model = provider.transcription("best"); let mut opts = options(mock_audio(), "audio/wav"); @@ -136,8 +144,12 @@ async fn should_extract_text() { let server = MockServer::start().await; mock_response(&server, &fixture_response(), &[]).await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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 result = model @@ -154,8 +166,12 @@ async fn should_extract_text() { #[tokio::test] async fn should_reject_streaming_model_for_non_streaming() { let server = MockServer::start().await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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("ink-2"); let result = model.do_generate(&options(mock_audio(), "audio/wav")).await; @@ -177,8 +193,12 @@ async fn should_include_response_data() { ) .await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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 result = model @@ -195,8 +215,12 @@ async fn should_use_real_date() { let server = MockServer::start().await; mock_response(&server, &fixture_response(), &[]).await; - let config = CartesiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CartesiaProvider::new(config); + 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 result = model diff --git a/aimux-providers/tests/cassette_exhaustive_test.rs b/aimux-providers/tests/cassette_exhaustive_test.rs index c234eb785..22c38b6ea 100644 --- a/aimux-providers/tests/cassette_exhaustive_test.rs +++ b/aimux-providers/tests/cassette_exhaustive_test.rs @@ -24,7 +24,9 @@ use serde_json::Value; use wiremock::{Mock, MockServer, Request, ResponseTemplate}; use aimux_core::generate::{GenerateTextOptions, generate_text, stream_text}; -use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::openai_compatible::{ + OpenAICompatibleProviderSettings, create_openai_compatible, +}; const CASSETTE_DIR: &str = "tests/cassettes"; @@ -195,7 +197,7 @@ async fn replay_single_cassette(cass: &Cassette) -> Result<(), String> { let model_id = extract_model(cass).unwrap_or_else(|| "gpt-4o".to_string()); - // Only test OpenAI-compatible providers via OpenAIProvider. + // Only test OpenAI-compatible providers via the OpenAI-compatible package. // Anthropic, Gemini, Cohere, Bedrock have their own providers but // most cassettes use /chat/completions path. let is_chat_completions = cass.req_path.ends_with("/chat/completions"); @@ -204,8 +206,14 @@ async fn replay_single_cassette(cass: &Cassette) -> Result<(), String> { return Ok(()); } - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(base_url)); - let model = provider.model(&model_id); + let provider = create_openai_compatible(OpenAICompatibleProviderSettings { + name: "compat".to_string(), + base_url, + api_key: Some("test-key".to_string()), + ..Default::default() + }) + .unwrap(); + let model = provider.chat(&model_id); let prompt = extract_prompt(cass); @@ -275,7 +283,7 @@ async fn replay_all_cassettes_exhaustive() { let total = cassettes.len(); assert!(total > 2000, "expected 2000+ cassettes, got {total}"); - // Filter to chat-completions cassettes (the only ones OpenAIProvider can handle) + // Filter to chat-completions cassettes (the only ones the compatible chat model handles) let chat_cassettes: Vec<&Cassette> = cassettes .iter() .filter(|c| c.req_path.ends_with("/chat/completions")) diff --git a/aimux-providers/tests/cassette_full_test.rs b/aimux-providers/tests/cassette_full_test.rs index 9cde165b4..5826bdb14 100644 --- a/aimux-providers/tests/cassette_full_test.rs +++ b/aimux-providers/tests/cassette_full_test.rs @@ -15,17 +15,54 @@ use futures::StreamExt; use wiremock::MockServer; use aimux_core::generate::{GenerateTextOptions, generate_text, stream_text}; +use aimux_core::language_model::LanguageModel; use aimux_core::stream_part::TextStreamPart; -use aimux_providers::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIProvider}; +use aimux_providers::deepseek::{DeepSeekProviderSettings, create_deepseek}; +use aimux_providers::groq::{GroqProviderSettings, create_groq}; +use aimux_providers::openai_compatible::{ + OpenAICompatibleProviderSettings, create_openai_compatible, +}; + +/// Which chat dialect a cassette directory is replayed through. +#[derive(Clone, Copy)] +enum Kind { + /// The generic OpenAI-compatible chat model. + Generic, + DeepSeek, + Groq, +} + +fn chat_model(kind: Kind, base_url: String, model_id: &str) -> std::sync::Arc { + let api_key = Some("test-key".to_string()); + match kind { + Kind::Generic => create_openai_compatible(OpenAICompatibleProviderSettings { + name: "compat".to_string(), + base_url, + api_key: Some("test-key".to_string()), + ..Default::default() + }) + .unwrap() + .call(model_id), + Kind::DeepSeek => create_deepseek(DeepSeekProviderSettings { + base_url: Some(base_url), + api_key, + ..Default::default() + }) + .unwrap() + .call(model_id), + Kind::Groq => create_groq(GroqProviderSettings { + base_url: Some(base_url), + api_key, + ..Default::default() + }) + .unwrap() + .call(model_id), + } +} /// Run generate_text + stream_text against a cassette directory. /// Hard asserts: must get a valid response with non-empty text and usage. -async fn replay_openai_compat( - cassette_dir: &str, - model_id: &str, - base_path: &str, - profile: OpenAICompatProfile, -) { +async fn replay_openai_compat(cassette_dir: &str, model_id: &str, base_path: &str, kind: Kind) { let server = MockServer::start().await; let n = replay::mount_cassettes(&server, cassette_dir).await; assert!(n > 0, "no cassettes loaded from {cassette_dir}"); @@ -36,15 +73,10 @@ async fn replay_openai_compat( format!("{}/{}", server.uri(), base_path) }; - let provider = OpenAIProvider::new( - OpenAIConfig::new("test-key") - .with_base_url(base_url) - .with_profile(profile), - ); - let model = provider.model(model_id); + let model = chat_model(kind, base_url, model_id); // ── Non-streaming: hard assert ── - let result = generate_text(&model, "Hello", GenerateTextOptions::default()) + let result = generate_text(&*model, "Hello", GenerateTextOptions::default()) .await .expect("{cassette_dir}: generate_text should succeed with cassette replay"); assert!( @@ -53,7 +85,7 @@ async fn replay_openai_compat( ); // ── Streaming: hard assert ── - let result = stream_text(&model, "Hello", GenerateTextOptions::default()) + let result = stream_text(&*model, "Hello", GenerateTextOptions::default()) .await .expect("{cassette_dir}: stream_text should succeed"); let mut stream = result.stream; @@ -78,13 +110,7 @@ async fn replay_openai_compat( #[tokio::test] async fn replay_openai() { - replay_openai_compat( - "tests/cassettes/openai", - "gpt-4o", - "v1", - OpenAICompatProfile::full(), - ) - .await; + replay_openai_compat("tests/cassettes/openai", "gpt-4o", "v1", Kind::Generic).await; } #[tokio::test] @@ -93,7 +119,7 @@ async fn replay_deepseek() { "tests/cassettes/deepseek", "deepseek-chat", "", - OpenAICompatProfile::deepseek(), + Kind::DeepSeek, ) .await; } @@ -104,7 +130,7 @@ async fn replay_groq() { "tests/cassettes/groq", "llama-3.3-70b-versatile", "openai/v1", - OpenAICompatProfile::groq(), + Kind::Groq, ) .await; } @@ -115,20 +141,14 @@ async fn replay_mistral() { "tests/cassettes/mistral", "ministral-8b-latest", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } #[tokio::test] async fn replay_perplexity() { - replay_openai_compat( - "tests/cassettes/perplexity", - "sonar", - "", - OpenAICompatProfile::full(), - ) - .await; + replay_openai_compat("tests/cassettes/perplexity", "sonar", "", Kind::Generic).await; } #[tokio::test] @@ -137,7 +157,7 @@ async fn replay_cerebras() { "tests/cassettes/cerebras", "llama3.3-70b", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -148,7 +168,7 @@ async fn replay_fireworks() { "tests/cassettes/fireworks", "llama-v3p1-8b-instruct", "inference/v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -159,7 +179,7 @@ async fn replay_togetherai() { "tests/cassettes/togetherai", "meta-llama/Llama-3.1-8B-Instruct-Turbo", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -170,20 +190,14 @@ async fn replay_moonshotai() { "tests/cassettes/moonshotai", "moonshot-v1-8k", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } #[tokio::test] async fn replay_copilot() { - replay_openai_compat( - "tests/cassettes/copilot", - "gpt-4o", - "", - OpenAICompatProfile::full(), - ) - .await; + replay_openai_compat("tests/cassettes/copilot", "gpt-4o", "", Kind::Generic).await; } #[tokio::test] @@ -192,7 +206,7 @@ async fn replay_baseten() { "tests/cassettes/baseten", "meta-llama/Llama-3.1-8B-Instruct", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -203,7 +217,7 @@ async fn replay_deepinfra() { "tests/cassettes/deepinfra", "meta-llama/Llama-3.1-8B-Instruct", "v1/openai", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -214,20 +228,14 @@ async fn replay_doubleword() { "tests/cassettes/doubleword", "Qwen/Qwen3.5-9B", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } #[tokio::test] async fn replay_github() { - replay_openai_compat( - "tests/cassettes/github", - "gpt-4o", - "", - OpenAICompatProfile::full(), - ) - .await; + replay_openai_compat("tests/cassettes/github", "gpt-4o", "", Kind::Generic).await; } #[tokio::test] @@ -236,7 +244,7 @@ async fn replay_llamafile() { "tests/cassettes/llamafile", "llama3.2:latest", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -247,7 +255,7 @@ async fn replay_lmstudio() { "tests/cassettes/lmstudio", "llama-3.2-3b-instruct", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -258,20 +266,14 @@ async fn replay_mistralrs() { "tests/cassettes/mistralrs", "Qwen/Qwen3-4B", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } #[tokio::test] async fn replay_ollama() { - replay_openai_compat( - "tests/cassettes/ollama", - "qwen3:4b", - "v1", - OpenAICompatProfile::full(), - ) - .await; + replay_openai_compat("tests/cassettes/ollama", "qwen3:4b", "v1", Kind::Generic).await; } #[tokio::test] @@ -280,7 +282,7 @@ async fn replay_sambanova() { "tests/cassettes/sambanova", "Meta-Llama-3.1-8B-Instruct", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -291,20 +293,14 @@ async fn replay_siliconflow() { "tests/cassettes/siliconflow", "Qwen/Qwen2.5-7B-Instruct", "v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } #[tokio::test] async fn replay_vercel() { - replay_openai_compat( - "tests/cassettes/vercel", - "gpt-4o", - "v1", - OpenAICompatProfile::full(), - ) - .await; + replay_openai_compat("tests/cassettes/vercel", "gpt-4o", "v1", Kind::Generic).await; } #[tokio::test] @@ -313,7 +309,7 @@ async fn replay_zai() { "tests/cassettes/zai", "glm-4.7", "api/paas/v4", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -324,7 +320,7 @@ async fn replay_alibaba() { "tests/cassettes/alibaba", "qwen-plus", "compatible-mode/v1", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -335,7 +331,7 @@ async fn replay_bytedance() { "tests/cassettes/bytedance", "doubao-pro-32k", "api/v3", - OpenAICompatProfile::full(), + Kind::Generic, ) .await; } @@ -373,16 +369,19 @@ async fn replay_chatgpt() {} #[tokio::test] async fn replay_anthropic() { - use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider}; + use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; let server = MockServer::start().await; let n = replay::mount_cassettes(&server, "tests/cassettes/anthropic").await; assert!(n > 0); - let provider = AnthropicProvider::new( - AnthropicConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())), - ); - let model = provider.model("claude-sonnet-4-6"); + let provider = create_anthropic(AnthropicProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap(); + let model = provider.messages("claude-sonnet-4-6"); let result = generate_text(&model, "Hello", GenerateTextOptions::default()) .await @@ -427,12 +426,14 @@ async fn replay_gemini() { // Use the OpenAI-compatible endpoint — model is in the body, // so replay can score and match correctly. - let provider = OpenAIProvider::new( - OpenAIConfig::new("test-key") - .with_base_url(format!("{}/v1beta/openai", server.uri())) - .with_provider("google"), - ); - let model = provider.model("gemini-2.5-pro-preview-05-06"); + let provider = create_openai_compatible(OpenAICompatibleProviderSettings { + name: "google".to_string(), + base_url: format!("{}/v1beta/openai", server.uri()), + api_key: Some("test-key".to_string()), + ..Default::default() + }) + .unwrap(); + let model = provider.chat("gemini-2.5-pro-preview-05-06"); let result = generate_text(&model, "Hello", GenerateTextOptions::default()) .await @@ -449,16 +450,19 @@ async fn replay_gemini() { #[tokio::test] async fn replay_cohere() { - use aimux_providers::cohere::{CohereConfig, CohereProvider}; + use aimux_providers::cohere::{CohereProviderSettings, create_cohere}; let server = MockServer::start().await; let n = replay::mount_cassettes(&server, "tests/cassettes/cohere").await; assert!(n > 0); - let provider = CohereProvider::new( - CohereConfig::new("test-key").with_base_url(format!("{}/v2", server.uri())), - ); - let model = provider.model("command-r-plus"); + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v2", server.uri())), + ..Default::default() + }) + .expect("valid settings"); + let model = provider.chat("command-r-plus"); let result = generate_text(&model, "Hello", GenerateTextOptions::default()) .await @@ -472,17 +476,21 @@ async fn replay_cohere() { #[tokio::test] async fn replay_bedrock() { - use aimux_providers::bedrock::{BedrockProvider, BedrockProviderConfig}; + use aimux_providers::bedrock::{AmazonBedrockProviderSettings, create_amazon_bedrock}; let server = MockServer::start().await; let n = replay::mount_cassettes(&server, "tests/cassettes/bedrock").await; assert!(n > 0); - let provider = BedrockProvider::new( - BedrockProviderConfig::new("test-key", "test-secret", "us-east-1") - .with_base_url(server.uri()), - ); - let model = provider.model("us.anthropic.claude-sonnet-4-20250514-v1:0"); + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + access_key_id: Some("test-key".to_string()), + secret_access_key: Some("test-secret".to_string()), + region: Some("us-east-1".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.chat("us.anthropic.claude-sonnet-4-20250514-v1:0"); let result = generate_text(&model, "Hello", GenerateTextOptions::default()) .await diff --git a/aimux-providers/tests/cassette_multimodal_test.rs b/aimux-providers/tests/cassette_multimodal_test.rs index 307fc724b..7fe92ba98 100644 --- a/aimux-providers/tests/cassette_multimodal_test.rs +++ b/aimux-providers/tests/cassette_multimodal_test.rs @@ -16,7 +16,20 @@ use aimux_core::files_model::{Files, UploadFileCallOptions, UploadFileData}; use aimux_core::image_model::ImageModel; use aimux_core::shared::FileBytes; use aimux_core::transcription_model::TranscriptionModel; -use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::openai::{OpenAIProvider, OpenAIProviderSettings, create_openai}; +use aimux_providers::openai_compatible::{ + OpenAICompatibleProviderSettings, create_openai_compatible, +}; + +/// The native OpenAI package pointed at the mock server. +fn native_provider(base_url: String) -> OpenAIProvider { + create_openai(OpenAIProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(base_url), + ..Default::default() + }) + .unwrap() +} const CASSETTE_DIR: &str = "tests/cassettes"; @@ -66,8 +79,8 @@ 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 = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(base_url)); - let model = provider.embedding_model("text-embedding-3-small"); + 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"); @@ -87,8 +100,8 @@ async fn cassette_openai_embedding_documents() { load_cassette("openai", "TestOpenAI.test_documents.json").expect("cassette should load"); let (_server, base_url) = mount_single(&cass).await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(base_url)); - let model = provider.embedding_model("text-embedding-3-small"); + let provider = native_provider(base_url); + let model = provider.embedding("text-embedding-3-small"); let opts = EmbeddingCallOptions { values: vec!["doc1".into(), "doc2".into()], @@ -112,8 +125,8 @@ async fn cassette_openai_embedding_error() { load_cassette("openai", "TestOpenAI.test_embed_error.json").expect("cassette should load"); let (_server, base_url) = mount_single(&cass).await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(base_url)); - let model = provider.embedding_model("nonexistent"); + let provider = native_provider(base_url); + let model = provider.embedding("nonexistent"); let opts = EmbeddingCallOptions::new("hello"); let result = model.do_embed(&opts).await; @@ -155,9 +168,14 @@ async fn cassette_cohere_embedding_query() { // Cohere base_url: the cassette path is /v2/embed, so base_url = server/v2 let base_url = format!("{}/v2", server.uri()); - use aimux_providers::cohere::{CohereConfig, CohereProvider}; - let provider = CohereProvider::new(CohereConfig::new("test-key").with_base_url(base_url)); - let model = provider.embedding_model("embed-english-v3.0"); + use aimux_providers::cohere::{CohereProviderSettings, create_cohere}; + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(base_url), + ..Default::default() + }) + .expect("valid settings"); + let model = provider.embedding("embed-english-v3.0"); let opts = EmbeddingCallOptions::new("hello"); let result = model.do_embed(&opts).await; @@ -207,7 +225,7 @@ async fn cassette_openai_files_upload() { .await; let base_url = format!("{}/v1", server.uri()); - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(base_url)); + let provider = native_provider(base_url); let files = provider.files(); let opts = UploadFileCallOptions::new( @@ -249,13 +267,15 @@ async fn cassette_xai_image_generation() { .mount(&server) .await; - // xAI is OpenAI-compatible — use OpenAI provider with xAI base_url + // xAI is OpenAI-compatible — use the compatible package with xAI's name let base_url = format!("{}/v1", server.uri()); - let provider = OpenAIProvider::new( - OpenAIConfig::new("test-key") - .with_base_url(base_url) - .with_provider("xai"), - ); + let provider = create_openai_compatible(OpenAICompatibleProviderSettings { + name: "xai".to_string(), + base_url, + api_key: Some("test-key".to_string()), + ..Default::default() + }) + .unwrap(); let model = provider.image("grok-2-image"); use aimux_core::image_model::ImageCallOptions; @@ -311,9 +331,9 @@ async fn cassette_transcription() { .mount(&server) .await; - // OpenRouter is OpenAI-compatible + // The transcription endpoint is OpenAI's wire format; OpenRouter serves it as is. let base_url = format!("{}/api/v1", server.uri()); - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(base_url)); + let provider = native_provider(base_url); let model = provider.transcription("whisper-1"); use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions}; diff --git a/aimux-providers/tests/cerebras_test.rs b/aimux-providers/tests/cerebras_test.rs index b58d5ad16..1139e1bbf 100644 --- a/aimux-providers/tests/cerebras_test.rs +++ b/aimux-providers/tests/cerebras_test.rs @@ -4,7 +4,7 @@ //! - `packages/cerebras/src/cerebras-provider.test.ts` //! - `packages/cerebras/src/cerebras-chat-language-model.test.ts` //! -//! Cerebras is a thin OpenAI-compatible wrapper over [`OpenAIProvider`]. The TS +//! Cerebras is a registry preset of the OpenAI-compatible package. The TS //! suite is mostly provider-configuration unit tests plus a `transformRequestBody` //! that renames assistant `reasoning_content` �?`reasoning`, and fixture-based //! finish-reason normalisation for structured-output tool calls. The Rust thin @@ -22,6 +22,7 @@ use serde_json::{Value, json}; use serial_test::serial; +use std::sync::Arc; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -29,7 +30,7 @@ use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; -use aimux_providers::{ProviderOptions, provider, provider_from_env}; +use aimux_providers::{PresetSettings, create_provider}; fn test_prompt() -> LanguageModelPrompt { vec![LanguageModelMessage::user_text("Hello")] @@ -54,28 +55,32 @@ fn text_completion_body() -> Value { }) } -fn make_provider(server: &MockServer) -> Box { - provider( +fn make_provider(server: &MockServer) -> Arc { + create_provider( "cerebras", - Some("test-api-key".to_string()), - "llama-3.3-70b", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-api-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("llama-3.3-70b") .expect("cerebras provider should build") } /// TS: `createCerebras()` produces a provider whose name is "cerebras". #[test] fn provider_builds_cerebras_model() { - let model = provider( + let model = create_provider( "cerebras", - Some("test-key".to_string()), - "llama-3.3-70b", - None, + PresetSettings { + api_key: Some("test-key".to_string().into()), + ..Default::default() + }, ) + .expect("provider should build") + .language_model("llama-3.3-70b") .expect("cerebras provider should build"); assert_eq!(model.model_id(), "llama-3.3-70b"); } @@ -111,15 +116,16 @@ async fn custom_api_key_used_in_auth_header() { .mount(&server) .await; - let model = provider( + let model = create_provider( "cerebras", - Some("my-custom-key".to_string()), - "llama-3.3-70b", - Some(ProviderOptions { + PresetSettings { + api_key: Some("my-custom-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("llama-3.3-70b") .expect("cerebras provider should build"); model @@ -140,15 +146,16 @@ async fn custom_headers_forwarded() { .mount(&server) .await; - let model = provider( + let model = create_provider( "cerebras", - Some("test-key".to_string()), - "llama-3.3-70b", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("llama-3.3-70b") .expect("cerebras provider should build"); let mut options = default_options(test_prompt()); @@ -174,15 +181,16 @@ async fn language_model_via_trait() { .mount(&server) .await; - let model = provider( + let model = create_provider( "cerebras", - Some("test-key".to_string()), - "llama-3.3-70b", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("llama-3.3-70b") .expect("provider should build a model"); model @@ -200,7 +208,9 @@ fn from_env_loads_cerebras_api_key() { std::env::set_var("CEREBRAS_API_KEY", "env-test-key"); } - let model = provider_from_env("cerebras", "llama-3.3-70b", None); + let model = create_provider("cerebras", PresetSettings::default()) + .expect("provider should build") + .language_model("llama-3.3-70b"); assert!(model.is_ok(), "from_env should succeed with env var set"); unsafe { @@ -213,15 +223,19 @@ fn from_env_loads_cerebras_api_key() { /// TS: without the env var, `createCerebras()` fails. #[serial] -#[test] -fn from_env_fails_without_env_var() { +#[tokio::test] +async fn from_env_fails_without_env_var() { let saved = std::env::var("CEREBRAS_API_KEY").ok(); unsafe { std::env::remove_var("CEREBRAS_API_KEY"); } - let model = provider_from_env("cerebras", "llama-3.3-70b", None); - assert!(model.is_err(), "from_env should fail without env var"); + let model = create_provider("cerebras", PresetSettings::default()) + .expect("provider should build") + .language_model("llama-3.3-70b") + .expect("provider should build a model"); + let result = model.do_generate(&default_options(test_prompt())).await; + assert!(result.is_err(), "from_env should fail without env var"); unsafe { if let Some(v) = saved { diff --git a/aimux-providers/tests/codex_test.rs b/aimux-providers/tests/codex_test.rs index 481f92096..7c354b56f 100644 --- a/aimux-providers/tests/codex_test.rs +++ b/aimux-providers/tests/codex_test.rs @@ -20,7 +20,25 @@ use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, StreamResult}; use aimux_core::stream_part::StreamPart; -use aimux_providers::{CODEX_OAUTH_TOKEN_URL, CodexConfig, CodexProvider, codex_refresh_at}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{ + CODEX_OAUTH_TOKEN_URL, CodexMode, CodexProvider, CodexProviderSettings, codex_refresh_at, + create_codex, +}; + +fn codex_provider( + mode: CodexMode, + base_url: impl Into, + originator: Option<&str>, +) -> CodexProvider { + create_codex(CodexProviderSettings { + mode, + base_url: Some(base_url.into()), + originator: originator.map(str::to_string), + ..Default::default() + }) + .expect("valid settings") +} // helpers @@ -152,8 +170,12 @@ async fn api_key_generates_text() { let server = MockServer::start().await; mock_json_response(&server, text_response_body()).await; - let config = CodexConfig::new("test-key").with_base_url(server.uri()); - let model = CodexProvider::new(config).model("gpt-5.2-codex"); + let provider = codex_provider( + CodexMode::ApiKey(Some(Resolvable::Value("test-key".to_string()))), + server.uri(), + None, + ); + let model = provider.responses("gpt-5.2-codex"); let result = model .do_generate(&default_options(test_prompt())) @@ -166,7 +188,7 @@ async fn api_key_generates_text() { } other => panic!("expected Text, got {other:?}"), } - assert_eq!(model.provider(), "codex"); + assert_eq!(model.provider(), "codex.responses"); assert_eq!(model.model_id(), "gpt-5.2-codex"); // Request shape: Responses endpoint, bearer auth, non-streaming body @@ -186,8 +208,12 @@ async fn api_key_streams_text() { let server = MockServer::start().await; mock_sse_response(&server, &sse_body(TEXT_STREAM_EVENTS)).await; - let config = CodexConfig::new("test-key").with_base_url(server.uri()); - let model = CodexProvider::new(config).model("gpt-5.2-codex"); + let provider = codex_provider( + CodexMode::ApiKey(Some(Resolvable::Value("test-key".to_string()))), + server.uri(), + None, + ); + let model = provider.responses("gpt-5.2-codex"); let result = model .do_stream(&default_options(test_prompt())) @@ -228,11 +254,15 @@ async fn subscription_generate_forces_stream_and_assembles_result() { .collect(); mock_sse_response(&server, &sse_body_strings(&events)).await; - let config = CodexConfig::subscription("account-token") - .with_base_url(server.uri()) - .with_chatgpt_account_id("acct_123") - .with_originator("aimux-test"); - let model = CodexProvider::new(config).model("gpt-5.2-codex"); + let provider = codex_provider( + CodexMode::ChatGptAccount { + token: Resolvable::Value("account-token".to_string()), + account_id: Some("acct_123".to_string()), + }, + server.uri(), + Some("aimux-test"), + ); + let model = provider.responses("gpt-5.2-codex"); let result = model .do_generate(&default_options(test_prompt())) @@ -283,9 +313,16 @@ async fn subscription_generate_continues_after_a_malformed_sse_frame() { ]; mock_sse_response(&server, &sse_body_strings(&events)).await; - let config = CodexConfig::subscription("account-token").with_base_url(server.uri()); - let result = CodexProvider::new(config) - .model("gpt-5.2-codex") + let provider = codex_provider( + CodexMode::ChatGptAccount { + token: Resolvable::Value("account-token".to_string()), + account_id: None, + }, + server.uri(), + None, + ); + let result = provider + .responses("gpt-5.2-codex") .do_generate(&default_options(test_prompt())) .await .expect("a later response.completed event should still be assembled"); @@ -301,9 +338,16 @@ async fn subscription_generate_returns_parse_error_without_a_completed_event() { let server = MockServer::start().await; mock_sse_response(&server, "data: {unparsable}\n\n").await; - let config = CodexConfig::subscription("account-token").with_base_url(server.uri()); - let error = CodexProvider::new(config) - .model("gpt-5.2-codex") + let provider = codex_provider( + CodexMode::ChatGptAccount { + token: Resolvable::Value("account-token".to_string()), + account_id: None, + }, + server.uri(), + None, + ); + let error = provider + .responses("gpt-5.2-codex") .do_generate(&default_options(test_prompt())) .await .unwrap_err(); @@ -316,8 +360,15 @@ async fn subscription_stream_forces_store_false() { let server = MockServer::start().await; mock_sse_response(&server, &sse_body(TEXT_STREAM_EVENTS)).await; - let config = CodexConfig::subscription("account-token").with_base_url(server.uri()); - let model = CodexProvider::new(config).model("gpt-5.2-codex"); + let provider = codex_provider( + CodexMode::ChatGptAccount { + token: Resolvable::Value("account-token".to_string()), + account_id: None, + }, + server.uri(), + None, + ); + let model = provider.responses("gpt-5.2-codex"); let result = model .do_stream(&default_options(test_prompt())) @@ -336,10 +387,15 @@ async fn subscription_user_headers_win_over_defaults() { let server = MockServer::start().await; mock_sse_response(&server, &sse_body(TEXT_STREAM_EVENTS)).await; - let config = CodexConfig::subscription("account-token") - .with_base_url(server.uri()) - .with_originator("library-default"); - let model = CodexProvider::new(config).model("gpt-5.2-codex"); + let provider = codex_provider( + CodexMode::ChatGptAccount { + token: Resolvable::Value("account-token".to_string()), + account_id: None, + }, + server.uri(), + Some("library-default"), + ); + let model = provider.responses("gpt-5.2-codex"); let options = CallOptions { headers: Some( @@ -371,8 +427,15 @@ async fn subscription_401_maps_to_token_expired() { .mount(&server) .await; - let config = CodexConfig::subscription("expired-token").with_base_url(server.uri()); - let model = CodexProvider::new(config).model("gpt-5.2-codex"); + let provider = codex_provider( + CodexMode::ChatGptAccount { + token: Resolvable::Value("expired-token".to_string()), + account_id: None, + }, + server.uri(), + None, + ); + let model = provider.responses("gpt-5.2-codex"); let options = default_options(test_prompt()); let gen_err = model diff --git a/aimux-providers/tests/cohere_embedding_test.rs b/aimux-providers/tests/cohere_embedding_test.rs index 393bdc5c5..5631903ab 100644 --- a/aimux-providers/tests/cohere_embedding_test.rs +++ b/aimux-providers/tests/cohere_embedding_test.rs @@ -9,8 +9,9 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; -use aimux_providers::{CohereConfig, CohereProvider}; +use aimux_core::embedding_model::EmbeddingCallOptions; +use aimux_core::provider::Provider; +use aimux_providers::{CohereProviderSettings, create_cohere}; const TEST_VALUES: &[&str] = &["sunny day at the beach", "rainy day in the city"]; @@ -58,9 +59,13 @@ async fn should_extract_embedding() { .mount(&server) .await; - let config = CohereConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CohereProvider::new(config); - let model = provider.embedding_model("embed-english-v3.0"); + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -97,9 +102,13 @@ async fn should_expose_raw_response() { .mount(&server) .await; - let config = CohereConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CohereProvider::new(config); - let model = provider.embedding_model("embed-english-v3.0"); + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -114,6 +123,10 @@ async fn should_expose_raw_response() { headers.get("test-header").map(std::string::String::as_str), Some("test-value") ); + assert_eq!( + result.response.unwrap().body, + Some(embedding_response_body()) + ); } /// TS: "should extract usage" @@ -126,9 +139,13 @@ async fn should_extract_usage() { .mount(&server) .await; - let config = CohereConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CohereProvider::new(config); - let model = provider.embedding_model("embed-english-v3.0"); + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -148,9 +165,13 @@ async fn should_pass_model_and_values() { .mount(&server) .await; - let config = CohereConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CohereProvider::new(config); - let model = provider.embedding_model("embed-english-v3.0"); + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); let _ = model .do_embed(&default_options(test_values())) @@ -179,9 +200,17 @@ async fn should_pass_headers() { .mount(&server) .await; - let config = CohereConfig::new("test-api-key").with_base_url(server.uri()); - let provider = CohereProvider::new(config); - let model = provider.embedding_model("embed-english-v3.0"); + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + headers: Some(HashMap::from([( + "Custom-Provider-Header".to_string(), + Some("provider-header-value".to_string()), + )])), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); let mut request_headers = HashMap::new(); request_headers.insert( @@ -203,6 +232,11 @@ async fn should_pass_headers() { let requests = server.received_requests().await.expect("requests recorded"); assert_eq!(requests.len(), 1); let h = &requests[0].headers; + assert_eq!( + h.get("custom-provider-header") + .and_then(|v| v.to_str().ok()), + Some("provider-header-value") + ); assert_eq!( h.get("authorization").and_then(|v| v.to_str().ok()), Some("Bearer test-api-key") diff --git a/aimux-providers/tests/common/mock_fetch.rs b/aimux-providers/tests/common/mock_fetch.rs new file mode 100644 index 000000000..1a084c614 --- /dev/null +++ b/aimux-providers/tests/common/mock_fetch.rs @@ -0,0 +1,217 @@ +//! A scripted [`Fetch`] transport and an environment-variable guard, for the +//! factory tests that replay recorded AI SDK requests. +//! +//! Included with `#[path = "common/mock_fetch.rs"] mod mock_fetch;`. + +#![allow(dead_code)] + +use std::collections::{BTreeMap, VecDeque}; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use bytes::Bytes; +use reqwest::StatusCode; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use serde_json::Value; + +use aimux_provider_utils::{Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse}; + +/// One canned HTTP response. +#[derive(Clone)] +pub struct Canned { + pub status: u16, + pub headers: Vec<(String, String)>, + pub body: Vec, +} + +impl Canned { + pub fn json(body: &Value) -> Self { + Self { + status: 200, + headers: vec![("content-type".into(), "application/json".into())], + body: serde_json::to_vec(body).unwrap(), + } + } +} + +/// One request as the transport saw it. +#[derive(Clone, Debug)] +pub struct Seen { + pub method: String, + pub url: String, + /// Lower-cased names (the `http` crate normalizes them). + pub headers: BTreeMap, + pub body: Vec, +} + +impl Seen { + pub fn json_body(&self) -> Value { + serde_json::from_slice(&self.body).expect("request body is JSON") + } +} + +#[derive(Default)] +pub struct MockFetch { + canned: Mutex>, + seen: Mutex>, +} + +impl MockFetch { + pub fn new(canned: Vec) -> Arc { + Arc::new(Self { + canned: Mutex::new(canned.into()), + seen: Mutex::default(), + }) + } + + pub fn seen(&self) -> Vec { + self.seen.lock().unwrap().clone() + } + + pub fn transport(self: &Arc) -> FetchFunction { + self.clone() + } +} + +#[async_trait] +impl Fetch for MockFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + self.seen.lock().unwrap().push(Seen { + method: request.method.to_string(), + url: request.url.to_string(), + headers: request + .headers + .iter() + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().expect("ASCII header").to_string(), + ) + }) + .collect(), + body: request.body.to_vec(), + }); + let canned = self + .canned + .lock() + .unwrap() + .pop_front() + .ok_or_else(|| FetchError::Other("no canned response left".into()))?; + let mut headers = HeaderMap::new(); + for (name, value) in &canned.headers { + headers.insert( + HeaderName::try_from(name.as_str()).unwrap(), + HeaderValue::try_from(value.as_str()).unwrap(), + ); + } + Ok(FetchResponse::from_bytes( + StatusCode::from_u16(canned.status).unwrap(), + headers, + request.url, + Bytes::from(canned.body), + )) + } +} + +/// A scripted [`Fetch`] that answers every request from a function of the +/// request, so a test can script "poll answers 503 forever" without counting +/// canned responses. Every request is recorded. +pub struct RouteFetch { + handler: Box Canned + Send + Sync>, + seen: Mutex>, +} + +impl RouteFetch { + pub fn new(handler: impl Fn(&Seen) -> Canned + Send + Sync + 'static) -> Arc { + Arc::new(Self { + handler: Box::new(handler), + seen: Mutex::default(), + }) + } + + pub fn seen(&self) -> Vec { + self.seen.lock().unwrap().clone() + } + + /// How many recorded requests had this method and URL path. + pub fn count(&self, method: &str, path: &str) -> usize { + self.seen() + .iter() + .filter(|seen| { + seen.method == method + && url::Url::parse(&seen.url).is_ok_and(|url| url.path() == path) + }) + .count() + } + + pub fn transport(self: &Arc) -> FetchFunction { + self.clone() + } +} + +#[async_trait] +impl Fetch for RouteFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + let seen = Seen { + method: request.method.to_string(), + url: request.url.to_string(), + headers: request + .headers + .iter() + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().expect("ASCII header").to_string(), + ) + }) + .collect(), + body: request.body.to_vec(), + }; + let canned = (self.handler)(&seen); + self.seen.lock().unwrap().push(seen); + let mut headers = HeaderMap::new(); + for (name, value) in &canned.headers { + headers.insert( + HeaderName::try_from(name.as_str()).unwrap(), + HeaderValue::try_from(value.as_str()).unwrap(), + ); + } + Ok(FetchResponse::from_bytes( + StatusCode::from_u16(canned.status).unwrap(), + headers, + request.url, + Bytes::from(canned.body), + )) + } +} + +/// Sets (or removes) an environment variable and restores it on drop. Tests +/// that use it are `#[serial]`. +pub struct EnvVar { + name: &'static str, + saved: Option, +} + +impl EnvVar { + pub fn set(name: &'static str, value: Option<&str>) -> Self { + let saved = std::env::var(name).ok(); + unsafe { + match value { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + Self { name, saved } + } +} + +impl Drop for EnvVar { + fn drop(&mut self) { + unsafe { + match &self.saved { + Some(value) => std::env::set_var(self.name, value), + None => std::env::remove_var(self.name), + } + } + } +} diff --git a/aimux-providers/tests/common/replay.rs b/aimux-providers/tests/common/replay.rs index 139ade247..e56bc1449 100644 --- a/aimux-providers/tests/common/replay.rs +++ b/aimux-providers/tests/common/replay.rs @@ -58,13 +58,18 @@ //! ```no_run //! # async fn example() { //! use wiremock::MockServer; -//! use aimux_providers::anthropic::AnthropicConfig; +//! use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; //! use reqwest::Client; //! //! let server = MockServer::start().await; //! common::replay::mount_cassettes(&server, "tests/cassettes/anthropic").await; //! -//! let config = AnthropicConfig::new("test-key").with_base_url(server.uri()); +//! let provider = create_anthropic(AnthropicProviderSettings { +//! api_key: Some("test-key".to_string()), +//! base_url: Some(format!("{}/v1", server.uri())), +//! ..Default::default() +//! }) +//! .unwrap(); //! // ... build the model, call do_generate / do_stream, assert on the parsed //! // result. The provider sees the cassette's real recorded response. //! # } diff --git a/aimux-providers/tests/config_snapshot_test.rs b/aimux-providers/tests/config_snapshot_test.rs deleted file mode 100644 index 54a1f49ed..000000000 --- a/aimux-providers/tests/config_snapshot_test.rs +++ /dev/null @@ -1,208 +0,0 @@ -//! RFC-0023 `config_snapshot` / `api_key_source` coverage (findings C2/C3/C6). -//! -//! These are pure construction-time checks (no HTTP): build a config the way -//! real callers do (`from_env` / `new` / registry `provider()`), then assert -//! the recorded `ProviderRecord` carries the right provider identity and -//! api-key *source* (never the plaintext key). - -use aimux_core::language_model::LanguageModel; - -use aimux_providers::provider; -use aimux_providers::{ - AnthropicConfig, AnthropicProvider, AzureConfig, AzureModel, CodexConfig, CodexProvider, - CohereConfig, CohereProvider, GoogleConfig, GoogleProvider, MistralConfig, MistralProvider, - XAIConfig, XAIProvider, -}; - -// ── C2: OpenAI-compatible chat path surfaces the registry provider name ───── - -/// A registry-built deepseek model records `provider == "deepseek"` (not the -/// old hardcoded `"openai"`). The Chat path must use `config.provider`. -#[test] -fn registry_chat_snapshot_records_real_provider() { - let model = provider("deepseek", Some("sk-test".into()), "deepseek-chat", None) - .expect("deepseek constructs via registry"); - let snap = model.config_snapshot(); - assert_eq!( - snap.provider, "deepseek", - "snapshot must carry the registry name" - ); - assert_eq!(snap.model_id, "deepseek-chat"); - assert_eq!(model.provider(), "deepseek"); - // Registry `provider(name, None, ..)` reads the env var → env source. - assert!( - snap.api_key_source == "explicit" || snap.api_key_source.starts_with("env:"), - "api_key_source should be explicit (key given) or env:, got {}", - snap.api_key_source - ); -} - -/// Explicit key → `explicit`; no plaintext leaks into the snapshot. -#[test] -fn registry_chat_explicit_key_source() { - let model = provider("groq", Some("sk-test".into()), "llama-3.3-70b", None).unwrap(); - let snap = model.config_snapshot(); - assert_eq!(snap.api_key_source, "explicit"); - let json = serde_json::to_string(&snap).unwrap(); - assert!(!json.contains("sk-test"), "plaintext key leaked: {json}"); -} - -// ── C3: native providers record env vs explicit api_key_source ────────────── - -/// Google `from_env` → `env:GOOGLE_GENERATIVE_AI_API_KEY`; `new` → `explicit`. -#[test] -fn google_api_key_source_env_vs_explicit() { - let explicit = GoogleConfig::new("g-explicit"); - let snap = GoogleProvider::new(explicit) - .model("gemini-2.0-flash") - .config_snapshot(); - assert_eq!(snap.provider, "google.generative-ai"); - assert_eq!(snap.api_key_source, "explicit"); - - // from_env requires the env var; set it so construction succeeds. - unsafe { std::env::set_var("GOOGLE_GENERATIVE_AI_API_KEY", "g-env") }; - let env_cfg = GoogleConfig::from_env().expect("env var set"); - unsafe { std::env::remove_var("GOOGLE_GENERATIVE_AI_API_KEY") }; - let snap = GoogleProvider::new(env_cfg) - .model("gemini-2.0-flash") - .config_snapshot(); - assert_eq!( - snap.api_key_source, "env:GOOGLE_GENERATIVE_AI_API_KEY", - "from_env must mark the env source" - ); - let json = serde_json::to_string(&snap).unwrap(); - assert!(!json.contains("g-env"), "plaintext key leaked: {json}"); -} - -/// Anthropic `from_env` → `env:ANTHROPIC_API_KEY`; `new` → `explicit`. Also -/// covers C6: the snapshot's `provider_options` is a ProviderOptions-shaped -/// object carrying `headers`/`api_version`/`max_retries`/`body_overrides`. -#[test] -fn anthropic_api_key_source_and_snapshot_shape() { - let explicit = AnthropicConfig::new("a-explicit"); - let snap = AnthropicProvider::new(explicit) - .model("claude-3-5-sonnet-20241022") - .config_snapshot(); - assert_eq!(snap.provider, "anthropic.messages"); - assert_eq!(snap.api_key_source, "explicit"); - - unsafe { std::env::set_var("ANTHROPIC_API_KEY", "a-env") }; - let env_cfg = AnthropicConfig::from_env().expect("env var set"); - unsafe { std::env::remove_var("ANTHROPIC_API_KEY") }; - let snap = AnthropicProvider::new(env_cfg) - .model("claude-3-5-sonnet-20241022") - .config_snapshot(); - assert_eq!(snap.api_key_source, "env:ANTHROPIC_API_KEY"); - - // C6: provider_options shape (not a bare headers map). - let opts = snap - .provider_options - .as_ref() - .expect("anthropic snapshot must carry provider_options"); - assert!(opts.get("headers").is_some(), "missing headers: {opts}"); - assert_eq!(opts["api_version"], "2023-06-01"); - assert!( - opts.get("max_retries").is_some(), - "missing max_retries: {opts}" - ); - assert!( - opts.get("body_overrides").is_some(), - "missing body_overrides: {opts}" - ); - - let json = serde_json::to_string(&snap).unwrap(); - assert!(!json.contains("a-env"), "plaintext key leaked: {json}"); -} - -/// Mistral `from_env` → `env:MISTRAL_API_KEY`. -#[test] -fn mistral_api_key_source_env() { - unsafe { std::env::set_var("MISTRAL_API_KEY", "m-env") }; - let cfg = MistralConfig::from_env().expect("env var set"); - unsafe { std::env::remove_var("MISTRAL_API_KEY") }; - let snap = MistralProvider::new(cfg) - .model("mistral-small-latest") - .config_snapshot(); - assert_eq!(snap.provider, "mistral"); - assert_eq!(snap.api_key_source, "env:MISTRAL_API_KEY"); -} - -/// Cohere `from_env` → `env:COHERE_API_KEY`. -#[test] -fn cohere_api_key_source_env() { - unsafe { std::env::set_var("COHERE_API_KEY", "c-env") }; - let cfg = CohereConfig::from_env().expect("env var set"); - unsafe { std::env::remove_var("COHERE_API_KEY") }; - let snap = CohereProvider::new(cfg) - .model("command-r-plus") - .config_snapshot(); - assert_eq!(snap.provider, "cohere"); - assert_eq!(snap.api_key_source, "env:COHERE_API_KEY"); -} - -/// Azure `from_env` → `env:AZURE_API_KEY`; explicit api-key → `explicit`. -#[test] -fn azure_api_key_source_env_vs_explicit() { - let cfg = AzureConfig::new() - .with_resource_name("my-resource") - .with_api_key("az-explicit"); - let snap = AzureModel::new("gpt-4o".into(), cfg).config_snapshot(); - assert_eq!(snap.provider, "azure"); - assert_eq!(snap.api_key_source, "explicit"); - - unsafe { - std::env::set_var("AZURE_API_KEY", "az-env"); - std::env::set_var("AZURE_RESOURCE_NAME", "res"); - } - let cfg = AzureConfig::from_env().expect("env var set"); - unsafe { - std::env::remove_var("AZURE_API_KEY"); - std::env::remove_var("AZURE_RESOURCE_NAME"); - } - let snap = AzureModel::new("gpt-4o".into(), cfg).config_snapshot(); - assert_eq!(snap.api_key_source, "env:AZURE_API_KEY"); - let json = serde_json::to_string(&snap).unwrap(); - assert!(!json.contains("az-env"), "plaintext key leaked: {json}"); -} - -/// Codex `from_env` → `env:CODEX_API_KEY`; `new` → `explicit`. -#[test] -fn codex_api_key_source_env_vs_explicit() { - let snap = CodexProvider::new(CodexConfig::new("sk-test")) - .model("gpt-5.2-codex") - .config_snapshot(); - assert_eq!(snap.provider, "codex"); - assert_eq!(snap.api_key_source, "explicit"); - - unsafe { std::env::set_var("CODEX_API_KEY", "cx-env") }; - let cfg = CodexConfig::from_env().expect("env var set"); - unsafe { std::env::remove_var("CODEX_API_KEY") }; - let snap = CodexProvider::new(cfg) - .model("gpt-5.2-codex") - .config_snapshot(); - assert_eq!(snap.api_key_source, "env:CODEX_API_KEY"); -} - -/// xAI `from_env` → `env:XAI_API_KEY`; `new` → `explicit`. M2b: `model()` must -/// preserve the credential source instead of reconstructing `XAIConfig::new` -/// (which previously dropped `api_key_source` → snapshot wrongly said -/// "explicit" even for env-built providers). -#[test] -fn xai_api_key_source_env_vs_explicit() { - let snap = XAIProvider::new(XAIConfig::new("x-explicit")) - .model("grok-2") - .config_snapshot(); - assert_eq!(snap.provider, "xai.chat"); - assert_eq!(snap.api_key_source, "explicit"); - - unsafe { std::env::set_var("XAI_API_KEY", "x-env") }; - let cfg = XAIConfig::from_env().expect("env var set"); - unsafe { std::env::remove_var("XAI_API_KEY") }; - let snap = XAIProvider::new(cfg).model("grok-2").config_snapshot(); - assert_eq!( - snap.api_key_source, "env:XAI_API_KEY", - "model() must preserve the env credential source from from_env" - ); - let json = serde_json::to_string(&snap).unwrap(); - assert!(!json.contains("x-env"), "plaintext key leaked: {json}"); -} diff --git a/aimux-providers/tests/conformance_test.rs b/aimux-providers/tests/conformance_test.rs index 5fe52f920..f4def891f 100644 --- a/aimux-providers/tests/conformance_test.rs +++ b/aimux-providers/tests/conformance_test.rs @@ -13,23 +13,56 @@ mod common; use common::replay::mount_cassettes; use futures::StreamExt; +use std::sync::Arc; use wiremock::MockServer; +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::result::GenerateContent; use aimux_core::stream_part::StreamPart; +use aimux_providers::anthropic::{AnthropicProvider, AnthropicProviderSettings, create_anthropic}; +use aimux_providers::openai::{OpenAIProvider, OpenAIProviderSettings, create_openai}; +use aimux_providers::openai_compatible::{ + OpenAICompatibleProvider, OpenAICompatibleProviderSettings, create_openai_compatible, +}; use aimux_providers::{ - AnthropicConfig, AnthropicProvider, BedrockProvider, BedrockProviderConfig, CohereConfig, - CohereProvider, GoogleConfig, GoogleProvider, HuggingFaceConfig, HuggingFaceProvider, - LlamafileConfig, LlamafileProvider, LmStudioConfig, LmStudioProvider, MistralConfig, - MistralProvider, MistralrsConfig, MistralrsProvider, OllamaConfig, OllamaProvider, - OpenAIConfig, OpenAIProvider, OpenRouterConfig, OpenRouterProvider, ProviderOptions, XAIConfig, - XAIProvider, provider, + AmazonBedrockProvider, AmazonBedrockProviderSettings, CohereProvider, CohereProviderSettings, + GoogleProvider, GoogleProviderSettings, MistralProvider, MistralProviderSettings, + PresetSettings, XAIProvider, XAIProviderSettings, create_amazon_bedrock, create_cohere, + create_google, create_mistral, create_provider, create_xai, }; +/// A registry preset pointed at `base_url`, as a language model. +fn registry_model( + name: &str, + api_key: String, + model_id: &str, + base_url: String, +) -> Result, AiMuxError> { + create_provider( + name, + PresetSettings { + api_key: Some(api_key.into()), + base_url: Some(base_url), + ..Default::default() + }, + )? + .language_model(model_id) +} + +/// The native OpenAI package pointed at the mock server. +fn native_openai(base_url: String) -> OpenAIProvider { + create_openai(OpenAIProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(base_url), + ..Default::default() + }) + .unwrap() +} + // 閳光偓閳光偓 helpers 閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓 fn test_prompt() -> LanguageModelPrompt { @@ -108,8 +141,12 @@ mod anthropic_conformance { use super::*; fn make_provider(server: &MockServer) -> AnthropicProvider { - let config = AnthropicConfig::new("test-key").with_base_url(server.uri()); - AnthropicProvider::new(config) + create_anthropic(AnthropicProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap() } /// doGenerate with a non-streaming cassette should return text content. @@ -119,7 +156,7 @@ mod anthropic_conformance { mount_cassettes(&server, "tests/cassettes/anthropic").await; let provider = make_provider(&server); - let model = provider.model("claude-sonnet-4-6"); + let model = provider.messages("claude-sonnet-4-6"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -153,7 +190,7 @@ mod anthropic_conformance { mount_cassettes(&server, "tests/cassettes/anthropic").await; let provider = make_provider(&server); - let model = provider.model("claude-sonnet-4-6"); + let model = provider.messages("claude-sonnet-4-6"); let result = model.do_stream(&default_options(test_prompt())).await; @@ -191,8 +228,7 @@ mod openai_conformance { fn make_provider(server: &MockServer) -> OpenAIProvider { // Cassettes are recorded against /v1/responses; the responses model // appends `/responses` to the base URL, so base must include `/v1`. - let config = OpenAIConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())); - OpenAIProvider::new(config) + native_openai(format!("{}/v1", server.uri())) } #[tokio::test] @@ -201,7 +237,7 @@ mod openai_conformance { mount_cassettes(&server, "tests/cassettes/openai").await; let provider = make_provider(&server); - let model = provider.responses_model("gpt-4o"); + let model = provider.responses("gpt-4o"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -232,7 +268,7 @@ mod openai_conformance { mount_cassettes(&server, "tests/cassettes/openai").await; let provider = make_provider(&server); - let model = provider.responses_model("gpt-4o"); + let model = provider.responses("gpt-4o"); let result = model.do_stream(&default_options(test_prompt())).await; @@ -265,15 +301,12 @@ mod openai_conformance { mod deepseek_conformance { use super::*; - fn make_provider(server: &MockServer) -> Box { - provider( + fn make_provider(server: &MockServer) -> Arc { + registry_model( "deepseek", - Some("test-key".to_string()), + "test-key".to_string(), "deepseek-chat", - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), + server.uri(), ) .expect("deepseek should construct from registry") } @@ -316,8 +349,12 @@ mod xai_conformance { fn make_provider(server: &MockServer) -> XAIProvider { // Cassettes are recorded against /v1/responses; the responses model // appends `/responses` to the base URL, so base must include `/v1`. - let config = XAIConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())); - XAIProvider::new(config) + create_xai(XAIProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .expect("valid settings") } #[tokio::test] @@ -326,7 +363,7 @@ mod xai_conformance { mount_cassettes(&server, "tests/cassettes/xai").await; let provider = make_provider(&server); - let model = provider.responses_model("grok-3-mini"); + let model = provider.responses("grok-3-mini"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -356,17 +393,14 @@ mod xai_conformance { mod groq_conformance { use super::*; - fn make_provider(server: &MockServer) -> Box { + fn make_provider(server: &MockServer) -> Arc { // Cassettes record paths under /openai/v1/chat/completions, matching // Groq's default base URL https://api.groq.com/openai/v1. - provider( + registry_model( "groq", - Some("test-key".to_string()), + "test-key".to_string(), "llama-3.3-70b-versatile", - Some(ProviderOptions { - base_url: Some(format!("{}/openai/v1", server.uri())), - ..Default::default() - }), + format!("{}/openai/v1", server.uri()), ) .expect("groq should construct from registry") } @@ -409,8 +443,12 @@ mod mistral_conformance { fn make_provider(server: &MockServer) -> MistralProvider { // Cassettes record paths under /v1/chat/completions, matching Mistral's // default base URL https://api.mistral.ai/v1. - let config = MistralConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())); - MistralProvider::new(config) + create_mistral(MistralProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .expect("valid settings") } #[tokio::test] @@ -419,7 +457,7 @@ mod mistral_conformance { mount_cassettes(&server, "tests/cassettes/mistral").await; let provider = make_provider(&server); - let model = provider.model("mistral-large-latest"); + let model = provider.chat("mistral-large-latest"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -449,17 +487,9 @@ mod mistral_conformance { mod perplexity_conformance { use super::*; - fn make_provider(server: &MockServer) -> Box { - provider( - "perplexity", - Some("test-key".to_string()), - "sonar", - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), - ) - .expect("perplexity should construct from registry") + fn make_provider(server: &MockServer) -> Arc { + registry_model("perplexity", "test-key".to_string(), "sonar", server.uri()) + .expect("perplexity should construct from registry") } #[tokio::test] @@ -498,9 +528,12 @@ mod gemini_conformance { use super::*; fn make_provider(server: &MockServer) -> GoogleProvider { - let config = - GoogleConfig::new("test-key").with_base_url(format!("{}/v1beta", server.uri())); - GoogleProvider::new(config) + create_google(GoogleProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1beta", server.uri())), + ..Default::default() + }) + .unwrap() } #[tokio::test] @@ -509,7 +542,7 @@ mod gemini_conformance { mount_cassettes(&server, "tests/cassettes/gemini").await; let provider = make_provider(&server); - let model = provider.model("gemini-2.5-flash"); + let model = provider.chat("gemini-2.5-flash"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -537,7 +570,7 @@ mod gemini_conformance { mount_cassettes(&server, "tests/cassettes/gemini").await; let provider = make_provider(&server); - let model = provider.model("gemini-2.5-flash"); + let model = provider.chat("gemini-2.5-flash"); let result = model.do_stream(&default_options(test_prompt())).await; @@ -566,12 +599,16 @@ mod gemini_conformance { mod openrouter_conformance { use super::*; - fn make_provider(server: &MockServer) -> OpenRouterProvider { + fn make_provider(server: &MockServer, model_id: &str) -> Arc { // Cassettes record paths under /api/v1/chat/completions, matching // OpenRouter's default base URL https://openrouter.ai/api/v1. - let config = - OpenRouterConfig::new("test-key").with_base_url(format!("{}/api/v1", server.uri())); - OpenRouterProvider::new(config) + registry_model( + "openrouter", + "test-key".to_string(), + model_id, + format!("{}/api/v1", server.uri()), + ) + .expect("registry provider should construct") } #[tokio::test] @@ -579,8 +616,7 @@ mod openrouter_conformance { let server = MockServer::start().await; mount_cassettes(&server, "tests/cassettes/openrouter").await; - let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o"); + let model = make_provider(&server, "openai/gpt-4o"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -610,20 +646,12 @@ mod openrouter_conformance { mod copilot_conformance { use super::*; - fn make_provider(server: &MockServer) -> Box { - // Copilot's base URL has no `/v1` prefix; the OpenAIProvider appends + fn make_provider(server: &MockServer) -> Arc { + // Copilot's base URL has no `/v1` prefix; the compatible chat model appends // `/chat/completions` directly, matching the cassette path // `/chat/completions`. - provider( - "copilot", - Some("test-key".to_string()), - "gpt-4o", - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), - ) - .expect("copilot should construct from registry") + registry_model("copilot", "test-key".to_string(), "gpt-4o", server.uri()) + .expect("copilot should construct from registry") } #[tokio::test] @@ -661,18 +689,15 @@ mod copilot_conformance { mod doubleword_conformance { use super::*; - fn make_provider(server: &MockServer) -> Box { - // Doubleword's base URL includes `/v1`; the OpenAIProvider appends + fn make_provider(server: &MockServer) -> Arc { + // Doubleword's base URL includes `/v1`; the compatible chat model appends // `/chat/completions`, so we point at `/v1` to match the // cassette path `/v1/chat/completions`. - provider( + registry_model( "doubleword", - Some("test-key".to_string()), + "test-key".to_string(), "Qwen/Qwen3.5-9B", - Some(ProviderOptions { - base_url: Some(format!("{}/v1", server.uri())), - ..Default::default() - }), + format!("{}/v1", server.uri()), ) .expect("doubleword should construct from registry") } @@ -712,12 +737,17 @@ mod doubleword_conformance { mod llamafile_conformance { use super::*; - fn make_provider(server: &MockServer) -> LlamafileProvider { - // llamafile's base URL includes `/v1`; the OpenAIProvider appends + fn make_provider(server: &MockServer, model_id: &str) -> Arc { + // llamafile's base URL includes `/v1`; the compatible chat model appends // `/chat/completions`, so we point at `/v1` to match the // cassette path `/v1/chat/completions`. - let config = LlamafileConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())); - LlamafileProvider::new(config) + registry_model( + "llamafile", + "test-key".to_string(), + model_id, + format!("{}/v1", server.uri()), + ) + .expect("registry provider should construct") } #[tokio::test] @@ -725,8 +755,7 @@ mod llamafile_conformance { let server = MockServer::start().await; mount_cassettes(&server, "tests/cassettes/llamafile").await; - let provider = make_provider(&server); - let model = provider.model("llama3.2:latest"); + let model = make_provider(&server, "llama3.2:latest"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -756,12 +785,17 @@ mod llamafile_conformance { mod mistralrs_conformance { use super::*; - fn make_provider(server: &MockServer) -> MistralrsProvider { - // mistral.rs's base URL includes `/v1`; the OpenAIProvider appends + fn make_provider(server: &MockServer, model_id: &str) -> Arc { + // mistral.rs's base URL includes `/v1`; the compatible chat model appends // `/chat/completions`, so we point at `/v1` to match the // cassette path `/v1/chat/completions`. - let config = MistralrsConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())); - MistralrsProvider::new(config) + registry_model( + "mistralrs", + "test-key".to_string(), + model_id, + format!("{}/v1", server.uri()), + ) + .expect("registry provider should construct") } #[tokio::test] @@ -769,8 +803,7 @@ mod mistralrs_conformance { let server = MockServer::start().await; mount_cassettes(&server, "tests/cassettes/mistralrs").await; - let provider = make_provider(&server); - let model = provider.model("Qwen/Qwen3-4B"); + let model = make_provider(&server, "Qwen/Qwen3-4B"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -800,12 +833,16 @@ mod mistralrs_conformance { mod bedrock_conformance { use super::*; - fn make_provider(server: &MockServer) -> BedrockProvider { + fn make_provider(server: &MockServer) -> AmazonBedrockProvider { // Bearer-token auth bypasses SigV4 signing so the mock server sees // plain requests — matching how the rig cassettes were recorded. - let config = BedrockProviderConfig::with_bearer_token("test-token", "us-east-1") - .with_base_url(server.uri()); - BedrockProvider::new(config) + create_amazon_bedrock(AmazonBedrockProviderSettings { + api_key: Some("test-token".to_string().into()), + region: Some("us-east-1".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap() } #[tokio::test] @@ -814,7 +851,7 @@ mod bedrock_conformance { mount_cassettes(&server, "tests/cassettes/bedrock").await; let provider = make_provider(&server); - let model = provider.model("amazon.nova-lite-v1:0"); + let model = provider.chat("amazon.nova-lite-v1:0"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -842,7 +879,7 @@ mod bedrock_conformance { mount_cassettes(&server, "tests/cassettes/bedrock").await; let provider = make_provider(&server); - let model = provider.model("amazon.nova-lite-v1:0"); + let model = provider.chat("amazon.nova-lite-v1:0"); let result = model.do_stream(&default_options(test_prompt())).await; @@ -872,17 +909,14 @@ mod bedrock_conformance { mod cerebras_conformance { use super::*; - fn make_provider(server: &MockServer) -> Box { + fn make_provider(server: &MockServer) -> Arc { // Cassettes record paths under /v1/chat/completions, matching Cerebras's // default base URL https://api.cerebras.ai/v1. - provider( + registry_model( "cerebras", - Some("test-key".to_string()), + "test-key".to_string(), "llama-3.3-70b", - Some(ProviderOptions { - base_url: Some(format!("{}/v1", server.uri())), - ..Default::default() - }), + format!("{}/v1", server.uri()), ) .expect("cerebras should construct from registry") } @@ -925,8 +959,12 @@ mod cohere_conformance { fn make_provider(server: &MockServer) -> CohereProvider { // Cassettes record paths under /v2/chat, matching Cohere's default base // URL https://api.cohere.com/v2; the model appends `/chat`. - let config = CohereConfig::new("test-key").with_base_url(format!("{}/v2", server.uri())); - CohereProvider::new(config) + create_cohere(CohereProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v2", server.uri())), + ..Default::default() + }) + .expect("valid settings") } #[tokio::test] @@ -935,7 +973,7 @@ mod cohere_conformance { mount_cassettes(&server, "tests/cassettes/cohere").await; let provider = make_provider(&server); - let model = provider.model("command-r-08-2024"); + let model = provider.chat("command-r-08-2024"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -965,13 +1003,19 @@ mod cohere_conformance { mod huggingface_conformance { use super::*; - fn make_provider(server: &MockServer) -> HuggingFaceProvider { + fn make_provider(server: &MockServer) -> OpenAICompatibleProvider { // Cassettes record paths under /together/v1/chat/completions (HF routes - // together-hosted models through this prefix). The OpenAI model appends - // `/chat/completions` to the base URL. - let config = HuggingFaceConfig::new("test-key") - .with_base_url(format!("{}/together/v1", server.uri())); - HuggingFaceProvider::new(config) + // together-hosted models through this prefix). The Hugging Face package + // serves the Responses API only, so the router's Chat Completions + // endpoint is reached through the OpenAI-compatible package, which + // appends `/chat/completions` to the base URL. + create_openai_compatible(OpenAICompatibleProviderSettings { + name: "huggingface".to_string(), + base_url: format!("{}/together/v1", server.uri()), + api_key: Some("test-key".to_string()), + ..Default::default() + }) + .expect("valid settings") } #[tokio::test] @@ -980,7 +1024,7 @@ mod huggingface_conformance { mount_cassettes(&server, "tests/cassettes/huggingface").await; let provider = make_provider(&server); - let model = provider.model("deepseek-ai/DeepSeek-R1"); + let model = provider.chat("deepseek-ai/DeepSeek-R1"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -1010,17 +1054,14 @@ mod huggingface_conformance { mod zai_conformance { use super::*; - fn make_provider(server: &MockServer) -> Box { + fn make_provider(server: &MockServer) -> Arc { // Cassettes record paths under /api/paas/v4/chat/completions, matching // Zai's base URL https://api.z.ai/api/paas/v4. - provider( + registry_model( "zai", - Some("test-key".to_string()), + "test-key".to_string(), "glm-4.7", - Some(ProviderOptions { - base_url: Some(format!("{}/api/paas/v4", server.uri())), - ..Default::default() - }), + format!("{}/api/paas/v4", server.uri()), ) .expect("zai should construct from registry") } @@ -1087,13 +1128,18 @@ mod zai_conformance { mod ollama_conformance { use super::*; - fn make_provider(server: &MockServer) -> OllamaProvider { + fn make_provider(server: &MockServer, model_id: &str) -> Arc { // pydantic-ai cassettes record paths under /v1/chat/completions, // matching Ollama's OpenAI-compatible endpoint. // (rig cassettes use the native /api/chat NDJSON endpoint and won't // be hit by this OpenAI-compatible provider — that's expected.) - let config = OllamaConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())); - OllamaProvider::new(config) + registry_model( + "ollama", + "test-key".to_string(), + model_id, + format!("{}/v1", server.uri()), + ) + .expect("registry provider should construct") } #[tokio::test] @@ -1101,8 +1147,7 @@ mod ollama_conformance { let server = MockServer::start().await; mount_cassettes(&server, "tests/cassettes/ollama").await; - let provider = make_provider(&server); - let model = provider.model("gpt-oss:20b"); + let model = make_provider(&server, "gpt-oss:20b"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -1135,9 +1180,7 @@ mod chatgpt_conformance { fn make_provider(server: &MockServer) -> OpenAIProvider { // Cassettes record paths under /backend-api/codex/responses; the // responses model appends `/responses` to the base URL. - let config = OpenAIConfig::new("test-key") - .with_base_url(format!("{}/backend-api/codex", server.uri())); - OpenAIProvider::new(config) + native_openai(format!("{}/backend-api/codex", server.uri())) } #[tokio::test] @@ -1146,7 +1189,7 @@ mod chatgpt_conformance { mount_cassettes(&server, "tests/cassettes/chatgpt").await; let provider = make_provider(&server); - let model = provider.responses_model("gpt-5.4"); + let model = provider.responses("gpt-5.4"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -1177,7 +1220,7 @@ mod chatgpt_conformance { mount_cassettes(&server, "tests/cassettes/chatgpt").await; let provider = make_provider(&server); - let model = provider.responses_model("gpt-5.4"); + let model = provider.responses("gpt-5.4"); let result = model.do_stream(&default_options(test_prompt())).await; @@ -1207,7 +1250,7 @@ mod chatgpt_conformance { // Thin-wrapper conformance (OpenAI-compatible providers without real cassettes) // // alibaba / baseten / bytedance / deepinfra / fireworks / moonshotai / -// togetherai / vercel are all thin wrappers around OpenAIProvider. Their +// togetherai / vercel are all registry presets of the OpenAI-compatible package. Their // responses are byte-for-byte OpenAI Chat Completions, so cassettes are // derived from real OpenAI recordings (path + model rewritten). See // scripts/generate_thin_wrapper_cassettes.py. @@ -1216,43 +1259,19 @@ mod chatgpt_conformance { /// Generates a conformance module for an OpenAI-compatible thin-wrapper provider. /// Each gets a generate + stream test. /// -/// Two arms: -/// - native arm (`$provider:ty, $config:ty`): providers with their own src/ -/// implementation (e.g. LmStudio) keep the config/provider construction. -/// - registry arm (`$name:literal`): the retired phase-4 shell types are built -/// through the registry-backed `provider(name, ...)` entry point. +/// Every provider is built through the registry-backed `create_provider(name, ...)` +/// entry point (a preset of `provider_registry.json`). macro_rules! thin_wrapper_conformance { - ($mod_name:ident, $provider:ty, $config:ty, $cassette_dir:expr, $base_url_prefix:expr, $model_id:expr) => { - mod $mod_name { - use super::*; - - fn make_provider(server: &MockServer) -> Box { - Box::new( - <$provider>::new(<$config>::new("test-key").with_base_url(format!( - "{}{}", - server.uri(), - $base_url_prefix - ))) - .model($model_id), - ) - } - - thin_wrapper_tests!($cassette_dir); - } - }; ($mod_name:ident, $name:literal, $cassette_dir:expr, $base_url_prefix:expr, $model_id:expr) => { mod $mod_name { use super::*; - fn make_provider(server: &MockServer) -> Box { - provider( + fn make_provider(server: &MockServer) -> Arc { + registry_model( $name, - Some("test-key".to_string()), + "test-key".to_string(), $model_id, - Some(ProviderOptions { - base_url: Some(format!("{}{}", server.uri(), $base_url_prefix)), - ..Default::default() - }), + format!("{}{}", server.uri(), $base_url_prefix), ) .expect("registry provider should construct") } @@ -1395,8 +1414,7 @@ thin_wrapper_conformance!( ); thin_wrapper_conformance!( lmstudio_conformance, - LmStudioProvider, - LmStudioConfig, + "lmstudio", "tests/cassettes/lmstudio", "/v1", "llama-3.2-3b-instruct" diff --git a/aimux-providers/tests/data_loss_regression_test.rs b/aimux-providers/tests/data_loss_regression_test.rs index 18147574f..cfbd5e165 100644 --- a/aimux-providers/tests/data_loss_regression_test.rs +++ b/aimux-providers/tests/data_loss_regression_test.rs @@ -22,7 +22,8 @@ //! against re-reads of the cassette (a re-read would pass even if both the //! producer and the assertion regressed together). -use aimux_core::tool::{RawToolCall, ToolResult}; +use aimux_core::tool::RawToolCall; +use aimux_core::tool::ToolResult; use std::collections::BTreeMap; use std::fs; use std::path::PathBuf; @@ -41,12 +42,16 @@ use aimux_core::stream_part::StreamPart; use aimux_core::tool::{ProviderTool, Tool}; use aimux_core::types::ProviderMetadata; -use aimux_providers::anthropic::AnthropicConfig; -use aimux_providers::anthropic::model::AnthropicModel; -use aimux_providers::bedrock::{BedrockAuth, BedrockConfig, BedrockModel, event_stream}; +use aimux_provider_utils::Resolvable; +use aimux_providers::anthropic::AnthropicMessagesModel; +use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; +use aimux_providers::bedrock::{ + AmazonBedrockProviderSettings, BedrockModel, create_amazon_bedrock, event_stream, +}; +use aimux_providers::openai::{OpenAIProviderSettings, create_openai}; use aimux_providers::{ - CohereConfig, CohereProvider, GoogleConfig, GoogleProvider, MistralConfig, MistralProvider, - OpenAIConfig, OpenAIProvider, + CohereProviderSettings, GoogleProvider, GoogleProviderSettings, MistralProviderSettings, + create_cohere, create_google, create_mistral, }; // ═══════════════════════════════════════════════════════════════════════════ @@ -329,7 +334,12 @@ const PNG_BASE64_PREFIX: &str = "iVBORw0KGgoAAAANSUhEUgAA"; const NANO_BANANA_BASE64_LEN: usize = 258_820; fn google_at(uri: &str) -> GoogleProvider { - GoogleProvider::new(GoogleConfig::new("test-api-key").with_base_url(format!("{uri}/v1beta"))) + create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(format!("{uri}/v1beta")), + ..Default::default() + }) + .unwrap() } #[tokio::test] @@ -339,7 +349,7 @@ async fn finding_1_gemini_inline_data_surfaces_as_file() { mount(&server, &c).await; let result = google_at(&server.uri()) - .model("gemini-2.5-flash-image") + .chat("gemini-2.5-flash-image") .do_generate(&opts()) .await .expect("do_generate should succeed"); @@ -372,7 +382,7 @@ async fn finding_1_gemini_image_and_text_output_both_survive() { mount(&server, &c).await; let result = google_at(&server.uri()) - .model("gemini-2.5-flash-image") + .chat("gemini-2.5-flash-image") .do_generate(&opts()) .await .expect("do_generate should succeed"); @@ -437,7 +447,7 @@ async fn finding_1_gemini_inline_data_streams_as_file_part() { let parts = collect( google_at(&server.uri()) - .model("gemini-2.5-flash-image") + .chat("gemini-2.5-flash-image") .do_stream(&opts()) .await .expect("do_stream should succeed"), @@ -484,7 +494,7 @@ async fn finding_3_gemini_thought_part_becomes_reasoning_not_text() { mount(&server, &c).await; let result = google_at(&server.uri()) - .model("gemini-3-pro-preview") + .chat("gemini-3-pro-preview") .do_generate(&opts()) .await .expect("do_generate should succeed"); @@ -525,7 +535,7 @@ async fn finding_5_gemini_thought_signature_reaches_provider_metadata() { mount(&server, &c).await; let result = google_at(&server.uri()) - .model("gemini-3-pro-preview") + .chat("gemini-3-pro-preview") .do_generate(&opts()) .await .expect("do_generate should succeed"); @@ -563,7 +573,7 @@ async fn finding_5_gemini_function_call_thought_signature_round_trips() { mount(&server, &c).await; let result = google_at(&server.uri()) - .model("gemini-2.5-flash") + .chat("gemini-2.5-flash") .do_generate(&opts()) .await .expect("do_generate should succeed"); @@ -593,7 +603,7 @@ async fn finding_12_gemini_code_execution_result_surfaced_with_matching_call_id( mount(&server, &c).await; let result = google_at(&server.uri()) - .model("gemini-3-flash-preview") + .chat("gemini-3-flash-preview") .do_generate(&opts_with_tools(vec![provider_tool( "google.code_execution", "code_execution", @@ -663,7 +673,7 @@ async fn finding_14_gemini_grounding_chunks_become_sources() { mount(&server, &c).await; let result = google_at(&server.uri()) - .model("gemini-2.5-pro") + .chat("gemini-2.5-pro") .do_generate(&opts_with_tools(vec![provider_tool( "google.google_search", "google_search", @@ -708,11 +718,14 @@ async fn finding_14_gemini_grounding_chunks_become_sources() { // camel-cased to match the upstream contract. // ═══════════════════════════════════════════════════════════════════════════ -fn anthropic_at(uri: &str) -> AnthropicModel { - AnthropicModel::new( - "claude-sonnet-4-0".to_string(), - AnthropicConfig::new("test-api-key").with_base_url(uri.to_string()), - ) +fn anthropic_at(uri: &str) -> AnthropicMessagesModel { + create_anthropic(AnthropicProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(format!("{uri}/v1")), + ..Default::default() + }) + .unwrap() + .messages("claude-sonnet-4-0") } #[tokio::test] @@ -1119,11 +1132,14 @@ async fn finding_25_cohere_citation_metadata_preserved_field_by_field() { let server = MockServer::start().await; mount(&server, &c).await; - let provider = CohereProvider::new( - CohereConfig::new("test-key").with_base_url(format!("{}/v2", server.uri())), - ); + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v2", server.uri())), + ..Default::default() + }) + .expect("valid settings"); let result = provider - .model("command-r-plus") + .chat("command-r-plus") .do_generate(&opts()) .await .expect("do_generate should succeed"); @@ -1176,11 +1192,14 @@ async fn finding_13_mistral_thinking_parts_become_reasoning_in_generate() { let server = MockServer::start().await; mount(&server, &c).await; - let provider = MistralProvider::new( - MistralConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())), - ); + let provider = create_mistral(MistralProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .expect("valid settings"); let result = provider - .model("mistral-small-latest") + .chat("mistral-small-latest") .do_generate(&opts()) .await .expect("do_generate should succeed"); @@ -1209,8 +1228,13 @@ async fn finding_13_mistral_thinking_parts_become_reasoning_in_generate() { // ═══════════════════════════════════════════════════════════════════════════ fn openai_responses_at(uri: &str, model: &str) -> impl LanguageModel { - OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(format!("{uri}/v1"))) - .responses_model(model) + create_openai(OpenAIProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{uri}/v1")), + ..Default::default() + }) + .unwrap() + .responses(model) } #[tokio::test] @@ -1313,15 +1337,14 @@ async fn finding_30_openai_responses_url_citation_annotation_becomes_source() { const BEDROCK_MODEL: &str = "us.anthropic.claude-sonnet-4-20250514-v1:0"; fn bedrock_at(uri: &str, model: &str) -> BedrockModel { - BedrockModel::new( - model.to_string(), - BedrockConfig { - base_url: uri.to_string(), - auth: BedrockAuth::BearerToken("test-token".to_string()), - retry_config: aimux_provider_utils::RetryConfig::default(), - api_key_source: None, - }, - ) + create_amazon_bedrock(AmazonBedrockProviderSettings { + api_key: Some(Resolvable::Value("test-token".to_string())), + region: Some("us-east-1".to_string()), + base_url: Some(uri.to_string()), + ..Default::default() + }) + .unwrap() + .chat(model) } /// The real signature recorded in @@ -1363,9 +1386,9 @@ 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_eq!( - m["bedrock"]["signature"], m["amazonBedrock"]["signature"], - "the signature is mirrored under both provider keys" + assert!( + m.get("bedrock").is_none(), + "the legacy provider key is not written" ); // The visible answer is separate from the reasoning. @@ -1457,7 +1480,7 @@ async fn finding_10_bedrock_stream_reasoning_signature_emitted_as_metadata_delta }) .expect("the ReasoningEnd must carry the signature"); assert_eq!(end_meta["amazonBedrock"]["signature"], json!(signature)); - assert_eq!(end_meta["bedrock"]["signature"], json!(signature)); + assert!(end_meta.get("bedrock").is_none()); } #[tokio::test] @@ -1510,9 +1533,9 @@ async fn finding_27_bedrock_stream_performance_config_and_service_tier_reach_fin meta["amazonBedrock"]["serviceTier"], json!({ "type": "flex" }) ); - assert_eq!( - meta["bedrock"], meta["amazonBedrock"], - "both provider keys carry the same payload" + assert!( + meta.get("bedrock").is_none(), + "the legacy provider key is not written" ); } @@ -1570,7 +1593,7 @@ async fn finding_27_bedrock_stream_guardrail_trace_reaches_finish() { json!(397), "nested metrics must not be flattened away" ); - assert_eq!(meta["bedrock"]["trace"], meta["amazonBedrock"]["trace"]); + assert!(meta.get("bedrock").is_none()); } #[tokio::test] @@ -1611,5 +1634,5 @@ async fn finding_26_bedrock_stream_stop_sequence_reaches_finish() { }) .expect("Finish must carry provider_metadata"); assert_eq!(meta["amazonBedrock"]["stopSequence"], json!("STOP")); - assert_eq!(meta["bedrock"]["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 b4517e30f..34a7930f4 100644 --- a/aimux-providers/tests/dataforseo_test.rs +++ b/aimux-providers/tests/dataforseo_test.rs @@ -11,10 +11,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{DataforseoConfig, DataforseoProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{DataforseoProvider, DataforseoProviderSettings, create_dataforseo}; const LOGIN: &str = "test-login"; const PASSWORD: &str = "test-password"; @@ -62,8 +62,13 @@ fn search_response_body() -> Value { } fn provider(server: &MockServer) -> DataforseoProvider { - let config = DataforseoConfig::new(LOGIN, PASSWORD).with_base_url(server.uri()); - DataforseoProvider::new(config) + let config = DataforseoProviderSettings { + login: Some(Resolvable::Value(LOGIN.to_string())), + password: Some(Resolvable::Value(PASSWORD.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_dataforseo(config).unwrap() } fn opts(query: &str, max_results: Option) -> SearchCallOptions { @@ -219,25 +224,3 @@ async fn status_401_maps_to_auth_error() { } // -- Provider trait ---------------------------------------------------------- - -#[tokio::test] -async fn provider_name_is_dataforseo() { - let config = DataforseoConfig::new(LOGIN, PASSWORD); - let provider = DataforseoProvider::new(config); - assert_eq!(provider.name(), "dataforseo"); -} - -#[test] -fn language_model_returns_unsupported_error() { - let config = DataforseoConfig::new(LOGIN, PASSWORD); - let provider = DataforseoProvider::new(config); - match provider.language_model("dataforseo-search") { - Err(AiMuxError::UnsupportedFunctionality(msg)) => { - assert!( - msg.contains("provider 'dataforseo' does not provide language models"), - "unexpected message: {msg}" - ); - } - _ => panic!("expected Unsupported error, got success or another error variant"), - } -} diff --git a/aimux-providers/tests/decision_recording_test.rs b/aimux-providers/tests/decision_recording_test.rs index ba18173f9..08cb9d732 100644 --- a/aimux-providers/tests/decision_recording_test.rs +++ b/aimux-providers/tests/decision_recording_test.rs @@ -42,9 +42,8 @@ fn response() -> Value { example().outcome.decision_result.unwrap()["response"]["body"].clone() } fn model(server: &MockServer) -> Box { - let mut config = + let config = JevConfig::new("live-key-secret").with_endpoint(format!("{}/v1/systemone", server.uri())); - config.retry_config.initial_delay = Duration::ZERO; JevProvider::new(config) .decision_model("jev-latest") .unwrap() @@ -82,6 +81,13 @@ async fn ring_records_one_complete_call_and_replays_exact_inputs_with_redaction( assert!(records[0].complete && records[0].transport_closed); assert_eq!(records[0].input.operation, RecordingOperation::Decision); assert!(records[0].input.prompt.is_empty()); + assert!(records[0].input.decision_capabilities.is_some()); + assert_eq!( + serde_json::to_value(&records[0].provider).unwrap(), + json!({ + "provider_id": "jev", "provider": "jev", "model_id": "jev-latest" + }) + ); assert_eq!(records[0].exchanges.len(), 1); assert_eq!(records[0].exchanges[0].attempt, 1); assert_eq!(records[0].exchanges[0].exchange_index, 1); @@ -110,7 +116,14 @@ async fn ring_records_one_complete_call_and_replays_exact_inputs_with_redaction( #[tokio::test] #[serial] async fn jsonl_flush_has_no_duplicate_placeholders_and_supports_offline_replay() { - let directory = std::env::temp_dir().join(recording::new_call_id()); + let directory = std::env::temp_dir().join(format!( + "aimux-decision-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); let recorder = Arc::new(JsonlRecorder::new(&directory)); recording::init_recording(Some(recorder.clone())); let _stop = StopRecording; @@ -224,7 +237,7 @@ async fn failed_response_and_timeout_have_complete_error_recordings() { #[tokio::test] #[serial] -async fn recorded_native_calls_replay_offline_and_provider_rebuild_preserves_endpoint() { +async fn recorded_native_calls_replay_offline_and_rebuild_uses_registered_endpoint() { recording::init_recording(None); let file = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) .join("tests/fixtures/jev_systemone_live.jsonl"); @@ -238,12 +251,17 @@ async fn recorded_native_calls_replay_offline_and_provider_rebuild_preserves_end ); let server = MockServer::start().await; common::replay::mount_recording(&server, &rec).await; - let mut snapshot = rec.provider.clone(); - snapshot.base_url = Some(format!("{}/v1/systemone", server.uri())); - snapshot.provider_options = - Some(json!({"headers": {"X-Routing": "billing", "Authorization": "[REDACTED]"}})); - let model = - aimux_providers::rebuild_decision_provider(&snapshot, Some("test-key")).unwrap(); + let mut config = + JevConfig::new("test-key").with_endpoint(format!("{}/v1/systemone", server.uri())); + config.headers.insert("X-Routing".into(), "billing".into()); + let registry = aimux_core::create_provider_registry( + std::collections::BTreeMap::from([( + "jev".into(), + Arc::new(JevProvider::new(config)) as Arc, + )]), + Default::default(), + ); + let model = aimux_providers::rebuild_decision_provider(&rec.provider, ®istry).unwrap(); replay_decision_with_model(&rec, model.as_ref()) .await .unwrap(); diff --git a/aimux-providers/tests/deepgram_transcription_test.rs b/aimux-providers/tests/deepgram_transcription_test.rs index 4d09cb343..d52f99e60 100644 --- a/aimux-providers/tests/deepgram_transcription_test.rs +++ b/aimux-providers/tests/deepgram_transcription_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Deepgram transcription model tests. +//! Rust translation of the Deepgram transcription model tests. //! //! Source: `reference/ai/packages/deepgram/src/deepgram-transcription-model.test.ts` @@ -9,7 +9,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; -use aimux_providers::{DeepgramConfig, DeepgramProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{DeepgramProviderSettings, create_deepgram}; fn mock_audio() -> Vec { vec![1u8, 2, 3, 4, 5] @@ -68,8 +69,12 @@ async fn should_pass_model_in_query() { let server = MockServer::start().await; mock_response(&server, &fixture_response()).await; - let config = DeepgramConfig::new("test-api-key").with_base_url(server.uri()); - let provider = DeepgramProvider::new(config); + 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"); model @@ -93,10 +98,13 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = DeepgramConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = DeepgramProvider::new(config); + let config = DeepgramProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_deepgram(config).unwrap(); let model = provider.transcription("nova-3"); let mut opts = options(mock_audio(), "audio/wav"); @@ -128,8 +136,12 @@ async fn should_extract_text() { let server = MockServer::start().await; mock_response(&server, &fixture_response()).await; - let config = DeepgramConfig::new("test-api-key").with_base_url(server.uri()); - let provider = DeepgramProvider::new(config); + 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 result = model @@ -153,8 +165,12 @@ async fn should_include_response_data() { ) .await; - let config = DeepgramConfig::new("test-api-key").with_base_url(server.uri()); - let provider = DeepgramProvider::new(config); + 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 result = model @@ -174,8 +190,12 @@ async fn should_use_real_date() { let server = MockServer::start().await; mock_response(&server, &fixture_response()).await; - let config = DeepgramConfig::new("test-api-key").with_base_url(server.uri()); - let provider = DeepgramProvider::new(config); + 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 result = model @@ -203,8 +223,12 @@ async fn should_return_none_language_when_not_detected() { ) .await; - let config = DeepgramConfig::new("test-api-key").with_base_url(server.uri()); - let provider = DeepgramProvider::new(config); + 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 result = model diff --git a/aimux-providers/tests/deepinfra_test.rs b/aimux-providers/tests/deepinfra_test.rs index a1d6e3799..54612580b 100644 --- a/aimux-providers/tests/deepinfra_test.rs +++ b/aimux-providers/tests/deepinfra_test.rs @@ -2,7 +2,7 @@ //! //! Translated from `packages/deepinfra/src/deepinfra-provider.test.ts`. //! -//! DeepInfra is a thin OpenAI-compatible wrapper over [`OpenAIProvider`]. The +//! DeepInfra is a registry preset of the OpenAI-compatible package. The //! TS suite is mostly provider-configuration unit tests plus image/completion //! model construction. The Rust wrapper models only the chat surface, so the //! image/completion model tests are not translated. The DeepInfra default base @@ -21,6 +21,7 @@ use serde_json::{Value, json}; use serial_test::serial; +use std::sync::Arc; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -28,7 +29,7 @@ use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; -use aimux_providers::{ProviderOptions, provider, provider_from_env}; +use aimux_providers::{PresetSettings, create_provider}; fn test_prompt() -> LanguageModelPrompt { vec![LanguageModelMessage::user_text("Hello")] @@ -53,28 +54,32 @@ fn text_completion_body() -> Value { }) } -fn make_provider(server: &MockServer) -> Box { - provider( +fn make_provider(server: &MockServer) -> Arc { + create_provider( "deepinfra", - Some("test-api-key".to_string()), - "meta-llama/Meta-Llama-3-70B-Instruct", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-api-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("meta-llama/Meta-Llama-3-70B-Instruct") .expect("deepinfra provider should build") } /// TS: `createDeepInfra()` produces a model for the "deepinfra" registry entry. #[test] fn provider_builds_deepinfra_model() { - let model = provider( + let model = create_provider( "deepinfra", - Some("test-key".to_string()), - "meta-llama/Meta-Llama-3-70B-Instruct", - None, + PresetSettings { + api_key: Some("test-key".to_string().into()), + ..Default::default() + }, ) + .expect("provider should build") + .language_model("meta-llama/Meta-Llama-3-70B-Instruct") .expect("deepinfra provider should build"); assert_eq!(model.model_id(), "meta-llama/Meta-Llama-3-70B-Instruct"); } @@ -110,15 +115,16 @@ async fn custom_api_key_used_in_auth_header() { .mount(&server) .await; - let model = provider( + let model = create_provider( "deepinfra", - Some("my-custom-key".to_string()), - "meta-llama/Meta-Llama-3-70B-Instruct", - Some(ProviderOptions { + PresetSettings { + api_key: Some("my-custom-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("meta-llama/Meta-Llama-3-70B-Instruct") .expect("deepinfra provider should build"); model @@ -139,15 +145,16 @@ async fn custom_headers_forwarded() { .mount(&server) .await; - let model = provider( + let model = create_provider( "deepinfra", - Some("test-key".to_string()), - "meta-llama/Meta-Llama-3-70B-Instruct", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("meta-llama/Meta-Llama-3-70B-Instruct") .expect("deepinfra provider should build"); let mut options = default_options(test_prompt()); @@ -190,7 +197,9 @@ fn from_env_loads_deepinfra_api_key() { std::env::set_var("DEEPINFRA_API_KEY", "env-test-key"); } - let model = provider_from_env("deepinfra", "meta-llama/Meta-Llama-3-70B-Instruct", None); + let model = create_provider("deepinfra", PresetSettings::default()) + .expect("provider should build") + .language_model("meta-llama/Meta-Llama-3-70B-Instruct"); assert!(model.is_ok(), "from_env should succeed with env var set"); unsafe { @@ -203,15 +212,19 @@ fn from_env_loads_deepinfra_api_key() { /// TS: without the env var, `createDeepInfra()` fails. #[serial] -#[test] -fn from_env_fails_without_env_var() { +#[tokio::test] +async fn from_env_fails_without_env_var() { let saved = std::env::var("DEEPINFRA_API_KEY").ok(); unsafe { std::env::remove_var("DEEPINFRA_API_KEY"); } - let model = provider_from_env("deepinfra", "meta-llama/Meta-Llama-3-70B-Instruct", None); - assert!(model.is_err(), "from_env should fail without env var"); + let model = create_provider("deepinfra", PresetSettings::default()) + .expect("provider should build") + .language_model("meta-llama/Meta-Llama-3-70B-Instruct") + .expect("provider should build a model"); + let result = model.do_generate(&default_options(test_prompt())).await; + assert!(result.is_err(), "from_env should fail without env var"); unsafe { if let Some(v) = saved { diff --git a/aimux-providers/tests/deepseek_chat_test.rs b/aimux-providers/tests/deepseek_chat_test.rs index 0ff6132c9..60ff24bd6 100644 --- a/aimux-providers/tests/deepseek_chat_test.rs +++ b/aimux-providers/tests/deepseek_chat_test.rs @@ -1,147 +1,179 @@ -//! DeepSeek basic chat tests, translated from the Vercel AI SDK TypeScript suite. +//! The DeepSeek chat model (`aimux_providers::deepseek`), ported from the AI SDK +//! `@ai-sdk/deepseek` 3.0.56 test suites: //! -//! Translation sources: -//! - `packages/deepseek/src/chat/deepseek-chat-language-model.test.ts` -//! - `describe('doGenerate') › describe('text')` — request body + text extraction -//! - `describe('doGenerate') › describe('tool call')` — tool call request body, -//! JSON response format (without/with schema), tool call extraction -//! - `describe('doStream') › describe('text')` — stream request body + text streaming -//! - `describe('doStream') › describe('tool call')` — tool call streaming -//! - Reasoning tests are in `deepseek_reasoning_test.rs` and are NOT duplicated here. -//! - `packages/deepseek/src/chat/convert-to-deepseek-chat-messages.test.ts` -//! - User messages (text, file warning, mediaType) -//! - Tool calls (stringify arguments, text output) -//! - Reasoning-related convert tests are excluded (covered by reasoning test file). -//! - `packages/deepseek/src/chat/deepseek-prepare-tools.test.ts` -//! - Strict mode pass-through (true, false, undefined, mixed) +//! - `chat/deepseek-chat-language-model.test.ts`: `model IDs`, `supportedUrls`, +//! `doGenerate` (`text`, `reasoning`, `logprobs`, `top-level reasoning`, +//! `tool call`, `json response format`, `assistant prefix completion`) and +//! `doStream` (stream errors, `text`, `reasoning`, `logprobs`, `tool call`, +//! `assistant prefix completion`). +//! - `chat/convert-to-deepseek-chat-messages.test.ts`: `message names`, +//! `user messages`, `tool calls`, `deepseek-v4 thinking mode` and +//! `assistant prefix completion`, observed through the request body. //! -//! HTTP is mocked with `wiremock` (a real loopback server), replacing the TS -//! `createTestServer`. Each test starts its own `MockServer` so parallel -//! `#[tokio::test]` runs do not collide. +//! 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). //! -//! # Mock URL note -//! -//! The real DeepSeek base URL is `https://api.deepseek.com/v1`, so production -//! requests hit `/v1/chat/completions`. In these tests `with_base_url` is -//! overridden with the mock server's root URI (no `/v1` suffix), so the -//! resulting request path is `/chat/completions`. +//! HTTP is a local `wiremock` server; every model is built with +//! `create_deepseek(..).chat(..)`. +use aimux_core::tool::RawToolCall; use futures::StreamExt; use serde_json::{Value, json}; -use wiremock::matchers::{method, path}; +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::{ - FilePart, LanguageModelMessage, LanguageModelPrompt, TextPart, UserPart, + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultOutput, ToolResultPart, UserPart, }; -use aimux_core::options::{CallOptions, Tool}; -use aimux_core::result::GenerateContent; +use aimux_core::options::{CallOptions, ResponseFormat, Tool}; +use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput}; +use aimux_core::shared::provider_namespace; use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::FunctionTool; -use aimux_core::types::FinishReasonUnified; +use aimux_core::types::{FinishReasonUnified, ReasoningEffort, ResponseMetadata, Warning}; +use aimux_providers::deepseek::{ + DeepSeekChatLanguageModel, DeepSeekProviderSettings, create_deepseek, +}; -use aimux_providers::openai::OpenAICompatProfile; -use aimux_providers::openai::convert::build_request_body_with_warnings; -use aimux_providers::{ProviderOptions, provider}; +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- -// ───────────────────────────────────────────────────────────────────────────── -// Shared helpers (mirrors deepseek_reasoning_test.rs) -// ───────────────────────────────────────────────────────────────────────────── +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, + }; + *provider_options = Some(serde_json::from_value(options).unwrap()); + message +} -/// The TS `TEST_PROMPT`: a single user text message "Hello". +/// The upstream `TEST_PROMPT`. fn test_prompt() -> LanguageModelPrompt { vec![LanguageModelMessage::user_text("Hello")] } -/// `CallOptions` with only `prompt` set (everything else default/None). -fn default_options(prompt: LanguageModelPrompt) -> CallOptions { +fn options() -> CallOptions { + CallOptions::new(test_prompt()) +} + +fn options_for(prompt: LanguageModelPrompt) -> CallOptions { CallOptions::new(prompt) } -/// Build a DeepSeek provider whose base URL points at the mock server. -fn make_provider(server: &MockServer, model_id: &str) -> Box { - provider( - "deepseek", - Some("test-api-key".to_string()), - model_id, - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() +fn with_provider_options(mut options: CallOptions, value: Value) -> CallOptions { + options.provider_options = Some(serde_json::from_value(value).unwrap()); + options +} + +fn weather_tool() -> FunctionTool { + FunctionTool { + name: "weather".to_string(), + description: None, + input_schema: json!({ + "type": "object", + "properties": { "location": { "type": "string" } }, + "required": ["location"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#" }), - ) - .expect("deepseek provider should build") + strict: None, + provider_options: None, + input_examples: None, + } } -/// Mount a JSON chat-completion response on `/chat/completions`. -async fn mock_json(server: &MockServer, body: Value) { +fn strict_tool(name: &str, strict: Option) -> Tool { + Tool::Function(FunctionTool { + name: name.to_string(), + description: None, + input_schema: json!({ "type": "object", "properties": {} }), + strict, + provider_options: None, + input_examples: None, + }) +} + +fn settings(server: &MockServer, base_path: &str) -> DeepSeekProviderSettings { + DeepSeekProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(format!("{}{base_path}", server.uri())), + ..Default::default() + } +} + +/// `provider.chat(model_id)` of the upstream tests. +fn chat(server: &MockServer, model_id: &str) -> DeepSeekChatLanguageModel { + create_deepseek(settings(server, "")) + .unwrap() + .chat(model_id) +} + +/// `betaProvider.chat(model_id)` of the upstream tests. +fn beta_chat(server: &MockServer, model_id: &str) -> DeepSeekChatLanguageModel { + create_deepseek(settings(server, "/beta")) + .unwrap() + .chat(model_id) +} + +async fn json_server(body: Value) -> MockServer { + let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/chat/completions")) + .and(path_regex(".*/chat/completions$")) .respond_with(ResponseTemplate::new(200).set_body_json(body)) - .mount(server) + .mount(&server) .await; + server +} + +fn data_event(chunk: &Value) -> String { + format!("data: {chunk}\n\n") } -/// Mount an SSE stream response on `/chat/completions`. -async fn mock_sse(server: &MockServer, body: String) { +async fn sse_server(chunks: Vec) -> MockServer { + let mut body = chunks.concat(); + body.push_str("data: [DONE]\n\n"); + let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/chat/completions")) + .and(path_regex(".*/chat/completions$")) .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") .set_body_string(body), ) - .mount(server) + .mount(&server) .await; + server } -/// Build a single SSE `data: \n\n` event string. -fn sse_event(json_str: &str) -> String { - format!("data: {json_str}\n\n") -} - -/// Concatenate SSE events and append the `[DONE]` sentinel. -fn sse_body(events: &[&str]) -> String { - let mut body = String::new(); - for event in events { - body.push_str(event); - } - body.push_str("data: [DONE]\n\n"); - body -} - -/// Collect every `StreamPart` from a `StreamResult` into a `Vec`. -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 -} - -/// Extract text deltas from a list of stream parts. -fn text_deltas(parts: &[StreamPart]) -> Vec { - parts +async fn request_paths(server: &MockServer) -> Vec { + server + .received_requests() + .await + .unwrap() .iter() - .filter_map(|p| match p { - StreamPart::TextDelta { delta, .. } => Some(delta.clone()), - _ => None, - }) + .map(|request| request.url.path().to_string()) .collect() } -// ───────────────────────────────────────────────────────────────────────────── -// Local fixtures -// ───────────────────────────────────────────────────────────────────────────── - -/// A minimal non-streaming text completion body returning a markdown response. -/// Abbreviated from the `deepseek-text` fixture. -fn text_completion_body() -> Value { +fn text_response() -> Value { json!({ "id": "00f10ecd-60b3-4707-b5db-e4bcadf7aea1", "object": "chat.completion", @@ -149,328 +181,355 @@ fn text_completion_body() -> Value { "model": "deepseek-chat", "choices": [{ "index": 0, - "message": { - "role": "assistant", - "content": "## **Holiday Name: Gratitude of Small Things Day (GST Day)**" - }, + "message": { "role": "assistant", "content": "Hello, World!" }, "logprobs": null, - "finish_reason": "length" + "finish_reason": "stop" }], "usage": { "prompt_tokens": 13, "completion_tokens": 300, "total_tokens": 313, - "prompt_tokens_details": { "cached_tokens": 0 } + "prompt_cache_hit_tokens": 0, + "prompt_cache_miss_tokens": 13 }, "system_fingerprint": "fp_eaab8d114b_prod0820_fp8_kvcache" }) } -// ════════════════════════════════════════════════════════════════════════════ -// doGenerate — text -// -// Translated from `deepseek-chat-language-model.test.ts` doGenerate › text. -// ════════════════════════════════════════════════════════════════════════════ +/// Generate against a canned text response and return the request body. +async fn generate_body(model_id: &str, options: &CallOptions) -> Value { + let server = json_server(text_response()).await; + chat(&server, model_id) + .do_generate(options) + .await + .unwrap() + .request + .and_then(|request| request.body) + .unwrap() +} -/// TS: doGenerate › text › "should extract text content" (line ~68). -/// -/// Verifies that the text content is extracted from the response and the -/// finish reason / usage are parsed correctly. -#[tokio::test] -async fn should_extract_text_content() { - let server = MockServer::start().await; - mock_json(&server, text_completion_body()).await; +async fn generate_result(model_id: &str, options: &CallOptions) -> GenerateResult { + let server = json_server(text_response()).await; + chat(&server, model_id).do_generate(options).await.unwrap() +} - let model = make_provider(&server, "deepseek-chat"); +/// The wire messages of a generate call. +async fn wire_messages(model_id: &str, prompt: LanguageModelPrompt) -> Vec { + generate_body(model_id, &options_for(prompt)).await["messages"] + .as_array() + .unwrap() + .clone() +} - let result = model - .do_generate(&default_options(test_prompt())) +/// Run a generate call that must fail before anything is sent. +async fn generate_error(model_id: &str, options: &CallOptions) -> String { + let server = json_server(text_response()).await; + let error = chat(&server, model_id) + .do_generate(options) .await - .expect("should succeed"); + .expect_err("the call must be rejected"); + assert!( + request_paths(&server).await.is_empty(), + "nothing may be fetched" + ); + error.to_string() +} - assert_eq!(result.content.len(), 1); - match &result.content[0] { - GenerateContent::Text { text, .. } => assert!( - text.contains("Gratitude of Small Things Day"), - "expected the fixture text, got: {text}" +fn warning_values(warnings: &[Warning]) -> Vec { + warnings + .iter() + .map(|warning| serde_json::to_value(warning).unwrap()) + .collect() +} + +fn deprecated(setting: &str) -> Warning { + Warning::Deprecated { + setting: setting.to_string(), + message: format!( + "{setting} is deprecated by DeepSeek and has been omitted. Remove {setting} from the request." ), - other => panic!("expected Text, got {other:?}"), } - assert_eq!(result.finish_reason.unified, FinishReasonUnified::Length); - assert_eq!(result.finish_reason.raw.as_deref(), Some("length")); - // Usage: prompt_tokens=13, completion_tokens=300 - assert_eq!(result.usage.input_tokens.total, Some(13)); - assert_eq!(result.usage.output_tokens.total, Some(300)); } -// ════════════════════════════════════════════════════════════════════════════ -// doGenerate — tool call -// -// Translated from `deepseek-chat-language-model.test.ts` doGenerate › tool call. -// ════════════════════════════════════════════════════════════════════════════ - -// ════════════════════════════════════════════════════════════════════════════ -// doGenerate — JSON response format -// -// Translated from `deepseek-chat-language-model.test.ts` doGenerate › tool call -// › "json response format" (lines ~378-568). -// -// TS behaviour: without `supportsStructuredOutputs`, DeepSeek uses -// `response_format: { type: "json_object" }` and injects a system message -// ("Return JSON." or "Return JSON that conforms to the following schema: ..."). -// ════════════════════════════════════════════════════════════════════════════ - -// ════════════════════════════════════════════════════════════════════════════ -// doStream — text -// -// Translated from `deepseek-chat-language-model.test.ts` doStream › text. -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: doStream › text › "should stream text" (line ~748). -/// -/// Verifies that text deltas are streamed as TextStart / TextDelta / TextEnd -/// parts, followed by a Finish part with usage. +fn thinking_unsupported(feature: &str) -> Warning { + Warning::Unsupported { + feature: feature.to_string(), + details: Some(format!( + "{feature} has no effect when DeepSeek thinking is enabled. Set providerOptions.deepseek.thinking.type to 'disabled' to use {feature}." + )), + } +} + +fn compatibility(feature: &str, details: &str) -> Warning { + Warning::Compatibility { + feature: feature.to_string(), + details: Some(details.to_string()), + } +} + +async fn collect(stream: aimux_core::result::StreamResult) -> Vec { + let mut parts = Vec::new(); + let mut stream = stream.stream; + while let Some(part) = stream.next().await { + parts.push(part.expect("stream part")); + } + parts +} + +fn text_chunks() -> Vec { + vec![ + data_event(&json!({ + "id": "c1", "object": "chat.completion.chunk", "created": 1764656316, + "model": "deepseek-chat", "system_fingerprint": "fp_eaab8d114b_prod0820_fp8_kvcache", + "choices": [{ "index": 0, "delta": { "role": "assistant", "content": "" }, "finish_reason": null }] + })), + data_event(&json!({ + "id": "c1", "object": "chat.completion.chunk", "created": 1764656316, + "model": "deepseek-chat", "system_fingerprint": "fp_eaab8d114b_prod0820_fp8_kvcache", + "choices": [{ "index": 0, "delta": { "content": "Hello" }, "finish_reason": null }] + })), + data_event(&json!({ + "id": "c1", "object": "chat.completion.chunk", "created": 1764656316, + "model": "deepseek-chat", "system_fingerprint": "fp_eaab8d114b_prod0820_fp8_kvcache", + "choices": [{ "index": 0, "delta": { "content": ", World!" }, "finish_reason": null }] + })), + data_event(&json!({ + "id": "c1", "object": "chat.completion.chunk", "created": 1764656316, + "model": "deepseek-chat", "system_fingerprint": "fp_eaab8d114b_prod0820_fp8_kvcache", + "choices": [{ "index": 0, "delta": { "content": "" }, "finish_reason": "stop" }], + "usage": { + "prompt_tokens": 13, "completion_tokens": 3, "total_tokens": 16, + "prompt_cache_hit_tokens": 5, "prompt_cache_miss_tokens": 8 + } + })), + ] +} + +fn finish_metadata(parts: &[StreamPart]) -> Value { + parts + .iter() + .find_map(|part| match part { + StreamPart::Finish { + provider_metadata, .. + } => provider_metadata + .as_ref() + .map(|metadata| serde_json::to_value(metadata).unwrap()), + _ => None, + }) + .expect("a finish part") +} + +fn stream_start_warnings(parts: &[StreamPart]) -> Vec { + match &parts[0] { + StreamPart::StreamStart { warnings } => warning_values(warnings), + other => panic!("expected stream-start, got {other:?}"), + } +} + +fn logprob(token: &str, logprob: f64, bytes: Value) -> Value { + json!({ + "token": token, "logprob": logprob, "bytes": bytes, + "top_logprobs": [{ "token": token, "logprob": logprob, "bytes": bytes }] + }) +} + +// =========================================================================== +// describe('model IDs') / describe('supportedUrls') +// =========================================================================== + #[tokio::test] -async fn should_stream_text() { +async fn should_forward_the_model_id() { + for model_id in ["deepseek-v4-flash", "deepseek-v4-pro"] { + assert_eq!(generate_body(model_id, &options()).await["model"], model_id); + } +} + +#[tokio::test] +async fn should_natively_support_http_image_urls() { let server = MockServer::start().await; - let body = sse_body(&[ - &sse_event( - r#"{"id":"x","object":"chat.completion.chunk","created":1,"model":"deepseek-chat","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}"#, - ), - &sse_event( - r#"{"id":"x","object":"chat.completion.chunk","created":1,"model":"deepseek-chat","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}"#, - ), - &sse_event( - r#"{"id":"x","object":"chat.completion.chunk","created":1,"model":"deepseek-chat","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}"#, - ), - &sse_event( - r#"{"id":"x","object":"chat.completion.chunk","created":1,"model":"deepseek-chat","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}"#, - ), - ]); - mock_sse(&server, body).await; + let urls = chat(&server, "deepseek-chat").supported_urls().await; + let patterns = &urls.0["image/*"]; + assert_eq!(patterns.len(), 1); + assert!(patterns[0].is_match("https://example.com/a.png")); + assert!(patterns[0].is_match("http://example.com/a.png")); + assert!(!patterns[0].is_match("data:image/png;base64,AAAA")); +} - let model = make_provider(&server, "deepseek-chat"); +// =========================================================================== +// describe('doGenerate') +// =========================================================================== - let result = model - .do_stream(&default_options(test_prompt())) +#[tokio::test] +async fn should_reject_a_response_without_choices() { + let server = json_server(json!({ + "id": "chatcmpl-empty", "object": "chat.completion", "created": 0, + "model": "deepseek-chat", "choices": [], + "usage": { "prompt_tokens": 1, "completion_tokens": 0, "total_tokens": 1 } + })) + .await; + let error = chat(&server, "deepseek-chat") + .do_generate(&options()) .await - .expect("do_stream should succeed"); - let parts = collect_stream(result).await; - - // Should have StreamStart, ResponseMetadata, TextStart, TextDelta x2, - // TextEnd, Finish. + .unwrap_err(); assert!( - parts - .iter() - .any(|p| matches!(p, StreamPart::StreamStart { .. })), - "expected StreamStart" + matches!(&error, AiMuxError::InvalidResponseData(m) if m == "Response did not contain any choices."), + "{error:?}" ); - assert!( - parts - .iter() - .any(|p| matches!(p, StreamPart::TextStart { .. })), - "expected TextStart" +} + +// ---- describe('text') ------------------------------------------------------ + +#[tokio::test] +async fn text_should_send_correct_request_body() { + let mut options = options_for(vec![ + LanguageModelMessage::System { + content: ("You are a helpful assistant.").into(), + provider_options: None, + }, + LanguageModelMessage::user_text("Hello"), + ]); + options.temperature = Some(0.5); + options.top_p = Some(0.3); + assert_eq!( + generate_body("deepseek-chat", &options).await, + json!({ + "messages": [ + { "content": "You are a helpful assistant.", "role": "system" }, + { "content": "Hello", "role": "user" } + ], + "model": "deepseek-chat", + "temperature": 0.5, + "top_p": 0.3 + }) ); - let deltas: Vec = text_deltas(&parts) - .into_iter() - .filter(|d| !d.is_empty()) - .collect(); +} + +#[tokio::test] +async fn text_should_omit_deprecated_and_ineffective_sampling_options_in_default_v4_thinking_mode() +{ + let mut options = options(); + options.temperature = Some(0.2); + options.top_p = Some(0.4); + options.frequency_penalty = Some(0.5); + options.presence_penalty = Some(0.6); + let result = generate_result("deepseek-v4-flash", &options).await; assert_eq!( - deltas, - vec!["Hello".to_string(), " world".to_string()], - "text deltas should match" + result.request.and_then(|request| request.body).unwrap(), + json!({ "model": "deepseek-v4-flash", "messages": [{ "role": "user", "content": "Hello" }] }) ); - assert!( - parts - .iter() - .any(|p| matches!(p, StreamPart::TextEnd { .. })), - "expected TextEnd" + assert_eq!( + warning_values(&result.warnings), + warning_values(&[ + deprecated("frequencyPenalty"), + deprecated("presencePenalty"), + thinking_unsupported("temperature"), + thinking_unsupported("topP"), + ]) ); - // The last part should be Finish. - match parts.last() { - Some(StreamPart::Finish { finish_reason, .. }) => { - assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); - } - other => panic!("expected Finish as last part, got {other:?}"), - } } -// ════════════════════════════════════════════════════════════════════════════ -// doStream — tool call -// -// Translated from `deepseek-chat-language-model.test.ts` doStream › tool call -// (line ~775). Uses an abbreviated chunk set without reasoning_content -// (reasoning streaming is tested in `deepseek_reasoning_test.rs`). -// ════════════════════════════════════════════════════════════════════════════ - -// ════════════════════════════════════════════════════════════════════════════ -// Message conversion — convert-to-deepseek-chat-messages.test.ts -// -// These tests call `build_request_body_with_warnings` directly (no HTTP mock -// needed) and assert on the `body["messages"]` field. -// -// Translated from `convert-to-deepseek-chat-messages.test.ts`. -// Reasoning-related convert tests are excluded (covered by reasoning test file). -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: user messages › "should convert messages with only a text part to a -/// string content" (line ~6). -/// -/// A user message with a single text part should produce -/// `{ role: "user", content: "Hello" }`. #[tokio::test] -async fn should_convert_user_text_part_to_string_content() { - let options = default_options(test_prompt()); - let result = build_request_body_with_warnings( - "deepseek-chat", - &options, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - +async fn text_should_preserve_supported_sampling_options_when_v4_thinking_is_disabled() { + let mut options = with_provider_options( + options(), + json!({ "deepseek": { "thinking": { "type": "disabled" } } }), + ); + options.temperature = Some(0.2); + options.top_p = Some(0.4); + options.frequency_penalty = Some(0.5); + options.presence_penalty = Some(0.6); + let result = generate_result("deepseek-v4-flash", &options).await; assert_eq!( - result.body["messages"], - json!([{ "role": "user", "content": "Hello" }]) + result.request.and_then(|request| request.body).unwrap(), + json!({ + "model": "deepseek-v4-flash", + "messages": [{ "role": "user", "content": "Hello" }], + "temperature": 0.2, + "top_p": 0.4, + "thinking": { "type": "disabled" } + }) + ); + assert_eq!( + warning_values(&result.warnings), + warning_values(&[ + deprecated("frequencyPenalty"), + deprecated("presencePenalty") + ]) ); - assert!(result.warnings.is_empty(), "expected no warnings"); } -/// TS: user messages › "should warn about unsupported file parts" (line ~31). -/// -/// A user message with a text part and a file part should produce only the -/// TS: file parts with image media type are converted to image_url (OpenAI format). -#[tokio::test] -async fn should_convert_image_file_parts_to_image_url() { - let prompt = vec![LanguageModelMessage::User { - content: vec![ - UserPart::Text(TextPart { - text: "Hello".into(), - provider_options: None, - }), - UserPart::File(FilePart { - data: FileData::Data { - data: FileBytes::Base64("AAECAw==".into()), +#[tokio::test] +async fn text_should_send_message_names() { + let name = |value: &str| json!({ "deepseek": { "name": value } }); + let messages = wire_messages( + "deepseek-chat", + vec![ + with_options( + LanguageModelMessage::System { + content: ("You are a helpful assistant.").into(), + provider_options: None, }, - media_type: "image/png".into(), - filename: None, - provider_options: None, - }), + name("guide"), + ), + with_options(LanguageModelMessage::user_text("Hello"), name("alice")), + with_options( + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: ("Hello, Alice.").into(), + provider_options: None, + })], + provider_options: None, + }, + name("assistant"), + ), + with_options( + LanguageModelMessage::user_text("How are you?"), + name("alice"), + ), ], - provider_options: None, - }]; - let options = default_options(prompt); - let result = build_request_body_with_warnings( - "deepseek-chat", - &options, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), ) - .unwrap(); + .await; + assert_eq!( + messages, + [ + json!({ "content": "You are a helpful assistant.", "name": "guide", "role": "system" }), + json!({ "content": "Hello", "name": "alice", "role": "user" }), + json!({ "content": "Hello, Alice.", "name": "assistant", "role": "assistant" }), + json!({ "content": "How are you?", "name": "alice", "role": "user" }), + ] + ); +} - // Image file parts are converted to image_url (OpenAI-compatible behaviour). - let messages_str = result.body["messages"].to_string(); - assert!( - messages_str.contains("image_url"), - "messages should contain image_url: {messages_str}" +#[tokio::test] +async fn text_should_pass_provider_options_user_id_as_user_id() { + let options = with_provider_options( + options(), + json!({ "deepseek": { "userId": "tenant_123-user" } }), ); - assert!( - messages_str.contains("data:image/png;base64,AAECAw=="), - "messages should contain base64 data URL: {messages_str}" + assert_eq!( + generate_body("deepseek-chat", &options).await["user_id"], + "tenant_123-user" ); } -/// TS: user messages › "should accept top-level-only mediaType without error -/// and not read it" (line ~71). -/// -/// A user message with a text part and a file part whose mediaType is just -/// "image" (top-level only) should still produce only the text content and -/// NOT include any image/mediaType data in the output. -#[tokio::test] -async fn should_accept_top_level_only_mediatype_without_error() { - let prompt = vec![LanguageModelMessage::User { - content: vec![ - UserPart::Text(TextPart { - text: "Hello".into(), - provider_options: None, - }), - UserPart::File(FilePart { - data: FileData::Data { - data: FileBytes::Base64("AAECAw==".into()), - }, - media_type: "image".into(), - filename: None, - provider_options: None, - }), - ], - provider_options: None, - }]; - let options = default_options(prompt); - let result = build_request_body_with_warnings( - "deepseek-chat", - &options, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - - // Top-level media type "image" is resolved to image_url (OpenAI-compatible). - let messages_str = result.body["messages"].to_string(); +#[tokio::test] +async fn text_should_omit_user_id_when_user_id_is_not_provided() { assert!( - messages_str.contains("image_url"), - "messages should contain image_url for top-level image media type: {messages_str}" + generate_body("deepseek-chat", &options()) + .await + .get("user_id") + .is_none() ); } -// ════════════════════════════════════════════════════════════════════════════ -// Tool preparation — deepseek-prepare-tools.test.ts -// -// These tests call `build_request_body_with_warnings` directly and assert on -// the `body["tools"]` field. -// -// Translated from `deepseek-prepare-tools.test.ts`. -// ════════════════════════════════════════════════════════════════════════════ - -/// TS: "should pass through strict mode when strict is true" (line ~5). -/// -/// A tool with `strict: true` should produce `"strict": true` in the -/// serialized tool definition. #[tokio::test] -async fn should_pass_through_strict_mode_when_strict_is_true() { - let tool = FunctionTool { - name: "testFunction".to_string(), - description: Some("A test function".to_string()), - input_schema: json!({ "type": "object", "properties": {} }), - strict: Some(true), - provider_options: None, - input_examples: None, - }; - let mut options = default_options(test_prompt()); - options.tools = Some(vec![Tool::from(tool)]); - let result = build_request_body_with_warnings( - "deepseek-chat", - &options, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - - assert_eq!( - result.body["tools"], - json!([{ - "type": "function", - "function": { - "description": "A test function", - "name": "testFunction", - "parameters": { "type": "object", "properties": {} }, - "strict": true - } - }]) - ); +async fn text_should_not_send_unknown_provider_options() { + // upstream: `deepseekLanguageModelChatOptions` is a zod object, which + // strips keys it does not declare. + let options = + with_provider_options(options(), json!({ "deepseek": { "user": "u1", "foo": 1 } })); + let body = generate_body("deepseek-chat", &options).await; + assert!(body.get("user").is_none() && body.get("foo").is_none()); } /// TS: "should pass through strict mode when strict is false" (line ~40). @@ -487,19 +546,12 @@ async fn should_pass_through_strict_mode_when_strict_is_false() { provider_options: None, input_examples: None, }; - let mut options = default_options(test_prompt()); + let mut options = options(); options.tools = Some(vec![Tool::from(tool)]); - let result = build_request_body_with_warnings( - "deepseek-chat", - &options, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); + let body = generate_body("deepseek-chat", &options).await; assert_eq!( - result.body["tools"], + body["tools"], json!([{ "type": "function", "function": { @@ -526,19 +578,12 @@ async fn should_not_include_strict_mode_when_strict_is_undefined() { provider_options: None, input_examples: None, }; - let mut options = default_options(test_prompt()); + let mut options = options(); options.tools = Some(vec![Tool::from(tool)]); - let result = build_request_body_with_warnings( - "deepseek-chat", - &options, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); + let body = generate_body("deepseek-chat", &options).await; assert_eq!( - result.body["tools"], + body["tools"], json!([{ "type": "function", "function": { @@ -550,70 +595,1955 @@ async fn should_not_include_strict_mode_when_strict_is_undefined() { ); // Verify strict is absent. assert!( - result.body["tools"][0]["function"].get("strict").is_none(), + body["tools"][0]["function"].get("strict").is_none(), "strict should be absent" ); } -/// TS: "should pass through strict mode for multiple tools with different strict -/// settings" (line ~108). -/// -/// Multiple tools with different strict settings should each preserve their -/// own strict value (or omit it when undefined). -#[tokio::test] -async fn should_pass_through_strict_mode_for_multiple_tools() { - let tools = vec![ - Tool::from(FunctionTool { - name: "strictTool".to_string(), - description: Some("A strict tool".to_string()), - input_schema: json!({ "type": "object", "properties": {} }), - strict: Some(true), - provider_options: None, - input_examples: None, - }), - Tool::from(FunctionTool { - name: "nonStrictTool".to_string(), - description: Some("A non-strict tool".to_string()), - input_schema: json!({ "type": "object", "properties": {} }), - strict: Some(false), - provider_options: None, - input_examples: None, - }), - Tool::from(FunctionTool { - name: "defaultTool".to_string(), - description: Some("A tool without strict setting".to_string()), - input_schema: json!({ "type": "object", "properties": {} }), - strict: None, - provider_options: None, - input_examples: None, - }), - ]; - let mut options = default_options(test_prompt()); - options.tools = Some(tools); - let result = build_request_body_with_warnings( - "deepseek-chat", - &options, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); +#[tokio::test] +async fn text_should_reject_strict_tools_on_the_standard_endpoint_before_fetching() { + let mut options = options(); + options.tools = Some(vec![strict_tool("getWeather", Some(true))]); + let message = generate_error("deepseek-chat", &options).await; + assert!( + message.contains("DeepSeek strict tool calls require a beta base URL ending in `/beta`."), + "{message}" + ); +} - let tools_arr = result.body["tools"].as_array().expect("tools array"); - assert_eq!(tools_arr.len(), 3); +#[tokio::test] +async fn text_should_send_all_strict_tools_on_the_beta_endpoint() { + let server = json_server(text_response()).await; + let mut options = options(); + options.tools = Some(vec![strict_tool("getWeather", Some(true))]); + let body = beta_chat(&server, "deepseek-chat") + .do_generate(&options) + .await + .unwrap() + .request + .and_then(|request| request.body) + .unwrap(); + assert_eq!(body["tools"][0]["type"], "function"); + assert_eq!(body["tools"][0]["function"]["name"], "getWeather"); + assert_eq!(body["tools"][0]["function"]["strict"], true); +} - // Tool 0: strict = true - assert_eq!(tools_arr[0]["function"]["name"], json!("strictTool")); - assert_eq!(tools_arr[0]["function"]["strict"], json!(true)); +#[tokio::test] +async fn text_should_reject_mixed_strict_tools_in_streaming_requests() { + let server = sse_server(text_chunks()).await; + let mut options = options(); + options.tools = Some(vec![ + strict_tool("strictTool", Some(true)), + strict_tool("nonStrictTool", None), + ]); + let Err(error) = beta_chat(&server, "deepseek-chat") + .do_stream(&options) + .await + else { + panic!("the call must be rejected"); + }; + assert!( + error.to_string().contains( + "DeepSeek strict mode requires every function tool in the request to set `strict: true`." + ), + "{error}" + ); + assert!(request_paths(&server).await.is_empty()); +} - // Tool 1: strict = false - assert_eq!(tools_arr[1]["function"]["name"], json!("nonStrictTool")); - assert_eq!(tools_arr[1]["function"]["strict"], json!(false)); +#[tokio::test] +async fn text_should_reject_invalid_user_id_before_making_an_api_request() { + for (user_id, expected) in [ + ("".to_string(), "userId must match /^[a-zA-Z0-9_-]+$/"), + ( + "contains space".to_string(), + "userId must match /^[a-zA-Z0-9_-]+$/", + ), + ( + "a".repeat(513), + "userId must be at most 512 characters long", + ), + ] { + let options = + with_provider_options(options(), json!({ "deepseek": { "userId": user_id } })); + let message = generate_error("deepseek-chat", &options).await; + assert!(message.contains("invalid provider options"), "{message}"); + assert!(message.contains(expected), "{message}"); + } +} - // Tool 2: strict absent - assert_eq!(tools_arr[2]["function"]["name"], json!("defaultTool")); +#[tokio::test] +async fn text_should_extract_text_content() { + let result = generate_result("deepseek-chat", &options()).await; + assert_eq!(result.content.len(), 1); assert!( - tools_arr[2]["function"].get("strict").is_none(), - "strict should be absent for defaultTool" + matches!(&result.content[0], GenerateContent::Text { text, .. } if text == "Hello, World!") + ); + 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(13)); + assert_eq!(result.usage.output_tokens.total, Some(300)); + assert_eq!( + result + .response + .as_ref() + .and_then(|response| response.id.as_deref()), + Some("00f10ecd-60b3-4707-b5db-e4bcadf7aea1") + ); + assert_eq!( + result + .response + .as_ref() + .and_then(|response| response.model_id.as_deref()), + Some("deepseek-chat") + ); +} + +#[tokio::test] +async fn text_should_report_usage_with_prompt_cache_tokens() { + // upstream convert-to-deepseek-usage.ts: cache hits are the cache-read tokens. + let mut body = text_response(); + body["usage"] = json!({ + "prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120, + "prompt_cache_hit_tokens": 70, "prompt_cache_miss_tokens": 30, + "completion_tokens_details": { "reasoning_tokens": 5 } + }); + let server = json_server(body).await; + let usage = chat(&server, "deepseek-chat") + .do_generate(&options()) + .await + .unwrap() + .usage; + assert_eq!(usage.input_tokens.total, Some(100)); + assert_eq!(usage.input_tokens.cache_read, Some(70)); + assert_eq!(usage.input_tokens.no_cache, Some(30)); + assert_eq!(usage.output_tokens.total, Some(20)); + assert_eq!(usage.output_tokens.reasoning, Some(5)); + assert_eq!(usage.output_tokens.text, Some(15)); + assert_eq!(usage.raw.unwrap()["prompt_cache_hit_tokens"], 70); +} + +#[tokio::test] +async fn text_should_report_the_provider_metadata() { + // upstream doGenerate: promptCacheHitTokens, promptCacheMissTokens, + // responseObject, choiceIndex, messageRole, systemFingerprint. + let result = generate_result("deepseek-chat", &options()).await; + assert_eq!( + serde_json::to_value(result.provider_metadata.unwrap()).unwrap(), + json!({ "deepseek": { + "promptCacheHitTokens": 0, + "promptCacheMissTokens": 13, + "responseObject": "chat.completion", + "choiceIndex": 0, + "messageRole": "assistant", + "systemFingerprint": "fp_eaab8d114b_prod0820_fp8_kvcache" + } }) + ); +} + +#[tokio::test] +async fn text_should_include_the_system_fingerprint_in_provider_metadata() { + let result = generate_result("deepseek-chat", &options()).await; + assert_eq!( + result.provider_metadata.unwrap()["deepseek"]["systemFingerprint"], + "fp_eaab8d114b_prod0820_fp8_kvcache" + ); +} + +#[tokio::test] +async fn text_should_tolerate_a_null_or_missing_system_fingerprint() { + for fingerprint in [Some(Value::Null), None] { + let mut body = text_response(); + match fingerprint { + Some(value) => body["system_fingerprint"] = value, + None => { + body.as_object_mut().unwrap().remove("system_fingerprint"); + } + } + let server = json_server(body).await; + let result = chat(&server, "deepseek-chat") + .do_generate(&options()) + .await + .unwrap(); + assert!( + result.provider_metadata.unwrap()["deepseek"] + .get("systemFingerprint") + .is_none() + ); + } +} + +// ---- describe('reasoning') --------------------------------------------------- + +fn reasoning_response() -> Value { + json!({ + "id": "r1", "object": "chat.completion", "created": 1764656316, "model": "deepseek-reasoner", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": "There are 3 r's in strawberry.", + "reasoning_content": "Let me count the letters." + }, + "finish_reason": "stop" + }], + "usage": { "prompt_tokens": 20, "completion_tokens": 30, "total_tokens": 50, + "completion_tokens_details": { "reasoning_tokens": 10 } } + }) +} + +#[tokio::test] +async fn reasoning_should_send_correct_request_body() { + let options = with_provider_options( + options_for(vec![LanguageModelMessage::user_text( + "How many \"r\"s are in the word \"strawberry\"?", + )]), + json!({ "deepseek": { "thinking": { "type": "enabled" } } }), + ); + assert_eq!( + generate_body("deepseek-reasoner", &options).await, + json!({ + "messages": [{ "content": "How many \"r\"s are in the word \"strawberry\"?", "role": "user" }], + "model": "deepseek-reasoner", + "thinking": { "type": "enabled" } + }) + ); +} + +#[tokio::test] +async fn reasoning_should_extract_reasoning_before_text() { + let server = json_server(reasoning_response()).await; + let result = chat(&server, "deepseek-reasoner") + .do_generate(&options()) + .await + .unwrap(); + assert_eq!(result.content.len(), 2); + assert!( + matches!(&result.content[0], GenerateContent::Reasoning(ReasoningOutput { text, .. }) if text == "Let me count the letters.") + ); + assert!( + matches!(&result.content[1], GenerateContent::Text { text, .. } if text == "There are 3 r's in strawberry.") + ); + assert_eq!(result.usage.output_tokens.reasoning, Some(10)); + assert_eq!(result.usage.output_tokens.text, Some(20)); +} + +#[tokio::test] +async fn reasoning_should_not_read_the_reasoning_field() { + // upstream deepseekChatResponseSchema only declares `reasoning_content`. + let mut body = reasoning_response(); + let message = body["choices"][0]["message"].as_object_mut().unwrap(); + let reasoning = message.remove("reasoning_content").unwrap(); + message.insert("reasoning".to_string(), reasoning); + let server = json_server(body).await; + let result = chat(&server, "deepseek-reasoner") + .do_generate(&options()) + .await + .unwrap(); + assert_eq!(result.content.len(), 1); + assert!(matches!(&result.content[0], GenerateContent::Text { .. })); +} + +// ---- describe('logprobs') ---------------------------------------------------- + +fn logprobs_response() -> Value { + let mut body = text_response(); + body["choices"][0]["message"] = json!({ + "role": "assistant", "content": "OK", "reasoning_content": "Reasoning" + }); + body["choices"][0]["logprobs"] = json!({ + "content": [logprob("OK", -0.00002467602, json!([79, 75]))], + "reasoning_content": [logprob("Reasoning", -0.1, Value::Null)] + }); + body +} + +#[tokio::test] +async fn logprobs_should_send_logprobs_provider_options() { + let options = with_provider_options( + options(), + json!({ "deepseek": { "logprobs": true, "topLogprobs": 1 } }), + ); + assert_eq!( + generate_body("deepseek-v4-flash", &options).await, + json!({ + "logprobs": true, + "messages": [{ "content": "Hello", "role": "user" }], + "model": "deepseek-v4-flash", + "top_logprobs": 1 + }) + ); +} + +#[tokio::test] +async fn logprobs_should_enable_logprobs_when_top_logprobs_is_set() { + let options = with_provider_options(options(), json!({ "deepseek": { "topLogprobs": 1 } })); + let body = generate_body("deepseek-v4-flash", &options).await; + assert_eq!(body["logprobs"], true); + assert_eq!(body["top_logprobs"], 1); +} + +#[tokio::test] +async fn logprobs_should_extract_content_and_reasoning_logprobs() { + let server = json_server(logprobs_response()).await; + let options = with_provider_options(options(), json!({ "deepseek": { "logprobs": true } })); + let result = chat(&server, "deepseek-v4-flash") + .do_generate(&options) + .await + .unwrap(); + assert_eq!( + result.provider_metadata.unwrap()["deepseek"]["logprobs"], + json!({ + "content": [logprob("OK", -0.00002467602, json!([79, 75]))], + "reasoning_content": [logprob("Reasoning", -0.1, Value::Null)] + }) + ); +} + +// ---- describe('top-level reasoning') ------------------------------------------- + +fn reasoning_options(effort: ReasoningEffort) -> CallOptions { + CallOptions { + reasoning: Some(effort), + ..options() + } +} + +#[tokio::test] +async fn top_level_reasoning_should_map_to_thinking_enabled() { + let body = generate_body( + "deepseek-reasoner", + &reasoning_options(ReasoningEffort::High), + ) + .await; + assert_eq!(body["thinking"], json!({ "type": "enabled" })); + assert_eq!(body["reasoning_effort"], "high"); +} + +#[tokio::test] +async fn top_level_reasoning_none_should_map_to_thinking_disabled() { + let body = generate_body( + "deepseek-reasoner", + &reasoning_options(ReasoningEffort::None), + ) + .await; + assert_eq!(body["thinking"], json!({ "type": "disabled" })); + assert!(body.get("reasoning_effort").is_none()); +} + +#[tokio::test] +async fn top_level_reasoning_none_should_keep_the_temperature() { + // thinking is disabled, so temperature applies and is not warned about. + let mut options = reasoning_options(ReasoningEffort::None); + options.temperature = Some(0.4); + let result = generate_result("deepseek-reasoner", &options).await; + assert_eq!( + result.request.and_then(|request| request.body).unwrap()["temperature"], + 0.4 + ); + assert!(result.warnings.is_empty()); +} + +#[tokio::test] +async fn top_level_reasoning_xhigh_should_map_to_reasoning_effort_max() { + let result = generate_result( + "deepseek-reasoner", + &reasoning_options(ReasoningEffort::Xhigh), + ) + .await; + assert_eq!( + result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + "max" + ); + assert!(warning_values(&result.warnings).contains(&serde_json::to_value(compatibility( + "reasoning", + "reasoning \"xhigh\" is not directly supported by this model. mapped to effort \"max\"." + )) + .unwrap())); +} + +#[tokio::test] +async fn top_level_reasoning_low_should_map_to_reasoning_effort_low_without_a_compatibility_warning() + { + let result = generate_result( + "deepseek-reasoner", + &reasoning_options(ReasoningEffort::Low), + ) + .await; + assert_eq!( + result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + "low" + ); + assert!( + !result + .warnings + .iter() + .any(|w| matches!(w, Warning::Compatibility { feature, .. } if feature == "reasoning")) + ); +} + +#[tokio::test] +async fn top_level_reasoning_medium_should_map_to_reasoning_effort_high_with_a_compatibility_warning() + { + let result = generate_result( + "deepseek-reasoner", + &reasoning_options(ReasoningEffort::Medium), + ) + .await; + assert_eq!( + result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + "high" + ); + assert_eq!( + warning_values(&result.warnings), + warning_values(&[compatibility( + "reasoning", + "reasoning \"medium\" is not directly supported by this model. mapped to effort \"high\"." + )]) + ); +} + +#[tokio::test] +async fn top_level_reasoning_minimal_should_map_to_reasoning_effort_low_with_compatibility_warning() +{ + let result = generate_result( + "deepseek-reasoner", + &reasoning_options(ReasoningEffort::Minimal), + ) + .await; + assert_eq!( + result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + "low" + ); + assert_eq!( + warning_values(&result.warnings), + warning_values(&[compatibility( + "reasoning", + "reasoning \"minimal\" is not directly supported by this model. mapped to effort \"low\"." + )]) + ); +} + +#[tokio::test] +async fn top_level_reasoning_should_map_provider_options_reasoning_effort() { + for (input, output, warns) in [ + ("low", "low", false), + ("medium", "high", true), + ("xhigh", "max", true), + ] { + let options = with_provider_options( + options(), + json!({ "deepseek": { "reasoningEffort": input } }), + ); + let result = generate_result("deepseek-reasoner", &options).await; + assert_eq!( + result.request.and_then(|request| request.body).unwrap()["reasoning_effort"], + output + ); + let expected = if warns { + vec![compatibility( + "reasoningEffort", + &format!( + "reasoningEffort \"{input}\" is not a canonical DeepSeek value. mapped to \"{output}\"." + ), + )] + } else { + vec![] + }; + assert_eq!(warning_values(&result.warnings), warning_values(&expected)); + } +} + +#[tokio::test] +async fn top_level_reasoning_should_map_legacy_thinking_type_adaptive_to_enabled_with_a_compatibility_warning() + { + let options = with_provider_options( + options(), + json!({ "deepseek": { "thinking": { "type": "adaptive" } } }), + ); + let result = generate_result("deepseek-reasoner", &options).await; + assert_eq!( + result.request.and_then(|request| request.body).unwrap()["thinking"], + json!({ "type": "enabled" }) + ); + assert!(warning_values(&result.warnings).contains(&serde_json::to_value(compatibility( + "thinking.type", + "thinking.type \"adaptive\" is not a canonical DeepSeek value. mapped to \"enabled\"." + )) + .unwrap())); +} + +#[tokio::test] +async fn top_level_reasoning_should_pass_provider_options_reasoning_effort() { + let options = with_provider_options( + options(), + json!({ "deepseek": { "reasoningEffort": "max" } }), + ); + let body = generate_body("deepseek-reasoner", &options).await; + assert_eq!(body["reasoning_effort"], "max"); + // When only reasoningEffort is set without thinking, thinking stays unset. + assert!(body.get("thinking").is_none()); +} + +#[tokio::test] +async fn top_level_reasoning_should_prefer_provider_options_thinking_over_top_level_reasoning() { + let options = with_provider_options( + reasoning_options(ReasoningEffort::None), + json!({ "deepseek": { "thinking": { "type": "enabled" } } }), + ); + assert_eq!( + generate_body("deepseek-reasoner", &options).await["thinking"], + json!({ "type": "enabled" }) + ); +} + +#[tokio::test] +async fn top_level_reasoning_should_prefer_provider_options_reasoning_effort_over_top_level_reasoning() + { + let options = with_provider_options( + reasoning_options(ReasoningEffort::High), + json!({ "deepseek": { "reasoningEffort": "max" } }), + ); + assert_eq!( + generate_body("deepseek-reasoner", &options).await["reasoning_effort"], + "max" + ); +} + +#[tokio::test] +async fn top_level_reasoning_should_not_set_thinking_when_reasoning_is_not_specified() { + assert!( + generate_body("deepseek-reasoner", &options()) + .await + .get("thinking") + .is_none() + ); +} + +#[tokio::test] +async fn top_level_reasoning_provider_default_sends_neither_thinking_nor_effort() { + let body = generate_body( + "deepseek-reasoner", + &reasoning_options(ReasoningEffort::ProviderDefault), + ) + .await; + assert!(body.get("thinking").is_none() && body.get("reasoning_effort").is_none()); +} + +// ---- describe('tool call') --------------------------------------------------- + +fn tool_call_response() -> Value { + json!({ + "id": "t1", "object": "chat.completion", "created": 1764656316, "model": "deepseek-reasoner", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", "content": "", + "reasoning_content": "I should look up the weather.", + "tool_calls": [{ + "id": "call_00_abc", "type": "function", + "function": { "name": "weather", "arguments": "{\"location\":\"San Francisco\"}" } + }] + }, + "finish_reason": "tool_calls" + }], + "usage": { "prompt_tokens": 50, "completion_tokens": 20, "total_tokens": 70 } + }) +} + +fn tool_options() -> CallOptions { + let mut options = with_provider_options( + options(), + json!({ "deepseek": { "thinking": { "type": "enabled" } } }), + ); + options.tools = Some(vec![Tool::Function(weather_tool())]); + options +} + +fn weather_wire_tool() -> Value { + json!({ + "type": "function", + "function": { + "name": "weather", + "parameters": { + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "properties": { "location": { "type": "string" } }, + "required": ["location"], + "type": "object" + } + } + }) +} + +#[tokio::test] +async fn tool_call_should_send_correct_request_body() { + assert_eq!( + generate_body("deepseek-reasoner", &tool_options()).await, + json!({ + "messages": [{ "content": "Hello", "role": "user" }], + "model": "deepseek-reasoner", + "thinking": { "type": "enabled" }, + "tools": [weather_wire_tool()] + }) + ); +} + +#[tokio::test] +async fn tool_call_should_extract_tool_call_content() { + let server = json_server(tool_call_response()).await; + let result = chat(&server, "deepseek-reasoner") + .do_generate(&tool_options()) + .await + .unwrap(); + assert!( + matches!(&result.content[0], GenerateContent::Reasoning(ReasoningOutput { text, .. }) if text == "I should look up the weather.") + ); + assert!(matches!( + &result.content[1], + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. }) + if tool_call_id == "call_00_abc" && tool_name == "weather" + && input == "{\"location\":\"San Francisco\"}" + )); + assert_eq!(result.content.len(), 2); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); + assert_eq!( + result.provider_metadata.unwrap()["deepseek"]["toolCallTypes"], + json!(["function"]) + ); +} + +#[tokio::test] +async fn tool_call_should_generate_an_id_when_the_server_sends_none() { + let mut body = tool_call_response(); + body["choices"][0]["message"]["tool_calls"][0]["id"] = Value::Null; + let server = json_server(body).await; + let result = chat(&server, "deepseek-reasoner") + .do_generate(&tool_options()) + .await + .unwrap(); + assert!(matches!( + &result.content[1], + GenerateContent::ToolCall(RawToolCall { tool_call_id, .. }) if !tool_call_id.is_empty() + )); +} + +// ---- describe('json response format') ------------------------------------------- + +fn json_options(schema: Option) -> CallOptions { + let mut options = tool_options(); + options.response_format = Some(ResponseFormat::Json { + schema, + name: None, + description: None, + }); + options +} + +#[tokio::test] +async fn json_response_format_should_send_correct_request_body_without_schema() { + assert_eq!( + generate_body("deepseek-reasoner", &json_options(None)).await, + json!({ + "messages": [ + { "content": "Return JSON.", "role": "system" }, + { "content": "Hello", "role": "user" } + ], + "model": "deepseek-reasoner", + "response_format": { "type": "json_object" }, + "thinking": { "type": "enabled" }, + "tools": [weather_wire_tool()] + }) + ); +} + +#[tokio::test] +async fn json_response_format_should_send_correct_request_body_with_schema() { + let schema = json!({ + "type": "object", + "properties": { "elements": { "type": "array", "items": { "type": "object", + "properties": { "location": { "type": "string" }, "temperature": { "type": "number" }, + "condition": { "type": "string" } }, + "required": ["location", "temperature", "condition"], "additionalProperties": false } } }, + "required": ["elements"], "additionalProperties": false, + "$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(); + assert_eq!( + body["messages"][0], + json!({ + "role": "system", + "content": format!("Return JSON that conforms to the following schema: {schema}") + }) + ); + assert_eq!( + body["messages"][1], + json!({ "role": "user", "content": "Hello" }) + ); + assert_eq!(body["response_format"], json!({ "type": "json_object" })); + // upstream: supportsStructuredOutputs is false, so the schema is injected + // into the system message with a compatibility warning. + assert_eq!( + warning_values(&result.warnings), + warning_values(&[compatibility( + "responseFormat JSON schema", + "JSON response schema is injected into the system message." + )]) + ); +} + +#[tokio::test] +async fn json_response_format_should_extract_text_content() { + let mut body = text_response(); + body["choices"][0]["message"]["content"] = json!("{\"elements\":[]}"); + let server = json_server(body).await; + let result = chat(&server, "deepseek-reasoner") + .do_generate(&json_options(None)) + .await + .unwrap(); + assert!( + matches!(&result.content[0], GenerateContent::Text { text, .. } if text == "{\"elements\":[]}") + ); +} + +// ---- describe('assistant prefix completion') -------------------------------------- + +fn prefix_prompt(prefix_options: Value) -> LanguageModelPrompt { + vec![ + LanguageModelMessage::user_text("Complete this sentence."), + with_options( + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: ("The answer is").into(), + provider_options: None, + })], + provider_options: None, + }, + json!({ "deepseek": prefix_options }), + ), + ] +} + +#[tokio::test] +async fn prefix_should_send_name_and_prefix_on_the_final_assistant_message() { + let server = json_server(text_response()).await; + let body = beta_chat(&server, "deepseek-chat") + .do_generate(&options_for(prefix_prompt( + json!({ "name": "assistant", "prefix": true }), + ))) + .await + .unwrap() + .request + .and_then(|request| request.body) + .unwrap(); + assert_eq!( + body["messages"], + json!([ + { "role": "user", "content": "Complete this sentence." }, + { "role": "assistant", "content": "The answer is", "name": "assistant", "prefix": true } + ]) + ); + assert_eq!(body["model"], "deepseek-chat"); +} + +#[tokio::test] +async fn prefix_should_reject_prefix_completion_with_the_default_base_url() { + let message = generate_error( + "deepseek-chat", + &options_for(vec![with_options( + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: ("The answer is").into(), + provider_options: None, + })], + provider_options: None, + }, + json!({ "deepseek": { "prefix": true } }), + )]), + ) + .await; + assert!( + message.contains( + "DeepSeek assistant prefix completion requires a beta base URL ending in `/beta`." + ), + "{message}" + ); +} + +// =========================================================================== +// describe('doStream') +// =========================================================================== + +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") + .do_stream(&options()) + .await + .unwrap(); + collect(stream) + .await + .into_iter() + .find_map(|part| match part { + StreamPart::Error { + error: AiMuxError::ApiCall(error), + } => Some(*error), + _ => None, + }) + .expect("an API call error part") +} + +/// `deepSeekErrorSchema` as the chat model's failed-response handler +/// (deepseek@3.0.56 `deepseek-chat-api-types.ts:117-124`, +/// `deepseek-chat-language-model.ts:213-217`): `code` may be a number. +#[tokio::test] +async fn a_failed_response_reads_deepseeks_error_schema() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(".*/chat/completions$")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ "error": { + "message": "Insufficient Balance", "type": "unknown_error", "param": null, "code": 4020 + } }))) + .mount(&server) + .await; + let Err(AiMuxError::ApiCall(error)) = + chat(&server, "deepseek-chat").do_generate(&options()).await + else { + panic!("expected an API call error"); + }; + assert_eq!(error.message, "Insufficient Balance"); + assert_eq!(error.provider_code.as_deref(), Some("4020")); + assert_eq!(error.status_code, Some(402)); +} + +#[tokio::test] +async fn stream_should_preserve_a_provider_error_envelope_in_stream_errors() { + let data = json!({ "error": { + "message": "Rate limit reached", "type": "rate_limit_error", "code": "rate_limit_exceeded" + } }); + let error = stream_error(&data).await; + assert_eq!(error.message, "Rate limit reached"); + assert_eq!(error.provider_code.as_deref(), Some("rate_limit_exceeded")); + assert_eq!(error.status_code, Some(429)); + assert!(error.is_retryable); + // upstream: `data` is the whole error envelope. + let envelope: Value = serde_json::from_str(error.response_body.as_deref().unwrap()).unwrap(); + assert_eq!(envelope, data); +} + +#[tokio::test] +async fn stream_should_classify_insufficient_quota_as_non_retryable() { + let data = json!({ "error": { + "message": "You exceeded your current quota.", "type": "rate_limit_error", + "code": "insufficient_quota" + } }); + let error = stream_error(&data).await; + assert_eq!(error.status_code, Some(429)); + assert!(!error.is_retryable); +} + +#[tokio::test] +async fn stream_should_preserve_the_provider_type_when_code_is_an_http_status() { + let data = json!({ "error": { + "message": "Rate limit reached", "type": "rate_limit_error", "code": "429" + } }); + let error = stream_error(&data).await; + assert_eq!(error.provider_code.as_deref(), Some("429")); + assert_eq!(error.status_code, Some(429)); + assert!(error.is_retryable); +} + +// ---- describe('text') ------------------------------------------------------------ + +async fn stream_parts( + model_id: &str, + options: &CallOptions, + chunks: Vec, +) -> (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(); + (body, collect(result).await) +} + +#[tokio::test] +async fn stream_text_should_send_model_id_settings_and_input() { + let mut options = options_for(vec![ + LanguageModelMessage::System { + content: ("You are a helpful assistant.").into(), + provider_options: None, + }, + LanguageModelMessage::user_text("Hello"), + ]); + options.temperature = Some(0.5); + options.top_p = Some(0.3); + let (body, _) = stream_parts("deepseek-chat", &options, text_chunks()).await; + assert_eq!( + body, + json!({ + "messages": [ + { "content": "You are a helpful assistant.", "role": "system" }, + { "content": "Hello", "role": "user" } + ], + "model": "deepseek-chat", + "stream": true, + "stream_options": { "include_usage": true }, + "temperature": 0.5, + "top_p": 0.3 + }) + ); +} + +#[tokio::test] +async fn stream_text_should_omit_deprecated_and_ineffective_sampling_options_in_default_v4_thinking_mode() + { + let mut options = options(); + options.temperature = Some(0.2); + options.top_p = Some(0.4); + options.frequency_penalty = Some(0.5); + options.presence_penalty = Some(0.6); + let (body, parts) = stream_parts("deepseek-v4-flash", &options, text_chunks()).await; + assert_eq!( + body, + json!({ + "model": "deepseek-v4-flash", + "messages": [{ "role": "user", "content": "Hello" }], + "stream": true, + "stream_options": { "include_usage": true } + }) + ); + assert_eq!( + stream_start_warnings(&parts), + warning_values(&[ + deprecated("frequencyPenalty"), + deprecated("presencePenalty"), + thinking_unsupported("temperature"), + thinking_unsupported("topP"), + ]) + ); +} + +#[tokio::test] +async fn stream_text_should_preserve_supported_sampling_options_when_v4_thinking_is_disabled() { + let mut options = with_provider_options( + options(), + json!({ "deepseek": { "thinking": { "type": "disabled" } } }), + ); + options.temperature = Some(0.2); + options.top_p = Some(0.4); + options.frequency_penalty = Some(0.5); + options.presence_penalty = Some(0.6); + let (body, parts) = stream_parts("deepseek-v4-flash", &options, text_chunks()).await; + assert_eq!( + body, + json!({ + "model": "deepseek-v4-flash", + "messages": [{ "role": "user", "content": "Hello" }], + "temperature": 0.2, + "top_p": 0.4, + "thinking": { "type": "disabled" }, + "stream": true, + "stream_options": { "include_usage": true } + }) + ); + assert_eq!( + stream_start_warnings(&parts), + warning_values(&[ + deprecated("frequencyPenalty"), + deprecated("presencePenalty") + ]) + ); +} + +#[tokio::test] +async fn stream_text_should_send_message_names() { + let name = |value: &str| json!({ "deepseek": { "name": value } }); + let options = options_for(vec![ + with_options( + LanguageModelMessage::System { + content: ("You are a helpful assistant.").into(), + provider_options: None, + }, + name("guide"), + ), + with_options(LanguageModelMessage::user_text("Hello"), name("alice")), + with_options( + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: ("Hello, Alice.").into(), + provider_options: None, + })], + provider_options: None, + }, + name("assistant"), + ), + ]); + let (body, _) = stream_parts("deepseek-chat", &options, text_chunks()).await; + assert_eq!( + body["messages"], + json!([ + { "content": "You are a helpful assistant.", "name": "guide", "role": "system" }, + { "content": "Hello", "name": "alice", "role": "user" }, + { "content": "Hello, Alice.", "name": "assistant", "role": "assistant" } + ]) + ); +} + +#[tokio::test] +async fn stream_text_should_pass_provider_options_user_id_as_user_id() { + let options = with_provider_options( + options(), + json!({ "deepseek": { "userId": "tenant_123-user" } }), + ); + let (body, _) = stream_parts("deepseek-chat", &options, text_chunks()).await; + assert_eq!(body["user_id"], "tenant_123-user"); + assert_eq!(body["stream"], true); +} + +#[tokio::test] +async fn stream_text_should_stream_text() { + let (_, parts) = stream_parts("deepseek-chat", &options(), text_chunks()).await; + assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); + assert!(matches!( + &parts[1], + StreamPart::ResponseMetadata(ResponseMetadata { id, model_id, .. }) + if id.as_deref() == Some("c1") && model_id.as_deref() == Some("deepseek-chat") + )); + let deltas: Vec<&str> = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextDelta { id, delta, .. } if id == "txt-0" => Some(delta.as_str()), + _ => None, + }) + .collect(); + assert_eq!(deltas, ["Hello", ", World!"]); + assert!( + parts + .iter() + .any(|p| matches!(p, StreamPart::TextStart { id, .. } if id == "txt-0")) + ); + assert!( + parts + .iter() + .any(|p| matches!(p, StreamPart::TextEnd { id, .. } if id == "txt-0")) + ); + match parts.last().unwrap() { + StreamPart::Finish { + finish_reason, + usage, + .. + } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(finish_reason.raw.as_deref(), Some("stop")); + // usage with prompt-cache tokens + assert_eq!(usage.input_tokens.total, Some(13)); + assert_eq!(usage.input_tokens.cache_read, Some(5)); + assert_eq!(usage.input_tokens.no_cache, Some(8)); + assert_eq!(usage.output_tokens.total, Some(3)); + } + other => panic!("expected finish, got {other:?}"), + } + assert_eq!( + finish_metadata(&parts), + json!({ "deepseek": { + "promptCacheHitTokens": 5, + "promptCacheMissTokens": 8, + "responseObject": "chat.completion.chunk", + "choiceIndex": 0, + "messageRole": "assistant", + "systemFingerprint": "fp_eaab8d114b_prod0820_fp8_kvcache" + } }) + ); +} + +#[tokio::test] +async fn stream_text_should_include_the_repeated_system_fingerprint_in_provider_metadata() { + let (_, parts) = stream_parts("deepseek-chat", &options(), text_chunks()).await; + assert_eq!( + finish_metadata(&parts)["deepseek"]["systemFingerprint"], + "fp_eaab8d114b_prod0820_fp8_kvcache" + ); +} + +#[tokio::test] +async fn stream_text_should_keep_the_latest_non_null_system_fingerprint() { + let chunks = vec![ + "data: {\"system_fingerprint\":\"fp_initial\",\"choices\":[{\"delta\":{\"content\":\"OK\"},\"finish_reason\":null}],\"usage\":null}\n\n".to_string(), + "data: {\"system_fingerprint\":null,\"choices\":[{\"delta\":{},\"finish_reason\":null}],\"usage\":null}\n\n".to_string(), + "data: {\"system_fingerprint\":\"fp_latest\",\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":null}\n\n".to_string(), + ]; + let (_, parts) = stream_parts("deepseek-chat", &options(), chunks).await; + assert_eq!( + finish_metadata(&parts)["deepseek"]["systemFingerprint"], + "fp_latest" + ); +} + +#[tokio::test] +async fn stream_text_should_tolerate_a_null_or_missing_system_fingerprint() { + for fingerprint in [Some(Value::Null), None] { + let mut chunk = json!({ + "choices": [{ "delta": { "content": "OK" }, "finish_reason": "stop" }], + "usage": null + }); + if let Some(value) = fingerprint { + chunk["system_fingerprint"] = value; + } + let (_, parts) = stream_parts("deepseek-chat", &options(), vec![data_event(&chunk)]).await; + assert!( + finish_metadata(&parts)["deepseek"] + .get("systemFingerprint") + .is_none() + ); + } +} + +// ---- describe('reasoning') ------------------------------------------------------------- + +#[tokio::test] +async fn stream_reasoning_should_keep_reasoning_active_when_deltas_include_empty_tool_calls() { + let chunk = |delta: Value, finish: Value| { + data_event(&json!({ + "id": "chatcmpl-test", "object": "chat.completion.chunk", "created": 1, + "model": "test-model", + "choices": [{ "index": 0, "delta": delta, "finish_reason": finish }] + })) + }; + let chunks = vec![ + chunk( + json!({ "role": "assistant", "content": "", "reasoning_content": "Think ", "tool_calls": [] }), + Value::Null, + ), + chunk( + json!({ "content": "", "reasoning_content": "more...", "tool_calls": [] }), + Value::Null, + ), + chunk( + json!({ "content": "Hello", "reasoning_content": "", "tool_calls": [] }), + json!("stop"), + ), + ]; + let (_, parts) = stream_parts("deepseek-reasoner", &options(), chunks).await; + let reasoning: Vec = parts + .iter() + .filter_map(|part| match part { + StreamPart::ReasoningStart { id, .. } => Some(format!("start {id}")), + StreamPart::ReasoningDelta { id, delta, .. } => Some(format!("delta {id} {delta}")), + StreamPart::ReasoningEnd { id, .. } => Some(format!("end {id}")), + _ => None, + }) + .collect(); + assert_eq!( + reasoning, + [ + "start reasoning-0", + "delta reasoning-0 Think ", + "delta reasoning-0 more...", + "end reasoning-0" + ] + ); +} + +#[tokio::test] +async fn stream_reasoning_should_map_legacy_thinking_and_generic_reasoning_to_canonical_request_values() + { + let options = with_provider_options( + reasoning_options(ReasoningEffort::Medium), + json!({ "deepseek": { "thinking": { "type": "adaptive" } } }), + ); + let (body, parts) = stream_parts("deepseek-reasoner", &options, text_chunks()).await; + assert_eq!(body["thinking"], json!({ "type": "enabled" })); + assert_eq!(body["reasoning_effort"], "high"); + assert_eq!( + stream_start_warnings(&parts), + warning_values(&[ + compatibility( + "thinking.type", + "thinking.type \"adaptive\" is not a canonical DeepSeek value. mapped to \"enabled\"." + ), + compatibility( + "reasoning", + "reasoning \"medium\" is not directly supported by this model. mapped to effort \"high\"." + ), + ]) + ); +} + +#[tokio::test] +async fn stream_reasoning_should_stream_reasoning_before_text() { + let chunk = |delta: Value, finish: Value| { + data_event(&json!({ + "id": "c", "object": "chat.completion.chunk", "created": 1, "model": "deepseek-reasoner", + "choices": [{ "index": 0, "delta": delta, "finish_reason": finish }] + })) + }; + let chunks = vec![ + chunk( + json!({ "role": "assistant", "content": null, "reasoning_content": "Let me " }), + Value::Null, + ), + chunk( + json!({ "content": null, "reasoning_content": "think." }), + Value::Null, + ), + chunk(json!({ "content": "Answer" }), Value::Null), + chunk(json!({ "content": "" }), json!("stop")), + ]; + let (_, parts) = stream_parts("deepseek-reasoner", &options(), chunks).await; + let order: Vec<&str> = parts + .iter() + .filter_map(|part| match part { + StreamPart::ReasoningStart { .. } => Some("reasoning-start"), + StreamPart::ReasoningDelta { .. } => Some("reasoning-delta"), + StreamPart::ReasoningEnd { .. } => Some("reasoning-end"), + StreamPart::TextStart { .. } => Some("text-start"), + StreamPart::TextDelta { .. } => Some("text-delta"), + StreamPart::TextEnd { .. } => Some("text-end"), + _ => None, + }) + .collect(); + assert_eq!( + order, + [ + "reasoning-start", + "reasoning-delta", + "reasoning-delta", + "text-start", + "reasoning-end", + "text-delta", + "text-end" + ] + ); +} + +// ---- describe('logprobs') ---------------------------------------------------------------- + +#[tokio::test] +async fn stream_logprobs_should_send_logprobs_provider_options_and_collect_streamed_logprobs() { + let chunk = |choice: Value, usage: Value| { + data_event(&json!({ + "id": "c", "object": "chat.completion.chunk", "created": 1, "model": "deepseek-v4-flash", + "choices": [choice], "usage": usage + })) + }; + let chunks = vec![ + chunk( + json!({ "index": 0, "delta": { "role": "assistant", "reasoning_content": "Reasoning" }, + "logprobs": { "reasoning_content": [logprob("Reasoning", -0.1, Value::Null)] }, + "finish_reason": null }), + Value::Null, + ), + chunk( + json!({ "index": 0, "delta": { "content": "OK" }, + "logprobs": { "content": [logprob("OK", -0.00002467602, json!([79, 75]))] }, + "finish_reason": "stop" }), + json!({ "prompt_tokens": 9, "completion_tokens": 2, "total_tokens": 11, + "prompt_cache_hit_tokens": 0, "prompt_cache_miss_tokens": 9 }), + ), + ]; + let options = with_provider_options( + options(), + json!({ "deepseek": { "logprobs": true, "topLogprobs": 1 } }), + ); + let (body, parts) = stream_parts("deepseek-v4-flash", &options, chunks).await; + assert_eq!(body["logprobs"], true); + assert_eq!(body["top_logprobs"], 1); + assert_eq!( + finish_metadata(&parts), + json!({ "deepseek": { + "choiceIndex": 0, + "logprobs": { + "content": [logprob("OK", -0.00002467602, json!([79, 75]))], + "reasoning_content": [logprob("Reasoning", -0.1, Value::Null)] + }, + "messageRole": "assistant", + "promptCacheHitTokens": 0, + "promptCacheMissTokens": 9, + "responseObject": "chat.completion.chunk" + } }) + ); +} + +// ---- describe('tool call') ------------------------------------------------------------------ + +#[tokio::test] +async fn stream_tool_call_should_stream_tool_call() { + let chunk = |delta: Value, finish: Value| { + data_event(&json!({ + "id": "c", "object": "chat.completion.chunk", "created": 1, "model": "deepseek-reasoner", + "choices": [{ "index": 0, "delta": delta, "finish_reason": finish }] + })) + }; + let chunks = vec![ + chunk( + json!({ "role": "assistant", "reasoning_content": "Look it up." }), + Value::Null, + ), + chunk( + json!({ "tool_calls": [{ "index": 0, "id": "call_00_abc", "type": "function", + "function": { "name": "weather", "arguments": "" } }] }), + Value::Null, + ), + chunk( + json!({ "tool_calls": [{ "index": 0, "function": { "arguments": "{\"location\":" } }] }), + Value::Null, + ), + chunk( + json!({ "tool_calls": [{ "index": 0, "function": { "arguments": "\"San Francisco\"}" } }] }), + Value::Null, + ), + chunk(json!({ "content": "" }), json!("tool_calls")), + ]; + let (_, parts) = stream_parts("deepseek-reasoner", &tool_options(), chunks).await; + assert!(parts.iter().any(|p| matches!( + p, + StreamPart::ToolInputStart { id, tool_name, .. } if id == "call_00_abc" && tool_name == "weather" + ))); + let input: String = parts + .iter() + .filter_map(|p| match p { + StreamPart::ToolInputDelta { id, delta, .. } if id == "call_00_abc" => { + Some(delta.as_str()) + } + _ => None, + }) + .collect(); + assert_eq!(input, "{\"location\":\"San Francisco\"}"); + assert!(parts.iter().any(|p| matches!( + p, + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. }) + if tool_call_id == "call_00_abc" && tool_name == "weather" + && input == "{\"location\":\"San Francisco\"}" + ))); + assert!( + parts + .iter() + .any(|p| matches!(p, StreamPart::ReasoningEnd { .. })) + ); + match parts.last().unwrap() { + StreamPart::Finish { finish_reason, .. } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::ToolCalls); + } + other => panic!("expected finish, got {other:?}"), + } + assert_eq!( + finish_metadata(&parts)["deepseek"]["toolCallTypes"], + json!(["function"]) + ); +} + +// ---- describe('assistant prefix completion') ----------------------------------------------------- + +#[tokio::test] +async fn stream_prefix_should_send_prefix_true_on_the_final_assistant_message() { + let server = sse_server(text_chunks()).await; + let result = beta_chat(&server, "deepseek-chat") + .do_stream(&options_for(prefix_prompt(json!({ "prefix": true })))) + .await + .unwrap(); + let body = result.request.and_then(|request| request.body).unwrap(); + assert_eq!( + body["messages"], + json!([ + { "role": "user", "content": "Complete this sentence." }, + { "role": "assistant", "content": "The answer is", "prefix": true } + ]) + ); + assert_eq!(body["model"], "deepseek-chat"); + assert_eq!(body["stream"], true); +} + +// =========================================================================== +// convert-to-deepseek-chat-messages.test.ts (through the request body) +// =========================================================================== + +fn png_data() -> UserPart { + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("AAECAw==".into()), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + }) +} + +fn file_with_options(mut part: UserPart, options: Value) -> UserPart { + if let UserPart::File(file) = &mut part { + file.provider_options = Some(serde_json::from_value(options).unwrap()); + } + part +} + +// ---- describe('message names') ------------------------------------------------------------- + +#[tokio::test] +async fn convert_should_ignore_a_name_on_a_tool_message_with_an_unsupported_warning() { + let options = options_for(vec![with_options( + LanguageModelMessage::Tool { + content: vec![ToolPart::ToolResult(ToolResultPart { + tool_call_id: ("call-1").into(), + output: ToolResultOutput::Text { + value: "sunny".into(), + provider_options: None, + }, + tool_name: "getWeather".into(), + provider_options: None, + })], + provider_options: None, + }, + json!({ "deepseek": { "name": "weather_tool" } }), + )]); + let result = generate_result("deepseek-chat", &options).await; + assert_eq!( + result.request.and_then(|request| request.body).unwrap()["messages"], + json!([{ "role": "tool", "tool_call_id": "call-1", "content": "sunny" }]) + ); + assert_eq!( + warning_values(&result.warnings), + warning_values(&[Warning::Unsupported { + feature: "message name on tool messages".to_string(), + details: None, + }]) + ); +} + +#[tokio::test] +async fn convert_should_reject_a_non_string_name() { + let options = options_for(vec![with_options( + LanguageModelMessage::user_text("Hello"), + json!({ "deepseek": { "name": 123 } }), + )]); + let message = generate_error("deepseek-chat", &options).await; + assert!(message.contains("invalid provider options"), "{message}"); +} + +// ---- describe('user messages') --------------------------------------------------------------- + +#[tokio::test] +async fn convert_should_convert_messages_with_only_a_text_part_to_a_string_content() { + let messages = wire_messages( + "deepseek-chat", + vec![LanguageModelMessage::user_text("Hello")], + ) + .await; + assert_eq!(messages, [json!({ "role": "user", "content": "Hello" })]); +} + +#[tokio::test] +async fn convert_should_convert_image_data_to_an_image_url_content_part() { + let messages = wire_messages( + "deepseek-chat", + vec![LanguageModelMessage::User { + content: vec![ + UserPart::Text(TextPart { + text: ("Hello").into(), + provider_options: None, + }), + png_data(), + ], + provider_options: None, + }], + ) + .await; + assert_eq!( + messages, + [json!({ "role": "user", "content": [ + { "type": "text", "text": "Hello" }, + { "type": "image_url", "image_url": { "url": "data:image/png;base64,AAECAw==" } } + ] })] + ); +} + +#[tokio::test] +async fn convert_should_convert_an_image_url_to_an_image_url_content_part() { + let messages = wire_messages( + "deepseek-chat", + vec![LanguageModelMessage::User { + content: vec![ + UserPart::Text(TextPart { + text: ("Hello").into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Url { + url: ("https://example.com/image.png").into(), + original_url: None, + }, + media_type: ("image/png").into(), + filename: None, + provider_options: None, + }), + ], + provider_options: None, + }], + ) + .await; + assert_eq!( + messages, + [json!({ "role": "user", "content": [ + { "type": "text", "text": "Hello" }, + { "type": "image_url", "image_url": { "url": "https://example.com/image.png" } } + ] })] + ); +} + +#[tokio::test] +async fn convert_should_pass_image_detail_to_image_url_content_parts() { + let messages = wire_messages( + "deepseek-v4-flash-vision-exp", + vec![LanguageModelMessage::User { + content: vec![file_with_options( + UserPart::File(FilePart { + data: FileData::Url { + url: ("https://example.com/image.webp").into(), + original_url: None, + }, + media_type: ("image/webp").into(), + filename: None, + provider_options: None, + }), + json!({ "deepseek": { "imageDetail": "low" } }), + )], + provider_options: None, + }], + ) + .await; + assert_eq!( + messages, + [json!({ "role": "user", "content": [ + { "type": "image_url", "image_url": { "url": "https://example.com/image.webp", "detail": "low" } } + ] })] + ); +} + +#[tokio::test] +async fn convert_should_convert_inline_image_data_to_file_data_and_preserve_its_filename() { + let part = UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("AAECAw==".to_string()), + }, + media_type: "image/jpg".to_string(), + filename: Some("sample.jpg".to_string()), + provider_options: Some( + provider_namespace("deepseek", json!({ "fileData": true })).unwrap(), + ), + }); + let messages = wire_messages( + "deepseek-v4-flash-vision-exp", + vec![LanguageModelMessage::User { + content: vec![part], + provider_options: None, + }], + ) + .await; + assert_eq!( + messages, + [json!({ "role": "user", "content": [ + { "type": "file", "file_data": "data:image/jpeg;base64,AAECAw==", "filename": "sample.jpg" } + ] })] + ); +} + +#[tokio::test] +async fn convert_should_reject_image_detail_together_with_file_data() { + let part = file_with_options( + png_data(), + json!({ "deepseek": { "fileData": true, "imageDetail": "high" } }), + ); + let message = generate_error( + "deepseek-v4-flash-vision-exp", + &options_for(vec![LanguageModelMessage::User { + content: vec![part], + provider_options: None, + }]), + ) + .await; + assert!( + message.contains("DeepSeek `imageDetail` cannot be combined with `fileData`."), + "{message}" + ); +} + +#[tokio::test] +async fn convert_should_reject_image_urls_longer_than_8192_characters() { + let url = format!("https://example.com/{}", "a".repeat(8192)); + let message = generate_error( + "deepseek-v4-flash-vision-exp", + &options_for(vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Url { + url, + original_url: None, + }, + media_type: ("image/png").into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]), + ) + .await; + assert!( + message.contains("DeepSeek image URLs must not exceed 8192 characters."), + "{message}" + ); +} + +#[tokio::test] +async fn convert_should_reject_unsupported_image_formats() { + let message = generate_error( + "deepseek-v4-flash-vision-exp", + &options_for(vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64(("AAECAw==").into()), + }, + media_type: ("image/svg+xml").into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]), + ) + .await; + assert!( + message.contains("DeepSeek supports JPEG, PNG, GIF, and WebP image inputs."), + "{message}" + ); +} + +#[tokio::test] +async fn convert_should_convert_an_image_provider_reference_to_a_file_content_part() { + let messages = wire_messages( + "deepseek-v4-flash-vision-exp", + vec![LanguageModelMessage::User { + content: vec![ + UserPart::Text(TextPart { + text: ("Hello").into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Reference { + reference: serde_json::from_value( + json!({ "deepseek": "file-api-deepseek", "openai": "file-openai" }), + ) + .unwrap(), + }, + media_type: ("image/png").into(), + filename: None, + provider_options: None, + }), + ], + provider_options: None, + }], + ) + .await; + assert_eq!( + messages, + [json!({ "role": "user", "content": [ + { "type": "text", "text": "Hello" }, + { "type": "file", "file_id": "file-api-deepseek" } + ] })] + ); +} + +#[tokio::test] +async fn convert_should_throw_when_an_image_reference_has_no_deepseek_identifier() { + let message = generate_error( + "deepseek-v4-flash-vision-exp", + &options_for(vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Reference { + reference: serde_json::from_value(json!({ "openai": "file-openai" })).unwrap(), + }, + media_type: ("image/png").into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]), + ) + .await; + assert!(message.contains("deepseek"), "{message}"); +} + +#[tokio::test] +async fn convert_should_warn_about_unsupported_non_image_file_parts() { + let result = generate_result( + "deepseek-chat", + &options_for(vec![LanguageModelMessage::User { + content: vec![ + UserPart::Text(TextPart { + text: ("Hello").into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64(("AAECAw==").into()), + }, + media_type: ("application/pdf").into(), + filename: None, + provider_options: None, + }), + ], + provider_options: None, + }]), + ) + .await; + assert_eq!( + result.request.and_then(|request| request.body).unwrap()["messages"], + json!([{ "role": "user", "content": "Hello" }]) + ); + assert_eq!( + warning_values(&result.warnings), + warning_values(&[Warning::Unsupported { + feature: "user message part type: file".to_string(), + details: None, + }]) + ); +} + +// ---- describe('tool calls') ------------------------------------------------------------------ + +fn tool_turn(reasoning: bool) -> Vec { + let mut assistant = vec![]; + if reasoning { + assistant.push(AssistantPart::Reasoning(ReasoningPart { + text: ("I think the tool will return the correct value.").into(), + provider_options: None, + })); + } + assistant.push(AssistantPart::ToolCall(ToolCallPart { + tool_call_id: ("quux").into(), + tool_name: ("thwomp").into(), + input: json!({ "foo": "bar123" }), + provider_executed: None, + provider_options: None, + })); + vec![ + LanguageModelMessage::Assistant { + content: assistant, + provider_options: None, + }, + LanguageModelMessage::Tool { + content: vec![ToolPart::ToolResult(ToolResultPart { + tool_call_id: ("quux").into(), + output: ToolResultOutput::Json { + value: json!({ "oof": "321rab" }), + provider_options: None, + }, + tool_name: "thwomp".into(), + provider_options: None, + })], + provider_options: None, + }, + ] +} + +fn wire_tool_call() -> Value { + json!([{ + "id": "quux", "type": "function", + "function": { "name": "thwomp", "arguments": "{\"foo\":\"bar123\"}" } + }]) +} + +#[tokio::test] +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"], + json!([ + // upstream sends the empty text as `content: ""`, not null. + { "role": "assistant", "content": "", "tool_calls": wire_tool_call() }, + { "role": "tool", "tool_call_id": "quux", "content": "{\"oof\":\"321rab\"}" } + ]) + ); + assert!(result.warnings.is_empty()); +} + +#[tokio::test] +async fn convert_should_handle_text_output_type_in_tool_results() { + let messages = wire_messages( + "deepseek-chat", + vec![ + LanguageModelMessage::Assistant { + content: vec![AssistantPart::ToolCall(ToolCallPart { + tool_call_id: ("call-1").into(), + tool_name: ("getWeather").into(), + input: json!({ "query": "weather" }), + provider_executed: None, + provider_options: None, + })], + provider_options: None, + }, + LanguageModelMessage::Tool { + content: vec![ToolPart::ToolResult(ToolResultPart { + tool_call_id: ("call-1").into(), + output: ToolResultOutput::Text { + value: "It is sunny today".into(), + provider_options: None, + }, + tool_name: "getWeather".into(), + provider_options: None, + })], + provider_options: None, + }, + ], + ) + .await; + assert_eq!( + messages[1], + json!({ "role": "tool", "tool_call_id": "call-1", "content": "It is sunny today" }) + ); +} + +#[tokio::test] +async fn convert_should_support_reasoning_content_in_tool_calls() { + let mut prompt = vec![LanguageModelMessage::user_text("Hello")]; + prompt.extend(tool_turn(true)); + let messages = wire_messages("deepseek-chat", prompt).await; + assert_eq!( + messages[1], + json!({ + "role": "assistant", "content": "", + "reasoning_content": "I think the tool will return the correct value.", + "tool_calls": wire_tool_call() + }) + ); +} + +#[tokio::test] +async fn convert_should_filter_out_reasoning_content_from_turns_before_the_last_user_message() { + let mut prompt = vec![LanguageModelMessage::user_text("Hello")]; + prompt.extend(tool_turn(true)); + prompt.push(LanguageModelMessage::user_text("Goodbye")); + let messages = wire_messages("deepseek-chat", prompt).await; + assert_eq!( + messages[1], + json!({ "role": "assistant", "content": "", "tool_calls": wire_tool_call() }) + ); +} + +// ---- describe('deepseek-v4 thinking mode') ------------------------------------------------------- + +#[tokio::test] +async fn convert_should_preserve_reasoning_content_from_prior_turns_for_deepseek_v4() { + let mut prompt = vec![LanguageModelMessage::user_text("Hello")]; + prompt.extend(tool_turn(true)); + prompt.push(LanguageModelMessage::user_text("Goodbye")); + let messages = wire_messages("deepseek-v4-pro", prompt).await; + assert_eq!( + messages[1], + json!({ + "role": "assistant", "content": "", + "reasoning_content": "I think the tool will return the correct value.", + "tool_calls": wire_tool_call() + }) + ); +} + +#[tokio::test] +async fn convert_should_preserve_reasoning_content_from_prior_turns_for_the_deepseek_flash_alias() { + let messages = wire_messages( + "deepseek-flash", + vec![ + LanguageModelMessage::user_text("Hello"), + LanguageModelMessage::Assistant { + content: vec![ + AssistantPart::Reasoning(ReasoningPart { + text: ("Prior-turn reasoning.").into(), + provider_options: None, + }), + AssistantPart::Text(TextPart { + text: ("Hi there").into(), + provider_options: None, + }), + ], + provider_options: None, + }, + LanguageModelMessage::user_text("Again"), + ], + ) + .await; + assert_eq!( + messages[1], + json!({ "role": "assistant", "content": "Hi there", "reasoning_content": "Prior-turn reasoning." }) + ); +} + +#[tokio::test] +async fn convert_should_back_fill_empty_reasoning_content_for_deepseek_v4_assistant_messages_with_no_reasoning_part() + { + let messages = wire_messages( + "deepseek-v4-pro", + vec![ + LanguageModelMessage::user_text("Hello"), + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: ("Hi there").into(), + provider_options: None, + })], + provider_options: None, + }, + LanguageModelMessage::user_text("Again"), + ], + ) + .await; + assert_eq!( + messages[1], + json!({ "role": "assistant", "content": "Hi there", "reasoning_content": "" }) + ); +} + +// ---- describe('assistant prefix completion') ----------------------------------------------------- + +#[tokio::test] +async fn convert_should_reject_prefix_completion_on_a_non_assistant_message() { + let message = generate_error( + "deepseek-chat", + &options_for(vec![with_options( + LanguageModelMessage::user_text("Hello"), + json!({ "deepseek": { "prefix": true } }), + )]), + ) + .await; + assert!( + message.contains( + "DeepSeek assistant prefix completion requires `prefix: true` on an assistant message." + ), + "{message}" + ); +} + +#[tokio::test] +async fn convert_should_reject_prefix_completion_on_a_non_final_assistant_message() { + let server = json_server(text_response()).await; + let prompt = vec![ + with_options( + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: ("The answer is").into(), + provider_options: None, + })], + provider_options: None, + }, + json!({ "deepseek": { "prefix": true } }), + ), + LanguageModelMessage::user_text("Continue"), + ]; + let error = beta_chat(&server, "deepseek-chat") + .do_generate(&options_for(prompt)) + .await + .unwrap_err(); + assert!( + error.to_string().contains( + "DeepSeek assistant prefix completion requires the prefixed assistant message to be the final message." + ), + "{error}" ); } diff --git a/aimux-providers/tests/deepseek_test.rs b/aimux-providers/tests/deepseek_test.rs new file mode 100644 index 000000000..396695a1c --- /dev/null +++ b/aimux-providers/tests/deepseek_test.rs @@ -0,0 +1,50 @@ +//! The DeepSeek provider (`aimux_providers::deepseek`): `create_deepseek` +//! settings and the recorded cassette replay. The chat +//! model's wire behaviour (ported from the `@ai-sdk/deepseek` tests) is in +//! `deepseek_chat_test.rs`. + +mod common; + +use futures::StreamExt; +use wiremock::MockServer; + +use aimux_core::generate::{GenerateTextOptions, generate_text, stream_text}; +use aimux_core::stream_part::TextStreamPart; +use aimux_providers::deepseek::{DeepSeekProvider, DeepSeekProviderSettings, create_deepseek}; + +fn deepseek_at(server: &MockServer) -> DeepSeekProvider { + create_deepseek(DeepSeekProviderSettings { + base_url: Some(server.uri()), + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap() +} + +/// The recorded DeepSeek exchanges (`tests/cassettes/deepseek`) replayed +/// through the package. +#[tokio::test] +async fn recorded_deepseek_cassettes_replay_through_the_package() { + let server = MockServer::start().await; + let n = common::replay::mount_cassettes(&server, "tests/cassettes/deepseek").await; + assert!(n > 0, "no deepseek cassettes"); + let model = deepseek_at(&server).chat("deepseek-chat"); + + let result = generate_text(&model, "Hello", GenerateTextOptions::default()) + .await + .expect("generate_text should succeed with cassette replay"); + assert!(!result.text().is_empty() || !result.tool_calls().is_empty()); + assert!(result.usage().input_tokens.is_some()); + + let result = stream_text(&model, "Hello", GenerateTextOptions::default()) + .await + .expect("stream_text should succeed"); + let mut stream = result.stream; + let mut finished = false; + while let Some(part) = stream.next().await { + if let TextStreamPart::Finish { .. } = part.expect("stream part") { + finished = true; + } + } + assert!(finished, "stream should finish"); +} diff --git a/aimux-providers/tests/download_ssrf_guard_test.rs b/aimux-providers/tests/download_ssrf_guard_test.rs index 133a95016..607e12ee6 100644 --- a/aimux-providers/tests/download_ssrf_guard_test.rs +++ b/aimux-providers/tests/download_ssrf_guard_test.rs @@ -14,7 +14,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::image_model::{ImageCallOptions, ImageModel}; -use aimux_providers::{RecraftConfig, RecraftProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{RecraftProviderSettings, create_recraft}; const PROMPT: &str = "A cute baby sea otter"; @@ -34,8 +35,12 @@ async fn mock_generations_with_url(server: &MockServer, url: &str) { } fn make_model(server: &MockServer) -> impl ImageModel { - let config = RecraftConfig::new("test-recraft-key").with_base_url(server.uri()); - RecraftProvider::new(config).image("recraftv3") + let config = RecraftProviderSettings { + api_key: Some(Resolvable::Value("test-recraft-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_recraft(config).unwrap().image("recraftv3") } #[tokio::test] diff --git a/aimux-providers/tests/e2e_test.rs b/aimux-providers/tests/e2e_test.rs index 09cbadd42..26c9832d4 100644 --- a/aimux-providers/tests/e2e_test.rs +++ b/aimux-providers/tests/e2e_test.rs @@ -11,12 +11,34 @@ use aimux_core::error::AiMuxError; use aimux_core::generate::{GenerateTextOptions, generate_text, stream_text}; use aimux_core::message::ModelMessage; use aimux_core::stream_part::TextStreamPart; -use aimux_core::tool::{FunctionTool, Tool, ToolChoice}; +use aimux_core::tool::{FunctionTool, Tool, ToolCall, ToolChoice}; use aimux_core::types::FinishReasonUnified; +use aimux_providers::anthropic::{AnthropicProvider, AnthropicProviderSettings, create_anthropic}; +use aimux_providers::openai::{OpenAIProvider, OpenAIProviderSettings, create_openai}; use aimux_providers::{ - AnthropicConfig, AnthropicProvider, CohereConfig, CohereProvider, GoogleConfig, GoogleProvider, - MistralConfig, MistralProvider, OpenAIConfig, OpenAIProvider, + CohereProviderSettings, GoogleProviderSettings, MistralProviderSettings, create_cohere, + create_google, create_mistral, }; + +/// The native Anthropic package pointed at a mock server. +fn anthropic_at(server_uri: String, key: &str) -> AnthropicProvider { + create_anthropic(AnthropicProviderSettings { + api_key: Some(key.to_string()), + base_url: Some(format!("{server_uri}/v1")), + ..Default::default() + }) + .expect("valid settings") +} + +/// The native OpenAI package pointed at a mock server. +fn openai_at(base_url: String, key: &str) -> OpenAIProvider { + create_openai(OpenAIProviderSettings { + api_key: Some(key.to_string()), + base_url: Some(base_url), + ..Default::default() + }) + .unwrap() +} use futures::StreamExt; use serde_json::json; use wiremock::{Mock, MockServer, ResponseTemplate, matchers::*}; @@ -53,8 +75,8 @@ async fn e2e_openai_generate_text() { .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gpt-4o"); + let provider = openai_at(server.uri(), "test-key"); + let model = provider.chat("gpt-4o"); let result = generate_text(&model, "What is Rust?", GenerateTextOptions::default()) .await @@ -88,8 +110,8 @@ async fn e2e_openai_stream_text() { .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gpt-4o"); + let provider = openai_at(server.uri(), "test-key"); + let model = provider.chat("gpt-4o"); let result = stream_text(&model, "Say hello", GenerateTextOptions::default()) .await @@ -141,8 +163,8 @@ async fn e2e_openai_generate_with_tools() { .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gpt-4o"); + let provider = openai_at(server.uri(), "test-key"); + let model = provider.chat("gpt-4o"); let result = generate_text( &model, @@ -187,8 +209,8 @@ async fn e2e_openai_error_401() { .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gpt-4o"); + let provider = openai_at(server.uri(), "test-key"); + let model = provider.chat("gpt-4o"); let result = generate_text(&model, "test", GenerateTextOptions::default()).await; @@ -220,9 +242,8 @@ async fn e2e_anthropic_generate_text() { .mount(&server) .await; - let provider = - AnthropicProvider::new(AnthropicConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("claude-3-5-sonnet-20241022"); + let provider = anthropic_at(server.uri(), "test-key"); + let model = provider.messages("claude-3-5-sonnet-20241022"); let result = generate_text( &model, @@ -269,9 +290,8 @@ async fn e2e_anthropic_stream_text() { .mount(&server) .await; - let provider = - AnthropicProvider::new(AnthropicConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("claude-3-5-sonnet-20241022"); + let provider = anthropic_at(server.uri(), "test-key"); + let model = provider.messages("claude-3-5-sonnet-20241022"); let result = stream_text( &model, @@ -297,9 +317,8 @@ async fn e2e_anthropic_error_429() { .mount(&server) .await; - let provider = - AnthropicProvider::new(AnthropicConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("claude-3-5-sonnet-20241022"); + let provider = anthropic_at(server.uri(), "test-key"); + let model = provider.messages("claude-3-5-sonnet-20241022"); let result = generate_text(&model, "test", GenerateTextOptions::default()).await; @@ -340,8 +359,8 @@ async fn e2e_openai_stream_part_sequence() { .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gpt-4o"); + let provider = openai_at(server.uri(), "test-key"); + let model = provider.chat("gpt-4o"); let result = stream_text(&model, "test", GenerateTextOptions::default()) .await @@ -404,9 +423,8 @@ async fn e2e_anthropic_stream_part_sequence() { .mount(&server) .await; - let provider = - AnthropicProvider::new(AnthropicConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("claude-3-5-sonnet-20241022"); + let provider = anthropic_at(server.uri(), "test-key"); + let model = provider.messages("claude-3-5-sonnet-20241022"); let result = stream_text(&model, "test", GenerateTextOptions::default()) .await @@ -453,8 +471,7 @@ async fn e2e_provider_interchangeability() { .mount(&openai_server) .await; - let openai = OpenAIProvider::new(OpenAIConfig::new("key").with_base_url(openai_server.uri())) - .model("gpt-4o"); + let openai = openai_at(openai_server.uri(), "key").chat("gpt-4o"); let openai_result = generate_text(&openai, "test", GenerateTextOptions::default()) .await @@ -474,9 +491,7 @@ async fn e2e_provider_interchangeability() { .mount(&anthropic_server) .await; - let anthropic = - AnthropicProvider::new(AnthropicConfig::new("key").with_base_url(anthropic_server.uri())) - .model("claude-3"); + let anthropic = anthropic_at(anthropic_server.uri(), "key").messages("claude-3"); let anthropic_result = generate_text(&anthropic, "test", GenerateTextOptions::default()) .await @@ -535,8 +550,8 @@ async fn e2e_openai_tool_call_round_trip() { .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gpt-4o"); + let provider = openai_at(server.uri(), "test-key"); + let model = provider.chat("gpt-4o"); let weather_tool = || { Tool::Function( @@ -651,8 +666,8 @@ async fn e2e_openai_multi_turn_dialog() { .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gpt-4o"); + let provider = openai_at(server.uri(), "test-key"); + let model = provider.chat("gpt-4o"); // Full multi-turn conversation: // system → user(q1) → assistant(a1) → user(q2 follow-up) @@ -723,8 +738,8 @@ async fn e2e_openai_tool_choice_required() { .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gpt-4o"); + let provider = openai_at(server.uri(), "test-key"); + let model = provider.chat("gpt-4o"); let weather_tool = Tool::Function( FunctionTool::new( @@ -791,8 +806,8 @@ async fn e2e_openai_stream_tool_calls() { .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gpt-4o"); + let provider = openai_at(server.uri(), "test-key"); + let model = provider.chat("gpt-4o"); let weather_tool = Tool::Function( FunctionTool::new( @@ -848,7 +863,7 @@ async fn e2e_openai_stream_tool_calls() { // If a complete ToolCall part exists, verify its fields let tool_call = parts.iter().find_map(|p| match p { - TextStreamPart::ToolCall(aimux_core::tool::ToolCall { + TextStreamPart::ToolCall(ToolCall { tool_name, input, .. }) => Some((tool_name, input)), _ => None, @@ -887,8 +902,13 @@ async fn e2e_google_generate_text() { .mount(&server) .await; - let provider = GoogleProvider::new(GoogleConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gemini-2.0-flash"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.chat("gemini-2.0-flash"); let result = generate_text(&model, "What is Rust?", GenerateTextOptions::default()) .await @@ -926,8 +946,13 @@ async fn e2e_google_stream_text() { .mount(&server) .await; - let provider = GoogleProvider::new(GoogleConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("gemini-2.0-flash"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.chat("gemini-2.0-flash"); let result = stream_text(&model, "Say hello", GenerateTextOptions::default()) .await @@ -968,8 +993,13 @@ async fn e2e_mistral_generate_text() { .mount(&server) .await; - let provider = MistralProvider::new(MistralConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("mistral-small-latest"); + let provider = create_mistral(MistralProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .expect("valid settings"); + let model = provider.chat("mistral-small-latest"); let result = generate_text(&model, "What is Rust?", GenerateTextOptions::default()) .await @@ -1010,8 +1040,13 @@ async fn e2e_mistral_stream_text() { .mount(&server) .await; - let provider = MistralProvider::new(MistralConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("mistral-small-latest"); + let provider = create_mistral(MistralProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .expect("valid settings"); + let model = provider.chat("mistral-small-latest"); let result = stream_text(&model, "Say hello", GenerateTextOptions::default()) .await @@ -1054,8 +1089,13 @@ async fn e2e_cohere_generate_text() { .mount(&server) .await; - let provider = CohereProvider::new(CohereConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("command-r-plus"); + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .expect("valid settings"); + let model = provider.chat("command-r-plus"); let result = generate_text(&model, "What is Rust?", GenerateTextOptions::default()) .await @@ -1104,8 +1144,13 @@ async fn e2e_cohere_stream_text() { .mount(&server) .await; - let provider = CohereProvider::new(CohereConfig::new("test-key").with_base_url(server.uri())); - let model = provider.model("command-r-plus"); + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .expect("valid settings"); + let model = provider.chat("command-r-plus"); let result = stream_text(&model, "Say hello", GenerateTextOptions::default()) .await diff --git a/aimux-providers/tests/elevenlabs_realtime_stream_test.rs b/aimux-providers/tests/elevenlabs_realtime_stream_test.rs index e3963235a..eecf99b98 100644 --- a/aimux-providers/tests/elevenlabs_realtime_stream_test.rs +++ b/aimux-providers/tests/elevenlabs_realtime_stream_test.rs @@ -10,6 +10,7 @@ //! (2026-09). Live-API smoke pending (no key at implementation time) — same //! posture as RFC-0028 D4. +use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -25,7 +26,26 @@ use aimux_core::transcription_model::{ AudioChunk, InputAudioFormat, TranscriptionModel, TranscriptionStreamOptions, TranscriptionStreamPart, }; -use aimux_providers::{ElevenLabsConfig, ElevenLabsProvider}; +use aimux_providers::{ElevenLabsProvider, ElevenLabsProviderSettings, create_elevenlabs}; + +fn test_provider( + api_key: &str, + base_url: impl Into, + headers: Option>, +) -> ElevenLabsProvider { + create_elevenlabs(ElevenLabsProviderSettings { + api_key: Some(api_key.to_string().into()), + base_url: Some(base_url.into()), + headers: headers.map(|headers| { + headers + .into_iter() + .map(|(name, value)| (name, Some(value))) + .collect() + }), + ..Default::default() + }) + .expect("valid settings") +} // ── helpers ───────────────────────────────────────────────────────────────── @@ -135,12 +155,7 @@ async fn start( } fn realtime_model(base_url: &str, model_id: &str) -> aimux_providers::ElevenLabsTranscriptionModel { - let config = ElevenLabsConfig { - api_key: "test-api-key".to_string(), - base_url: base_url.to_string(), - headers: None, - }; - ElevenLabsProvider::new(config).transcription(model_id) + test_provider("test-api-key", base_url, None).transcription(model_id) } fn stream_options( diff --git a/aimux-providers/tests/elevenlabs_speech_test.rs b/aimux-providers/tests/elevenlabs_speech_test.rs index ea7bd2ecf..5c73793ab 100644 --- a/aimux-providers/tests/elevenlabs_speech_test.rs +++ b/aimux-providers/tests/elevenlabs_speech_test.rs @@ -9,13 +9,34 @@ //! Note: the TS test mounts at `https://api.elevenlabs.io/v1/text-to-speech/*`. //! The Rust mock server uses a wildcard path matcher to capture any voice id. +use std::collections::HashMap; + use serde_json::{Value, json}; use wiremock::matchers::{method, path_regex}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::shared::Warning; use aimux_core::speech_model::{SpeechCallOptions, SpeechModel}; -use aimux_providers::{ElevenLabsConfig, ElevenLabsProvider}; +use aimux_providers::{ElevenLabsProvider, ElevenLabsProviderSettings, create_elevenlabs}; + +fn test_provider( + api_key: &str, + base_url: impl Into, + headers: Option>, +) -> ElevenLabsProvider { + create_elevenlabs(ElevenLabsProviderSettings { + api_key: Some(api_key.to_string().into()), + base_url: Some(base_url.into()), + headers: headers.map(|headers| { + headers + .into_iter() + .map(|(name, value)| (name, Some(value))) + .collect() + }), + ..Default::default() + }) + .expect("valid settings") +} // 鈹€鈹€ helpers 鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€ @@ -51,8 +72,7 @@ async fn should_generate_speech_with_required_parameters() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.speech("eleven_multilingual_v2"); let mut options = speech_options("Hello, world!"); @@ -80,8 +100,7 @@ async fn should_handle_custom_output_format() { let server = MockServer::start().await; mock_audio_response(&server, "pcm").await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.speech("eleven_multilingual_v2"); let mut options = speech_options("Hello, world!"); @@ -108,8 +127,7 @@ async fn should_handle_language_parameter() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.speech("eleven_multilingual_v2"); let mut options = speech_options("Hola, mundo!"); @@ -137,8 +155,7 @@ async fn should_handle_speed_parameter_in_voice_settings() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.speech("eleven_multilingual_v2"); let mut options = speech_options("Hello, world!"); @@ -160,8 +177,7 @@ async fn should_warn_about_unsupported_instructions_parameter() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.speech("eleven_multilingual_v2"); let mut options = speech_options("Hello, world!"); @@ -191,8 +207,7 @@ async fn should_pass_provider_specific_options() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.speech("eleven_multilingual_v2"); let mut options = speech_options("Hello, world!"); @@ -239,8 +254,7 @@ async fn should_include_api_key_header() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.speech("eleven_multilingual_v2"); let mut options = speech_options("Hello, world!"); diff --git a/aimux-providers/tests/elevenlabs_transcription_test.rs b/aimux-providers/tests/elevenlabs_transcription_test.rs index 9729e4c48..accf16f4f 100644 --- a/aimux-providers/tests/elevenlabs_transcription_test.rs +++ b/aimux-providers/tests/elevenlabs_transcription_test.rs @@ -8,7 +8,26 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; -use aimux_providers::{ElevenLabsConfig, ElevenLabsProvider}; +use aimux_providers::{ElevenLabsProvider, ElevenLabsProviderSettings, create_elevenlabs}; + +fn test_provider( + api_key: &str, + base_url: impl Into, + headers: Option>, +) -> ElevenLabsProvider { + create_elevenlabs(ElevenLabsProviderSettings { + api_key: Some(api_key.to_string().into()), + base_url: Some(base_url.into()), + headers: headers.map(|headers| { + headers + .into_iter() + .map(|(name, value)| (name, Some(value))) + .collect() + }), + ..Default::default() + }) + .expect("valid settings") +} fn mock_audio() -> Vec { vec![1u8, 2, 3] @@ -75,8 +94,7 @@ async fn should_pass_model_id() { let server = MockServer::start().await; mock_response(&server, &fixture_response(), &[]).await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.transcription("scribe_v1"); model @@ -100,10 +118,7 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = ElevenLabsConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), Some(ph)); let model = provider.transcription("scribe_v1"); let mut opts = options(mock_audio(), "audio/wav"); @@ -134,8 +149,7 @@ async fn should_extract_text() { let server = MockServer::start().await; mock_response(&server, &fixture_response(), &[]).await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.transcription("scribe_v1"); let result = model @@ -160,8 +174,7 @@ async fn should_include_response_data() { ) .await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.transcription("scribe_v1"); let result = model @@ -178,8 +191,7 @@ async fn should_use_real_date() { let server = MockServer::start().await; mock_response(&server, &fixture_response(), &[]).await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.transcription("scribe_v1"); let result = model @@ -201,8 +213,7 @@ async fn should_handle_no_words() { ) .await; - let config = ElevenLabsConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ElevenLabsProvider::new(config); + let provider = test_provider("test-api-key", server.uri(), None); let model = provider.transcription("scribe_v1"); let result = model diff --git a/aimux-providers/tests/exa_ai_test.rs b/aimux-providers/tests/exa_ai_test.rs index 8c8a56e8b..f64627180 100644 --- a/aimux-providers/tests/exa_ai_test.rs +++ b/aimux-providers/tests/exa_ai_test.rs @@ -6,22 +6,25 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::provider::Provider; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{ExaAiConfig, ExaAiProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{ExaAiProvider, ExaAiProviderSettings, create_exa_ai}; fn make_provider(server: &MockServer) -> ExaAiProvider { - let config = ExaAiConfig::new("test-api-key").with_base_url(server.uri()); - ExaAiProvider::new(config) -} - -#[test] -fn provider_name_is_exa_ai() { - let provider = ExaAiProvider::new(ExaAiConfig::new("test-key")); - assert_eq!(provider.name(), "exa_ai"); + let config = ExaAiProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_exa_ai(config).unwrap() } #[test] fn language_model_returns_unsupported() { - let provider = ExaAiProvider::new(ExaAiConfig::new("test-key")); + let provider = create_exa_ai(ExaAiProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + ..Default::default() + }) + .unwrap(); assert!(provider.language_model("any").is_err()); } @@ -64,8 +67,12 @@ async fn uses_x_api_key_header() { .mount(&server) .await; - let config = ExaAiConfig::new("my-key").with_base_url(server.uri()); - let provider = ExaAiProvider::new(config); + let config = ExaAiProviderSettings { + api_key: Some(Resolvable::Value("my-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_exa_ai(config).unwrap(); let model = provider.search_model(); model .do_search(&SearchCallOptions::new("test")) diff --git a/aimux-providers/tests/fal_image_test.rs b/aimux-providers/tests/fal_image_test.rs index 3628ed656..5eb823ca3 100644 --- a/aimux-providers/tests/fal_image_test.rs +++ b/aimux-providers/tests/fal_image_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the fal image model tests. +//! Rust translation of the fal image model tests. //! Source: `reference/ai/packages/fal/src/fal-image-model.test.ts` use serde_json::{Value, json}; @@ -6,7 +6,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; -use aimux_providers::{FalConfig, FalProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{FalProviderSettings, create_fal}; const PROMPT: &str = "A cute baby sea otter"; @@ -35,8 +36,12 @@ fn options(prompt: &str) -> ImageCallOptions { async fn should_extract_generated_images() { let server = MockServer::start().await; mock_fal(&server, fal_response(&server.uri())).await; - let config = FalConfig::new("test-key").with_base_url(server.uri()); - let model = FalProvider::new(config).image("fal-ai/flux/schnell"); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_fal(config).unwrap().image("fal-ai/flux/schnell"); let result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { ImageOutputs::Binary(imgs) => { @@ -51,8 +56,12 @@ async fn should_extract_generated_images() { async fn should_pass_prompt_and_n() { let server = MockServer::start().await; mock_fal(&server, fal_response(&server.uri())).await; - let config = FalConfig::new("test-key").with_base_url(server.uri()); - let model = FalProvider::new(config).image("fal-ai/flux/schnell"); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_fal(config).unwrap().image("fal-ai/flux/schnell"); let mut opts = options(PROMPT); opts.n = 1; model.do_generate(&opts).await.unwrap(); @@ -66,8 +75,12 @@ async fn should_pass_prompt_and_n() { async fn should_pass_size() { let server = MockServer::start().await; mock_fal(&server, fal_response(&server.uri())).await; - let config = FalConfig::new("test-key").with_base_url(server.uri()); - let model = FalProvider::new(config).image("fal-ai/flux/schnell"); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_fal(config).unwrap().image("fal-ai/flux/schnell"); let mut opts = options(PROMPT); opts.n = 1; opts.size = Some(aimux_core::shared::Size::new(1024, 1024)); @@ -82,8 +95,12 @@ async fn should_pass_size() { async fn should_pass_aspect_ratio() { let server = MockServer::start().await; mock_fal(&server, fal_response(&server.uri())).await; - let config = FalConfig::new("test-key").with_base_url(server.uri()); - let model = FalProvider::new(config).image("fal-ai/flux/schnell"); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_fal(config).unwrap().image("fal-ai/flux/schnell"); let mut opts = options(PROMPT); opts.n = 1; opts.aspect_ratio = Some(aimux_core::shared::AspectRatio::new(16, 9)); @@ -97,8 +114,12 @@ async fn should_pass_aspect_ratio() { async fn should_pass_seed() { let server = MockServer::start().await; mock_fal(&server, fal_response(&server.uri())).await; - let config = FalConfig::new("test-key").with_base_url(server.uri()); - let model = FalProvider::new(config).image("fal-ai/flux/schnell"); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_fal(config).unwrap().image("fal-ai/flux/schnell"); let mut opts = options(PROMPT); opts.n = 1; opts.seed = Some(42); @@ -112,8 +133,12 @@ async fn should_pass_seed() { async fn should_pass_auth_headers() { let server = MockServer::start().await; mock_fal(&server, fal_response(&server.uri())).await; - let config = FalConfig::new("test-key").with_base_url(server.uri()); - let model = FalProvider::new(config).image("fal-ai/flux/schnell"); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_fal(config).unwrap().image("fal-ai/flux/schnell"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); assert_eq!( @@ -131,8 +156,12 @@ async fn should_pass_auth_headers() { async fn should_forward_provider_options() { let server = MockServer::start().await; mock_fal(&server, fal_response(&server.uri())).await; - let config = FalConfig::new("test-key").with_base_url(server.uri()); - let model = FalProvider::new(config).image("fal-ai/flux/schnell"); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_fal(config).unwrap().image("fal-ai/flux/schnell"); let mut opts = options(PROMPT); opts.n = 1; opts.provider_options.insert( @@ -155,8 +184,12 @@ async fn should_handle_single_image_response() { let server = MockServer::start().await; let body = json!({ "image": { "url": format!("{}/image.png", server.uri()) } }); mock_fal(&server, body).await; - let config = FalConfig::new("test-key").with_base_url(server.uri()); - let model = FalProvider::new(config).image("fal-ai/flux/schnell"); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_fal(config).unwrap().image("fal-ai/flux/schnell"); let result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { ImageOutputs::Binary(imgs) => assert_eq!(imgs.len(), 1), diff --git a/aimux-providers/tests/fal_transcription_test.rs b/aimux-providers/tests/fal_transcription_test.rs index ca56295ce..8ff7c6051 100644 --- a/aimux-providers/tests/fal_transcription_test.rs +++ b/aimux-providers/tests/fal_transcription_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Fal transcription model tests. +//! Rust translation of the Fal transcription model tests. //! Source: `reference/ai/packages/fal/src/fal-transcription-model.test.ts` use std::collections::HashMap; @@ -8,7 +8,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; -use aimux_providers::{FalConfig, FalProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{FalProviderSettings, create_fal}; fn mock_audio() -> Vec { vec![1u8, 2, 3] @@ -57,8 +58,12 @@ async fn should_pass_model_and_audio() { let server = MockServer::start().await; mock_queue_and_result(&server, &fixture_result(), &[]).await; - let config = FalConfig::new("test-api-key").with_base_url(server.uri()); - let provider = FalProvider::new(config); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_fal(config).unwrap(); let model = provider.transcription("wizper"); model @@ -89,10 +94,13 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = FalConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = FalProvider::new(config); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_fal(config).unwrap(); let model = provider.transcription("wizper"); let mut opts = options(mock_audio(), "audio/wav"); @@ -124,8 +132,12 @@ async fn should_extract_text() { let server = MockServer::start().await; mock_queue_and_result(&server, &fixture_result(), &[]).await; - let config = FalConfig::new("test-api-key").with_base_url(server.uri()); - let provider = FalProvider::new(config); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_fal(config).unwrap(); let model = provider.transcription("wizper"); let result = model @@ -149,8 +161,12 @@ async fn should_include_response_data() { ) .await; - let config = FalConfig::new("test-api-key").with_base_url(server.uri()); - let provider = FalProvider::new(config); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_fal(config).unwrap(); let model = provider.transcription("wizper"); let result = model @@ -167,8 +183,12 @@ async fn should_use_real_date() { let server = MockServer::start().await; mock_queue_and_result(&server, &fixture_result(), &[]).await; - let config = FalConfig::new("test-api-key").with_base_url(server.uri()); - let provider = FalProvider::new(config); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_fal(config).unwrap(); let model = provider.transcription("wizper"); let result = model diff --git a/aimux-providers/tests/fal_video_test.rs b/aimux-providers/tests/fal_video_test.rs index dd7de0c57..081f347a3 100644 --- a/aimux-providers/tests/fal_video_test.rs +++ b/aimux-providers/tests/fal_video_test.rs @@ -2,13 +2,16 @@ //! Source: `reference/ai/packages/fal/src/fal-video-model.test.ts` use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +use aimux_core::download::{Download, DownloadRequest, DownloadedFile}; use aimux_core::video_model::{VideoCallOptions, generate_video}; -use aimux_providers::{FalConfig, FalProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{FalProviderSettings, create_fal}; fn fast_poll() -> Option { Some(aimux_core::video_model::VideoPollOptions { @@ -20,9 +23,27 @@ fn fast_poll() -> Option { fn options(prompt: &str) -> VideoCallOptions { let mut o = VideoCallOptions::new(prompt); o.poll = fast_poll(); + o.download = Some(recording_download(Arc::default())); o } +/// The bytes the test download returns for every video URL. +const VIDEO_BYTES: &[u8] = b"video bytes"; + +/// A download that fetches nothing: it records each URL `generate_video` +/// asks for and returns [`VIDEO_BYTES`]. +fn recording_download(requests: Arc>>) -> Download { + Download::new(move |request: DownloadRequest| { + requests.lock().unwrap().push(request.url); + async { + Ok(DownloadedFile { + data: VIDEO_BYTES.to_vec(), + media_type: None, + }) + } + }) +} + async fn mock_queue_and_result(server: &MockServer, result: &Value) { Mock::given(method("POST")) .and(path("/fal-ai/kling-video")) @@ -42,8 +63,12 @@ async fn should_generate_video() { let result = json!({"video": {"url": "https://cdn.fal.run/video.mp4"}}); mock_queue_and_result(&server, &result).await; - let config = FalConfig::new("test-api-key").with_base_url(server.uri()); - let provider = FalProvider::new(config); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_fal(config).unwrap(); let model = provider.video("fal-ai/kling-video"); let r = generate_video(&model, options("A cat playing")) @@ -58,8 +83,12 @@ async fn should_pass_prompt() { let result = json!({"video": {"url": "https://cdn.fal.run/video.mp4"}}); mock_queue_and_result(&server, &result).await; - let config = FalConfig::new("test-api-key").with_base_url(server.uri()); - let provider = FalProvider::new(config); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_fal(config).unwrap(); let model = provider.video("fal-ai/kling-video"); generate_video(&model, options("A cat playing")) @@ -79,10 +108,13 @@ async fn should_pass_headers() { let mut ph = HashMap::new(); ph.insert("Custom-Header".to_string(), "val".to_string()); - let config = FalConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = FalProvider::new(config); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_fal(config).unwrap(); let model = provider.video("fal-ai/kling-video"); let mut opts = options("test"); @@ -105,8 +137,12 @@ async fn should_include_response_data() { let result = json!({"video": {"url": "https://cdn.fal.run/video.mp4"}}); mock_queue_and_result(&server, &result).await; - let config = FalConfig::new("test-api-key").with_base_url(server.uri()); - let provider = FalProvider::new(config); + let config = FalProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_fal(config).unwrap(); let model = provider.video("fal-ai/kling-video"); let r = generate_video(&model, options("test")).await.unwrap(); diff --git a/aimux-providers/tests/firecrawl_test.rs b/aimux-providers/tests/firecrawl_test.rs index 894458b1b..b890897b9 100644 --- a/aimux-providers/tests/firecrawl_test.rs +++ b/aimux-providers/tests/firecrawl_test.rs @@ -6,22 +6,25 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::provider::Provider; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{FirecrawlConfig, FirecrawlProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{FirecrawlProvider, FirecrawlProviderSettings, create_firecrawl}; fn make_provider(server: &MockServer) -> FirecrawlProvider { - let config = FirecrawlConfig::new("test-api-key").with_base_url(server.uri()); - FirecrawlProvider::new(config) -} - -#[test] -fn provider_name_is_firecrawl() { - let provider = FirecrawlProvider::new(FirecrawlConfig::new("test-key")); - assert_eq!(provider.name(), "firecrawl"); + let config = FirecrawlProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_firecrawl(config).unwrap() } #[test] fn language_model_returns_unsupported() { - let provider = FirecrawlProvider::new(FirecrawlConfig::new("test-key")); + let provider = create_firecrawl(FirecrawlProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + ..Default::default() + }) + .unwrap(); assert!(provider.language_model("any").is_err()); } @@ -69,8 +72,12 @@ async fn uses_bearer_auth_header() { .mount(&server) .await; - let config = FirecrawlConfig::new("my-key").with_base_url(server.uri()); - let provider = FirecrawlProvider::new(config); + let config = FirecrawlProviderSettings { + api_key: Some(Resolvable::Value("my-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_firecrawl(config).unwrap(); let model = provider.search_model(); model .do_search(&SearchCallOptions::new("test")) diff --git a/aimux-providers/tests/fireworks_test.rs b/aimux-providers/tests/fireworks_test.rs index 7261fd73c..d47bcf143 100644 --- a/aimux-providers/tests/fireworks_test.rs +++ b/aimux-providers/tests/fireworks_test.rs @@ -2,7 +2,7 @@ //! //! Translated from `packages/fireworks/src/fireworks-provider.test.ts`. //! -//! Fireworks is a thin OpenAI-compatible wrapper over [`OpenAIProvider`]. The +//! Fireworks is a registry preset of the OpenAI-compatible package. The //! TS suite is mostly provider-configuration unit tests plus a //! `transformRequestBody` that snake-cases `thinking.budgetTokens`, //! `reasoningHistory`, `promptCacheKey`, `serviceTier` and remaps @@ -25,6 +25,7 @@ use serde_json::{Value, json}; use serial_test::serial; +use std::sync::Arc; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -32,7 +33,7 @@ use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; -use aimux_providers::{ProviderOptions, provider, provider_from_env}; +use aimux_providers::{PresetSettings, create_provider}; fn test_prompt() -> LanguageModelPrompt { vec![LanguageModelMessage::user_text("Hello")] @@ -57,28 +58,32 @@ fn text_completion_body() -> Value { }) } -fn make_provider(server: &MockServer) -> Box { - provider( +fn make_provider(server: &MockServer) -> Arc { + create_provider( "fireworks", - Some("test-api-key".to_string()), - "accounts/fireworks/models/llama-v3p1-70b-instruct", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-api-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("accounts/fireworks/models/llama-v3p1-70b-instruct") .expect("fireworks provider should build") } /// TS: `createFireworks()` produces a provider whose name is "fireworks". #[test] fn provider_builds_fireworks_model() { - let model = provider( + let model = create_provider( "fireworks", - Some("test-key".to_string()), - "accounts/fireworks/models/llama-v3p1-70b-instruct", - None, + PresetSettings { + api_key: Some("test-key".to_string().into()), + ..Default::default() + }, ) + .expect("provider should build") + .language_model("accounts/fireworks/models/llama-v3p1-70b-instruct") .expect("fireworks provider should build"); assert_eq!( model.model_id(), @@ -117,15 +122,16 @@ async fn custom_api_key_used_in_auth_header() { .mount(&server) .await; - let model = provider( + let model = create_provider( "fireworks", - Some("my-custom-key".to_string()), - "accounts/fireworks/models/llama-v3p1-70b-instruct", - Some(ProviderOptions { + PresetSettings { + api_key: Some("my-custom-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("accounts/fireworks/models/llama-v3p1-70b-instruct") .expect("fireworks provider should build"); model @@ -146,15 +152,16 @@ async fn custom_headers_forwarded() { .mount(&server) .await; - let model = provider( + let model = create_provider( "fireworks", - Some("test-key".to_string()), - "accounts/fireworks/models/llama-v3p1-70b-instruct", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("accounts/fireworks/models/llama-v3p1-70b-instruct") .expect("fireworks provider should build"); let mut options = default_options(test_prompt()); @@ -180,15 +187,16 @@ async fn language_model_via_trait() { .mount(&server) .await; - let model = provider( + let model = create_provider( "fireworks", - Some("test-key".to_string()), - "accounts/fireworks/models/llama-v3p1-70b-instruct", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("accounts/fireworks/models/llama-v3p1-70b-instruct") .expect("provider should build a model"); model @@ -206,11 +214,9 @@ fn from_env_loads_fireworks_api_key() { std::env::set_var("FIREWORKS_API_KEY", "env-test-key"); } - let model = provider_from_env( - "fireworks", - "accounts/fireworks/models/llama-v3p1-70b-instruct", - None, - ); + let model = create_provider("fireworks", PresetSettings::default()) + .expect("provider should build") + .language_model("accounts/fireworks/models/llama-v3p1-70b-instruct"); assert!(model.is_ok(), "from_env should succeed with env var set"); unsafe { @@ -223,19 +229,19 @@ fn from_env_loads_fireworks_api_key() { /// TS: without the env var, `createFireworks()` fails. #[serial] -#[test] -fn from_env_fails_without_env_var() { +#[tokio::test] +async fn from_env_fails_without_env_var() { let saved = std::env::var("FIREWORKS_API_KEY").ok(); unsafe { std::env::remove_var("FIREWORKS_API_KEY"); } - let model = provider_from_env( - "fireworks", - "accounts/fireworks/models/llama-v3p1-70b-instruct", - None, - ); - assert!(model.is_err(), "from_env should fail without env var"); + let model = create_provider("fireworks", PresetSettings::default()) + .expect("provider should build") + .language_model("accounts/fireworks/models/llama-v3p1-70b-instruct") + .expect("provider should build a model"); + let result = model.do_generate(&default_options(test_prompt())).await; + assert!(result.is_err(), "from_env should fail without env var"); unsafe { if let Some(v) = saved { diff --git a/aimux-providers/tests/fixtures/jev_systemone_live.jsonl b/aimux-providers/tests/fixtures/jev_systemone_live.jsonl index 08c8490fa..d7d7b39a8 100644 --- a/aimux-providers/tests/fixtures/jev_systemone_live.jsonl +++ b/aimux-providers/tests/fixtures/jev_systemone_live.jsonl @@ -1,2 +1,2 @@ -{"schema":2,"call_id":"call-1791293057374663321-0","recorded_at":"2026-10-06T13:24:17.375Z","input":{"operation":"decision","prompt":[],"options":{"state":"I was billed twice and payments are blocked. Please fix this urgently.","questions":[{"type":"boolean","id":"urgent","instructions":"Does the customer request urgent action?"},{"type":"choice","id":"department","instructions":"Choose the responsible department.","options":[{"label":"billing"},{"label":"technical"},{"label":"sales"}]},{"type":"score","id":"severity","instructions":"Rate the impact of the issue.","levels":["Minor inconvenience","Significant issue","Payments completely blocked"]}],"max_retries":0,"timeout":{"total_ms":30000,"step_ms":null,"first_chunk_ms":null,"chunk_ms":null}}},"provider":{"provider":"jev","model_id":"jev-latest","base_url":"https://api.typesafe.ai/v1/systemone","api_key_source":"explicit","profile":{"decision_capabilities":{"rounding":{"probability_decimals":2,"score_decimals":2},"probability_source":"native","supports_boolean":true,"supports_choice":true,"supports_score":true,"returns_distributions":true,"max_questions":null,"min_choices":1,"max_choices":255,"max_score_levels":10}},"provider_options":null},"exchanges":[{"attempt":1,"exchange_index":1,"request":{"method":"POST","url":"https://api.typesafe.ai/v1/systemone","headers":[["Authorization","[REDACTED]"]],"body":"{\"model\":\"jev-latest\",\"state\":\"I was billed twice and payments are blocked. Please fix this urgently.\",\"questions\":{\"urgent\":{\"type\":\"noul\",\"instructions\":\"Does the customer request urgent action?\"},\"department\":{\"type\":\"choice\",\"instructions\":\"Choose the responsible department.\",\"criteria\":{\"billing\":null,\"technical\":null,\"sales\":null}},\"severity\":{\"type\":\"score\",\"instructions\":\"Rate the impact of the issue.\",\"criteria\":[\"Minor inconvenience\",\"Significant issue\",\"Payments completely blocked\"]}}}"},"response":{"status":200,"headers":[["date","Tue, 06 Oct 2026 13:24:17 GMT"],["content-type","application/json"],["content-length","440"],["connection","keep-alive"],["server","cloudflare"],["x-typesafe-request-id","req_01a11162eaba71f5b6825f0466bca790"],["cf-cache-status","DYNAMIC"],["set-cookie","[REDACTED]"],["cf-ray","a464fc4aaee03760-LAX"],["alt-svc","h3=\":443\"; ma=86400"]],"body":"{\"model\":\"jev-1.13.0\",\"answers\":{\"urgent\":{\"type\":\"noul\",\"noul\":0.99},\"department\":{\"type\":\"choice\",\"choice\":\"billing\",\"confidence\":1.0,\"probabilities\":{\"billing\":1.0,\"technical\":0.0,\"sales\":0.0}},\"severity\":{\"type\":\"score\",\"score\":2.0,\"confidence\":1.0,\"legend\":{\"0\":\"Minor inconvenience\",\"1\":\"Significant issue\",\"2\":\"Payments completely blocked\"},\"probabilities\":{\"0\":0.0,\"1\":0.0,\"2\":1.0}}},\"usage\":{\"input_tokens\":365,\"output_tokens\":68}}","stream_chunks":1,"ttfb_ms":486},"timing":{"latency_ms":486,"ttfb_ms":486},"error":null,"finalized":true}],"outcome":{"decision_result":{"rounding":{"probability_decimals":2,"score_decimals":2},"answers":{"department":{"type":"choice","selected":"billing","probabilities":{"billing":1.0,"sales":0.0,"technical":0.0},"confidence":1.0},"severity":{"type":"score","expected_value":2.0,"levels":["Minor inconvenience","Significant issue","Payments completely blocked"],"probabilities":[0.0,0.0,1.0],"confidence":1.0},"urgent":{"type":"boolean","probability_true":0.99}},"provider":"jev","model":"jev-1.13.0","model_version":null,"probability_source":"native","usage":{"input_tokens":{"total":365},"output_tokens":{"total":68},"raw":{"input_tokens":365,"output_tokens":68}},"latency_ms":null,"provider_metadata":{"jev":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.99},"department":{"type":"choice","choice":"billing","confidence":1.0,"probabilities":{"billing":1.0,"technical":0.0,"sales":0.0}},"severity":{"type":"score","score":2.0,"confidence":1.0,"legend":{"0":"Minor inconvenience","1":"Significant issue","2":"Payments completely blocked"},"probabilities":{"0":0.0,"1":0.0,"2":1.0}}},"usage":{"input_tokens":365,"output_tokens":68}}},"response":{"headers":{"connection":"keep-alive","date":"Tue, 06 Oct 2026 13:24:17 GMT","set-cookie":"[REDACTED]","alt-svc":"h3=\":443\"; ma=86400","x-typesafe-request-id":"req_01a11162eaba71f5b6825f0466bca790","cf-cache-status":"DYNAMIC","server":"cloudflare","content-length":"440","content-type":"application/json","cf-ray":"a464fc4aaee03760-LAX"},"body":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.99},"department":{"type":"choice","choice":"billing","confidence":1.0,"probabilities":{"billing":1.0,"technical":0.0,"sales":0.0}},"severity":{"type":"score","score":2.0,"confidence":1.0,"legend":{"0":"Minor inconvenience","1":"Significant issue","2":"Payments completely blocked"},"probabilities":{"0":0.0,"1":0.0,"2":1.0}}},"usage":{"input_tokens":365,"output_tokens":68}}}},"status":"success","finish_reason":null,"error":null,"error_value":null,"usage":{"input_tokens":{"total":365},"output_tokens":{"total":68},"raw":{"input_tokens":365,"output_tokens":68}}},"complete":true,"transport_closed":true} -{"schema":2,"call_id":"call-1791293057868486266-1","recorded_at":"2026-10-06T13:24:17.868Z","input":{"operation":"decision","prompt":[],"options":{"state":{"ticket":"Payments blocked"},"questions":[{"type":"boolean","id":"urgent","instructions":{"question":"Is action urgent?","context":["payments"]},"criteria":{"true":{"includes":["blocked"]},"false":["routine","can wait"]}},{"type":"choice","id":"department","instructions":["Route the ticket",{"priority":"payments"}],"options":[{"label":"billing","description":{"includes":["payments"],"examples":["payout failure"]}},{"label":"technical"}]},{"type":"score","id":"severity","instructions":"Rate severity","levels":[{"label":"low","includes":["minor inconvenience"]},["blocking","No workaround"]]}],"max_retries":0,"timeout":{"total_ms":30000,"step_ms":null,"first_chunk_ms":null,"chunk_ms":null}}},"provider":{"provider":"jev","model_id":"jev-latest","base_url":"https://api.typesafe.ai/v1/systemone","api_key_source":"explicit","profile":{"decision_capabilities":{"rounding":{"probability_decimals":2,"score_decimals":2},"probability_source":"native","supports_boolean":true,"supports_choice":true,"supports_score":true,"returns_distributions":true,"max_questions":null,"min_choices":1,"max_choices":255,"max_score_levels":10}},"provider_options":null},"exchanges":[{"attempt":1,"exchange_index":1,"request":{"method":"POST","url":"https://api.typesafe.ai/v1/systemone","headers":[["Authorization","[REDACTED]"]],"body":"{\"model\":\"jev-latest\",\"state\":{\"ticket\":\"Payments blocked\"},\"questions\":{\"urgent\":{\"type\":\"noul\",\"instructions\":{\"question\":\"Is action urgent?\",\"context\":[\"payments\"]},\"criteria\":{\"true\":{\"includes\":[\"blocked\"]},\"false\":[\"routine\",\"can wait\"]}},\"department\":{\"type\":\"choice\",\"instructions\":[\"Route the ticket\",{\"priority\":\"payments\"}],\"criteria\":{\"billing\":{\"includes\":[\"payments\"],\"examples\":[\"payout failure\"]},\"technical\":null}},\"severity\":{\"type\":\"score\",\"instructions\":\"Rate severity\",\"criteria\":[{\"label\":\"low\",\"includes\":[\"minor inconvenience\"]},[\"blocking\",\"No workaround\"]]}}}"},"response":{"status":200,"headers":[["date","Tue, 06 Oct 2026 13:24:18 GMT"],["content-type","application/json"],["content-length","431"],["connection","keep-alive"],["server","cloudflare"],["set-cookie","[REDACTED]"],["x-typesafe-request-id","req_01a11162ebc076ee849d8e67ae47f547"],["cf-cache-status","DYNAMIC"],["cf-ray","a464fc4c4b1c3760-LAX"],["alt-svc","h3=\":443\"; ma=86400"]],"body":"{\"model\":\"jev-1.13.0\",\"answers\":{\"urgent\":{\"type\":\"noul\",\"noul\":0.92},\"department\":{\"type\":\"choice\",\"choice\":\"billing\",\"confidence\":0.98,\"probabilities\":{\"technical\":0.01,\"billing\":0.99}},\"severity\":{\"type\":\"score\",\"score\":0.99,\"confidence\":0.98,\"legend\":{\"0\":{\"label\":\"low\",\"includes\":[\"minor inconvenience\"]},\"1\":[\"blocking\",\"No workaround\"]},\"probabilities\":{\"0\":0.01,\"1\":0.99}}},\"usage\":{\"input_tokens\":448,\"output_tokens\":61}}","stream_chunks":1,"ttfb_ms":257},"timing":{"latency_ms":257,"ttfb_ms":257},"error":null,"finalized":true}],"outcome":{"decision_result":{"rounding":{"probability_decimals":2,"score_decimals":2},"answers":{"department":{"type":"choice","selected":"billing","probabilities":{"billing":0.99,"technical":0.01},"confidence":0.98},"severity":{"type":"score","expected_value":0.99,"levels":[{"label":"low","includes":["minor inconvenience"]},["blocking","No workaround"]],"probabilities":[0.01,0.99],"confidence":0.98},"urgent":{"type":"boolean","probability_true":0.92}},"provider":"jev","model":"jev-1.13.0","model_version":null,"probability_source":"native","usage":{"input_tokens":{"total":448},"output_tokens":{"total":61},"raw":{"input_tokens":448,"output_tokens":61}},"latency_ms":null,"provider_metadata":{"jev":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.92},"department":{"type":"choice","choice":"billing","confidence":0.98,"probabilities":{"technical":0.01,"billing":0.99}},"severity":{"type":"score","score":0.99,"confidence":0.98,"legend":{"0":{"label":"low","includes":["minor inconvenience"]},"1":["blocking","No workaround"]},"probabilities":{"0":0.01,"1":0.99}}},"usage":{"input_tokens":448,"output_tokens":61}}},"response":{"headers":{"date":"Tue, 06 Oct 2026 13:24:18 GMT","content-length":"431","set-cookie":"[REDACTED]","connection":"keep-alive","server":"cloudflare","cf-ray":"a464fc4c4b1c3760-LAX","alt-svc":"h3=\":443\"; ma=86400","x-typesafe-request-id":"req_01a11162ebc076ee849d8e67ae47f547","cf-cache-status":"DYNAMIC","content-type":"application/json"},"body":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.92},"department":{"type":"choice","choice":"billing","confidence":0.98,"probabilities":{"technical":0.01,"billing":0.99}},"severity":{"type":"score","score":0.99,"confidence":0.98,"legend":{"0":{"label":"low","includes":["minor inconvenience"]},"1":["blocking","No workaround"]},"probabilities":{"0":0.01,"1":0.99}}},"usage":{"input_tokens":448,"output_tokens":61}}}},"status":"success","finish_reason":null,"error":null,"error_value":null,"usage":{"input_tokens":{"total":448},"output_tokens":{"total":61},"raw":{"input_tokens":448,"output_tokens":61}}},"complete":true,"transport_closed":true} +{"schema":2,"call_id":"call-1791293057374663321-0","recorded_at":"2026-10-06T13:24:17.375Z","input":{"operation":"decision","prompt":[],"options":{"state":"I was billed twice and payments are blocked. Please fix this urgently.","questions":[{"type":"boolean","id":"urgent","instructions":"Does the customer request urgent action?"},{"type":"choice","id":"department","instructions":"Choose the responsible department.","options":[{"label":"billing"},{"label":"technical"},{"label":"sales"}]},{"type":"score","id":"severity","instructions":"Rate the impact of the issue.","levels":["Minor inconvenience","Significant issue","Payments completely blocked"]}],"max_retries":0,"timeout":{"total_ms":30000,"step_ms":null,"first_chunk_ms":null,"chunk_ms":null}},"decision_capabilities":{"rounding":{"probability_decimals":2,"score_decimals":2},"probability_source":"native","supports_boolean":true,"supports_choice":true,"supports_score":true,"returns_distributions":true,"max_questions":null,"min_choices":1,"max_choices":255,"max_score_levels":10}},"provider":{"provider":"jev","model_id":"jev-latest","provider_id":"jev"},"exchanges":[{"attempt":1,"exchange_index":1,"request":{"method":"POST","url":"https://api.typesafe.ai/v1/systemone","headers":[["Authorization","[REDACTED]"]],"body":"{\"model\":\"jev-latest\",\"state\":\"I was billed twice and payments are blocked. Please fix this urgently.\",\"questions\":{\"urgent\":{\"type\":\"noul\",\"instructions\":\"Does the customer request urgent action?\"},\"department\":{\"type\":\"choice\",\"instructions\":\"Choose the responsible department.\",\"criteria\":{\"billing\":null,\"technical\":null,\"sales\":null}},\"severity\":{\"type\":\"score\",\"instructions\":\"Rate the impact of the issue.\",\"criteria\":[\"Minor inconvenience\",\"Significant issue\",\"Payments completely blocked\"]}}}"},"response":{"status":200,"headers":[["date","Tue, 06 Oct 2026 13:24:17 GMT"],["content-type","application/json"],["content-length","440"],["connection","keep-alive"],["server","cloudflare"],["x-typesafe-request-id","req_01a11162eaba71f5b6825f0466bca790"],["cf-cache-status","DYNAMIC"],["set-cookie","[REDACTED]"],["cf-ray","a464fc4aaee03760-LAX"],["alt-svc","h3=\":443\"; ma=86400"]],"body":"{\"model\":\"jev-1.13.0\",\"answers\":{\"urgent\":{\"type\":\"noul\",\"noul\":0.99},\"department\":{\"type\":\"choice\",\"choice\":\"billing\",\"confidence\":1.0,\"probabilities\":{\"billing\":1.0,\"technical\":0.0,\"sales\":0.0}},\"severity\":{\"type\":\"score\",\"score\":2.0,\"confidence\":1.0,\"legend\":{\"0\":\"Minor inconvenience\",\"1\":\"Significant issue\",\"2\":\"Payments completely blocked\"},\"probabilities\":{\"0\":0.0,\"1\":0.0,\"2\":1.0}}},\"usage\":{\"input_tokens\":365,\"output_tokens\":68}}","stream_chunks":1,"ttfb_ms":486},"timing":{"latency_ms":486,"ttfb_ms":486},"error":null,"finalized":true}],"outcome":{"decision_result":{"rounding":{"probability_decimals":2,"score_decimals":2},"answers":{"department":{"type":"choice","selected":"billing","probabilities":{"billing":1.0,"sales":0.0,"technical":0.0},"confidence":1.0},"severity":{"type":"score","expected_value":2.0,"levels":["Minor inconvenience","Significant issue","Payments completely blocked"],"probabilities":[0.0,0.0,1.0],"confidence":1.0},"urgent":{"type":"boolean","probability_true":0.99}},"provider":"jev","model":"jev-1.13.0","model_version":null,"probability_source":"native","usage":{"input_tokens":{"total":365},"output_tokens":{"total":68},"raw":{"input_tokens":365,"output_tokens":68}},"latency_ms":null,"provider_metadata":{"jev":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.99},"department":{"type":"choice","choice":"billing","confidence":1.0,"probabilities":{"billing":1.0,"technical":0.0,"sales":0.0}},"severity":{"type":"score","score":2.0,"confidence":1.0,"legend":{"0":"Minor inconvenience","1":"Significant issue","2":"Payments completely blocked"},"probabilities":{"0":0.0,"1":0.0,"2":1.0}}},"usage":{"input_tokens":365,"output_tokens":68}}},"response":{"headers":{"connection":"keep-alive","date":"Tue, 06 Oct 2026 13:24:17 GMT","set-cookie":"[REDACTED]","alt-svc":"h3=\":443\"; ma=86400","x-typesafe-request-id":"req_01a11162eaba71f5b6825f0466bca790","cf-cache-status":"DYNAMIC","server":"cloudflare","content-length":"440","content-type":"application/json","cf-ray":"a464fc4aaee03760-LAX"},"body":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.99},"department":{"type":"choice","choice":"billing","confidence":1.0,"probabilities":{"billing":1.0,"technical":0.0,"sales":0.0}},"severity":{"type":"score","score":2.0,"confidence":1.0,"legend":{"0":"Minor inconvenience","1":"Significant issue","2":"Payments completely blocked"},"probabilities":{"0":0.0,"1":0.0,"2":1.0}}},"usage":{"input_tokens":365,"output_tokens":68}}}},"status":"success","finish_reason":null,"error":null,"error_value":null,"usage":{"input_tokens":{"total":365},"output_tokens":{"total":68},"raw":{"input_tokens":365,"output_tokens":68}}},"complete":true,"transport_closed":true} +{"schema":2,"call_id":"call-1791293057868486266-1","recorded_at":"2026-10-06T13:24:17.868Z","input":{"operation":"decision","prompt":[],"options":{"state":{"ticket":"Payments blocked"},"questions":[{"type":"boolean","id":"urgent","instructions":{"question":"Is action urgent?","context":["payments"]},"criteria":{"true":{"includes":["blocked"]},"false":["routine","can wait"]}},{"type":"choice","id":"department","instructions":["Route the ticket",{"priority":"payments"}],"options":[{"label":"billing","description":{"includes":["payments"],"examples":["payout failure"]}},{"label":"technical"}]},{"type":"score","id":"severity","instructions":"Rate severity","levels":[{"label":"low","includes":["minor inconvenience"]},["blocking","No workaround"]]}],"max_retries":0,"timeout":{"total_ms":30000,"step_ms":null,"first_chunk_ms":null,"chunk_ms":null}},"decision_capabilities":{"rounding":{"probability_decimals":2,"score_decimals":2},"probability_source":"native","supports_boolean":true,"supports_choice":true,"supports_score":true,"returns_distributions":true,"max_questions":null,"min_choices":1,"max_choices":255,"max_score_levels":10}},"provider":{"provider":"jev","model_id":"jev-latest","provider_id":"jev"},"exchanges":[{"attempt":1,"exchange_index":1,"request":{"method":"POST","url":"https://api.typesafe.ai/v1/systemone","headers":[["Authorization","[REDACTED]"]],"body":"{\"model\":\"jev-latest\",\"state\":{\"ticket\":\"Payments blocked\"},\"questions\":{\"urgent\":{\"type\":\"noul\",\"instructions\":{\"question\":\"Is action urgent?\",\"context\":[\"payments\"]},\"criteria\":{\"true\":{\"includes\":[\"blocked\"]},\"false\":[\"routine\",\"can wait\"]}},\"department\":{\"type\":\"choice\",\"instructions\":[\"Route the ticket\",{\"priority\":\"payments\"}],\"criteria\":{\"billing\":{\"includes\":[\"payments\"],\"examples\":[\"payout failure\"]},\"technical\":null}},\"severity\":{\"type\":\"score\",\"instructions\":\"Rate severity\",\"criteria\":[{\"label\":\"low\",\"includes\":[\"minor inconvenience\"]},[\"blocking\",\"No workaround\"]]}}}"},"response":{"status":200,"headers":[["date","Tue, 06 Oct 2026 13:24:18 GMT"],["content-type","application/json"],["content-length","431"],["connection","keep-alive"],["server","cloudflare"],["set-cookie","[REDACTED]"],["x-typesafe-request-id","req_01a11162ebc076ee849d8e67ae47f547"],["cf-cache-status","DYNAMIC"],["cf-ray","a464fc4c4b1c3760-LAX"],["alt-svc","h3=\":443\"; ma=86400"]],"body":"{\"model\":\"jev-1.13.0\",\"answers\":{\"urgent\":{\"type\":\"noul\",\"noul\":0.92},\"department\":{\"type\":\"choice\",\"choice\":\"billing\",\"confidence\":0.98,\"probabilities\":{\"technical\":0.01,\"billing\":0.99}},\"severity\":{\"type\":\"score\",\"score\":0.99,\"confidence\":0.98,\"legend\":{\"0\":{\"label\":\"low\",\"includes\":[\"minor inconvenience\"]},\"1\":[\"blocking\",\"No workaround\"]},\"probabilities\":{\"0\":0.01,\"1\":0.99}}},\"usage\":{\"input_tokens\":448,\"output_tokens\":61}}","stream_chunks":1,"ttfb_ms":257},"timing":{"latency_ms":257,"ttfb_ms":257},"error":null,"finalized":true}],"outcome":{"decision_result":{"rounding":{"probability_decimals":2,"score_decimals":2},"answers":{"department":{"type":"choice","selected":"billing","probabilities":{"billing":0.99,"technical":0.01},"confidence":0.98},"severity":{"type":"score","expected_value":0.99,"levels":[{"label":"low","includes":["minor inconvenience"]},["blocking","No workaround"]],"probabilities":[0.01,0.99],"confidence":0.98},"urgent":{"type":"boolean","probability_true":0.92}},"provider":"jev","model":"jev-1.13.0","model_version":null,"probability_source":"native","usage":{"input_tokens":{"total":448},"output_tokens":{"total":61},"raw":{"input_tokens":448,"output_tokens":61}},"latency_ms":null,"provider_metadata":{"jev":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.92},"department":{"type":"choice","choice":"billing","confidence":0.98,"probabilities":{"technical":0.01,"billing":0.99}},"severity":{"type":"score","score":0.99,"confidence":0.98,"legend":{"0":{"label":"low","includes":["minor inconvenience"]},"1":["blocking","No workaround"]},"probabilities":{"0":0.01,"1":0.99}}},"usage":{"input_tokens":448,"output_tokens":61}}},"response":{"headers":{"date":"Tue, 06 Oct 2026 13:24:18 GMT","content-length":"431","set-cookie":"[REDACTED]","connection":"keep-alive","server":"cloudflare","cf-ray":"a464fc4c4b1c3760-LAX","alt-svc":"h3=\":443\"; ma=86400","x-typesafe-request-id":"req_01a11162ebc076ee849d8e67ae47f547","cf-cache-status":"DYNAMIC","content-type":"application/json"},"body":{"model":"jev-1.13.0","answers":{"urgent":{"type":"noul","noul":0.92},"department":{"type":"choice","choice":"billing","confidence":0.98,"probabilities":{"technical":0.01,"billing":0.99}},"severity":{"type":"score","score":0.99,"confidence":0.98,"legend":{"0":{"label":"low","includes":["minor inconvenience"]},"1":["blocking","No workaround"]},"probabilities":{"0":0.01,"1":0.99}}},"usage":{"input_tokens":448,"output_tokens":61}}}},"status":"success","finish_reason":null,"error":null,"error_value":null,"usage":{"input_tokens":{"total":448},"output_tokens":{"total":61},"raw":{"input_tokens":448,"output_tokens":61}}},"complete":true,"transport_closed":true} diff --git a/aimux-providers/tests/fixtures/jev_systemone_live.md b/aimux-providers/tests/fixtures/jev_systemone_live.md index 463030716..229632cfe 100644 --- a/aimux-providers/tests/fixtures/jev_systemone_live.md +++ b/aimux-providers/tests/fixtures/jev_systemone_live.md @@ -11,6 +11,8 @@ description and structured Score levels. Each exercises Boolean, Choice and Score. Inputs are synthetic test data. Authorization is automatically redacted. These replace the earlier manually collected JSON fixture. Standard Recording -schema 2 includes the actual HTTP exchange, timing, capability snapshot and +The recording includes the actual HTTP exchange, timing, capability snapshot and normalized result. Offline provider tests replay the wire exchange via the shared replay helper; Core and the CLI replay the normalized result. + +The V4 layers rewrite the provider section to the identity they record (`provider_id`, `provider`, `model_id`) and move the decision capabilities under `input`; the envelope stays schema 2. Recorded requests, responses and normalized outcomes are unchanged. The file as master recorded it is kept as `aimux-core/tests/fixtures/master_decision_recordings.jsonl`. diff --git a/aimux-providers/tests/fixtures/openai_realtime_transcription.json b/aimux-providers/tests/fixtures/openai_realtime_transcription.json new file mode 100644 index 000000000..cb724280f --- /dev/null +++ b/aimux-providers/tests/fixtures/openai_realtime_transcription.json @@ -0,0 +1,10 @@ +{ + "description": "Server events for one gpt-realtime-whisper transcription session (OpenAI realtime API wire shape, RFC-0028 §3.2), replayed after the client commits its audio buffer.", + "events": [ + {"type": "session.created", "session": {"id": "sess_1", "type": "transcription"}}, + {"type": "input_audio_buffer.committed", "item_id": "item_1"}, + {"type": "conversation.item.input_audio_transcription.delta", "item_id": "item_1", "delta": "Hello"}, + {"type": "conversation.item.input_audio_transcription.delta", "item_id": "item_1", "delta": " world"}, + {"type": "conversation.item.input_audio_transcription.completed", "item_id": "item_1", "transcript": "Hello world"} + ] +} diff --git a/aimux-providers/tests/gladia_transcription_test.rs b/aimux-providers/tests/gladia_transcription_test.rs index ab4ee57ea..dfc75706a 100644 --- a/aimux-providers/tests/gladia_transcription_test.rs +++ b/aimux-providers/tests/gladia_transcription_test.rs @@ -8,7 +8,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; -use aimux_providers::{GladiaConfig, GladiaProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{GladiaProviderSettings, create_gladia}; fn mock_audio() -> Vec { vec![1u8, 2, 3] @@ -70,8 +71,12 @@ async fn should_extract_text() { let server = MockServer::start().await; mock_all(&server, &fixture_result(), &[]).await; - let config = GladiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GladiaProvider::new(config); + let config = GladiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_gladia(config).unwrap(); let model = provider.transcription("default"); let result = model @@ -95,10 +100,13 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = GladiaConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = GladiaProvider::new(config); + let config = GladiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_gladia(config).unwrap(); let model = provider.transcription("default"); let mut opts = options(mock_audio(), "audio/wav"); @@ -129,8 +137,12 @@ async fn should_include_response_data() { let server = MockServer::start().await; mock_all(&server, &fixture_result(), &[("x-request-id", "test-req")]).await; - let config = GladiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GladiaProvider::new(config); + let config = GladiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_gladia(config).unwrap(); let model = provider.transcription("default"); let result = model @@ -147,8 +159,12 @@ async fn should_use_real_date() { let server = MockServer::start().await; mock_all(&server, &fixture_result(), &[]).await; - let config = GladiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GladiaProvider::new(config); + let config = GladiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_gladia(config).unwrap(); let model = provider.transcription("default"); let result = model @@ -165,8 +181,12 @@ async fn should_include_provider_metadata() { let server = MockServer::start().await; mock_all(&server, &fixture_result(), &[]).await; - let config = GladiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GladiaProvider::new(config); + let config = GladiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_gladia(config).unwrap(); let model = provider.transcription("default"); let result = model diff --git a/aimux-providers/tests/google_embedding_test.rs b/aimux-providers/tests/google_embedding_test.rs index a14328afc..d2b56c574 100644 --- a/aimux-providers/tests/google_embedding_test.rs +++ b/aimux-providers/tests/google_embedding_test.rs @@ -10,8 +10,9 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; -use aimux_providers::{GoogleConfig, GoogleProvider}; +use aimux_core::embedding_model::EmbeddingCallOptions; +use aimux_core::provider::Provider; +use aimux_providers::{GoogleProviderSettings, create_google}; const TEST_VALUES: &[&str] = &["sunny day at the beach", "rainy day in the city"]; @@ -57,9 +58,13 @@ async fn should_extract_embedding() { .mount(&server) .await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.embedding_model("gemini-embedding-001"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("gemini-embedding-001").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -82,9 +87,13 @@ async fn should_expose_raw_response() { .mount(&server) .await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.embedding_model("gemini-embedding-001"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("gemini-embedding-001").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -111,9 +120,13 @@ async fn should_pass_model_and_values() { .mount(&server) .await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.embedding_model("gemini-embedding-001"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("gemini-embedding-001").unwrap(); let _ = model .do_embed(&default_options(test_values())) @@ -146,9 +159,13 @@ async fn should_pass_headers() { .mount(&server) .await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.embedding_model("gemini-embedding-001"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("gemini-embedding-001").unwrap(); let mut request_headers = HashMap::new(); request_headers.insert( @@ -190,9 +207,13 @@ async fn should_use_batch_endpoint() { .mount(&server) .await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.embedding_model("gemini-embedding-001"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("gemini-embedding-001").unwrap(); let _ = model .do_embed(&default_options(test_values())) @@ -213,9 +234,13 @@ async fn should_use_single_endpoint() { .mount(&server) .await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.embedding_model("gemini-embedding-001"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("gemini-embedding-001").unwrap(); let result = model .do_embed(&default_options(vec![TEST_VALUES[0].to_string()])) @@ -232,9 +257,12 @@ async fn should_use_single_endpoint() { /// TS: "should expose the Google batch embedding API limit" #[test] fn should_expose_max_embeddings_per_call() { - let config = GoogleConfig::new("test-api-key"); - let provider = GoogleProvider::new(config); - let model = provider.embedding_model("gemini-embedding-001"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("gemini-embedding-001").unwrap(); assert_eq!(model.max_embeddings_per_call(), Some(100)); } diff --git a/aimux-providers/tests/google_factory_test.rs b/aimux-providers/tests/google_factory_test.rs new file mode 100644 index 000000000..3cc6feee8 --- /dev/null +++ b/aimux-providers/tests/google_factory_test.rs @@ -0,0 +1,459 @@ +//! `create_google` against the recorded AI SDK behavior. +//! +//! `fixtures/aisdk/google/*.json` hold what `@ai-sdk/google`'s `createGoogle` +//! sent over a mocked fetch for a fixed input: URL, headers, body and the +//! model's `provider` string. Each fixture is replayed here through an +//! injected [`Fetch`] (the factory's `fetch` setting), so the test also proves +//! the transport is the one the settings name. The expected request comes from +//! the fixture; the inputs are rebuilt from the fixture's recorded `sdk.input`. +//! +//! The recorded `x-goog-api-key` is redacted, so the key the case used is +//! substituted back. The recorded SDK and runtime user-agent identifiers are +//! replaced with the pinned provider package's user-agent suffix. + +#[path = "common/mock_fetch.rs"] +mod mock_fetch; + +use aimux_core::tool::RawToolCall; +use std::collections::BTreeMap; +use std::path::PathBuf; +use std::sync::Arc; + +use futures::StreamExt; +use serde_json::{Value, json}; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; +use aimux_core::options::CallOptions; +use aimux_core::result::{GenerateContent, GenerateResult}; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{FunctionTool, LanguageModelV4ToolChoice, Tool}; +use aimux_core::types::FinishReasonUnified; +use aimux_providers::google::{GoogleProviderSettings, create_google}; + +use mock_fetch::{Canned, MockFetch, Seen}; + +const KEY: &str = "sk-test-fixture"; + +fn settings(mock: &Arc) -> GoogleProviderSettings { + GoogleProviderSettings { + api_key: Some(KEY.to_string()), + fetch: Some(mock.transport()), + ..Default::default() + } +} + +// ── fixtures ───────────────────────────────────────────────────────────────── + +/// Every fixture under `fixtures/aisdk/google/`; the directory listing is +/// compared to this so a new fixture cannot go unreplayed. +const FIXTURES: [&str; 5] = [ + "embedding-basic", + "generate-basic", + "generate-provider-options", + "generate-stream-basic", + "generate-tools", +]; + +fn fixtures_dir() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../fixtures/aisdk/google") +} + +struct Fixture { + name: &'static str, + json: Value, +} + +impl Fixture { + fn load(name: &'static str) -> Self { + assert!(FIXTURES.contains(&name), "unlisted fixture {name}"); + let path = fixtures_dir().join(format!("{name}.json")); + let text = std::fs::read_to_string(&path) + .unwrap_or_else(|e| panic!("cannot read {}: {e}", path.display())); + Self { + name, + json: serde_json::from_str(&text).expect("fixture is JSON"), + } + } + + fn provider_string(&self) -> &str { + self.json["model"]["provider"].as_str().unwrap() + } + + fn model_id(&self) -> &str { + self.json["model"]["modelId"].as_str().unwrap() + } + + fn sdk_input(&self) -> &Value { + &self.json["sdk"]["input"] + } + + /// The transport answering with the recorded response. + fn mock(&self) -> Arc { + let response = &self.json["response"]; + let body = match &response["body"] { + Value::String(text) => text.clone().into_bytes(), + other => serde_json::to_vec(other).unwrap(), + }; + MockFetch::new(vec![Canned { + status: response["status"].as_u64().unwrap() as u16, + headers: response["headers"] + .as_object() + .unwrap() + .iter() + .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string())) + .collect(), + body, + }]) + } + + /// The recorded request must equal what went out: method, URL, headers + /// (names case-insensitively, with the pinned user-agent) and body. + fn assert_request(&self, seen: &Seen) { + let recorded = &self.json["request"]; + assert_eq!( + seen.method, + recorded["method"].as_str().unwrap(), + "{}", + self.name + ); + assert_eq!(seen.url, recorded["url"].as_str().unwrap(), "{}", self.name); + + let mut expected: BTreeMap = recorded["headers"] + .as_object() + .unwrap() + .iter() + .map(|(k, v)| (k.to_ascii_lowercase(), v.as_str().unwrap().to_string())) + .collect(); + expected.insert( + "user-agent".to_string(), + "ai-sdk-google/4.0.85 ai-sdk-provider-utils/5.0.51".to_string(), + ); + assert_eq!( + expected.get("x-goog-api-key").map(String::as_str), + Some(""), + "{}: the recording redacts the key", + self.name + ); + expected.insert("x-goog-api-key".to_string(), KEY.to_string()); + assert_eq!(seen.headers, expected, "{}: headers", self.name); + + assert_eq!(seen.json_body(), recorded["body"], "{}: body", self.name); + } + + /// The call options the recording was made with. + fn call_options(&self) -> CallOptions { + call_options_from(self.sdk_input()) + } +} + +/// Rebuild aimux call options from a recorded `sdk.input` (the subset the +/// Google fixtures use). +fn call_options_from(input: &Value) -> CallOptions { + let mut prompt: LanguageModelPrompt = Vec::new(); + if let Some(system) = input.get("system").and_then(Value::as_str) { + prompt.push(LanguageModelMessage::System { + content: system.to_string(), + provider_options: None, + }); + } + if let Some(text) = input.get("prompt").and_then(Value::as_str) { + prompt.push(LanguageModelMessage::user_text(text)); + } + let mut options = CallOptions::new(prompt); + if let Some(tools) = input.get("tools").and_then(Value::as_object) { + options.tools = Some( + tools + .iter() + .map(|(name, tool)| { + let mut function = FunctionTool::new(name.clone(), tool["inputSchema"].clone()); + if let Some(description) = tool["description"].as_str() { + function = function.with_description(description); + } + Tool::Function(function) + }) + .collect(), + ); + } + options.tool_choice = match input.get("toolChoice").and_then(Value::as_str) { + Some("required") => Some(LanguageModelV4ToolChoice::Required), + Some("none") => Some(LanguageModelV4ToolChoice::None), + Some("auto") => Some(LanguageModelV4ToolChoice::Auto), + None => options + .tools + .as_ref() + .filter(|tools| !tools.is_empty()) + .map(|_| LanguageModelV4ToolChoice::Auto), + Some(other) => panic!("unmapped toolChoice {other}"), + }; + options.provider_options = input + .get("providerOptions") + .map(|value| serde_json::from_value(value.clone()).unwrap()); + options +} + +fn generated_text(result: &GenerateResult) -> String { + result + .content + .iter() + .filter_map(|part| match part { + GenerateContent::Text { text, .. } => Some(text.as_str()), + _ => None, + }) + .collect() +} + +fn unified_finish(recorded: &str) -> FinishReasonUnified { + match recorded { + "stop" => FinishReasonUnified::Stop, + "tool-calls" => FinishReasonUnified::ToolCalls, + other => panic!("unmapped finish reason {other}"), + } +} + +/// Result fields every recorded `generateText` case pins. +fn assert_generate_result(fixture: &Fixture, result: &GenerateResult) { + let recorded = &fixture.json["result"]; + assert_eq!( + generated_text(result), + recorded["text"].as_str().unwrap(), + "{}: text", + fixture.name + ); + assert_eq!( + result.response.as_ref().and_then(|r| r.id.as_deref()), + recorded["response"]["id"].as_str(), + "{}: 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 _), + "{}: input tokens", + fixture.name + ); + assert_eq!( + result.usage.output_tokens.total, + recorded["usage"]["outputTokens"].as_u64().map(|v| v as _), + "{}: output tokens", + fixture.name + ); + assert_eq!( + result.finish_reason.unified, + unified_finish(recorded["finishReason"].as_str().unwrap()), + "{}: finish", + fixture.name + ); + assert_eq!( + result.finish_reason.raw.as_deref(), + recorded["rawFinishReason"].as_str(), + "{}: raw finish", + fixture.name + ); + // The usage metadata comes back under the Gemini namespace. + assert_eq!( + result.provider_metadata.as_ref().unwrap()["google"]["usageMetadata"], + recorded["providerMetadata"]["google"]["usageMetadata"], + "{}: usage metadata", + fixture.name + ); +} + +// ── the fixtures ───────────────────────────────────────────────────────────── + +#[test] +fn every_fixture_is_replayed() { + let mut on_disk: Vec = std::fs::read_dir(fixtures_dir()) + .expect("fixtures/aisdk/google exists") + .map(|entry| { + let name = entry.unwrap().file_name().into_string().unwrap(); + name.strip_suffix(".json") + .unwrap_or_else(|| panic!("unexpected file {name}")) + .to_string() + }) + .collect(); + on_disk.sort(); + assert_eq!( + on_disk, FIXTURES, + "add a test for the new fixture and list it" + ); +} + +/// Replay a fixture whose request is a plain `generateText`. +async fn replay_generate(name: &'static str) -> (Fixture, Arc, GenerateResult) { + let fixture = Fixture::load(name); + let mock = fixture.mock(); + let provider = create_google(settings(&mock)).unwrap(); + let model = provider.chat(fixture.model_id()); + assert_eq!( + model.provider(), + fixture.provider_string(), + "{name}: provider()" + ); + assert_eq!(model.model_id(), fixture.model_id()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1, "{name}: exactly one request"); + fixture.assert_request(&seen[0]); + assert_generate_result(&fixture, &result); + (fixture, mock, result) +} + +#[tokio::test] +async fn generate_basic() { + replay_generate("generate-basic").await; +} + +#[tokio::test] +async fn generate_provider_options_are_read_from_the_google_namespace() { + let (fixture, _, _) = replay_generate("generate-provider-options").await; + // The recorded body carries both option groups the case set. + let body = &fixture.json["request"]["body"]; + assert_eq!( + body["generationConfig"]["thinkingConfig"]["thinkingBudget"], + json!(1024) + ); + assert_eq!( + body["safetySettings"][0]["threshold"], + json!("BLOCK_ONLY_HIGH") + ); +} + +#[tokio::test] +async fn generate_tools() { + let (fixture, _, result) = replay_generate("generate-tools").await; + let recorded = &fixture.json["result"]["content"][0]; + let call = result + .content + .iter() + .find_map(|part| match part { + GenerateContent::ToolCall(RawToolCall { + tool_name, input, .. + }) => Some((tool_name.clone(), input.clone())), + _ => None, + }) + .expect("a tool call"); + assert_eq!(call.0, recorded["toolName"].as_str().unwrap()); + assert_eq!( + serde_json::from_str::(&call.1).unwrap(), + recorded["input"] + ); +} + +#[tokio::test] +async fn generate_stream_basic() { + let fixture = Fixture::load("generate-stream-basic"); + let mock = fixture.mock(); + let model = create_google(settings(&mock)) + .unwrap() + .chat(fixture.model_id()); + assert_eq!(model.provider(), fixture.provider_string()); + + let result = model.do_stream(&fixture.call_options()).await.unwrap(); + let parts: Vec = result + .stream + .map(|part| part.expect("no stream error")) + .collect() + .await; + + // The recording's user-agent names the provider package; the stream case + // is checked like the others. + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + fixture.assert_request(&seen[0]); + + let text: String = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextDelta { delta, .. } => Some(delta.as_str()), + _ => None, + }) + .collect(); + assert_eq!(text, "Hello world"); + let recorded = &fixture.json["result"]; + let finish = parts.iter().find_map(|part| match part { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => Some((finish_reason, usage, provider_metadata)), + _ => None, + }); + let (reason, usage, metadata) = finish.expect("a finish part"); + assert_eq!( + reason.unified, + unified_finish(recorded["finishReason"].as_str().unwrap()) + ); + assert_eq!(reason.raw.as_deref(), Some("STOP")); + assert_eq!( + usage.input_tokens.total, + recorded["usage"]["inputTokens"].as_u64().map(|v| v as _) + ); + assert_eq!( + usage.output_tokens.total, + recorded["usage"]["outputTokens"].as_u64().map(|v| v as _) + ); + assert_eq!( + metadata.as_ref().unwrap()["google"]["usageMetadata"], + recorded["parts"] + .as_array() + .unwrap() + .iter() + .find(|part| part["type"] == "finish-step") + .unwrap()["providerMetadata"]["google"]["usageMetadata"] + ); +} + +#[tokio::test] +async fn embedding_basic() { + let fixture = Fixture::load("embedding-basic"); + let mock = fixture.mock(); + let model = create_google(settings(&mock)) + .unwrap() + .embedding(fixture.model_id()); + assert_eq!(model.provider(), fixture.provider_string()); + + let values: Vec = fixture.sdk_input()["values"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap().to_string()) + .collect(); + let result = model + .do_embed(&EmbeddingCallOptions { + values, + abort_signal: None, + max_retries: None, + timeout: None, + provider_options: None, + headers: None, + }) + .await + .unwrap(); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + fixture.assert_request(&seen[0]); + let recorded: Vec> = fixture.json["result"]["embeddings"] + .as_array() + .unwrap() + .iter() + .map(|row| { + row.as_array() + .unwrap() + .iter() + .map(|v| v.as_f64().unwrap() as f32) + .collect() + }) + .collect(); + assert_eq!(result.embeddings, recorded); +} diff --git a/aimux-providers/tests/google_files_test.rs b/aimux-providers/tests/google_files_test.rs index e18b5e667..77008b61c 100644 --- a/aimux-providers/tests/google_files_test.rs +++ b/aimux-providers/tests/google_files_test.rs @@ -1,7 +1,7 @@ //! Rust translations of the Google Files provider tests. //! //! Source: `reference/ai/packages/google/src/google-files.test.ts` -//! (21 test cases). +//! (16 test cases). //! //! Google uses a resumable upload protocol with three phases: //! 1. POST `/upload/v1beta/files` (init) - returns `x-goog-upload-url` header. @@ -12,10 +12,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::files_model::{Files, UploadFileCallOptions, UploadFileData}; -use aimux_core::retry::RetryConfig; -use aimux_core::shared::{FileBytes, provider_namespace}; -use aimux_providers::{GoogleConfig, GoogleProvider}; +use aimux_core::files_model::{UploadFileCallOptions, UploadFileData}; +use aimux_core::provider::Provider; +use aimux_core::shared::FileBytes; +use aimux_providers::{GoogleProvider, GoogleProviderSettings, create_google}; // -- helpers ----------------------------------------------------------------- @@ -35,17 +35,12 @@ fn default_file_resource() -> Value { } fn provider(server: &MockServer) -> GoogleProvider { - provider_with_retries(server, RetryConfig::default().max_retries) -} - -fn provider_with_retries(server: &MockServer, max_retries: u32) -> GoogleProvider { - let config = GoogleConfig::new("test-api-key") - .with_base_url(format!("{}/v1beta", server.uri())) - .with_retry_config(RetryConfig { - max_retries, - ..Default::default() - }); - GoogleProvider::new(config) + create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(format!("{}/v1beta", server.uri())), + ..Default::default() + }) + .unwrap() } fn upload_options() -> UploadFileCallOptions { @@ -86,8 +81,12 @@ async fn mount_success_mocks(server: &MockServer) { #[tokio::test] async fn should_expose_correct_provider_and_specification_version() { - let provider = GoogleProvider::new(GoogleConfig::new("test-api-key")); - let files = provider.files(); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap(); + let files = Provider::files(&provider).unwrap(); assert_eq!(files.provider(), "google.generative-ai"); assert_eq!(files.specification_version(), "v4"); } @@ -100,7 +99,7 @@ async fn should_send_correct_headers_for_resumable_upload_initiation() { mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let opts = UploadFileCallOptions { data: UploadFileData::Data { @@ -157,7 +156,7 @@ async fn should_omit_display_name_from_body_when_not_provided() { mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); files.upload_file(&upload_options()).await.unwrap(); @@ -174,7 +173,7 @@ async fn should_send_file_data_to_upload_url_with_correct_headers() { mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let opts = UploadFileCallOptions { data: UploadFileData::Data { @@ -217,7 +216,7 @@ async fn should_return_provider_reference_with_google_key_set_to_file_uri() { mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await.unwrap(); @@ -233,7 +232,7 @@ async fn should_return_empty_warnings_when_filename_not_provided() { mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await.unwrap(); @@ -246,7 +245,7 @@ async fn should_return_unsupported_warning_when_filename_provided() { mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let mut opts = upload_options(); opts.filename = Some("test.pdf".to_string()); @@ -267,7 +266,7 @@ async fn should_return_provider_metadata_with_file_details() { mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await.unwrap(); @@ -294,7 +293,7 @@ async fn should_handle_base64_string_data() { mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); // "hello" encoded as base64. let opts = UploadFileCallOptions { @@ -321,7 +320,7 @@ async fn should_not_poll_when_file_is_immediately_active() { mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); files.upload_file(&upload_options()).await.unwrap(); @@ -355,7 +354,7 @@ async fn should_throw_when_file_state_is_failed() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await; @@ -380,8 +379,8 @@ async fn should_throw_when_initiation_request_fails() { .await; // Only the message text matters here; skip the retries so it stays fast. - let provider = provider_with_retries(&server, 0); - let files = provider.files(); + let provider = provider(&server); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await; @@ -406,7 +405,7 @@ async fn should_throw_when_no_upload_url_returned() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await; @@ -436,8 +435,8 @@ async fn should_throw_when_upload_request_fails() { .await; // Only the message text matters here; skip the retries so it stays fast. - let provider = provider_with_retries(&server, 0); - let files = provider.files(); + let provider = provider(&server); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await; @@ -451,71 +450,19 @@ async fn should_throw_when_upload_request_fails() { // -- provider options tests -------------------------------------------------- -#[tokio::test] -async fn should_accept_valid_provider_options() { - let server = MockServer::start().await; - mount_success_mocks(&server).await; - - let provider = provider(&server); - let files = provider.files(); - - let po = provider_namespace( - "google", - json!({ "displayName": "test", "pollIntervalMs": 5000, "pollTimeoutMs": 60000 }), - ) - .unwrap(); - let opts = UploadFileCallOptions { - data: UploadFileData::Data { - data: FileBytes::Binary(vec![1]), - }, - media_type: "text/plain".to_string(), - filename: None, - provider_options: Some(po), - abort_signal: None, - }; - - let result = files.upload_file(&opts).await.unwrap(); - - assert!(result.provider_reference.contains_key("google")); -} - #[tokio::test] async fn should_work_without_provider_options() { let server = MockServer::start().await; mount_success_mocks(&server).await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await.unwrap(); assert!(result.provider_reference.contains_key("google")); } -#[tokio::test] -async fn should_pass_through_unknown_properties() { - let server = MockServer::start().await; - mount_success_mocks(&server).await; - - let provider = provider(&server); - let files = provider.files(); - - let po = provider_namespace("google", json!({ "customField": "custom-value" })).unwrap(); - let opts = UploadFileCallOptions { - data: UploadFileData::Data { - data: FileBytes::Binary(vec![1]), - }, - media_type: "text/plain".to_string(), - filename: None, - provider_options: Some(po), - abort_signal: None, - }; - - let result = files.upload_file(&opts).await.unwrap(); - - assert!(result.provider_reference.contains_key("google")); -} - // -- response metadata tests ------------------------------------------------- #[tokio::test] @@ -543,16 +490,16 @@ async fn should_omit_optional_fields_from_metadata_when_not_present() { .await; let provider = provider(&server); - let files = provider.files(); + let files = Provider::files(&provider).unwrap(); let result = files.upload_file(&upload_options()).await.unwrap(); let metadata = result.provider_metadata.expect("metadata"); let google = metadata.get("google").expect("google metadata"); assert_eq!(google["name"], "files/minimal"); - assert_eq!(google["displayName"], Value::Null); + assert!(google.get("displayName").is_none()); assert_eq!(google["mimeType"], "text/plain"); - assert_eq!(google["sizeBytes"], Value::Null); + assert!(google.get("sizeBytes").is_none()); assert_eq!(google["state"], "ACTIVE"); assert_eq!( google["uri"], @@ -562,81 +509,3 @@ async fn should_omit_optional_fields_from_metadata_when_not_present() { assert!(google.get("expirationTime").is_none()); assert!(google.get("sha256Hash").is_none()); } - -/// A transient 503 on the upload stage followed by 200 must succeed, and -/// must not re-run the init stage — a retried upload reuses the -/// already-minted `upload_url` instead of requesting a new one. -#[tokio::test] -async fn transient_upload_failure_is_retried_without_re_initiating() { - let server = MockServer::start().await; - let upload_url = format!("{}/resume", server.uri()); - - let init_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let init_observed = std::sync::Arc::clone(&init_attempts); - Mock::given(method("POST")) - .and(path("/upload/v1beta/files")) - .respond_with(move |_: &wiremock::Request| { - init_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - ResponseTemplate::new(200).insert_header("x-goog-upload-url", upload_url.as_str()) - }) - .mount(&server) - .await; - - let upload_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let upload_observed = std::sync::Arc::clone(&upload_attempts); - Mock::given(method("POST")) - .and(path("/resume")) - .respond_with(move |_: &wiremock::Request| { - if upload_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { - ResponseTemplate::new(503) - .insert_header("retry-after-ms", "0") - .set_body_json(json!({"error": {"message": "try again"}})) - } else { - ResponseTemplate::new(200).set_body_json(json!({ "file": default_file_resource() })) - } - }) - .mount(&server) - .await; - - let provider = provider_with_retries(&server, 1); - let files = provider.files(); - - let result = files.upload_file(&upload_options()).await.unwrap(); - - assert_eq!( - result.provider_reference.get("google"), - Some(&"https://generativelanguage.googleapis.com/v1beta/files/abc123".to_string()) - ); - assert_eq!( - init_attempts.load(std::sync::atomic::Ordering::SeqCst), - 1, - "the init stage must not be replayed by an upload-stage retry" - ); - assert_eq!(upload_attempts.load(std::sync::atomic::Ordering::SeqCst), 2); -} - -// ════════════════════════════════════════════════════════════════════════════ -// get_model_path (TS: get-model-path.test.ts) -// ════════════════════════════════════════════════════════════════════════════ - -mod get_model_path_tests { - use aimux_providers::google::utils::get_model_path; - - #[test] - fn pass_through_for_models_slash() { - assert_eq!(get_model_path("models/some-model"), "models/some-model"); - } - - #[test] - fn pass_through_for_tuned_models_slash() { - assert_eq!( - get_model_path("tunedModels/some-model"), - "tunedModels/some-model" - ); - } - - #[test] - fn add_prefix_to_models_without_slash() { - assert_eq!(get_model_path("some-model"), "models/some-model"); - } -} diff --git a/aimux-providers/tests/google_image_test.rs b/aimux-providers/tests/google_image_test.rs index 6292f6011..f0bd2f18c 100644 --- a/aimux-providers/tests/google_image_test.rs +++ b/aimux-providers/tests/google_image_test.rs @@ -1,8 +1,6 @@ //! Rust translation of the Google image model tests. //! //! Source: `reference/ai/packages/google/src/google-image-model.test.ts` -//! -//! Tests both Imagen (non-gemini) and Gemini image model paths. use std::collections::HashMap; @@ -13,23 +11,15 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, }; +use aimux_core::provider::Provider; use aimux_core::shared::{AspectRatio, Size}; use aimux_core::types::Warning; -use aimux_providers::{GoogleConfig, GoogleImageSettings, GoogleProvider}; +use aimux_providers::{GoogleImageSettings, GoogleProviderSettings, create_google}; // ── helpers ───────────────────────────────────────────────────────────────── const PROMPT: &str = "A cute baby sea otter"; -fn imagen_response_body() -> Value { - json!({ - "predictions": [ - { "bytesBase64Encoded": "base64-image-1" }, - { "bytesBase64Encoded": "base64-image-2" }, - ] - }) -} - fn gemini_response_body() -> Value { json!({ "candidates": [{ @@ -52,26 +42,6 @@ fn gemini_response_body() -> Value { }) } -async fn mock_imagen_response(server: &MockServer, body: Value) { - mock_imagen_response_with_headers(server, body, &[]).await; -} - -async fn mock_imagen_response_with_headers( - server: &MockServer, - body: Value, - headers: &[(&str, &str)], -) { - let mut template = ResponseTemplate::new(200).set_body_json(body); - for (k, v) in headers { - template = template.insert_header(*k, *v); - } - Mock::given(method("POST")) - .and(path("/models/imagen-3.0-generate-002:predict")) - .respond_with(template) - .mount(server) - .await; -} - async fn mock_gemini_response(server: &MockServer, body: Value) { mock_gemini_response_with_headers(server, body, &[]).await; } @@ -104,21 +74,25 @@ fn base64_file(media_type: &str, b64: &str) -> ImageFile { } // ════════════════════════════════════════════════════════════════════════════ -// Imagen tests +// Image model tests // ════════════════════════════════════════════════════════════════════════════ /// TS: "should pass headers" #[tokio::test] async fn should_pass_headers() { let server = MockServer::start().await; - mock_imagen_response(&server, imagen_response_body()).await; + mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("imagen-3.0-generate-002"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options(PROMPT); - opts.n = 2; + opts.n = 1; let mut req_headers = HashMap::new(); req_headers.insert( "Custom-Request-Header".to_string(), @@ -144,10 +118,13 @@ async fn should_pass_headers() { /// TS: "should respect maxImagesPerCall setting" #[tokio::test] async fn should_respect_max_images_per_call_setting() { - let config = GoogleConfig::new("test-api-key"); - let provider = GoogleProvider::new(config); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap(); let model = provider.image_with_settings( - "imagen-3.0-generate-002", + "gemini-2.5-flash-image", GoogleImageSettings { max_images_per_call: Some(2), }, @@ -158,33 +135,37 @@ async fn should_respect_max_images_per_call_setting() { /// TS: "should use default maxImagesPerCall when not specified" #[tokio::test] async fn should_use_default_max_images_per_call_when_not_specified() { - let config = GoogleConfig::new("test-api-key"); - let provider = GoogleProvider::new(config); - let model = provider.image("imagen-3.0-generate-002"); - assert_eq!(model.max_images_per_call(), Some(4)); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); + assert_eq!(model.max_images_per_call(), Some(1)); } /// TS: "should extract the generated images" #[tokio::test] async fn should_extract_the_generated_images() { let server = MockServer::start().await; - mock_imagen_response(&server, imagen_response_body()).await; + mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("imagen-3.0-generate-002"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options(PROMPT); - opts.n = 2; + opts.n = 1; let result = model.do_generate(&opts).await.unwrap(); match &result.images { ImageOutputs::Base64(imgs) => { - assert_eq!( - imgs, - &["base64-image-1".to_string(), "base64-image-2".to_string()] - ); + assert_eq!(imgs, &["base64-generated-image".to_string()]); } _ => panic!("expected Base64 images"), } @@ -194,11 +175,15 @@ async fn should_extract_the_generated_images() { #[tokio::test] async fn sends_aspect_ratio_in_the_request() { let server = MockServer::start().await; - mock_imagen_response(&server, imagen_response_body()).await; + mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("imagen-3.0-generate-002"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("test prompt"); opts.n = 1; @@ -208,20 +193,26 @@ async fn sends_aspect_ratio_in_the_request() { let requests = server.received_requests().await.expect("requests recorded"); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!(body["instances"][0]["prompt"], "test prompt"); - assert_eq!(body["parameters"]["sampleCount"], 1); - assert_eq!(body["parameters"]["aspectRatio"], "16:9"); + assert_eq!(body["contents"][0]["parts"][0]["text"], "test prompt"); + assert_eq!( + body["generationConfig"]["imageConfig"]["aspectRatio"], + "16:9" + ); } /// TS: "should return warnings for unsupported settings" #[tokio::test] async fn should_return_warnings_for_unsupported_settings() { let server = MockServer::start().await; - mock_imagen_response(&server, imagen_response_body()).await; + mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("imagen-3.0-generate-002"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -231,7 +222,7 @@ async fn should_return_warnings_for_unsupported_settings() { let result = model.do_generate(&opts).await.unwrap(); - assert_eq!(result.warnings.len(), 2); + assert_eq!(result.warnings.len(), 1); match &result.warnings[0] { Warning::Unsupported { feature, details } => { assert_eq!(feature, "size"); @@ -242,25 +233,15 @@ async fn should_return_warnings_for_unsupported_settings() { } _ => panic!("expected Unsupported warning for size"), } - match &result.warnings[1] { - Warning::Unsupported { feature, details } => { - assert_eq!(feature, "seed"); - assert_eq!( - details.as_deref(), - Some("This model does not support the `seed` option through this provider.") - ); - } - _ => panic!("expected Unsupported warning for seed"), - } } /// TS: "should include response data with timestamp, modelId and headers" #[tokio::test] async fn should_include_response_data_with_timestamp_model_id_and_headers() { let server = MockServer::start().await; - mock_imagen_response_with_headers( + mock_gemini_response_with_headers( &server, - imagen_response_body(), + gemini_response_body(), &[ ("request-id", "test-request-id"), ("x-goog-quota-remaining", "123"), @@ -268,9 +249,13 @@ async fn should_include_response_data_with_timestamp_model_id_and_headers() { ) .await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("imagen-3.0-generate-002"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -280,7 +265,7 @@ async fn should_include_response_data_with_timestamp_model_id_and_headers() { assert!(result.response.timestamp.is_some()); assert_eq!( result.response.model_id.as_deref(), - Some("imagen-3.0-generate-002") + Some("gemini-2.5-flash-image") ); let headers = result.response.headers.as_ref().unwrap(); assert_eq!( @@ -295,63 +280,19 @@ async fn should_include_response_data_with_timestamp_model_id_and_headers() { ); } -/// TS: "should use real date when no custom date provider is specified" -#[tokio::test] -async fn should_use_real_date_when_no_custom_date_provider_is_specified() { - let server = MockServer::start().await; - mock_imagen_response(&server, imagen_response_body()).await; - - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("imagen-3.0-generate-002"); - - let mut opts = options(PROMPT); - opts.n = 2; - - let before = chrono::Utc::now(); - let result = model.do_generate(&opts).await.unwrap(); - let after = chrono::Utc::now(); - - let ts = result.response.timestamp.as_ref().unwrap(); - let parsed = chrono::DateTime::parse_from_rfc3339(ts).unwrap(); - let parsed_utc = parsed.with_timezone(&chrono::Utc); - assert!(parsed_utc >= before); - assert!(parsed_utc <= after); - assert_eq!( - result.response.model_id.as_deref(), - Some("imagen-3.0-generate-002") - ); -} - -/// TS: "should throw error when files are provided" -#[tokio::test] -async fn should_throw_error_when_files_are_provided() { - let server = MockServer::start().await; - mock_imagen_response(&server, imagen_response_body()).await; - - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("imagen-3.0-generate-002"); - - let mut opts = options("Edit this image"); - opts.n = 1; - opts.files = Some(vec![base64_file("image/png", "base64-source-image")]); - - let result = model.do_generate(&opts).await; - assert!(result.is_err()); - let err = result.unwrap_err().to_string(); - assert!(err.contains("does not support image editing with Imagen models")); -} - /// TS: "should throw error when mask is provided" #[tokio::test] async fn should_throw_error_when_mask_is_provided() { let server = MockServer::start().await; - mock_imagen_response(&server, imagen_response_body()).await; + mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("imagen-3.0-generate-002"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("Edit this image"); opts.n = 1; @@ -360,7 +301,7 @@ async fn should_throw_error_when_mask_is_provided() { let result = model.do_generate(&opts).await; assert!(result.is_err()); let err = result.unwrap_err().to_string(); - assert!(err.contains("does not support image editing with masks")); + assert!(err.contains("do not support mask-based image editing")); } // ════════════════════════════════════════════════════════════════════════════ @@ -370,17 +311,23 @@ async fn should_throw_error_when_mask_is_provided() { /// TS: "should return 10 for Gemini image models by default" #[tokio::test] async fn gemini_should_return_10_for_max_images_per_call() { - let config = GoogleConfig::new("test-api-key"); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); - assert_eq!(model.max_images_per_call(), Some(10)); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); + assert_eq!(model.max_images_per_call(), Some(1)); } /// TS: "should respect custom maxImagesPerCall setting" #[tokio::test] async fn gemini_should_respect_custom_max_images_per_call() { - let config = GoogleConfig::new("test-api-key"); - let provider = GoogleProvider::new(config); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap(); let model = provider.image_with_settings( "gemini-2.5-flash-image", GoogleImageSettings { @@ -396,9 +343,13 @@ async fn gemini_should_extract_the_generated_image() { let server = MockServer::start().await; mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A beautiful sunset"); opts.n = 1; @@ -419,9 +370,13 @@ async fn gemini_should_send_correct_request_body() { let server = MockServer::start().await; mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A beautiful sunset"); opts.n = 1; @@ -446,9 +401,13 @@ async fn gemini_should_pass_aspect_ratio_via_image_config() { let server = MockServer::start().await; mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A beautiful sunset"); opts.n = 1; @@ -470,9 +429,13 @@ async fn gemini_should_support_21_9_aspect_ratio() { let server = MockServer::start().await; mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A cinematic landscape"); opts.n = 1; @@ -494,9 +457,13 @@ async fn gemini_should_pass_seed_in_generation_config() { let server = MockServer::start().await; mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A beautiful sunset"); opts.n = 1; @@ -537,9 +504,13 @@ async fn gemini_should_include_usage_in_response() { ) .await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A beautiful sunset"); opts.n = 1; @@ -558,9 +529,13 @@ async fn gemini_should_return_warning_for_size() { let server = MockServer::start().await; mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A beautiful sunset"); opts.n = 1; @@ -574,35 +549,19 @@ async fn gemini_should_return_warning_for_size() { })); } -/// TS: "should not send a tools field when googleSearch is not set" -#[tokio::test] -async fn gemini_should_not_send_tools_when_no_google_search() { - let server = MockServer::start().await; - mock_gemini_response(&server, gemini_response_body()).await; - - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); - - let mut opts = options("A beautiful sunset"); - opts.n = 1; - - model.do_generate(&opts).await.unwrap(); - - let requests = server.received_requests().await.expect("requests recorded"); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert!(body.get("tools").is_none()); -} - /// TS: "should include input images in request for editing" #[tokio::test] async fn gemini_should_include_input_images_for_editing() { let server = MockServer::start().await; mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("Add a hat to this cat"); opts.n = 1; @@ -621,34 +580,19 @@ async fn gemini_should_include_input_images_for_editing() { ); } -/// TS: "should throw error when n > 1" -#[tokio::test] -async fn gemini_should_throw_error_when_n_gt_1() { - let server = MockServer::start().await; - mock_gemini_response(&server, gemini_response_body()).await; - - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); - - let mut opts = options("A beautiful sunset"); - opts.n = 2; - - let result = model.do_generate(&opts).await; - assert!(result.is_err()); - let err = result.unwrap_err().to_string(); - assert!(err.contains("do not support generating a set number of images per call")); -} - /// TS: "should throw error when mask is provided" #[tokio::test] async fn gemini_should_throw_error_when_mask_is_provided() { let server = MockServer::start().await; mock_gemini_response(&server, gemini_response_body()).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("Edit this image"); opts.n = 1; diff --git a/aimux-providers/tests/google_pse_test.rs b/aimux-providers/tests/google_pse_test.rs index efeadf59f..760e6ec73 100644 --- a/aimux-providers/tests/google_pse_test.rs +++ b/aimux-providers/tests/google_pse_test.rs @@ -10,10 +10,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path, query_param}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{GooglePseConfig, GooglePseProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{GooglePseProvider, GooglePseProviderSettings, create_google_pse}; const API_KEY: &str = "test-google-key"; const CX: &str = "test-cx-id"; @@ -39,10 +39,13 @@ fn search_body() -> Value { } fn provider(server: &MockServer) -> GooglePseProvider { - let config = GooglePseConfig::new(API_KEY) - .with_cx(CX) - .with_base_url(format!("{}/customsearch/v1", server.uri())); - GooglePseProvider::new(config) + let config = GooglePseProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + cx: Some(CX.to_string()), + base_url: Some(format!("{}/customsearch/v1", server.uri()).to_string()), + ..Default::default() + }; + create_google_pse(config).unwrap() } fn opts(query: &str) -> SearchCallOptions { @@ -158,9 +161,12 @@ async fn do_search_forwards_max_results_as_num() { async fn missing_cx_returns_invalid_argument_error() { let server = MockServer::start().await; // Config without cx; no provider_options either. - let config = - GooglePseConfig::new(API_KEY).with_base_url(format!("{}/customsearch/v1", server.uri())); - let provider = GooglePseProvider::new(config); + let config = GooglePseProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(format!("{}/customsearch/v1", server.uri()).to_string()), + ..Default::default() + }; + let provider = create_google_pse(config).unwrap(); let model = provider.search_model(); let result = model.do_search(&opts("rust language")).await; @@ -220,33 +226,15 @@ async fn status_403_maps_to_provider_error() { // -- Provider trait ---------------------------------------------------------- -#[tokio::test] -async fn provider_name_is_google_pse() { - let config = GooglePseConfig::new(API_KEY).with_cx(CX); - let provider = GooglePseProvider::new(config); - assert_eq!(provider.name(), "google_pse"); -} - -#[test] -fn language_model_returns_unsupported_error() { - let config = GooglePseConfig::new(API_KEY).with_cx(CX); - let provider = GooglePseProvider::new(config); - match provider.language_model("google-pse-search") { - Err(AiMuxError::UnsupportedFunctionality(msg)) => { - assert!( - msg.contains("provider 'google_pse' does not provide language models"), - "unexpected message: {msg}" - ); - } - _ => panic!("expected Unsupported error, got success or another error variant"), - } -} - #[test] fn model_id_is_google_pse_search() { - let config = GooglePseConfig::new(API_KEY).with_cx(CX); - let provider = GooglePseProvider::new(config); + let config = GooglePseProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + cx: Some(CX.to_string()), + ..Default::default() + }; + let provider = create_google_pse(config).unwrap(); let model = provider.search_model(); assert_eq!(model.model_id(), "google-pse-search"); - assert_eq!(model.provider(), "google_pse"); + assert_eq!(model.provider(), "google_pse.search"); } diff --git a/aimux-providers/tests/google_vertex_image_test.rs b/aimux-providers/tests/google_vertex_image_test.rs index 0140319a4..9880c8181 100644 --- a/aimux-providers/tests/google_vertex_image_test.rs +++ b/aimux-providers/tests/google_vertex_image_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Google Vertex AI image model tests. +//! Rust translation of the Google Vertex AI image model tests. //! //! Source: `reference/ai/packages/google-vertex/src/google-vertex-image-model.test.ts` @@ -6,18 +6,14 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::image_model::{ - ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, -}; +use aimux_core::image_model::{ImageCallOptions, ImageFile, ImageFileData, ImageOutputs}; +use aimux_core::provider::Provider; use aimux_core::shared::{AspectRatio, Size}; -use aimux_providers::{VertexProvider, VertexProviderConfig}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{VertexProviderSettings, create_google_vertex}; const PROMPT: &str = "A cute baby sea otter"; -fn imagen_response() -> Value { - json!({ "predictions": [{ "bytesBase64Encoded": "base64-image-1", "mimeType": "image/png" }] }) -} - fn gemini_response() -> Value { json!({ "candidates": [{ @@ -28,13 +24,12 @@ fn gemini_response() -> Value { }) } -fn config(server: &MockServer) -> VertexProviderConfig { - VertexProviderConfig::with_api_key("test-api-key").with_base_url(server.uri()) -} - -fn config_bearer(server: &MockServer) -> VertexProviderConfig { - VertexProviderConfig::new("test-token", "test-project", "us-central1") - .with_base_url(server.uri()) +fn config(server: &MockServer) -> VertexProviderSettings { + VertexProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri()), + ..Default::default() + } } fn options(prompt: &str) -> ImageCallOptions { @@ -45,15 +40,15 @@ fn options(prompt: &str) -> ImageCallOptions { async fn imagen_should_extract_generated_images() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/models/imagen-4.0-generate-001:predict")) - .respond_with(ResponseTemplate::new(200).set_body_json(imagen_response())) + .and(path("/models/gemini-2.5-flash-image:generateContent")) + .respond_with(ResponseTemplate::new(200).set_body_json(gemini_response())) .mount(&server) .await; - let provider = VertexProvider::new(config(&server)); - let model = provider.image("imagen-4.0-generate-001"); + let provider = create_google_vertex(config(&server)).unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { - ImageOutputs::Base64(i) => assert_eq!(i, ["base64-image-1"]), + ImageOutputs::Base64(i) => assert_eq!(i, ["base64-generated-image"]), _ => panic!("expected Base64"), } } @@ -62,32 +57,34 @@ async fn imagen_should_extract_generated_images() { async fn imagen_should_send_aspect_ratio() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/models/imagen-4.0-generate-001:predict")) - .respond_with(ResponseTemplate::new(200).set_body_json(imagen_response())) + .and(path("/models/gemini-2.5-flash-image:generateContent")) + .respond_with(ResponseTemplate::new(200).set_body_json(gemini_response())) .mount(&server) .await; - let provider = VertexProvider::new(config(&server)); - let model = provider.image("imagen-4.0-generate-001"); + let provider = create_google_vertex(config(&server)).unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("test prompt"); opts.n = 1; opts.aspect_ratio = Some(AspectRatio::new(16, 9)); model.do_generate(&opts).await.unwrap(); let reqs = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&reqs[0].body).unwrap(); - assert_eq!(body["parameters"]["aspectRatio"], "16:9"); - assert_eq!(body["parameters"]["sampleCount"], 1); + assert_eq!( + body["generationConfig"]["imageConfig"]["aspectRatio"], + "16:9" + ); } #[tokio::test] async fn imagen_should_warn_for_size() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/models/imagen-4.0-generate-001:predict")) - .respond_with(ResponseTemplate::new(200).set_body_json(imagen_response())) + .and(path("/models/gemini-2.5-flash-image:generateContent")) + .respond_with(ResponseTemplate::new(200).set_body_json(gemini_response())) .mount(&server) .await; - let provider = VertexProvider::new(config(&server)); - let model = provider.image("imagen-4.0-generate-001"); + let provider = create_google_vertex(config(&server)).unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options(PROMPT); opts.n = 1; opts.size = Some(Size::new(1024, 1024)); @@ -99,12 +96,12 @@ async fn imagen_should_warn_for_size() { async fn imagen_should_support_editing() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/models/imagen-4.0-generate-001:predict")) - .respond_with(ResponseTemplate::new(200).set_body_json(imagen_response())) + .and(path("/models/gemini-2.5-flash-image:generateContent")) + .respond_with(ResponseTemplate::new(200).set_body_json(gemini_response())) .mount(&server) .await; - let provider = VertexProvider::new(config(&server)); - let model = provider.image("imagen-4.0-generate-001"); + let provider = create_google_vertex(config(&server)).unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("Edit this"); opts.n = 1; opts.files = Some(vec![ImageFile::File { @@ -114,10 +111,9 @@ async fn imagen_should_support_editing() { model.do_generate(&opts).await.unwrap(); let reqs = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&reqs[0].body).unwrap(); - assert!(body["instances"][0]["referenceImages"].is_array()); assert_eq!( - body["parameters"]["editMode"], - "EDIT_MODE_INPAINT_INSERTION" + body["contents"][0]["parts"][1], + json!({ "inlineData": { "mimeType": "image/png", "data": "base64-src" } }) ); } @@ -129,8 +125,8 @@ async fn gemini_should_extract_image() { .respond_with(ResponseTemplate::new(200).set_body_json(gemini_response())) .mount(&server) .await; - let provider = VertexProvider::new(config(&server)); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google_vertex(config(&server)).unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A sunset"); opts.n = 1; let result = model.do_generate(&opts).await.unwrap(); @@ -148,8 +144,8 @@ async fn gemini_should_send_response_modalities() { .respond_with(ResponseTemplate::new(200).set_body_json(gemini_response())) .mount(&server) .await; - let provider = VertexProvider::new(config(&server)); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google_vertex(config(&server)).unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A sunset"); opts.n = 1; model.do_generate(&opts).await.unwrap(); @@ -169,8 +165,8 @@ async fn gemini_should_pass_aspect_ratio() { .respond_with(ResponseTemplate::new(200).set_body_json(gemini_response())) .mount(&server) .await; - let provider = VertexProvider::new(config(&server)); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google_vertex(config(&server)).unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("A sunset"); opts.n = 1; opts.aspect_ratio = Some(AspectRatio::new(16, 9)); @@ -183,21 +179,11 @@ async fn gemini_should_pass_aspect_ratio() { ); } -#[tokio::test] -async fn gemini_should_throw_for_n_gt_1() { - let server = MockServer::start().await; - let provider = VertexProvider::new(config(&server)); - let model = provider.image("gemini-2.5-flash-image"); - let mut opts = options("A sunset"); - opts.n = 2; - assert!(model.do_generate(&opts).await.is_err()); -} - #[tokio::test] async fn gemini_should_throw_for_mask() { let server = MockServer::start().await; - let provider = VertexProvider::new(config(&server)); - let model = provider.image("gemini-2.5-flash-image"); + let provider = create_google_vertex(config(&server)).unwrap(); + let model = provider.image_model("gemini-2.5-flash-image").unwrap(); let mut opts = options("Edit"); opts.n = 1; opts.mask = Some(ImageFile::File { @@ -207,29 +193,6 @@ async fn gemini_should_throw_for_mask() { assert!(model.do_generate(&opts).await.is_err()); } -#[tokio::test] -async fn should_use_bearer_token_auth() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/models/imagen-4.0-generate-001:predict")) - .respond_with(ResponseTemplate::new(200).set_body_json(imagen_response())) - .mount(&server) - .await; - let provider = VertexProvider::new(config_bearer(&server)); - let model = provider.image("imagen-4.0-generate-001"); - model.do_generate(&options(PROMPT)).await.unwrap(); - let reqs = server.received_requests().await.unwrap(); - assert!( - reqs[0] - .headers - .get("authorization") - .unwrap() - .to_str() - .unwrap() - .starts_with("Bearer ") - ); -} - // Helper trait for warning feature access trait WarningFeature { fn feature(&self) -> &str; diff --git a/aimux-providers/tests/google_video_test.rs b/aimux-providers/tests/google_video_test.rs index a965f646d..23afeae57 100644 --- a/aimux-providers/tests/google_video_test.rs +++ b/aimux-providers/tests/google_video_test.rs @@ -1,12 +1,12 @@ //! Rust translation of the Google video model tests. //! Source: `reference/ai/packages/google/src/google-video-model.test.ts` -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; -use aimux_core::AiMuxError; +use aimux_core::download::{Download, DownloadRequest, DownloadedFile}; +use aimux_core::provider::Provider; use aimux_core::video_model::{VideoCallOptions, generate_video}; -use aimux_providers::{GoogleConfig, GoogleProvider}; +use aimux_providers::{GoogleProviderSettings, create_google}; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -21,14 +21,25 @@ fn fast_poll() -> Option { fn options(p: &str) -> VideoCallOptions { let mut o = VideoCallOptions::new(p); o.poll = fast_poll(); + o.download = Some(recording_download(Arc::default())); o } -fn request_header<'a>(request: &'a wiremock::Request, name: &str) -> Option<&'a str> { - request - .headers - .get(name) - .and_then(|value| value.to_str().ok()) +/// The bytes the test download returns for every video URL. +const VIDEO_BYTES: &[u8] = b"video bytes"; + +/// A download that fetches nothing: it records each URL `generate_video` +/// asks for and returns [`VIDEO_BYTES`]. +fn recording_download(requests: Arc>>) -> Download { + Download::new(move |request: DownloadRequest| { + requests.lock().unwrap().push(request.url); + async { + Ok(DownloadedFile { + data: VIDEO_BYTES.to_vec(), + media_type: None, + }) + } + }) } async fn mock_predict_and_poll(server: &MockServer, result: &Value) { @@ -51,247 +62,58 @@ async fn mock_predict_and_poll(server: &MockServer, result: &Value) { #[tokio::test] async fn should_generate_video() { let server = MockServer::start().await; - let result = - json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); + let result = json!({"done": true, "response": {"generateVideoResponse": {"generatedSamples": [{"video": {"uri": format!("{}/files/video-123.mp4", server.uri())}}]}}}); mock_predict_and_poll(&server, &result).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.video("veo-3.0-generate-001"); - let r = generate_video(&model, options("A cat")).await.unwrap(); - assert_eq!(r.videos.len(), 1); -} - -#[tokio::test] -async fn poll_retry_does_not_submit_a_second_generation() { - let server = MockServer::start().await; - - Mock::given(method("POST")) - .and(path("/models/veo-3.0-generate-001:predictLongRunning")) - .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"name": "operations/test-op"})), - ) - .mount(&server) - .await; - - let poll_attempts = Arc::new(AtomicUsize::new(0)); - let responder_attempts = Arc::clone(&poll_attempts); - Mock::given(method("GET")) - .and(path("/operations/test-op")) - .respond_with(move |_: &wiremock::Request| { - if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 { - ResponseTemplate::new(503) - .insert_header("retry-after-ms", "0") - .set_body_json(json!({ - "error": {"message": "try again", "status": "UNAVAILABLE"} - })) - } else { - ResponseTemplate::new(200).set_body_json(json!({ - "done": true, - "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]} - })) - } - }) - .mount(&server) - .await; - - let config = GoogleConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_retry_config(aimux_provider_utils::RetryConfig { - max_retries: 1, - ..Default::default() - }); - let provider = GoogleProvider::new(config); - let model = provider.video("veo-3.0-generate-001"); - let mut opts = options("A cat"); - opts.max_retries = Some(1); - - let result = generate_video(&model, opts).await.unwrap(); - assert_eq!(result.videos.len(), 1); - assert_eq!(poll_attempts.load(Ordering::SeqCst), 2); - - let requests = server.received_requests().await.unwrap(); - assert_eq!( - requests - .iter() - .filter(|request| request.url.path().ends_with(":predictLongRunning")) - .count(), - 1 - ); - assert!( - requests - .iter() - .find(|request| request.url.path().ends_with(":predictLongRunning")) - .and_then(|request| request_header(request, "idempotency-key")) - .is_some(), - "the start request should receive a generated idempotency key" - ); - assert!( - requests - .iter() - .filter(|request| request.url.path() == "/operations/test-op") - .all(|request| request_header(request, "idempotency-key").is_none()), - "the generated start key must not leak to status requests" - ); -} - -#[tokio::test] -async fn start_retry_reuses_the_same_idempotency_key() { - let server = MockServer::start().await; - let start_attempts = Arc::new(AtomicUsize::new(0)); - let responder_attempts = Arc::clone(&start_attempts); - - Mock::given(method("POST")) - .and(path("/models/veo-3.0-generate-001:predictLongRunning")) - .respond_with(move |_: &wiremock::Request| { - if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 { - ResponseTemplate::new(503) - .insert_header("retry-after-ms", "0") - .set_body_json(json!({ - "error": {"message": "try again", "status": "UNAVAILABLE"} - })) - } else { - ResponseTemplate::new(200).set_body_json(json!({"name": "operations/test-op"})) - } - }) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/operations/test-op")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "done": true, - "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]} - }))) - .mount(&server) - .await; - - let config = GoogleConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_retry_config(aimux_provider_utils::RetryConfig { - max_retries: 1, - ..Default::default() - }); - let provider = GoogleProvider::new(config); - let model = provider.video("veo-3.0-generate-001"); - let mut opts = options("A cat"); - opts.max_retries = Some(1); - - generate_video(&model, opts).await.unwrap(); - - let requests = server.received_requests().await.unwrap(); - let starts: Vec<_> = requests - .iter() - .filter(|request| request.url.path().ends_with(":predictLongRunning")) - .collect(); - assert_eq!(starts.len(), 2); - let first_key = request_header(starts[0], "idempotency-key"); - assert!(first_key.is_some()); - assert_eq!(first_key, request_header(starts[1], "idempotency-key")); - assert!( - requests - .iter() - .filter(|request| request.url.path() == "/operations/test-op") - .all(|request| request_header(request, "idempotency-key").is_none()) - ); -} - -#[tokio::test] -async fn poll_deadline_reached_during_delay_does_not_issue_status_request() { - let server = MockServer::start().await; - let result = - json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); - mock_predict_and_poll(&server, &result).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.video("veo-3.0-generate-001"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider + .video_model("veo-3.0-generate-001") + .unwrap() + .unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); let mut opts = options("A cat"); - opts.poll = Some(aimux_core::video_model::VideoPollOptions { - interval_ms: Some(100), - timeout_ms: Some(10), - }); - - let error = generate_video(&model, opts).await.unwrap_err(); - - assert!(matches!(error, AiMuxError::Timeout(_))); - let requests = server.received_requests().await.unwrap(); + opts.download = Some(recording_download(requests.clone())); + let r = generate_video(model.as_ref(), opts).await.unwrap(); + // `generate_video` downloads the returned URL, as the AI SDK does. assert_eq!( - requests - .iter() - .filter(|request| request.url.path().ends_with(":predictLongRunning")) - .count(), - 1 - ); - assert_eq!( - requests - .iter() - .filter(|request| request.url.path() == "/operations/test-op") - .count(), - 0 - ); -} - -#[tokio::test] -async fn poll_retry_exhaustion_does_not_submit_a_second_generation() { - let server = MockServer::start().await; - - Mock::given(method("POST")) - .and(path("/models/veo-3.0-generate-001:predictLongRunning")) - .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"name": "operations/test-op"})), - ) - .mount(&server) - .await; - - Mock::given(method("GET")) - .and(path("/operations/test-op")) - .respond_with( - ResponseTemplate::new(503) - .insert_header("retry-after-ms", "0") - .set_body_json(json!({ - "error": {"message": "still unavailable", "status": "UNAVAILABLE"} - })), - ) - .mount(&server) - .await; - - let config = GoogleConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_retry_config(aimux_provider_utils::RetryConfig { - max_retries: 1, - ..Default::default() - }); - let provider = GoogleProvider::new(config); - let model = provider.video("veo-3.0-generate-001"); - - let error = generate_video(&model, options("A cat")).await.unwrap_err(); - assert!(matches!(error, AiMuxError::Retry(_))); - - let requests = server.received_requests().await.unwrap(); - assert_eq!( - requests - .iter() - .filter(|request| request.url.path().ends_with(":predictLongRunning")) - .count(), - 1 - ); - assert_eq!( - requests - .iter() - .filter(|request| request.url.path() == "/operations/test-op") - .count(), - 2 + *requests.lock().unwrap(), + [format!( + "{}/files/video-123.mp4?key=test-api-key", + server.uri() + )] ); + assert_eq!(r.videos.len(), 1); + match &r.videos[0] { + aimux_core::video_model::VideoData::Binary { data, media_type } => { + assert_eq!(data, VIDEO_BYTES); + assert_eq!(media_type, "video/mp4"); + } + _ => panic!("expected the downloaded video"), + } } #[tokio::test] async fn should_pass_prompt() { let server = MockServer::start().await; - let result = - json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); + let result = json!({"done": true, "response": {"generateVideoResponse": {"generatedSamples": [{"video": {"uri": "https://generativelanguage.googleapis.com/files/video-123.mp4"}}]}}}); mock_predict_and_poll(&server, &result).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.video("veo-3.0-generate-001"); - generate_video(&model, options("A cat")).await.unwrap(); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider + .video_model("veo-3.0-generate-001") + .unwrap() + .unwrap(); + generate_video(model.as_ref(), options("A cat")) + .await + .unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); assert_eq!(body["instances"][0]["prompt"], "A cat"); @@ -300,16 +122,22 @@ async fn should_pass_prompt() { #[tokio::test] async fn should_pass_aspect_ratio_and_duration() { let server = MockServer::start().await; - let result = - json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); + let result = json!({"done": true, "response": {"generateVideoResponse": {"generatedSamples": [{"video": {"uri": "https://generativelanguage.googleapis.com/files/video-123.mp4"}}]}}}); mock_predict_and_poll(&server, &result).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.video("veo-3.0-generate-001"); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider + .video_model("veo-3.0-generate-001") + .unwrap() + .unwrap(); let mut opts = options("test"); opts.aspect_ratio = Some(aimux_core::shared::AspectRatio::new(16, 9)); opts.duration = Some(5); - generate_video(&model, opts).await.unwrap(); + generate_video(model.as_ref(), opts).await.unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); assert_eq!(body["parameters"]["aspectRatio"], "16:9"); @@ -319,13 +147,21 @@ async fn should_pass_aspect_ratio_and_duration() { #[tokio::test] async fn should_include_response_data() { let server = MockServer::start().await; - let result = - json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); + let result = json!({"done": true, "response": {"generateVideoResponse": {"generatedSamples": [{"video": {"uri": "https://generativelanguage.googleapis.com/files/video-123.mp4"}}]}}}); mock_predict_and_poll(&server, &result).await; - let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); - let provider = GoogleProvider::new(config); - let model = provider.video("veo-3.0-generate-001"); - let r = generate_video(&model, options("test")).await.unwrap(); + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider + .video_model("veo-3.0-generate-001") + .unwrap() + .unwrap(); + let r = generate_video(model.as_ref(), options("test")) + .await + .unwrap(); assert!(r.response.timestamp.is_some()); assert_eq!( r.response.model_id, diff --git a/aimux-providers/tests/groq_test.rs b/aimux-providers/tests/groq_test.rs index 46279b3e4..6f8635ad3 100644 --- a/aimux-providers/tests/groq_test.rs +++ b/aimux-providers/tests/groq_test.rs @@ -12,11 +12,8 @@ //! - `packages/groq/src/groq-chat-language-model-options.test.ts` → Zod schema //! validation (TypeScript-specific; not directly translatable) //! -//! Groq is an OpenAI-compatible thin wrapper. The Rust implementation reuses -//! `OpenAIProvider` but sets `provider = "groq"` in the config, which enables -//! groq-specific behaviour in the shared request builder (provider-options key, -//! browser_search tool, stream_options, etc.). Reasoning effort is a direct -//! passthrough — no vendor normalization. +//! Every test builds the model through the Groq package +//! (`create_groq(..).chat(model)`), so it exercises `GroqChatLanguageModel`. use aimux_core::tool::RawToolCall; use futures::StreamExt; @@ -31,12 +28,12 @@ use aimux_core::language_model_message::{ }; use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat, Tool}; use aimux_core::result::{GenerateContent, ReasoningOutput}; -use aimux_core::shared::{FileBytes, FileData}; +use aimux_core::shared::{FileBytes, FileData, provider_namespace}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::FunctionTool; use aimux_core::types::{FinishReasonUnified, ReasoningEffort}; -use aimux_providers::{ProviderOptions, provider}; +use aimux_providers::groq::{GroqChatLanguageModel, GroqProviderSettings, create_groq, groq}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -50,18 +47,27 @@ fn default_options(prompt: LanguageModelPrompt) -> CallOptions { CallOptions::new(prompt) } -/// Build a provider pointed at the mock server. -fn make_provider(server: &MockServer) -> Box { - provider( - "groq", - Some("test-api-key".to_string()), - "gemma2-9b-it", - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), - ) +/// A Groq chat model pointed at the mock server. +fn model_at(server: &MockServer, model_id: &str) -> GroqChatLanguageModel { + create_groq(GroqProviderSettings { + base_url: Some(server.uri()), + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) .expect("groq provider should build") + .chat(model_id) +} + +fn browser_search_model_ids() -> Vec<&'static str> { + include_str!("../src/groq/browser_search_models.rs") + .split('"') + .filter(|value| value.starts_with("openai/")) + .collect() +} + +/// Build the default test model pointed at the mock server. +fn make_provider(server: &MockServer) -> GroqChatLanguageModel { + model_at(server, "gemma2-9b-it") } /// A standard non-streaming chat-completion JSON body returning "Hello, World!". @@ -307,11 +313,11 @@ mod convert_messages { LanguageModelMessage::Tool { content: vec![ToolPart::ToolResult(ToolResultPart { tool_call_id: "quux".into(), - tool_name: "thwomp".into(), output: ToolResultOutput::Json { value: json!({"oof":"321rab"}), provider_options: None, }, + tool_name: "thwomp".into(), provider_options: None, })], provider_options: None, @@ -420,14 +426,19 @@ mod convert_messages { })], provider_options: None, }]; - // Upstream rejects file parts with provider references outright. + // upstream convert-to-groq-chat-messages.ts: a provider reference + // throws UnsupportedFunctionalityError. let result = model.do_generate(&default_options(prompt)).await; assert!(result.is_err()); let err = result.unwrap_err(); + assert!(matches!( + err, + aimux_core::AiMuxError::UnsupportedFunctionality(_) + )); assert!( err.to_string() .contains("file parts with provider references"), - "error should mention 'file parts with provider references': {err}" + "error should name the unsupported functionality: {err}" ); } } @@ -440,8 +451,6 @@ mod convert_usage { use super::*; /// TS: "should convert basic usage without token details" - /// Note: Rust OpenAI converter returns Some(0) for missing cache_read and - /// reasoning (matching TS OpenAI), while TS Groq returns undefined. #[tokio::test] async fn basic_usage() { let server = MockServer::start().await; @@ -731,6 +740,35 @@ mod prepare_tools { assert_eq!(body["tools"][0]["function"]["parameters"]["type"], "object"); } + /// TS: "should add warnings for unsupported provider-defined tools" + #[tokio::test] + async fn unsupported_provider_tool_warning() { + let server = MockServer::start().await; + mock_json(&server, text_completion_body()).await; + + let model = make_provider(&server); + + let tool = Tool::Provider(aimux_core::tool::ProviderTool { + id: "some.unsupported_tool".parse().unwrap(), + name: "unsupported_tool".to_string(), + args: serde_json::Map::new(), + supports_deferred_results: None, + }); + let options = CallOptions { + tools: Some(vec![tool]), + ..default_options(test_prompt()) + }; + let result = model.do_generate(&options).await.unwrap(); + + // The unsupported tool should produce a warning. + assert!(result.warnings.iter().any(|w| match w { + aimux_core::types::Warning::Unsupported { feature, .. } => { + feature.contains("some.unsupported_tool") + } + _ => false, + })); + } + /// TS: "should pass through strict mode when strict is true" #[tokio::test] async fn strict_mode_true() { @@ -820,6 +858,7 @@ mod prepare_tools { model.do_generate(&options).await.unwrap(); let body = first_request_body(&server).await; + // upstream groq-prepare-tools.ts: `case 'auto': tool_choice: 'auto'` assert_eq!(body["tool_choice"], "auto"); } @@ -885,6 +924,117 @@ mod prepare_tools { assert_eq!(body["tool_choice"]["type"], "function"); assert_eq!(body["tool_choice"]["function"]["name"], "testFunction"); } + + /// TS: "should handle browser search tool with supported model" + #[tokio::test] + async fn browser_search_supported_model() { + let server = MockServer::start().await; + mock_json(&server, text_completion_body()).await; + + let model = model_at(&server, browser_search_model_ids()[1]); + + let tool = Tool::Provider(aimux_core::tool::ProviderTool { + id: "groq.browser_search".parse().unwrap(), + name: "browser_search".to_string(), + args: serde_json::Map::new(), + supports_deferred_results: None, + }); + let options = CallOptions { + tools: Some(vec![tool]), + ..default_options(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert_eq!(body["tools"][0]["type"], "browser_search"); + } + + /// TS: "should warn when browser search is used with unsupported model" + #[tokio::test] + async fn browser_search_unsupported_model() { + let server = MockServer::start().await; + mock_json(&server, text_completion_body()).await; + + let model = make_provider(&server); + + let tool = Tool::Provider(aimux_core::tool::ProviderTool { + id: "groq.browser_search".parse().unwrap(), + name: "browser_search".to_string(), + args: serde_json::Map::new(), + supports_deferred_results: None, + }); + let options = CallOptions { + tools: Some(vec![tool]), + ..default_options(test_prompt()) + }; + let result = model.do_generate(&options).await.unwrap(); + + assert!(result.warnings.iter().any(|w| match w { + aimux_core::types::Warning::Unsupported { + feature, details, .. + } => { + feature.contains("groq.browser_search") + && details.as_ref().is_some_and(|d| d.contains("gemma2-9b-it")) + } + _ => false, + })); + } + + /// TS: "should handle mixed tools with model validation" + #[tokio::test] + async fn mixed_tools_with_browser_search() { + let server = MockServer::start().await; + mock_json(&server, text_completion_body()).await; + + let model = model_at(&server, browser_search_model_ids()[0]); + + let func_tool = FunctionTool::new("test-tool", json!({"type":"object","properties":{}})) + .with_description("A test tool"); + let browser_tool = Tool::Provider(aimux_core::tool::ProviderTool { + id: "groq.browser_search".parse().unwrap(), + name: "browser_search".to_string(), + args: serde_json::Map::new(), + supports_deferred_results: None, + }); + let options = CallOptions { + tools: Some(vec![Tool::from(func_tool), browser_tool]), + ..default_options(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert_eq!(body["tools"][0]["type"], "function"); + assert_eq!(body["tools"][0]["function"]["name"], "test-tool"); + assert_eq!(body["tools"][1]["type"], "browser_search"); + } + + /// TS: "should validate all browser search supported models" + #[tokio::test] + async fn browser_search_all_supported_models() { + let model_ids = browser_search_model_ids(); + assert_eq!(model_ids.len(), 2); + for model_id in model_ids { + let server = MockServer::start().await; + mock_json(&server, text_completion_body()).await; + + let model = model_at(&server, model_id); + + let tool = Tool::Provider(aimux_core::tool::ProviderTool { + id: "groq.browser_search".parse().unwrap(), + name: "browser_search".to_string(), + args: serde_json::Map::new(), + supports_deferred_results: None, + }); + let options = CallOptions { + tools: Some(vec![tool]), + ..default_options(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert_eq!(body["tools"][0]["type"], "browser_search"); + } + } } // ════════════════════════════════════════════════════════════════════════════ @@ -894,6 +1044,82 @@ mod prepare_tools { mod do_generate { use super::*; + async fn failed_call(status: u16, body: Value) -> aimux_core::error::ApiCallError { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(status).set_body_json(body)) + .mount(&server) + .await; + match make_provider(&server) + .do_generate(&default_options(test_prompt())) + .await + { + Err(aimux_core::AiMuxError::ApiCall(error)) => *error, + other => panic!("expected an API call error, got {other:?}"), + } + } + + /// `groqFailedResponseHandler` (groq@4.0.52 `groq-error.ts:4-16`): the + /// message is `error.message`, and `error.type` is the provider code. + #[tokio::test] + async fn a_failed_response_reads_groqs_error_schema() { + let error = failed_call( + 400, + json!({ "error": { + "message": "The model `x` does not exist", + "type": "invalid_request_error", + "code": "model_not_found" + } }), + ) + .await; + assert_eq!(error.message, "The model `x` does not exist"); + assert_eq!( + error.provider_code.as_deref(), + Some("invalid_request_error") + ); + assert_eq!(error.status_code, Some(400)); + assert!(!error.is_retryable); + } + + /// A body outside the schema falls back to the status text, as + /// `createJsonErrorResponseHandler` does when the schema does not parse. + #[tokio::test] + async fn a_failed_response_outside_the_schema_uses_the_status_text() { + let error = failed_call(503, json!({ "message": "top-level only" })).await; + assert_eq!(error.message, "Service Unavailable"); + assert_eq!(error.provider_code, None); + assert!(error.is_retryable); + } + + /// upstream: "should reject a response without choices" + #[tokio::test] + async fn rejects_a_response_without_choices() { + let server = MockServer::start().await; + mock_json( + &server, + json!({ + "id": "chatcmpl-empty", + "object": "chat.completion", + "created": 1711115037, + "model": "gemma2-9b-it", + "choices": [], + "usage": {"prompt_tokens": 4, "total_tokens": 4, "completion_tokens": 0} + }), + ) + .await; + let model = make_provider(&server); + + let error = model + .do_generate(&default_options(test_prompt())) + .await + .unwrap_err(); + assert!( + matches!(&error, aimux_core::AiMuxError::InvalidResponseData(m) if m == "Response did not contain any choices."), + "{error:?}" + ); + } + /// TS: "should extract text content" #[tokio::test] async fn extracts_text_content() { @@ -1000,9 +1226,9 @@ mod do_generate { assert_eq!(body["reasoning_effort"], "high"); } - /// v3: minimal 直传(不再归一到 low) + /// upstream: "should coerce top-level reasoning minimal to low" #[tokio::test] - async fn reasoning_effort_minimal_passthrough() { + async fn reasoning_effort_minimal_coerced_to_low() { let server = MockServer::start().await; mock_json(&server, groq_text_body()).await; @@ -1015,12 +1241,12 @@ mod do_generate { model.do_generate(&options).await.unwrap(); let body = first_request_body(&server).await; - assert_eq!(body["reasoning_effort"], "minimal"); + assert_eq!(body["reasoning_effort"], "low"); } - /// v3: xhigh 直传(不再归一到 high) + /// upstream: "should coerce top-level reasoning xhigh to high" #[tokio::test] - async fn reasoning_effort_xhigh_passthrough() { + async fn reasoning_effort_xhigh_coerced_to_high() { let server = MockServer::start().await; mock_json(&server, groq_text_body()).await; @@ -1033,25 +1259,71 @@ mod do_generate { model.do_generate(&options).await.unwrap(); let body = first_request_body(&server).await; - assert_eq!(body["reasoning_effort"], "xhigh"); + assert_eq!(body["reasoning_effort"], "high"); } - /// v3: none 直传为 "none"(不再跳过) + /// upstream: "should map top-level reasoning none to reasoning_effort for Qwen 3.6" #[tokio::test] - async fn reasoning_none_passthrough() { + async fn reasoning_none_maps_for_qwen() { let server = MockServer::start().await; mock_json(&server, groq_text_body()).await; - let model = make_provider(&server); + let model = model_at(&server, "qwen/qwen3.6-27b"); let options = CallOptions { reasoning: Some(ReasoningEffort::None), ..default_options(test_prompt()) }; - model.do_generate(&options).await.unwrap(); + let result = model.do_generate(&options).await.unwrap(); let body = first_request_body(&server).await; assert_eq!(body["reasoning_effort"], "none"); + assert!(result.warnings.is_empty()); + } + + /// upstream: "should omit unsupported top-level reasoning none and warn" + #[tokio::test] + async fn reasoning_none_omitted_and_warns_for_other_models() { + let server = MockServer::start().await; + mock_json(&server, groq_text_body()).await; + + let model = model_at(&server, browser_search_model_ids()[1]); + + let options = CallOptions { + reasoning: Some(ReasoningEffort::None), + ..default_options(test_prompt()) + }; + let result = model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert!(body.get("reasoning_effort").is_none()); + assert!(result.warnings.iter().any(|w| matches!( + w, + aimux_core::types::Warning::Unsupported { feature, details } + if feature == "reasoning" + && details.as_deref() + == Some("reasoning \"none\" is not supported by this model.") + ))); + } + + /// TS: "should prefer providerOptions reasoningEffort over top-level reasoning" + #[tokio::test] + async fn provider_option_reasoning_effort_preferred() { + let server = MockServer::start().await; + mock_json(&server, groq_text_body()).await; + + let model = make_provider(&server); + + let provider_opts = provider_namespace("groq", json!({"reasoningEffort": "high"})); + let options = CallOptions { + reasoning: Some(ReasoningEffort::Medium), + provider_options: Some(provider_opts.unwrap()), + ..default_options(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert_eq!(body["reasoning_effort"], "high"); } /// TS: "should extract usage" @@ -1192,6 +1464,72 @@ mod do_generate { assert_eq!(result.finish_reason.raw.as_deref(), Some("eos")); } + /// TS: "should pass provider options" (reasoningFormat, user, parallelToolCalls) + #[tokio::test] + async fn pass_provider_options() { + let server = MockServer::start().await; + mock_json(&server, groq_text_body()).await; + + let model = make_provider(&server); + + let provider_opts = provider_namespace( + "groq", + json!({ + "reasoningFormat": "hidden", + "user": "test-user-id", + "parallelToolCalls": false + }), + ); + let options = CallOptions { + provider_options: Some(provider_opts.unwrap()), + ..default_options(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert_eq!(body["reasoning_format"], "hidden"); + assert_eq!(body["user"], "test-user-id"); + assert_eq!(body["parallel_tool_calls"], false); + } + + /// TS: "should pass serviceTier provider option" + #[tokio::test] + async fn pass_service_tier_flex() { + let server = MockServer::start().await; + mock_json(&server, groq_text_body()).await; + + let model = make_provider(&server); + + let provider_opts = provider_namespace("groq", json!({"serviceTier": "flex"})); + let options = CallOptions { + provider_options: Some(provider_opts.unwrap()), + ..default_options(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert_eq!(body["service_tier"], "flex"); + } + + /// TS: "should pass performance serviceTier provider option" + #[tokio::test] + async fn pass_service_tier_performance() { + let server = MockServer::start().await; + mock_json(&server, groq_text_body()).await; + + let model = make_provider(&server); + + let provider_opts = provider_namespace("groq", json!({"serviceTier": "performance"})); + let options = CallOptions { + provider_options: Some(provider_opts.unwrap()), + ..default_options(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert_eq!(body["service_tier"], "performance"); + } + /// TS: "should pass tools and toolChoice" #[tokio::test] async fn pass_tools_and_tool_choice() { @@ -1255,6 +1593,92 @@ mod do_generate { assert_eq!(body["response_format"]["json_schema"]["name"], "test-name"); assert_eq!(body["response_format"]["json_schema"]["strict"], true); } + + /// TS: "should pass response format as json_object when structuredOutputs + /// explicitly disabled" + #[tokio::test] + async fn response_format_json_object_when_disabled() { + let server = MockServer::start().await; + mock_json(&server, groq_text_body()).await; + + let model = make_provider(&server); + + let provider_opts = provider_namespace("groq", json!({"structuredOutputs": false})); + let options = CallOptions { + provider_options: Some(provider_opts.unwrap()), + response_format: Some(ResponseFormat::Json { + schema: Some(json!({ + "type": "object", + "properties": { "value": { "type": "string" } }, + "required": ["value"], + "additionalProperties": false, + })), + name: Some("test-name".to_string()), + description: Some("test description".to_string()), + }), + ..default_options(test_prompt()) + }; + let result = model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert_eq!(body["response_format"]["type"], "json_object"); + + // Should have a warning about structuredOutputs + assert!(result.warnings.iter().any(|w| match w { + aimux_core::types::Warning::Unsupported { feature, .. } => feature == "responseFormat", + _ => false, + })); + } + + /// TS: "should send strict: false when strictJsonSchema is explicitly disabled" + #[tokio::test] + async fn strict_json_schema_false() { + let server = MockServer::start().await; + mock_json(&server, groq_text_body()).await; + + let model = make_provider(&server); + + let provider_opts = provider_namespace("groq", json!({"strictJsonSchema": false})); + let options = CallOptions { + provider_options: Some(provider_opts.unwrap()), + response_format: Some(ResponseFormat::Json { + schema: Some(json!({ + "type": "object", + "properties": { "value": { "type": "string" } }, + "required": ["value"], + "additionalProperties": false, + })), + name: Some("test-name".to_string()), + description: Some("test description".to_string()), + }), + ..default_options(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let body = first_request_body(&server).await; + assert_eq!(body["response_format"]["json_schema"]["strict"], false); + } + + /// TS: "should send request body" (request.body) + #[tokio::test] + async fn request_body_string() { + let server = MockServer::start().await; + mock_json(&server, groq_text_body()).await; + + let model = make_provider(&server); + + let result = model + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + + let request_body = result + .request + .and_then(|request| request.body) + .expect("should have request body"); + assert_eq!(request_body["model"], "gemma2-9b-it"); + assert_eq!(request_body["messages"][0]["content"], "Hello"); + } } // ════════════════════════════════════════════════════════════════════════════ @@ -1428,7 +1852,7 @@ mod do_stream { let (id, name, input) = tool_call.expect("should have tool call"); assert_eq!(id, "call_abc"); assert_eq!(name, "test-tool"); - assert_eq!(input, r#"{"value":"Sparkle Day"}"#); + assert_eq!(input, Value::String(r#"{"value":"Sparkle Day"}"#.into())); // Should have tool-calls finish reason let finish = parts.iter().find_map(|p| match p { @@ -1450,10 +1874,55 @@ mod do_stream { let model = make_provider(&server); let result = model.do_stream(&default_options(test_prompt())).await; - // Groq error in first chunk should reject the promise + // 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()); } + /// upstream: "should handle error stream parts" (an error chunk after the + /// first event, which stays in the stream) + #[tokio::test] + async fn error_chunk_is_an_error_part_with_status() { + let server = MockServer::start().await; + let body = sse_body(&[ + &sse_event( + r#"{"id":"c","model":"m","choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":null}]}"#, + ), + &sse_event(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 + .unwrap(); + let parts = collect_stream(result).await; + + let error = parts + .iter() + .find_map(|p| match p { + StreamPart::Error { error } => Some(error), + _ => None, + }) + .expect("an error part"); + match error { + aimux_core::AiMuxError::ApiCall(call) => { + assert_eq!(call.message, "Rate limit reached"); + assert_eq!(call.status_code, Some(429)); + assert!(error.is_retryable()); + } + other => panic!("expected ApiCall, got {other:?}"), + } + match parts.last().unwrap() { + StreamPart::Finish { finish_reason, .. } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Error); + } + other => panic!("expected Finish, got {other:?}"), + } + } + /// TS: "should stream tool call that is sent in one chunk" #[tokio::test] async fn streams_tool_call_one_chunk() { @@ -1485,7 +1954,10 @@ mod do_stream { StreamPart::ToolCall(RawToolCall { input, .. }) => Some(input.clone()), _ => None, }); - assert_eq!(tool_call.unwrap(), r#"{"value":"Sparkle Day"}"#); + assert_eq!( + tool_call.unwrap(), + Value::String(r#"{"value":"Sparkle Day"}"#.into()) + ); } /// TS: "should stream usage from x_groq.usage" @@ -1523,6 +1995,127 @@ mod do_stream { assert_eq!(raw["prompt_tokens"], json!(10)); assert_eq!(raw["total_tokens"], json!(15)); } + + /// upstream: "should keep reasoning active when deltas include empty tool calls" + #[tokio::test] + async fn keeps_reasoning_active_with_empty_tool_calls() { + let server = MockServer::start().await; + let body = sse_body(&[ + &sse_event( + r#"{"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test-model","choices":[{"index":0,"delta":{"role":"assistant","content":"","reasoning":"Think ","tool_calls":[]},"finish_reason":null}]}"#, + ), + &sse_event( + r#"{"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test-model","choices":[{"index":0,"delta":{"content":"","reasoning":"more...","tool_calls":[]},"finish_reason":null}]}"#, + ), + &sse_event( + r#"{"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"test-model","choices":[{"index":0,"delta":{"content":"Hello","reasoning":"","tool_calls":[]},"finish_reason":"stop"}]}"#, + ), + ]); + mock_sse(&server, body).await; + let model = make_provider(&server); + + let result = model + .do_stream(&default_options(test_prompt())) + .await + .unwrap(); + let parts = collect_stream(result).await; + let reasoning: Vec = parts + .iter() + .filter_map(|p| match p { + StreamPart::ReasoningStart { id, .. } => Some(format!("start {id}")), + StreamPart::ReasoningDelta { id, delta, .. } => Some(format!("delta {id} {delta}")), + StreamPart::ReasoningEnd { id, .. } => Some(format!("end {id}")), + _ => None, + }) + .collect(); + assert_eq!( + reasoning, + [ + "start reasoning-0", + "delta reasoning-0 Think ", + "delta reasoning-0 more...", + "end reasoning-0" + ] + ); + } + + /// upstream: "should handle unparsable stream parts" + #[tokio::test] + async fn unparsable_chunk_is_an_error_part_and_finish_reason_error() { + let server = MockServer::start().await; + mock_sse( + &server, + "data: {unparsable}\n\ndata: [DONE]\n\n".to_string(), + ) + .await; + let model = make_provider(&server); + + let options = CallOptions { + include_raw_chunks: Some(true), + ..default_options(test_prompt()) + }; + let result = model.do_stream(&options).await.unwrap(); + let parts = collect_stream(result).await; + + assert!(matches!(parts[0], StreamPart::StreamStart { .. })); + assert!(matches!(&parts[1], StreamPart::Raw { raw_value } if raw_value.is_null())); + assert!(matches!(parts[2], StreamPart::Error { .. }), "{parts:?}"); + match &parts[3] { + StreamPart::Finish { + finish_reason, + usage, + .. + } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Error); + assert_eq!(finish_reason.raw, None); + assert_eq!(usage.input_tokens.total, None); + } + other => panic!("expected Finish, got {other:?}"), + } + assert_eq!(parts.len(), 4); + } + + /// upstream: "should stream raw chunks when includeRawChunks is true" + #[tokio::test] + async fn streams_raw_chunks() { + let server = MockServer::start().await; + let chunks = [ + r#"{"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gemma2-9b-it","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}"#, + r#"{"id":"chatcmpl-456","object":"chat.completion.chunk","created":1234567890,"model":"gemma2-9b-it","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}"#, + r#"{"id":"chatcmpl-789","object":"chat.completion.chunk","created":1234567890,"model":"gemma2-9b-it","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"x_groq":{"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}}"#, + ]; + mock_sse( + &server, + sse_body( + &chunks + .map(sse_event) + .iter() + .map(String::as_str) + .collect::>(), + ), + ) + .await; + let model = make_provider(&server); + + let options = CallOptions { + include_raw_chunks: Some(true), + ..default_options(test_prompt()) + }; + let result = model.do_stream(&options).await.unwrap(); + let parts = collect_stream(result).await; + + let raw: Vec<&Value> = parts + .iter() + .filter_map(|p| match p { + StreamPart::Raw { raw_value } => Some(raw_value), + _ => None, + }) + .collect(); + assert_eq!(raw.len(), 3); + for (raw, chunk) in raw.iter().zip(chunks) { + assert_eq!(**raw, serde_json::from_str::(chunk).unwrap()); + } + } } // ════════════════════════════════════════════════════════════════════════════ @@ -1577,3 +2170,196 @@ mod auth { assert_eq!(auth, "Bearer test-api-key"); } } + +// ════════════════════════════════════════════════════════════════════════════ +// The package: identity, namespace, metadata, usage, cassettes +// ════════════════════════════════════════════════════════════════════════════ + +mod package { + use super::*; + + use serial_test::serial; + + use aimux_core::AiMuxError; + use aimux_core::provider::Provider; + + fn groq_at(server: &MockServer) -> aimux_providers::groq::GroqProvider { + create_groq(GroqProviderSettings { + base_url: Some(server.uri()), + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap() + } + + fn options_with(provider_options: Value) -> CallOptions { + let mut options = default_options(test_prompt()); + options.provider_options = Some(serde_json::from_value(provider_options).unwrap()); + options + } + + #[test] + fn every_entry_point_reports_groq_chat() { + assert_eq!( + groq().chat("llama-3.3-70b-versatile").provider(), + "groq.chat" + ); + assert_eq!(groq().call("m").provider(), "groq.chat"); + assert_eq!(groq().language_model("m").unwrap().provider(), "groq.chat"); + // The registry entry point and the preset build the same identity. + let model = + aimux_providers::create_provider("groq", aimux_providers::PresetSettings::default()) + .unwrap() + .language_model("m") + .unwrap(); + assert_eq!(model.provider(), "groq.chat"); + } + + #[test] + fn embedding_and_image_models_are_no_such_model() { + let provider = groq(); + for result in [ + provider.embedding_model("e").map(|_| ()), + provider.image_model("i").map(|_| ()), + ] { + assert!(matches!(result, Err(AiMuxError::NoSuchModel { .. }))); + } + } + + #[tokio::test] + async fn only_the_groq_namespace_is_read() { + let server = MockServer::start().await; + mock_json(&server, text_completion_body()).await; + let model = groq_at(&server).chat("gemma2-9b-it"); + + let result = model + .do_generate(&options_with(json!({ + "groq": {"user": "u1", "reasoningFormat": "parsed", "serviceTier": "flex", + "parallelToolCalls": false}, + "openai": {"user": "ignored"}, + "openaiCompatible": {"reasoningEffort": "low"}, + }))) + .await + .unwrap(); + + let body = result.request.and_then(|request| request.body).unwrap(); + assert_eq!(body["user"], "u1"); + assert_eq!(body["reasoning_format"], "parsed"); + assert_eq!(body["service_tier"], "flex"); + assert_eq!(body["parallel_tool_calls"], json!(false)); + // upstream: parseProviderOptions({ provider: 'groq' }) reads that namespace only + assert!(body.get("reasoning_effort").is_none()); + } + + #[tokio::test] + async fn streaming_asks_for_no_stream_options_and_reports_x_groq_usage() { + let server = MockServer::start().await; + let body = sse_body(&[ + &sse_event( + r#"{"id":"c","model":"m","choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":null}]}"#, + ), + &sse_event( + r#"{"id":"c","model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"x_groq":{"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,"prompt_tokens_details":{"cached_tokens":4}}}}"#, + ), + ]); + mock_sse(&server, body).await; + let model = groq_at(&server).chat("m"); + + let result = model + .do_stream(&default_options(test_prompt())) + .await + .unwrap(); + assert!( + result + .request + .as_ref() + .and_then(|request| request.body.as_ref()) + .unwrap() + .get("stream_options") + .is_none() + ); + let parts = collect_stream(result).await; + + let (usage, metadata) = parts + .iter() + .find_map(|p| match p { + StreamPart::Finish { + usage, + provider_metadata, + .. + } => Some((usage.clone(), provider_metadata.clone())), + _ => None, + }) + .expect("a finish part"); + assert_eq!(usage.input_tokens.total, Some(10)); + assert_eq!(usage.input_tokens.cache_read, Some(4)); + assert_eq!(usage.input_tokens.no_cache, Some(6)); + assert_eq!(usage.output_tokens.total, Some(5)); + // upstream: the finish part carries no providerMetadata + assert!(metadata.is_none()); + } + + #[tokio::test] + async fn the_token_limit_is_max_tokens() { + let server = MockServer::start().await; + mock_json(&server, text_completion_body()).await; + let model = groq_at(&server).chat("gemma2-9b-it"); + let mut options = default_options(test_prompt()); + 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(); + // upstream: getArgs sends `max_tokens: maxOutputTokens` + assert_eq!(body["max_tokens"], 64); + assert!(body.get("max_completion_tokens").is_none()); + assert!(body.get("top_k").is_none(), "Groq has no top_k"); + assert!( + result + .warnings + .iter() + .any(|w| matches!(w, aimux_core::types::Warning::Unsupported { feature, .. } if feature == "topK")) + ); + } + + #[serial] + #[tokio::test] + async fn the_key_is_read_from_the_environment_per_request_not_at_creation() { + let saved = std::env::var("GROQ_API_KEY").ok(); + unsafe { std::env::remove_var("GROQ_API_KEY") }; + let server = MockServer::start().await; + mock_json(&server, text_completion_body()).await; + let provider = create_groq(GroqProviderSettings { + base_url: Some(server.uri()), + ..Default::default() + }) + .expect("no key is read at creation"); + let model = provider.chat("m"); + + let error = model + .do_generate(&default_options(test_prompt())) + .await + .unwrap_err(); + assert!( + matches!(&error, AiMuxError::LoadApiKey { env_var, .. } if env_var == "GROQ_API_KEY"), + "{error:?}" + ); + assert!(server.received_requests().await.unwrap().is_empty()); + + unsafe { std::env::set_var("GROQ_API_KEY", "env-key") }; + model + .do_generate(&default_options(test_prompt())) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests[0].headers.get("authorization").unwrap(), + "Bearer env-key" + ); + unsafe { + match saved { + Some(v) => std::env::set_var("GROQ_API_KEY", v), + None => std::env::remove_var("GROQ_API_KEY"), + } + } + } +} diff --git a/aimux-providers/tests/huggingface_responses_test.rs b/aimux-providers/tests/huggingface_responses_test.rs index e7168f230..3568a330e 100644 --- a/aimux-providers/tests/huggingface_responses_test.rs +++ b/aimux-providers/tests/huggingface_responses_test.rs @@ -16,7 +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, ToolResult}; +use aimux_core::tool::RawToolCall; +use aimux_core::tool::ToolResult; use futures::StreamExt; use serde_json::{Value, json}; @@ -27,16 +28,18 @@ use aimux_core::error::AiMuxError; use aimux_core::generate::{GenerateTextOptions, generate_text}; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{ - FilePart, LanguageModelMessage, LanguageModelPrompt, UserPart, + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, ReasoningPart, TextPart, + ToolCallPart, ToolPart, ToolResultOutput, ToolResultPart, UserPart, }; -use aimux_core::options::CallOptions; -use aimux_core::result::{GenerateContent, Source}; +use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat, Tool}; +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; -use aimux_providers::{HuggingFaceConfig, HuggingFaceProvider}; +use aimux_providers::{HuggingFaceProvider, HuggingFaceProviderSettings, create_huggingface}; // ───────────────────────────────────────────────────────────────────────────── // Shared helpers @@ -54,8 +57,12 @@ fn default_options(prompt: LanguageModelPrompt) -> CallOptions { /// Build a Hugging Face provider whose base URL points at the mock server. fn make_provider(server: &MockServer) -> HuggingFaceProvider { - let config = HuggingFaceConfig::new("APIKEY").with_base_url(server.uri()); - HuggingFaceProvider::new(config) + create_huggingface(HuggingFaceProviderSettings { + api_key: Some("APIKEY".to_string().into()), + base_url: Some(server.uri()), + ..Default::default() + }) + .expect("valid settings") } /// Mount a JSON response on `/responses`. @@ -158,7 +165,7 @@ async fn should_generate_text() { mock_json(&server, basic_response_body()).await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_generate(&default_options(test_prompt())) @@ -181,7 +188,7 @@ async fn should_extract_usage() { mock_json(&server, basic_response_body()).await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_generate(&default_options(test_prompt())) @@ -238,7 +245,7 @@ async fn should_extract_text_from_output_array_when_output_text_missing() { .await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_generate(&default_options(test_prompt())) @@ -284,7 +291,7 @@ async fn should_handle_missing_usage_gracefully() { .await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_generate(&default_options(test_prompt())) @@ -301,6 +308,48 @@ async fn should_handle_missing_usage_gracefully() { assert_eq!(result.usage.output_tokens.reasoning, None); } +/// TS: doGenerate › basic text response › "should send model id, settings, and +/// input" +#[tokio::test] +async fn should_send_model_id_settings_and_input() { + let server = MockServer::start().await; + mock_json(&server, basic_response_body()).await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); + + let prompt = vec![ + LanguageModelMessage::System { + content: "You are a helpful assistant.".into(), + provider_options: None, + }, + LanguageModelMessage::user_text("Hello"), + ]; + let mut options = default_options(prompt); + options.temperature = Some(0.5); + options.top_p = Some(0.3); + options.max_output_tokens = Some(100); + + let result = model.do_generate(&options).await.expect("should succeed"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + + assert_eq!(body["model"], json!("deepseek-ai/DeepSeek-V3-0324")); + assert_eq!(body["temperature"], json!(0.5)); + assert_eq!(body["top_p"], json!(0.3)); + assert_eq!(body["max_output_tokens"], json!(100)); + assert_eq!(body["stream"], json!(false)); + assert_eq!( + body["input"], + json!([ + { "role": "system", "content": "You are a helpful assistant." }, + { "role": "user", "content": [{ "type": "input_text", "text": "Hello" }] } + ]) + ); +} + /// TS: doGenerate › basic text response › "should handle unsupported settings /// with warnings" #[tokio::test] @@ -309,7 +358,7 @@ async fn should_handle_unsupported_settings_with_warnings() { mock_json(&server, basic_response_body()).await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let mut options = default_options(test_prompt()); options.top_k = Some(10.0); @@ -388,7 +437,7 @@ async fn should_generate_text_and_sources_from_annotations() { .await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_generate(&default_options(test_prompt())) @@ -476,7 +525,7 @@ async fn should_handle_mcp_tools_with_annotations() { .await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_generate(&default_options(test_prompt())) @@ -496,7 +545,10 @@ async fn should_handle_mcp_tools_with_annotations() { }) => { assert_eq!(tool_call_id, "mcp_search_test"); assert_eq!(tool_name, "search"); - assert_eq!(input, r#"{"query": "San Francisco tech events"}"#); + assert_eq!( + input, + &Value::String(r#"{"query": "San Francisco tech events"}"#.into()) + ); } other => panic!("expected ToolCall at [0], got {other:?}"), } @@ -569,7 +621,7 @@ async fn should_stream_text_deltas() { mock_sse(&server, chunks).await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_stream(&default_options(test_prompt())) @@ -650,7 +702,7 @@ async fn should_handle_streaming_without_usage() { mock_sse(&server, chunks).await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_stream(&default_options(test_prompt())) @@ -690,7 +742,7 @@ async fn should_handle_non_message_item_types() { mock_sse(&server, chunks).await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_stream(&default_options(test_prompt())) @@ -734,7 +786,7 @@ async fn should_handle_streaming_errors() { mock_sse(&server, chunks).await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_stream(&default_options(test_prompt())) @@ -762,6 +814,36 @@ async fn should_handle_streaming_errors() { ))); } +/// TS: doStream › "should send correct streaming request" +#[tokio::test] +async fn should_send_correct_streaming_request() { + let server = MockServer::start().await; + let chunks = sse_body(&[sse_event( + r#"{"type":"response.completed","response":{"id":"resp_test","status":"completed"},"sequence_number":1}"#, + )]); + mock_sse(&server, chunks).await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); + + let mut options = default_options(test_prompt()); + options.temperature = Some(0.7); + + let result = model.do_stream(&options).await.expect("should succeed"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + + assert_eq!(body["model"], json!("deepseek-ai/DeepSeek-V3-0324")); + assert_eq!(body["temperature"], json!(0.7)); + assert_eq!(body["stream"], json!(true)); + assert_eq!( + body["input"], + json!([{ "role": "user", "content": [{ "type": "input_text", "text": "Hello" }] }]) + ); +} + // ════════════════════════════════════════════════════════════════════════════ // Message conversion // ════════════════════════════════════════════════════════════════════════════ @@ -794,6 +876,51 @@ async fn stub_empty_response(server: &MockServer) { .await; } +/// TS: message conversion › "should convert user messages with images" +#[tokio::test] +async fn should_convert_user_messages_with_images() { + let server = MockServer::start().await; + stub_empty_response(&server).await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); + + let prompt = vec![LanguageModelMessage::User { + content: vec![ + UserPart::Text(TextPart { + text: "What do you see?".into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Base64("AQIDBA==".into()), + }, + media_type: "image/jpeg".into(), + filename: None, + provider_options: None, + }), + ], + provider_options: None, + }]; + + let result = model + .do_generate(&default_options(prompt)) + .await + .expect("ok"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + + assert_eq!( + body["input"][0]["content"], + json!([ + { "text": "What do you see?", "type": "input_text" }, + { "image_url": "data:image/jpeg;base64,AQIDBA==", "type": "input_image" } + ]) + ); +} + /// TS: message conversion › "should throw for file parts with provider /// references" #[tokio::test] @@ -802,7 +929,7 @@ async fn should_throw_for_file_parts_with_provider_references() { stub_empty_response(&server).await; let provider = make_provider(&server); - let model = provider.responses_model("Qwen/Qwen2.5-VL-32B-Instruct"); + let model = provider.responses("Qwen/Qwen2.5-VL-32B-Instruct"); let prompt = vec![LanguageModelMessage::User { content: vec![UserPart::File(FilePart { @@ -827,6 +954,139 @@ async fn should_throw_for_file_parts_with_provider_references() { ); } +/// TS: message conversion › "should handle assistant messages" +#[tokio::test] +async fn should_handle_assistant_messages() { + let server = MockServer::start().await; + stub_empty_response(&server).await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); + + let prompt = vec![ + LanguageModelMessage::user_text("Hello"), + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: "Hi there!".into(), + provider_options: None, + })], + provider_options: None, + }, + LanguageModelMessage::user_text("How are you?"), + ]; + + let result = model + .do_generate(&default_options(prompt)) + .await + .expect("ok"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + + assert_eq!( + body["input"], + json!([ + { "content": [{ "text": "Hello", "type": "input_text" }], "role": "user" }, + { "content": [{ "text": "Hi there!", "type": "output_text" }], "role": "assistant" }, + { "content": [{ "text": "How are you?", "type": "input_text" }], "role": "user" } + ]) + ); +} + +/// TS: message conversion › "should warn about unsupported assistant content +/// types" — tool-call, tool-result, and reasoning in assistant messages +/// produce NO warnings (they are silently skipped). +#[tokio::test] +async fn should_not_warn_about_assistant_content_types() { + let server = MockServer::start().await; + stub_empty_response(&server).await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); + + let prompt = vec![LanguageModelMessage::Assistant { + content: vec![ + AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "test".into(), + tool_name: "test".into(), + input: json!({}), + provider_executed: None, + provider_options: None, + }), + AssistantPart::ToolResult(ToolResultPart { + tool_call_id: "test".into(), + tool_name: "test".into(), + output: ToolResultOutput::Text { + value: "test".into(), + provider_options: None, + }, + + provider_options: None, + }), + AssistantPart::Reasoning(ReasoningPart { + text: "thinking...".into(), + + provider_options: None, + }), + ], + provider_options: None, + }]; + + let result = model + .do_generate(&default_options(prompt)) + .await + .expect("ok"); + + // TS expects no warnings (tool calls/results/reasoning are silently + // skipped, reasoning is included as output_text). + assert!( + result.warnings.is_empty(), + "expected no warnings, got: {:?}", + result.warnings + ); +} + +/// TS: message conversion › "should warn about tool messages" +#[tokio::test] +async fn should_warn_about_tool_messages() { + let server = MockServer::start().await; + stub_empty_response(&server).await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); + + let prompt = vec![LanguageModelMessage::Tool { + content: vec![ToolPart::ToolResult(ToolResultPart { + tool_call_id: "test".into(), + tool_name: "test".into(), + output: ToolResultOutput::Text { + value: "test".into(), + provider_options: None, + }, + + provider_options: None, + })], + provider_options: None, + }]; + + let result = model + .do_generate(&default_options(prompt)) + .await + .expect("ok"); + + let features: Vec = result + .warnings + .iter() + .filter_map(|w| match w { + aimux_core::types::Warning::Unsupported { feature, .. } => Some(feature.clone()), + _ => None, + }) + .collect(); + + assert_eq!(features, vec!["tool messages"]); +} + // ════════════════════════════════════════════════════════════════════════════ // Tool calls // ════════════════════════════════════════════════════════════════════════════ @@ -878,7 +1138,7 @@ async fn should_handle_function_call_tool_responses() { .await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_generate(&default_options(test_prompt())) @@ -897,7 +1157,7 @@ async fn should_handle_function_call_tool_responses() { }) => { assert_eq!(tool_call_id, "call_123"); assert_eq!(tool_name, "getWeather"); - assert_eq!(input, r#"{"location": "New York"}"#); + assert_eq!(input, &Value::String(r#"{"location": "New York"}"#.into())); } other => panic!("expected ToolCall at [0], got {other:?}"), } @@ -936,7 +1196,7 @@ async fn should_stream_tool_calls() { mock_sse(&server, chunks).await; let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); let result = model .do_stream(&default_options(test_prompt())) @@ -979,7 +1239,10 @@ async fn should_stream_tool_calls() { }) => { assert_eq!(tool_call_id, "call_456"); assert_eq!(tool_name, "calculator"); - assert_eq!(input, r#"{"operation": "add", "a": 5, "b": 3}"#); + assert_eq!( + input, + &Value::String(r#"{"operation": "add", "a": 5, "b": 3}"#.into()) + ); } other => panic!("expected ToolCall, got {other:?}"), } @@ -1012,18 +1275,566 @@ async fn should_stream_tool_calls() { // Structured output // ════════════════════════════════════════════════════════════════════════════ +/// TS: structured output › "should send text.format for structured output" +#[tokio::test] +async fn should_send_text_format_for_structured_output() { + let server = MockServer::start().await; + mock_json( + &server, + json!({ + "id": "resp_structured", + "model": "moonshotai/Kimi-K2-Instruct", + "object": "response", + "created_at": 1741257730, + "status": "completed", + "error": null, + "instructions": null, + "max_output_tokens": null, + "metadata": null, + "tool_choice": "auto", + "tools": [], + "temperature": 1.0, + "top_p": 1.0, + "incomplete_details": null, + "usage": null, + "output": [ + { + "id": "msg_structured", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { "type": "output_text", "text": "{\"name\": \"John Doe\", \"age\": 30}" } + ] + } + ], + "output_text": null + }), + ) + .await; + + let provider = make_provider(&server); + let model = provider.responses("moonshotai/Kimi-K2-Instruct"); + + let mut options = default_options(test_prompt()); + options.response_format = Some(ResponseFormat::Json { + schema: Some(json!({ + "type": "object", + "properties": { + "name": { "type": "string" }, + "age": { "type": "number" } + }, + "required": ["name", "age"] + })), + name: None, + description: None, + }); + + let result = model.do_generate(&options).await.expect("should succeed"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + + assert_eq!( + body["text"]["format"], + json!({ + "type": "json_schema", + "strict": false, + "name": "response", + "schema": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "age": { "type": "number" } + }, + "required": ["name", "age"] + } + }) + ); + // description should NOT be present (it was None). + assert!( + body["text"]["format"].get("description").is_none(), + "description should be absent when None" + ); +} + +/// TS: structured output › "should handle structured output with custom name +/// and description" +#[tokio::test] +async fn should_handle_structured_output_with_custom_name_and_description() { + let server = MockServer::start().await; + mock_json( + &server, + json!({ + "id": "resp_structured", + "model": "moonshotai/Kimi-K2-Instruct", + "object": "response", + "created_at": 1741257730, + "status": "completed", + "error": null, + "instructions": null, + "max_output_tokens": null, + "metadata": null, + "tool_choice": "auto", + "tools": [], + "temperature": 1.0, + "top_p": 1.0, + "incomplete_details": null, + "usage": null, + "output": [], + "output_text": "{}" + }), + ) + .await; + + let provider = make_provider(&server); + let model = provider.responses("moonshotai/Kimi-K2-Instruct"); + + let mut options = default_options(test_prompt()); + options.response_format = Some(ResponseFormat::Json { + schema: Some(json!({ "type": "object", "properties": { "name": { "type": "string" } } })), + name: Some("person_profile".to_string()), + description: Some("A person profile with basic information".to_string()), + }); + + let result = model.do_generate(&options).await.expect("should succeed"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + + assert_eq!(body["text"]["format"]["name"], json!("person_profile")); + assert_eq!( + body["text"]["format"]["description"], + json!("A person profile with basic information") + ); +} + // ════════════════════════════════════════════════════════════════════════════ // Reasoning // ════════════════════════════════════════════════════════════════════════════ +/// TS: reasoning › "should handle reasoning content in responses" +#[tokio::test] +async fn should_handle_reasoning_content_in_responses() { + let server = MockServer::start().await; + mock_json( + &server, + json!({ + "id": "resp_reasoning", + "model": "deepseek-ai/DeepSeek-R1", + "object": "response", + "created_at": 1741257730, + "status": "completed", + "error": null, + "instructions": null, + "max_output_tokens": null, + "metadata": null, + "tool_choice": "auto", + "tools": [], + "temperature": 1.0, + "top_p": 1.0, + "incomplete_details": null, + "usage": { "input_tokens": 10, "output_tokens": 50, "total_tokens": 60 }, + "output": [ + { + "id": "reasoning_1", + "type": "reasoning", + "content": [ + { "type": "reasoning_text", "text": "Let me think about this problem step by step..." } + ] + }, + { + "id": "msg_after_reasoning", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { "type": "output_text", "text": "The answer is 42." } + ] + } + ], + "output_text": null + }), + ) + .await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-R1"); + + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("should succeed"); + + assert_eq!(result.content.len(), 2); + match &result.content[0] { + GenerateContent::Reasoning(ReasoningOutput { + text, + provider_metadata, + }) => { + assert_eq!(text, "Let me think about this problem step by step..."); + assert_eq!( + provider_metadata.as_ref(), + Some( + &aimux_core::shared::provider_namespace( + "huggingface", + json!({ "itemId": "reasoning_1" }) + ) + .expect("provider metadata object") + ) + ); + } + other => panic!("expected Reasoning at [0], got {other:?}"), + } + match &result.content[1] { + GenerateContent::Text { text, .. } => assert_eq!(text, "The answer is 42."), + other => panic!("expected Text at [1], got {other:?}"), + } +} + +/// TS: reasoning › "should stream reasoning content" +#[tokio::test] +async fn should_stream_reasoning_content() { + let server = MockServer::start().await; + let chunks = sse_body(&[ + sse_event( + r#"{"type":"response.created","response":{"id":"resp_reasoning_stream","object":"response","created_at":1741269019,"status":"in_progress","model":"deepseek-ai/DeepSeek-R1"}}"#, + ), + sse_event( + r#"{"type":"response.output_item.added","output_index":0,"item":{"id":"reasoning_stream","type":"reasoning"},"sequence_number":1}"#, + ), + sse_event( + r#"{"type":"response.reasoning_text.delta","item_id":"reasoning_stream","output_index":0,"content_index":0,"delta":"Thinking about","sequence_number":2}"#, + ), + sse_event( + r#"{"type":"response.reasoning_text.delta","item_id":"reasoning_stream","output_index":0,"content_index":0,"delta":" the problem...","sequence_number":3}"#, + ), + sse_event( + r#"{"type":"response.reasoning_text.done","item_id":"reasoning_stream","output_index":0,"content_index":0,"text":"Thinking about the problem...","sequence_number":4}"#, + ), + sse_event( + r#"{"type":"response.output_item.done","output_index":0,"item":{"id":"reasoning_stream","type":"reasoning","content":[{"type":"reasoning_text","text":"Thinking about the problem..."}]},"sequence_number":5}"#, + ), + sse_event( + r#"{"type":"response.output_item.added","output_index":1,"item":{"id":"msg_stream","type":"message","role":"assistant"},"sequence_number":6}"#, + ), + sse_event( + r#"{"type":"response.output_text.delta","item_id":"msg_stream","output_index":1,"content_index":0,"delta":"The solution is","sequence_number":7}"#, + ), + sse_event( + r#"{"type":"response.output_text.delta","item_id":"msg_stream","output_index":1,"content_index":0,"delta":" simple.","sequence_number":8}"#, + ), + sse_event( + r#"{"type":"response.output_item.done","output_index":1,"item":{"id":"msg_stream","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"The solution is simple."}]},"sequence_number":9}"#, + ), + sse_event( + r#"{"type":"response.completed","response":{"id":"resp_reasoning_stream","status":"completed","usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}},"sequence_number":10}"#, + ), + ]); + mock_sse(&server, chunks).await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-R1"); + + let result = model + .do_stream(&default_options(test_prompt())) + .await + .expect("should succeed"); + + let parts = collect_stream(result).await; + + // Expected: stream-start, response-metadata, reasoning-start, + // reasoning-delta×2, reasoning-end, text-start, text-delta×2, text-end, + // finish + assert_eq!(parts.len(), 11); + + assert!(matches!(&parts[0], StreamPart::StreamStart { .. })); + + match &parts[1] { + StreamPart::ResponseMetadata(ResponseMetadata { id, model_id, .. }) => { + assert_eq!(id.as_deref(), Some("resp_reasoning_stream")); + assert_eq!(model_id.as_deref(), Some("deepseek-ai/DeepSeek-R1")); + } + other => panic!("expected ResponseMetadata, got {other:?}"), + } + + match &parts[2] { + StreamPart::ReasoningStart { + id, + provider_metadata, + } => { + assert_eq!(id, "reasoning_stream"); + assert_eq!( + provider_metadata.as_ref(), + Some( + &aimux_core::shared::provider_namespace( + "huggingface", + json!({ "itemId": "reasoning_stream" }) + ) + .expect("provider metadata object") + ) + ); + } + other => panic!("expected ReasoningStart, got {other:?}"), + } + match &parts[3] { + StreamPart::ReasoningDelta { id, delta, .. } => { + assert_eq!(id, "reasoning_stream"); + assert_eq!(delta, "Thinking about"); + } + other => panic!("expected ReasoningDelta, got {other:?}"), + } + match &parts[4] { + StreamPart::ReasoningDelta { id, delta, .. } => { + assert_eq!(id, "reasoning_stream"); + assert_eq!(delta, " the problem..."); + } + other => panic!("expected ReasoningDelta, got {other:?}"), + } + match &parts[5] { + StreamPart::ReasoningEnd { id, .. } => assert_eq!(id, "reasoning_stream"), + other => panic!("expected ReasoningEnd, got {other:?}"), + } + match &parts[6] { + StreamPart::TextStart { id, .. } => assert_eq!(id, "msg_stream"), + other => panic!("expected TextStart, got {other:?}"), + } + match &parts[7] { + StreamPart::TextDelta { id, delta, .. } => { + assert_eq!(id, "msg_stream"); + assert_eq!(delta, "The solution is"); + } + other => panic!("expected TextDelta, got {other:?}"), + } + match &parts[8] { + StreamPart::TextDelta { id, delta, .. } => { + assert_eq!(id, "msg_stream"); + assert_eq!(delta, " simple."); + } + other => panic!("expected TextDelta, got {other:?}"), + } + match &parts[9] { + StreamPart::TextEnd { id, .. } => assert_eq!(id, "msg_stream"), + other => panic!("expected TextEnd, got {other:?}"), + } + match &parts[10] { + StreamPart::Finish { + finish_reason, + usage, + .. + } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(usage.input_tokens.total, Some(10)); + assert_eq!(usage.output_tokens.total, Some(20)); + } + other => panic!("expected Finish, got {other:?}"), + } +} + // ════════════════════════════════════════════════════════════════════════════ // Provider options // ════════════════════════════════════════════════════════════════════════════ +/// TS: provider options › "should send provider-specific options" +#[tokio::test] +async fn should_send_provider_specific_options() { + let server = MockServer::start().await; + mock_json( + &server, + json!({ + "id": "resp_provider_options", + "model": "deepseek-ai/DeepSeek-V3-0324", + "object": "response", + "created_at": 1741257730, + "status": "completed", + "error": null, + "instructions": null, + "max_output_tokens": null, + "metadata": null, + "tool_choice": "auto", + "tools": [], + "temperature": 1.0, + "top_p": 1.0, + "incomplete_details": null, + "usage": null, + "output": [], + "output_text": "Test" + }), + ) + .await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); + + let mut options = default_options(test_prompt()); + let mut po = aimux_core::shared::SharedProviderOptions::new(); + po.insert( + "huggingface".to_string(), + serde_json::from_value(json!({ + "metadata": { "key": "value" }, + "instructions": "Be concise", + "strictJsonSchema": true + })) + .unwrap(), + ); + options.provider_options = Some(po); + options.response_format = Some(ResponseFormat::Json { + schema: Some(json!({ "type": "object" })), + name: None, + description: None, + }); + + let result = model.do_generate(&options).await.expect("should succeed"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + + assert_eq!(body["metadata"], json!({ "key": "value" })); + assert_eq!(body["instructions"], json!("Be concise")); + assert_eq!(body["text"]["format"]["strict"], json!(true)); +} + // ════════════════════════════════════════════════════════════════════════════ // Tool preparation // ════════════════════════════════════════════════════════════════════════════ +/// TS: tool preparation › "should prepare tools correctly" +#[tokio::test] +async fn should_prepare_tools_correctly() { + let server = MockServer::start().await; + mock_json( + &server, + json!({ + "id": "resp_tools", + "model": "deepseek-ai/DeepSeek-V3-0324", + "object": "response", + "created_at": 1741257730, + "status": "completed", + "error": null, + "instructions": null, + "max_output_tokens": null, + "metadata": null, + "tool_choice": "auto", + "tools": [], + "temperature": 1.0, + "top_p": 1.0, + "incomplete_details": null, + "usage": null, + "output": [], + "output_text": "Test" + }), + ) + .await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); + + let mut options = default_options(test_prompt()); + options.tools = Some(vec![Tool::from( + FunctionTool::new( + "getWeather", + json!({ + "type": "object", + "properties": { "location": { "type": "string" } }, + "required": ["location"] + }), + ) + .with_description("Get weather information"), + )]); + options.tool_choice = Some(LanguageModelV4ToolChoice::Tool { + tool_name: "getWeather".to_string(), + }); + + let result = model.do_generate(&options).await.expect("should succeed"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + + assert_eq!( + body["tools"], + json!([{ + "type": "function", + "name": "getWeather", + "description": "Get weather information", + "parameters": { + "type": "object", + "properties": { "location": { "type": "string" } }, + "required": ["location"] + } + }]) + ); + assert_eq!( + body["tool_choice"], + json!({ "type": "function", "function": { "name": "getWeather" } }) + ); +} + +/// TS: tool preparation › "should handle auto and required tool choices" +#[tokio::test] +async fn should_handle_auto_and_required_tool_choices() { + let server = MockServer::start().await; + mock_json( + &server, + json!({ + "id": "resp_tools", + "model": "deepseek-ai/DeepSeek-V3-0324", + "object": "response", + "created_at": 1741257730, + "status": "completed", + "error": null, + "instructions": null, + "max_output_tokens": null, + "metadata": null, + "tool_choice": "auto", + "tools": [], + "temperature": 1.0, + "top_p": 1.0, + "incomplete_details": null, + "usage": null, + "output": [], + "output_text": "Test" + }), + ) + .await; + + let provider = make_provider(&server); + let model = provider.responses("deepseek-ai/DeepSeek-V3-0324"); + + // Test auto + let mut options = default_options(test_prompt()); + options.tools = Some(vec![Tool::from(FunctionTool::new( + "test", + json!({ "type": "object" }), + ))]); + options.tool_choice = Some(LanguageModelV4ToolChoice::Auto); + + let result = model.do_generate(&options).await.expect("should succeed"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + assert_eq!(body["tool_choice"], json!("auto")); + + // Test required + let mut options = default_options(test_prompt()); + options.tools = Some(vec![Tool::from(FunctionTool::new( + "test", + json!({ "type": "object" }), + ))]); + options.tool_choice = Some(LanguageModelV4ToolChoice::Required); + + let result = model.do_generate(&options).await.expect("should succeed"); + let body = result + .request + .and_then(|request| request.body) + .expect("body"); + assert_eq!(body["tool_choice"], json!("required")); +} + // ════════════════════════════════════════════════════════════════════════════ // Top-level-only media type resolution // ════════════════════════════════════════════════════════════════════════════ @@ -1088,6 +1899,36 @@ fn detects_image_subtype_from_inline_bytes_for_top_level_image() { ); } +/// TS: "passes through URL source for top-level-only image" +#[test] +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, + }, + media_type: "image".into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]; + + let (input, warnings) = + convert_to_huggingface_responses_messages(&prompt).expect("should succeed"); + + assert!(warnings.is_empty()); + let content = &input[0]["content"][0]; + assert_eq!( + content, + &json!({ + "type": "input_image", + "image_url": "https://example.com/x.png" + }) + ); +} + /// TS: "normalizes image/* wildcard via detection" #[test] fn normalizes_image_wildcard_via_detection() { diff --git a/aimux-providers/tests/hume_speech_test.rs b/aimux-providers/tests/hume_speech_test.rs index 28d22f16f..f14cfa8cb 100644 --- a/aimux-providers/tests/hume_speech_test.rs +++ b/aimux-providers/tests/hume_speech_test.rs @@ -18,7 +18,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::speech_model::{AudioData, SpeechCallOptions, SpeechModel}; -use aimux_providers::{HumeConfig, HumeProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{HumeProviderSettings, create_hume}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -71,8 +72,12 @@ async fn should_pass_the_model_and_text() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = HumeConfig::new("test-api-key").with_base_url(server.uri()); - let provider = HumeProvider::new(config); + let config = HumeProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_hume(config).unwrap(); let model = provider.speech(); model @@ -106,10 +111,18 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = HumeConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(provider_headers); - let provider = HumeProvider::new(config); + let config = HumeProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some( + (provider_headers) + .into_iter() + .map(|(k, v)| (k, Some(v))) + .collect(), + ), + ..Default::default() + }; + let provider = create_hume(config).unwrap(); let model = provider.speech(); let mut options = speech_options("Hello from the AI SDK!"); @@ -142,8 +155,12 @@ async fn should_pass_options() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = HumeConfig::new("test-api-key").with_base_url(server.uri()); - let provider = HumeProvider::new(config); + let config = HumeProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_hume(config).unwrap(); let model = provider.speech(); let mut options = speech_options("Hello from the AI SDK!"); @@ -176,8 +193,12 @@ async fn should_return_audio_data_with_correct_content_type() { ) .await; - let config = HumeConfig::new("test-api-key").with_base_url(server.uri()); - let provider = HumeProvider::new(config); + let config = HumeProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_hume(config).unwrap(); let model = provider.speech(); let mut options = speech_options("Hello from the AI SDK!"); @@ -202,8 +223,12 @@ async fn should_include_response_data_with_timestamp_modelid_and_headers() { ) .await; - let config = HumeConfig::new("test-api-key").with_base_url(server.uri()); - let provider = HumeProvider::new(config); + let config = HumeProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_hume(config).unwrap(); let model = provider.speech(); let result = model @@ -231,8 +256,12 @@ async fn should_use_real_date_when_no_custom_date_provider_is_specified() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = HumeConfig::new("test-api-key").with_base_url(server.uri()); - let provider = HumeProvider::new(config); + let config = HumeProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_hume(config).unwrap(); let model = provider.speech(); let result = model @@ -251,8 +280,12 @@ async fn should_handle_different_audio_formats() { let server = MockServer::start().await; mock_audio_response(&server, format).await; - let config = HumeConfig::new("test-api-key").with_base_url(server.uri()); - let provider = HumeProvider::new(config); + let config = HumeProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_hume(config).unwrap(); let model = provider.speech(); let mut options = speech_options("Hello from the AI SDK!"); @@ -273,8 +306,12 @@ async fn should_include_warnings_if_any_are_generated() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = HumeConfig::new("test-api-key").with_base_url(server.uri()); - let provider = HumeProvider::new(config); + let config = HumeProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_hume(config).unwrap(); let model = provider.speech(); let result = model diff --git a/aimux-providers/tests/jev_test.rs b/aimux-providers/tests/jev_test.rs index 652b861a5..9e7cf98f7 100644 --- a/aimux-providers/tests/jev_test.rs +++ b/aimux-providers/tests/jev_test.rs @@ -1,6 +1,5 @@ //! Offline contract tests using official examples and recorded live responses. //! No live credentials or external requests are needed to replay them. -use std::time::Duration; use aimux_core::decision_model::*; use aimux_core::{AiMuxError, Provider}; @@ -153,9 +152,7 @@ fn options() -> DecisionCallOptions { } fn model(server: &MockServer) -> Box { - let mut config = - JevConfig::new("test-key").with_endpoint(format!("{}/v1/systemone", server.uri())); - config.retry_config.initial_delay = Duration::ZERO; + let config = JevConfig::new("test-key").with_endpoint(format!("{}/v1/systemone", server.uri())); JevProvider::new(config) .decision_model("jev-latest") .unwrap() @@ -390,7 +387,14 @@ async fn official_provider_failures_use_standard_retry_rules_and_preserve_codes( #[test] fn non_decision_provider_returns_unsupported() { - let provider = aimux_providers::provider_handle("deepseek", Some("test".into()), None).unwrap(); + let provider = aimux_providers::create_provider( + "deepseek", + aimux_providers::PresetSettings { + api_key: Some("test".to_string().into()), + ..Default::default() + }, + ) + .unwrap(); assert!(matches!( provider.decision_model("gpt-test"), Err(AiMuxError::UnsupportedFunctionality(_)) diff --git a/aimux-providers/tests/jina_ai_test.rs b/aimux-providers/tests/jina_ai_test.rs index bccf0b61a..46270f6af 100644 --- a/aimux-providers/tests/jina_ai_test.rs +++ b/aimux-providers/tests/jina_ai_test.rs @@ -10,10 +10,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::reranking_model::{RerankingCallOptions, RerankingDocuments, RerankingModel}; -use aimux_providers::{JinaAiConfig, JinaAiProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{JinaAiProvider, JinaAiProviderSettings, create_jina_ai}; const API_KEY: &str = "test-jina-key"; const MODEL: &str = "jina-reranker-v2-base-multilingual"; @@ -49,8 +49,12 @@ fn rerank_response_body() -> Value { } fn provider(server: &MockServer) -> JinaAiProvider { - let config = JinaAiConfig::new(API_KEY).with_base_url(server.uri()); - JinaAiProvider::new(config) + let config = JinaAiProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_jina_ai(config).unwrap() } fn text_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { @@ -276,25 +280,3 @@ async fn status_401_maps_to_auth_error() { } // -- Provider trait ---------------------------------------------------------- - -#[tokio::test] -async fn provider_name_is_jina_ai() { - let config = JinaAiConfig::new(API_KEY); - let provider = JinaAiProvider::new(config); - assert_eq!(provider.name(), "jina_ai"); -} - -#[test] -fn language_model_returns_unsupported_error() { - let config = JinaAiConfig::new(API_KEY); - let provider = JinaAiProvider::new(config); - match provider.language_model(MODEL) { - Err(AiMuxError::UnsupportedFunctionality(msg)) => { - assert!( - msg.contains("provider 'jina_ai' does not provide language models"), - "unexpected message: {msg}" - ); - } - _ => panic!("expected Unsupported error, got success or another error variant"), - } -} diff --git a/aimux-providers/tests/klingai_video_test.rs b/aimux-providers/tests/klingai_video_test.rs index 2b08ddde6..82375e007 100644 --- a/aimux-providers/tests/klingai_video_test.rs +++ b/aimux-providers/tests/klingai_video_test.rs @@ -2,13 +2,16 @@ //! Source: `reference/ai/packages/klingai/src/klingai-video-model.test.ts` use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +use aimux_core::download::{Download, DownloadRequest, DownloadedFile}; use aimux_core::video_model::{VideoCallOptions, VideoModel, generate_video}; -use aimux_providers::{KlingAIConfig, KlingAIProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{KlingAIProviderSettings, create_klingai}; fn fast_poll() -> Option { Some(aimux_core::video_model::VideoPollOptions { @@ -20,9 +23,27 @@ fn fast_poll() -> Option { fn options(prompt: &str) -> VideoCallOptions { let mut o = VideoCallOptions::new(prompt); o.poll = fast_poll(); + o.download = Some(recording_download(Arc::default())); o } +/// The bytes the test download returns for every video URL. +const VIDEO_BYTES: &[u8] = b"video bytes"; + +/// A download that fetches nothing: it records each URL `generate_video` +/// asks for and returns [`VIDEO_BYTES`]. +fn recording_download(requests: Arc>>) -> Download { + Download::new(move |request: DownloadRequest| { + requests.lock().unwrap().push(request.url); + async { + Ok(DownloadedFile { + data: VIDEO_BYTES.to_vec(), + media_type: None, + }) + } + }) +} + async fn mock_task_and_result( server: &MockServer, result_videos: &[Value], @@ -67,20 +88,30 @@ async fn should_generate_video_from_prompt() { let videos = vec![json!({"url": "https://cdn.klingai.com/video.mp4"})]; mock_task_and_result(&server, &videos, &[]).await; - let config = KlingAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = KlingAIProvider::new(config); + let config = KlingAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_klingai(config).unwrap(); let model = provider.video("kling-v2.1-master-t2v"); - let result = generate_video(&model, options("A cat playing")) - .await - .unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let mut opts = options("A cat playing"); + opts.download = Some(recording_download(requests.clone())); + let result = generate_video(&model, opts).await.unwrap(); + // `generate_video` downloads the returned URL, as the AI SDK does. + assert_eq!( + *requests.lock().unwrap(), + ["https://cdn.klingai.com/video.mp4"] + ); assert_eq!(result.videos.len(), 1); match &result.videos[0] { - aimux_core::video_model::VideoData::Url { url, .. } => { - assert_eq!(url, "https://cdn.klingai.com/video.mp4"); + aimux_core::video_model::VideoData::Binary { data, .. } => { + assert_eq!(data, VIDEO_BYTES); } - v => panic!("expected Url, got: {v:?}"), + v => panic!("expected the downloaded video, got: {v:?}"), } } @@ -90,8 +121,12 @@ async fn should_pass_model_name_and_prompt() { let videos = vec![json!({"url": "https://cdn.klingai.com/video.mp4"})]; mock_task_and_result(&server, &videos, &[]).await; - let config = KlingAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = KlingAIProvider::new(config); + let config = KlingAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_klingai(config).unwrap(); let model = provider.video("kling-v2.1-master-t2v"); generate_video(&model, options("A cat playing")) @@ -112,10 +147,13 @@ async fn should_pass_headers() { let mut ph = HashMap::new(); ph.insert("Custom-Provider-Header".to_string(), "val".to_string()); - let config = KlingAIConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = KlingAIProvider::new(config); + let config = KlingAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_klingai(config).unwrap(); let model = provider.video("kling-v2.1-master-t2v"); let mut opts = options("test"); @@ -138,8 +176,12 @@ async fn should_pass_seed_duration_and_aspect_ratio() { let videos = vec![json!({"url": "https://cdn.klingai.com/video.mp4"})]; mock_task_and_result(&server, &videos, &[]).await; - let config = KlingAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = KlingAIProvider::new(config); + let config = KlingAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_klingai(config).unwrap(); let model = provider.video("kling-v2.1-master-t2v"); let mut opts = options("test"); @@ -162,8 +204,12 @@ async fn should_include_response_data() { let videos = vec![json!({"url": "https://cdn.klingai.com/video.mp4"})]; mock_task_and_result(&server, &videos, &[("x-request-id", "test-req")]).await; - let config = KlingAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = KlingAIProvider::new(config); + let config = KlingAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_klingai(config).unwrap(); let model = provider.video("kling-v2.1-master-t2v"); let result = generate_video(&model, options("test")).await.unwrap(); @@ -177,8 +223,11 @@ async fn should_include_response_data() { #[tokio::test] async fn should_return_max_videos_per_call() { - let config = KlingAIConfig::new("test-api-key"); - let provider = KlingAIProvider::new(config); + let config = KlingAIProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + ..Default::default() + }; + let provider = create_klingai(config).unwrap(); let model = provider.video("kling-v2.1-master-t2v"); assert_eq!(model.max_videos_per_call(), Some(1)); } diff --git a/aimux-providers/tests/linkup_test.rs b/aimux-providers/tests/linkup_test.rs index 73f25281b..6b71e5f07 100644 --- a/aimux-providers/tests/linkup_test.rs +++ b/aimux-providers/tests/linkup_test.rs @@ -10,10 +10,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{LinkupConfig, LinkupProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{LinkupProvider, LinkupProviderSettings, create_linkup}; const API_KEY: &str = "test-linkup-key"; @@ -52,8 +52,12 @@ fn sourced_answer_body() -> Value { } fn provider(server: &MockServer) -> LinkupProvider { - let config = LinkupConfig::new(API_KEY).with_base_url(server.uri()); - LinkupProvider::new(config) + let config = LinkupProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_linkup(config).unwrap() } fn opts(query: &str) -> SearchCallOptions { @@ -231,33 +235,14 @@ async fn status_403_maps_to_provider_error() { // -- Provider trait ---------------------------------------------------------- -#[tokio::test] -async fn provider_name_is_linkup() { - let config = LinkupConfig::new(API_KEY); - let provider = LinkupProvider::new(config); - assert_eq!(provider.name(), "linkup"); -} - -#[test] -fn language_model_returns_unsupported_error() { - let config = LinkupConfig::new(API_KEY); - let provider = LinkupProvider::new(config); - match provider.language_model("linkup-search") { - Err(AiMuxError::UnsupportedFunctionality(msg)) => { - assert!( - msg.contains("provider 'linkup' does not provide language models"), - "unexpected message: {msg}" - ); - } - _ => panic!("expected Unsupported error, got success or another error variant"), - } -} - #[test] fn model_id_is_linkup_search() { - let config = LinkupConfig::new(API_KEY); - let provider = LinkupProvider::new(config); + let config = LinkupProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + ..Default::default() + }; + let provider = create_linkup(config).unwrap(); let model = provider.search_model(); assert_eq!(model.model_id(), "linkup-search"); - assert_eq!(model.provider(), "linkup"); + assert_eq!(model.provider(), "linkup.search"); } diff --git a/aimux-providers/tests/list_models_test.rs b/aimux-providers/tests/list_models_test.rs index 67a7326de..6b5aedcdf 100644 --- a/aimux-providers/tests/list_models_test.rs +++ b/aimux-providers/tests/list_models_test.rs @@ -1,4 +1,4 @@ -//! Integration tests for `Provider::list_models` (RFC-0027). +//! Integration tests for `ProviderDiscovery::list_models` (RFC-0027). //! //! `list_models` returns `RuntimeModel` (provider's official data only). //! Community catalogue (`get_model_specs`) is tested separately — the two are @@ -14,10 +14,10 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::model_catalogue::RuntimeModel; -use aimux_core::provider::Provider; +use aimux_core::provider::ProviderDiscovery; use aimux_providers::catalogue; -use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; -use aimux_providers::{ProviderOptions, provider_handle}; +use aimux_providers::openai::{OpenAIProviderSettings, create_openai}; +use aimux_providers::{PresetSettings, create_provider}; async fn mount_cassette_file(server: &MockServer, cassette_path: &Path) -> String { let text = std::fs::read_to_string(cassette_path) @@ -62,8 +62,12 @@ async fn openai_provider_list_models() { let cassette = Path::new("tests/cassettes/openai/list_models_smoke.json"); let recorded_path = mount_cassette_file(&server, cassette).await; let base_url = base_url_for(&server.uri(), &recorded_path); - let config = OpenAIConfig::new("test-key").with_base_url(base_url); - let provider = OpenAIProvider::new(config); + let provider = create_openai(OpenAIProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(base_url), + ..Default::default() + }) + .unwrap(); let models: Vec = provider.list_models().await.unwrap(); assert!(models.iter().any(|m| m.id == "gpt-4o")); assert!(models.iter().all(|m| m.owned_by.is_some())); @@ -71,16 +75,21 @@ async fn openai_provider_list_models() { #[tokio::test] #[serial] -async fn deepseek_provider_list_models_via_handle() { +async fn deepseek_provider_list_models_via_discovery_handle() { let server = MockServer::start().await; let cassette = Path::new("tests/cassettes/deepseek/list_models_smoke.json"); let recorded_path = mount_cassette_file(&server, cassette).await; let base_url = base_url_for(&server.uri(), &recorded_path); - let opts = ProviderOptions { - base_url: Some(base_url), - ..Default::default() - }; - let handle = provider_handle("deepseek", Some("test-key".into()), Some(opts)).unwrap(); + let provider = create_provider( + "deepseek", + PresetSettings { + api_key: Some("test-key".to_string().into()), + base_url: Some(base_url), + ..Default::default() + }, + ) + .unwrap(); + let handle = provider.discovery().expect("deepseek lists its models"); let models = handle.list_models().await.unwrap(); assert_eq!(models.len(), 2); assert!(models.iter().any(|m| m.id == "deepseek-v4-flash")); @@ -89,17 +98,22 @@ async fn deepseek_provider_list_models_via_handle() { #[tokio::test] #[serial] -async fn provider_handle_then_language_model() { +async fn discovery_then_language_model() { let server = MockServer::start().await; let cassette = Path::new("tests/cassettes/deepseek/list_models_smoke.json"); let recorded_path = mount_cassette_file(&server, cassette).await; let base_url = base_url_for(&server.uri(), &recorded_path); - let opts = ProviderOptions { - base_url: Some(base_url), - ..Default::default() - }; - let handle = provider_handle("deepseek", Some("test-key".into()), Some(opts)).unwrap(); - let models = handle.list_models().await.unwrap(); + let handle = create_provider( + "deepseek", + PresetSettings { + api_key: Some("test-key".to_string().into()), + base_url: Some(base_url), + ..Default::default() + }, + ) + .unwrap(); + let discovery = handle.discovery().expect("deepseek lists its models"); + let models = discovery.list_models().await.unwrap(); let first_id = models[0].id.clone(); let _model = handle.language_model(&first_id).unwrap(); } @@ -111,9 +125,13 @@ async fn anthropic_list_models() { let cassette = Path::new("tests/cassettes/anthropic/list_models_smoke.json"); let _ = mount_cassette_file(&server, cassette).await; let base_url = server.uri().trim_end_matches('/').to_string(); - use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider}; - let config = AnthropicConfig::new("test-key").with_base_url(&base_url); - let provider = AnthropicProvider::new(config); + use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; + let provider = create_anthropic(AnthropicProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{base_url}/v1")), + ..Default::default() + }) + .unwrap(); let models: Vec = provider.list_models().await.unwrap(); assert!(!models.is_empty()); assert!(models.iter().all(|m| !m.id.is_empty())); @@ -132,36 +150,18 @@ async fn google_list_models() { .strip_suffix("/models") .unwrap_or(&recorded_path) ); - use aimux_providers::google::{GoogleConfig, GoogleProvider}; - let config = GoogleConfig::new("test-key").with_base_url(&base_url); - let provider = GoogleProvider::new(config); + use aimux_providers::google::{GoogleProviderSettings, create_google}; + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(base_url), + ..Default::default() + }) + .unwrap(); let models: Vec = provider.list_models().await.unwrap(); assert!(!models.is_empty()); assert!(models.iter().all(|m| !m.id.starts_with("models/"))); } -#[tokio::test] -#[serial] -async fn ollama_list_models_via_delegate_macro() { - let server = MockServer::start().await; - let body = r#"{"data":[{"id":"llama3.2","object":"model","owned_by":"ollama"},{"id":"qwen3:4b","object":"model","owned_by":"ollama"}]}"#; - Mock::given(method("GET")) - .and(path("/v1/models")) - .respond_with( - ResponseTemplate::new(200) - .append_header("content-type", "application/json") - .set_body_bytes(body.as_bytes().to_vec()), - ) - .mount(&server) - .await; - use aimux_providers::ollama::{OllamaConfig, OllamaProvider}; - let config = OllamaConfig::new("ollama").with_base_url(format!("{}/v1", server.uri())); - let provider = OllamaProvider::new(config); - let models: Vec = provider.list_models().await.unwrap(); - assert_eq!(models.len(), 2); - assert!(models.iter().any(|m| m.id == "llama3.2")); -} - #[tokio::test] #[serial] async fn list_models_malformed_response() { @@ -171,8 +171,12 @@ async fn list_models_malformed_response() { .respond_with(ResponseTemplate::new(200).set_body_bytes(b"not json at all".to_vec())) .mount(&server) .await; - let config = OpenAIConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())); - let provider = OpenAIProvider::new(config); + let provider = create_openai(OpenAIProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap(); let err = provider.list_models().await.unwrap_err(); assert!(matches!( err, @@ -191,8 +195,12 @@ async fn list_models_empty_data() { .respond_with(ResponseTemplate::new(200).set_body_bytes(br#"{"data":[]}"#.to_vec())) .mount(&server) .await; - let config = OpenAIConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())); - let provider = OpenAIProvider::new(config); + let provider = create_openai(OpenAIProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap(); let models: Vec = provider.list_models().await.unwrap(); assert!(models.is_empty()); } @@ -204,53 +212,23 @@ async fn list_models_http_error() { Mock::given(method("GET")) .and(path("/v1/models")) .respond_with( - ResponseTemplate::new(500).set_body_bytes(br#"{"error":"internal"}"#.to_vec()), + ResponseTemplate::new(500) + // A zero retry hint keeps the default retries instant. + .insert_header("retry-after-ms", "0") + .set_body_bytes(br#"{"error":"internal"}"#.to_vec()), ) .mount(&server) .await; - let config = OpenAIConfig::new("test-key") - .with_base_url(format!("{}/v1", server.uri())) - .with_retry_config(aimux_provider_utils::RetryConfig { - max_retries: 0, - ..Default::default() - }); - let provider = OpenAIProvider::new(config); + let provider = create_openai(OpenAIProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap(); let err = provider.list_models().await.unwrap_err(); assert!(!err.to_string().is_empty()); } -#[test] -fn provider_handle_unknown_name() { - match aimux_providers::provider_handle("no-such-provider", Some("k".into()), None) { - Err(e) => assert!(matches!(e, aimux_core::AiMuxError::NoSuchProvider { .. })), - Ok(_) => panic!("unknown provider should fail"), - } -} - -#[tokio::test] -async fn unsupported_provider_returns_unsupported() { - struct StubProvider; - impl Provider for StubProvider { - fn name(&self) -> &str { - "stub" - } - fn language_model( - &self, - _model_id: &str, - ) -> Result, aimux_core::AiMuxError> - { - Err(aimux_core::AiMuxError::UnsupportedFunctionality( - "none".into(), - )) - } - } - let err = StubProvider.list_models().await.unwrap_err(); - assert!(matches!( - err, - aimux_core::AiMuxError::UnsupportedFunctionality(_) - )); -} - #[test] fn catalogue_normalize_and_lookup() { assert_eq!( diff --git a/aimux-providers/tests/lmnt_speech_test.rs b/aimux-providers/tests/lmnt_speech_test.rs index aecc864bc..66108659b 100644 --- a/aimux-providers/tests/lmnt_speech_test.rs +++ b/aimux-providers/tests/lmnt_speech_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the LMNT speech (TTS) model tests. +//! Rust translation of the LMNT speech (TTS) model tests. //! //! Source: `reference/ai/packages/lmnt/src/lmnt-speech-model.test.ts` //! @@ -13,7 +13,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::speech_model::{AudioData, SpeechCallOptions, SpeechModel}; -use aimux_providers::{LMNTConfig, LMNTProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{LMNTProviderSettings, create_lmnt}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -66,8 +67,12 @@ async fn should_pass_the_model_and_text() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = LMNTConfig::new("test-api-key").with_base_url(server.uri()); - let provider = LMNTProvider::new(config); + let config = LMNTProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_lmnt(config).unwrap(); let model = provider.speech("aurora"); model @@ -93,10 +98,18 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = LMNTConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(provider_headers); - let provider = LMNTProvider::new(config); + let config = LMNTProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some( + (provider_headers) + .into_iter() + .map(|(k, v)| (k, Some(v))) + .collect(), + ), + ..Default::default() + }; + let provider = create_lmnt(config).unwrap(); let model = provider.speech("aurora"); let mut options = speech_options("Hello from the AI SDK!"); @@ -129,8 +142,12 @@ async fn should_pass_options() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = LMNTConfig::new("test-api-key").with_base_url(server.uri()); - let provider = LMNTProvider::new(config); + let config = LMNTProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_lmnt(config).unwrap(); let model = provider.speech("aurora"); let mut options = speech_options("Hello from the AI SDK!"); @@ -163,8 +180,12 @@ async fn should_return_audio_data_with_correct_content_type() { ) .await; - let config = LMNTConfig::new("test-api-key").with_base_url(server.uri()); - let provider = LMNTProvider::new(config); + let config = LMNTProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_lmnt(config).unwrap(); let model = provider.speech("aurora"); let mut options = speech_options("Hello from the AI SDK!"); @@ -189,8 +210,12 @@ async fn should_include_response_data_with_timestamp_modelid_and_headers() { ) .await; - let config = LMNTConfig::new("test-api-key").with_base_url(server.uri()); - let provider = LMNTProvider::new(config); + let config = LMNTProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_lmnt(config).unwrap(); let model = provider.speech("aurora"); let result = model @@ -218,8 +243,12 @@ async fn should_use_real_date_when_no_custom_date_provider_is_specified() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = LMNTConfig::new("test-api-key").with_base_url(server.uri()); - let provider = LMNTProvider::new(config); + let config = LMNTProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_lmnt(config).unwrap(); let model = provider.speech("aurora"); let result = model @@ -241,8 +270,12 @@ async fn should_handle_different_audio_formats() { let server = MockServer::start().await; mock_audio_response(&server, format).await; - let config = LMNTConfig::new("test-api-key").with_base_url(server.uri()); - let provider = LMNTProvider::new(config); + let config = LMNTProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_lmnt(config).unwrap(); let model = provider.speech("aurora"); let mut options = speech_options("Hello from the AI SDK!"); @@ -268,8 +301,12 @@ async fn should_include_warnings_if_any_are_generated() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = LMNTConfig::new("test-api-key").with_base_url(server.uri()); - let provider = LMNTProvider::new(config); + let config = LMNTProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_lmnt(config).unwrap(); let model = provider.speech("aurora"); let result = model diff --git a/aimux-providers/tests/luma_image_test.rs b/aimux-providers/tests/luma_image_test.rs index 699033239..2eabfb65c 100644 --- a/aimux-providers/tests/luma_image_test.rs +++ b/aimux-providers/tests/luma_image_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Luma image model tests. +//! Rust translation of the Luma image model tests. //! Source: `reference/ai/packages/luma/src/luma-image-model.test.ts` use serde_json::{Value, json}; @@ -6,7 +6,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::image_model::{ImageCallOptions, ImageFile, ImageModel, ImageOutputs}; -use aimux_providers::{LumaConfig, LumaProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{LumaProviderSettings, create_luma}; const PROMPT: &str = "A beautiful sunset"; @@ -41,8 +42,12 @@ fn options(prompt: &str) -> ImageCallOptions { async fn should_extract_generated_image() { let server = MockServer::start().await; mock_luma(&server).await; - let config = LumaConfig::new("test-key").with_base_url(server.uri()); - let model = LumaProvider::new(config).image("ray2"); + let config = LumaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_luma(config).unwrap().image("ray2"); let result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { ImageOutputs::Binary(imgs) => { @@ -57,8 +62,12 @@ async fn should_extract_generated_image() { async fn should_pass_prompt_and_model() { let server = MockServer::start().await; mock_luma(&server).await; - let config = LumaConfig::new("test-key").with_base_url(server.uri()); - let model = LumaProvider::new(config).image("ray2"); + let config = LumaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_luma(config).unwrap().image("ray2"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&reqs[0].body).unwrap(); @@ -70,8 +79,12 @@ async fn should_pass_prompt_and_model() { async fn should_pass_aspect_ratio() { let server = MockServer::start().await; mock_luma(&server).await; - let config = LumaConfig::new("test-key").with_base_url(server.uri()); - let model = LumaProvider::new(config).image("ray2"); + let config = LumaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_luma(config).unwrap().image("ray2"); let mut opts = options(PROMPT); opts.aspect_ratio = Some(aimux_core::shared::AspectRatio::new(16, 9)); model.do_generate(&opts).await.unwrap(); @@ -84,8 +97,12 @@ async fn should_pass_aspect_ratio() { async fn should_warn_for_seed() { let server = MockServer::start().await; mock_luma(&server).await; - let config = LumaConfig::new("test-key").with_base_url(server.uri()); - let model = LumaProvider::new(config).image("ray2"); + let config = LumaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_luma(config).unwrap().image("ray2"); let mut opts = options(PROMPT); opts.seed = Some(42); let result = model.do_generate(&opts).await.unwrap(); @@ -96,8 +113,12 @@ async fn should_warn_for_seed() { async fn should_warn_for_size() { let server = MockServer::start().await; mock_luma(&server).await; - let config = LumaConfig::new("test-key").with_base_url(server.uri()); - let model = LumaProvider::new(config).image("ray2"); + let config = LumaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_luma(config).unwrap().image("ray2"); let mut opts = options(PROMPT); opts.size = Some(aimux_core::shared::Size::new(1024, 1024)); let result = model.do_generate(&opts).await.unwrap(); @@ -108,8 +129,12 @@ async fn should_warn_for_size() { async fn should_pass_auth_headers() { let server = MockServer::start().await; mock_luma(&server).await; - let config = LumaConfig::new("test-key").with_base_url(server.uri()); - let model = LumaProvider::new(config).image("ray2"); + let config = LumaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_luma(config).unwrap().image("ray2"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); assert_eq!( @@ -127,8 +152,12 @@ async fn should_pass_auth_headers() { async fn should_support_image_editing_with_urls() { let server = MockServer::start().await; mock_luma(&server).await; - let config = LumaConfig::new("test-key").with_base_url(server.uri()); - let model = LumaProvider::new(config).image("ray2"); + let config = LumaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_luma(config).unwrap().image("ray2"); let mut opts = options("Edit this"); opts.n = 1; opts.files = Some(vec![ImageFile::Url { @@ -146,8 +175,12 @@ async fn should_support_image_editing_with_urls() { #[tokio::test] async fn should_throw_for_mask() { let server = MockServer::start().await; - let config = LumaConfig::new("test-key").with_base_url(server.uri()); - let model = LumaProvider::new(config).image("ray2"); + let config = LumaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_luma(config).unwrap().image("ray2"); let mut opts = options("Edit"); opts.n = 1; opts.mask = Some(ImageFile::Url { diff --git a/aimux-providers/tests/mistral_embedding_test.rs b/aimux-providers/tests/mistral_embedding_test.rs index 66eed80be..4755a5a39 100644 --- a/aimux-providers/tests/mistral_embedding_test.rs +++ b/aimux-providers/tests/mistral_embedding_test.rs @@ -9,10 +9,9 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; -use aimux_core::tool::FunctionTool; -use aimux_providers::mistral::convert::prepare_tools; -use aimux_providers::{MistralConfig, MistralProvider}; +use aimux_core::embedding_model::EmbeddingCallOptions; +use aimux_core::provider::Provider; +use aimux_providers::{MistralProviderSettings, create_mistral}; const TEST_VALUES: &[&str] = &["sunny day at the beach", "rainy day in the city"]; @@ -63,9 +62,13 @@ async fn should_extract_embedding() { .mount(&server) .await; - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.embedding_model("mistral-embed"); + let provider = create_mistral(MistralProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("mistral-embed").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -89,9 +92,13 @@ async fn should_extract_usage() { .mount(&server) .await; - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.embedding_model("mistral-embed"); + let provider = create_mistral(MistralProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("mistral-embed").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -115,9 +122,13 @@ async fn should_expose_raw_response() { .mount(&server) .await; - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.embedding_model("mistral-embed"); + let provider = create_mistral(MistralProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("mistral-embed").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -132,6 +143,10 @@ async fn should_expose_raw_response() { headers.get("test-header").map(std::string::String::as_str), Some("test-value") ); + assert_eq!( + result.response.unwrap().body, + Some(embedding_response_body(&dummy_embeddings(), None)) + ); } /// TS: "should pass the model and the values" @@ -147,9 +162,13 @@ async fn should_pass_model_and_values() { .mount(&server) .await; - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.embedding_model("mistral-embed"); + let provider = create_mistral(MistralProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("mistral-embed").unwrap(); let _ = model .do_embed(&default_options(test_values())) @@ -180,9 +199,17 @@ async fn should_pass_headers() { .mount(&server) .await; - let config = MistralConfig::new("test-api-key").with_base_url(server.uri()); - let provider = MistralProvider::new(config); - let model = provider.embedding_model("mistral-embed"); + let provider = create_mistral(MistralProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + headers: Some(HashMap::from([( + "Custom-Provider-Header".to_string(), + Some("provider-header-value".to_string()), + )])), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("mistral-embed").unwrap(); let mut request_headers = HashMap::new(); request_headers.insert( @@ -204,6 +231,11 @@ async fn should_pass_headers() { let requests = server.received_requests().await.expect("requests recorded"); assert_eq!(requests.len(), 1); let h = &requests[0].headers; + assert_eq!( + h.get("custom-provider-header") + .and_then(|v| v.to_str().ok()), + Some("provider-header-value") + ); assert_eq!( h.get("authorization").and_then(|v| v.to_str().ok()), Some("Bearer test-api-key") @@ -213,13 +245,3 @@ async fn should_pass_headers() { Some("request-header-value") ); } - -/// TS: "should include strict field when set" -#[test] -fn prepare_tools_function_tool_with_strict() { - let tool = FunctionTool::new("strictTool", json!({ "type": "object", "properties": {} })) - .with_strict(true); - let result = prepare_tools(&Some(vec![tool]), None); - let tools = result.tools.expect("tools should be Some"); - assert_eq!(tools[0]["function"]["strict"], json!(true)); -} diff --git a/aimux-providers/tests/moonshotai_test.rs b/aimux-providers/tests/moonshotai_test.rs index 0760eb048..0895f28be 100644 --- a/aimux-providers/tests/moonshotai_test.rs +++ b/aimux-providers/tests/moonshotai_test.rs @@ -4,7 +4,7 @@ //! - `packages/moonshotai/src/moonshotai-provider.test.ts` //! - `packages/moonshotai/src/convert-moonshotai-chat-usage.test.ts` //! -//! Moonshot AI is a thin OpenAI-compatible wrapper over [`OpenAIProvider`]. +//! Moonshot AI is a registry preset of the OpenAI-compatible package. //! The TS suite has a custom usage converter that recognises a **top-level** //! `cached_tokens` field (Moonshot's wire format) and prioritises it over the //! nested `prompt_tokens_details.cached_tokens`. The shared OpenAI @@ -20,6 +20,7 @@ use futures::StreamExt; use serde_json::{Value, json}; use serial_test::serial; +use std::sync::Arc; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -30,7 +31,7 @@ use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; use aimux_core::types::FinishReasonUnified; -use aimux_providers::{ProviderOptions, provider, provider_from_env}; +use aimux_providers::{PresetSettings, create_provider}; // ── shared helpers ─────────────────────────────────────────────────────────── @@ -83,16 +84,17 @@ async fn collect_stream(result: aimux_core::result::StreamResult) -> Vec Box { - provider( +fn make_provider(server: &MockServer) -> Arc { + create_provider( "moonshotai", - Some("test-api-key".to_string()), - "moonshot-v1-8k", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-api-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("moonshot-v1-8k") .expect("moonshotai provider should build") } @@ -121,12 +123,15 @@ async fn usage_for(usage: Value) -> aimux_core::types::Usage { /// TS: `createMoonshotAI()` produces a provider whose name is "moonshotai". #[test] fn provider_builds_moonshotai_model() { - let model = provider( + let model = create_provider( "moonshotai", - Some("test-key".to_string()), - "moonshot-v1-8k", - None, + PresetSettings { + api_key: Some("test-key".to_string().into()), + ..Default::default() + }, ) + .expect("provider should build") + .language_model("moonshot-v1-8k") .expect("moonshotai provider should build"); assert_eq!(model.model_id(), "moonshot-v1-8k"); } @@ -146,15 +151,16 @@ async fn custom_api_key_used_in_auth_header() { .mount(&server) .await; - let model = provider( + let model = create_provider( "moonshotai", - Some("my-custom-key".to_string()), - "moonshot-v1-8k", - Some(ProviderOptions { + PresetSettings { + api_key: Some("my-custom-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("moonshot-v1-8k") .expect("moonshotai provider should build"); model @@ -178,15 +184,16 @@ async fn custom_headers_forwarded() { .mount(&server) .await; - let model = provider( + let model = create_provider( "moonshotai", - Some("test-key".to_string()), - "moonshot-v1-8k", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("moonshot-v1-8k") .expect("moonshotai provider should build"); let mut options = default_options(test_prompt()); @@ -216,15 +223,16 @@ async fn language_model_via_trait() { .mount(&server) .await; - let model = provider( + let model = create_provider( "moonshotai", - Some("test-key".to_string()), - "moonshot-v1-8k", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("moonshot-v1-8k") .expect("provider should build a model"); model @@ -242,7 +250,9 @@ fn from_env_loads_moonshot_api_key() { std::env::set_var("MOONSHOT_API_KEY", "env-test-key"); } - let model = provider_from_env("moonshotai", "moonshot-v1-8k", None); + let model = create_provider("moonshotai", PresetSettings::default()) + .expect("provider should build") + .language_model("moonshot-v1-8k"); assert!(model.is_ok(), "from_env should succeed with env var set"); unsafe { @@ -255,15 +265,19 @@ fn from_env_loads_moonshot_api_key() { /// TS: without the env var, `createMoonshotAI()` fails. #[serial] -#[test] -fn from_env_fails_without_env_var() { +#[tokio::test] +async fn from_env_fails_without_env_var() { let saved = std::env::var("MOONSHOT_API_KEY").ok(); unsafe { std::env::remove_var("MOONSHOT_API_KEY"); } - let model = provider_from_env("moonshotai", "moonshot-v1-8k", None); - assert!(model.is_err(), "from_env should fail without env var"); + let model = create_provider("moonshotai", PresetSettings::default()) + .expect("provider should build") + .language_model("moonshot-v1-8k") + .expect("provider should build a model"); + let result = model.do_generate(&default_options(test_prompt())).await; + assert!(result.is_err(), "from_env should fail without env var"); unsafe { if let Some(v) = saved { diff --git a/aimux-providers/tests/open_responses_test.rs b/aimux-providers/tests/open_responses_test.rs index 7b6ccbc3f..b1d40e462 100644 --- a/aimux-providers/tests/open_responses_test.rs +++ b/aimux-providers/tests/open_responses_test.rs @@ -10,7 +10,6 @@ //! `#[tokio::test]` runs do not collide. use aimux_core::tool::RawToolCall; -use std::collections::HashMap; use futures::StreamExt; use serde_json::{Value, json}; @@ -19,7 +18,8 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{ - AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, TextPart, UserPart, + AssistantPart, FilePart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, + ToolPart, ToolResultContent, ToolResultOutput, ToolResultPart, UserPart, }; use aimux_core::options::{CallOptions, LanguageModelV4ToolChoice, ResponseFormat, Tool}; use aimux_core::result::{GenerateContent, ReasoningOutput}; @@ -29,8 +29,8 @@ use aimux_core::tool::FunctionTool; use aimux_core::types::{FinishReasonUnified, ReasoningEffort}; use aimux_providers::open_responses::{ - OpenResponsesConfig, OpenResponsesModel, OpenResponsesProvider, - convert_to_open_responses_input, map_open_responses_finish_reason, + OpenResponsesModel, OpenResponsesProviderSettings, convert_to_open_responses_input, + create_open_responses, map_open_responses_finish_reason, }; // ============================================================================ @@ -47,15 +47,15 @@ fn default_options(prompt: LanguageModelPrompt) -> CallOptions { CallOptions::new(prompt) } -/// Build an Open Responses config whose URL points at the mock server. -fn make_config(server: &MockServer) -> OpenResponsesConfig { - OpenResponsesConfig::new("lmstudio", "lmstudio", server.uri()) +/// Build Open Responses settings whose base URL points at the mock server. +fn make_settings(server: &MockServer) -> OpenResponsesProviderSettings { + OpenResponsesProviderSettings::new("lmstudio", server.uri()) } /// Create a model pointing at the mock server. fn make_model(server: &MockServer, model_id: &str) -> OpenResponsesModel { - let provider = OpenResponsesProvider::new(make_config(server)); - provider.model(model_id) + let provider = create_open_responses(make_settings(server)).expect("valid settings"); + provider.responses(model_id) } /// Mount a JSON response on any path. @@ -163,6 +163,40 @@ fn lmstudio_basic_json() -> Value { }) } +/// The `lmstudio-tool-call.1.json` fixture. +fn lmstudio_tool_call_json() -> Value { + json!({ + "id": "resp_930de53bd4b5933673481fa630f3dc5f58027a2c67598a2a", + "object": "response", + "created_at": 1769005553, + "status": "completed", + "incomplete_details": null, + "model": "mistralai/ministral-3-14b-reasoning", + "output": [ + { + "id": "fc_ru0kcno9erlzp8573yub", + "call_id": "call_2866856768160095", + "type": "function_call", + "name": "weather", + "arguments": "{\"location\":\"San Francisco\"}", + "status": "completed" + } + ], + "error": null, + "usage": { + "input_tokens": 1189, + "output_tokens": 11, + "total_tokens": 1200, + "input_tokens_details": { + "cached_tokens": 891 + }, + "output_tokens_details": { + "reasoning_tokens": 0 + } + } + }) +} + /// The `openai-pdf-input-file.1.json` fixture. fn openai_pdf_json() -> Value { json!({ @@ -171,7 +205,7 @@ fn openai_pdf_json() -> Value { "created_at": 1775566496, "status": "completed", "incomplete_details": null, - "model": "gpt-4.1-nano-2025-04-14", + "model": "test-pdf-model", "output": [ { "id": "msg_048edf44633e41ae0069d4fea0d1a08194af1e491c093df1d9", @@ -407,6 +441,31 @@ mod convert_tests { assert!(warnings.is_empty()); } + #[test] + fn convert_image_file_url_to_input_image() { + let prompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Url { + url: "https://example.com/image.png".into(), + original_url: None, + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([{ + "type": "message", + "role": "user", + "content": [{"type": "input_image", "image_url": "https://example.com/image.png"}] + }]) + ); + } + #[test] fn convert_pdf_file_base64_to_input_file() { let prompt = vec![LanguageModelMessage::User { @@ -441,6 +500,32 @@ mod convert_tests { assert!(warnings.is_empty()); } + #[test] + fn convert_pdf_file_url_to_input_file() { + let prompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Url { + url: "https://example.com/document.pdf".into(), + original_url: None, + }, + media_type: "application/pdf".into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]; + let (input, _, warnings) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([{ + "type": "message", + "role": "user", + "content": [{"type": "input_file", "file_url": "https://example.com/document.pdf"}] + }]) + ); + assert!(warnings.is_empty()); + } + #[test] fn convert_file_with_custom_filename() { let prompt = vec![LanguageModelMessage::User { @@ -509,43 +594,441 @@ mod convert_tests { let (input, _, _) = convert_to_open_responses_input(&prompt); assert_eq!( input, - json!([{ - "type": "message", - "role": "assistant", - "content": [ - {"type": "output_text", "text": "First part"}, - {"type": "output_text", "text": "Second part"} - ] - }]) + json!([{ + "type": "message", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "First part"}, + {"type": "output_text", "text": "Second part"} + ] + }]) + ); + } + + // -- Assistant messages with tool calls -- + + #[test] + fn convert_assistant_single_tool_call() { + let prompt = vec![LanguageModelMessage::Assistant { + content: vec![AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "call_123".into(), + tool_name: "get_weather".into(), + input: json!({"location": "San Francisco"}), + provider_executed: None, + provider_options: None, + })], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([{ + "type": "function_call", + "call_id": "call_123", + "name": "get_weather", + "arguments": "{\"location\":\"San Francisco\"}" + }]) + ); + } + + #[test] + fn convert_assistant_tool_call_string_input() { + let prompt = vec![LanguageModelMessage::Assistant { + content: vec![AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "call_124".into(), + tool_name: "get_weather".into(), + input: Value::String("{\"location\":\"Berlin\"}".to_string()), + provider_executed: None, + provider_options: None, + })], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([{ + "type": "function_call", + "call_id": "call_124", + "name": "get_weather", + "arguments": "{\"location\":\"Berlin\"}" + }]) + ); + } + + #[test] + fn convert_assistant_text_and_tool_call() { + let prompt = vec![LanguageModelMessage::Assistant { + content: vec![ + AssistantPart::Text(TextPart { + text: "Let me check the weather for you.".into(), + provider_options: None, + }), + AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "call_456".into(), + tool_name: "get_weather".into(), + input: json!({"location": "New York"}), + provider_executed: None, + provider_options: None, + }), + ], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Let me check the weather for you."}] + }, + { + "type": "function_call", + "call_id": "call_456", + "name": "get_weather", + "arguments": "{\"location\":\"New York\"}" + } + ]) + ); + } + + #[test] + fn convert_assistant_multiple_tool_calls() { + let prompt = vec![LanguageModelMessage::Assistant { + content: vec![ + AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "call_001".into(), + tool_name: "get_weather".into(), + input: json!({"location": "Paris"}), + provider_executed: None, + provider_options: None, + }), + AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "call_002".into(), + tool_name: "get_time".into(), + input: json!({"timezone": "Europe/Paris"}), + provider_executed: None, + provider_options: None, + }), + ], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([ + { + "type": "function_call", + "call_id": "call_001", + "name": "get_weather", + "arguments": "{\"location\":\"Paris\"}" + }, + { + "type": "function_call", + "call_id": "call_002", + "name": "get_time", + "arguments": "{\"timezone\":\"Europe/Paris\"}" + } + ]) + ); + } + + // -- Tool messages -- + + #[test] + fn convert_tool_message_json_output() { + 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(), + + provider_options: None, + })], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([{ + "type": "function_call_output", + "call_id": "call_123", + "output": "{\"temperature\":72,\"condition\":\"sunny\"}" + }]) + ); + } + + #[test] + fn convert_tool_message_text_output() { + 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(), + + provider_options: None, + })], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([{ + "type": "function_call_output", + "call_id": "call_456", + "output": "Search results: Found 5 items" + }]) + ); + } + + #[test] + fn convert_tool_message_error_text_output() { + 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(), + + provider_options: None, + })], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([{ + "type": "function_call_output", + "call_id": "call_789", + "output": "API request failed: timeout" + }]) + ); + } + + #[test] + fn convert_tool_message_execution_denied_output() { + 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(), + + provider_options: None, + })], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([{ + "type": "function_call_output", + "call_id": "call_denied", + "output": "User declined the action" + }]) + ); + } + + #[test] + fn convert_tool_message_content_output_text() { + 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(), + + provider_options: None, + })], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([{ + "type": "function_call_output", + "call_id": "call_content", + "output": [ + {"type": "input_text", "text": "First result"}, + {"type": "input_text", "text": "Second result"} + ] + }]) + ); + } + + #[test] + fn convert_tool_message_multiple_results() { + let prompt = vec![LanguageModelMessage::Tool { + 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(), + + 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(), + + provider_options: None, + }), + ], + provider_options: None, + }]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([ + { + "type": "function_call_output", + "call_id": "call_001", + "output": "{\"temp\":72}" + }, + { + "type": "function_call_output", + "call_id": "call_002", + "output": "3:00 PM" + } + ]) + ); + } + + // -- Message chains -- + + #[test] + fn convert_user_assistant_user_chain() { + let prompt = vec![ + LanguageModelMessage::user_text("What is the capital of France?"), + LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: "The capital of France is Paris.".into(), + provider_options: None, + })], + provider_options: None, + }, + LanguageModelMessage::user_text("And what about Germany?"), + ]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "The capital of France is Paris."}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "And what about Germany?"}]} + ]) + ); + } + + #[test] + fn convert_user_assistant_tool_tool_chain() { + let prompt = vec![ + LanguageModelMessage::user_text("What is the weather in Tokyo?"), + LanguageModelMessage::Assistant { + content: vec![AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "call_weather".into(), + tool_name: "get_weather".into(), + input: json!({"location": "Tokyo"}), + provider_executed: None, + provider_options: None, + })], + provider_options: None, + }, + 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(), + + provider_options: None, + })], + provider_options: None, + }, + ]; + let (input, _, _) = convert_to_open_responses_input(&prompt); + assert_eq!( + input, + json!([ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is the weather in Tokyo?"}]}, + {"type": "function_call", "call_id": "call_weather", "name": "get_weather", "arguments": "{\"location\":\"Tokyo\"}"}, + {"type": "function_call_output", "call_id": "call_weather", "output": "{\"temperature\":25,\"condition\":\"cloudy\"}"} + ]) ); } - // -- Assistant messages with tool calls -- - - // -- Tool messages -- - - // -- Message chains -- - #[test] - fn convert_user_assistant_user_chain() { + fn convert_tool_roundtrip_with_followup_assistant() { let prompt = vec![ - LanguageModelMessage::user_text("What is the capital of France?"), + LanguageModelMessage::user_text("What is the weather in Tokyo?"), + LanguageModelMessage::Assistant { + content: vec![AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "call_weather".into(), + tool_name: "get_weather".into(), + input: Value::String("{\"location\":\"Tokyo\"}".to_string()), + provider_executed: None, + provider_options: None, + })], + provider_options: None, + }, + 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(), + + provider_options: None, + })], + provider_options: None, + }, LanguageModelMessage::Assistant { content: vec![AssistantPart::Text(TextPart { - text: "The capital of France is Paris.".into(), + text: "It is 25 C and cloudy in Tokyo.".into(), provider_options: None, })], provider_options: None, }, - LanguageModelMessage::user_text("And what about Germany?"), ]; let (input, _, _) = convert_to_open_responses_input(&prompt); assert_eq!( input, json!([ - {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]}, - {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "The capital of France is Paris."}]}, - {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "And what about Germany?"}]} + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is the weather in Tokyo?"}]}, + {"type": "function_call", "call_id": "call_weather", "name": "get_weather", "arguments": "{\"location\":\"Tokyo\"}"}, + {"type": "function_call_output", "call_id": "call_weather", "output": "{\"temperature\":25,\"condition\":\"cloudy\"}"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "It is 25 C and cloudy in Tokyo."}]} ]) ); } @@ -961,10 +1444,269 @@ 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"), + ) + } + + #[tokio::test] + async fn provider_options_reasoning_summary_detailed() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_basic_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + provider_options: lmstudio_opts(json!({"reasoningSummary": "detailed"})), + ..CallOptions::new(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let request = &server.received_requests().await.unwrap()[0]; + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["reasoning"], json!({"summary": "detailed"})); + } + + #[tokio::test] + async fn provider_options_combines_effort_with_summary() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_basic_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + reasoning: Some(ReasoningEffort::High), + provider_options: lmstudio_opts(json!({"reasoningSummary": "auto"})), + ..CallOptions::new(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let request = &server.received_requests().await.unwrap()[0]; + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!( + body["reasoning"], + json!({"effort": "high", "summary": "auto"}) + ); + } + + #[tokio::test] + async fn provider_options_reasoning_summary_concise() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_basic_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + provider_options: lmstudio_opts(json!({"reasoningSummary": "concise"})), + ..CallOptions::new(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let request = &server.received_requests().await.unwrap()[0]; + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["reasoning"], json!({"summary": "concise"})); + } + + #[tokio::test] + async fn provider_options_no_reasoning_fields() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_basic_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + provider_options: lmstudio_opts(json!({})), + ..CallOptions::new(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let request = &server.received_requests().await.unwrap()[0]; + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert!(body.get("reasoning").is_none()); + } + // -- Tool call parsing tests -- + #[tokio::test] + async fn parse_tool_call_from_response() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_tool_call_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + tools: Some(vec![Tool::from( + FunctionTool::new( + "weather", + json!({ + "type": "object", + "properties": { + "location": {"type": "string", "description": "The location to get the weather for"} + }, + "required": ["location"] + }), + ) + .with_description("Get the weather in a location"), + )]), + tool_choice: Some(LanguageModelV4ToolChoice::Required), + ..CallOptions::new(test_prompt()) + }; + let result = model.do_generate(&options).await.unwrap(); + + assert_eq!(result.content.len(), 1); + match &result.content[0] { + GenerateContent::ToolCall(RawToolCall { + tool_call_id, + tool_name, + input, + .. + }) => { + assert_eq!(tool_call_id, "call_2866856768160095"); + assert_eq!(tool_name, "weather"); + assert_eq!( + input, + &Value::String(r#"{"location":"San Francisco"}"#.into()) + ); + } + other => panic!("expected ToolCall, got {other:?}"), + } + } + + #[tokio::test] + async fn tool_call_finish_reason() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_tool_call_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + tools: Some(vec![Tool::from( + FunctionTool::new("weather", json!({"type": "object"})) + .with_description("Get the weather in a location"), + )]), + tool_choice: Some(LanguageModelV4ToolChoice::Required), + ..CallOptions::new(test_prompt()) + }; + let result = model.do_generate(&options).await.unwrap(); + + assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); + assert_eq!(result.finish_reason.raw, None); + } + + #[tokio::test] + async fn tool_call_usage() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_tool_call_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + tools: Some(vec![Tool::from( + FunctionTool::new("weather", json!({"type": "object"})) + .with_description("Get the weather in a location"), + )]), + tool_choice: Some(LanguageModelV4ToolChoice::Required), + ..CallOptions::new(test_prompt()) + }; + let result = model.do_generate(&options).await.unwrap(); + + assert_eq!(result.usage.input_tokens.total, Some(1189)); + assert_eq!(result.usage.input_tokens.cache_read, Some(891)); + assert_eq!(result.usage.input_tokens.no_cache, Some(298)); + 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)); + } + // -- Tool choice tests -- + fn test_tool() -> Tool { + Tool::from( + FunctionTool::new( + "get_weather", + json!({ + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"] + }), + ) + .with_description("Get the current weather"), + ) + } + + #[tokio::test] + async fn tool_choice_auto() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_basic_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + tools: Some(vec![test_tool()]), + tool_choice: Some(LanguageModelV4ToolChoice::Auto), + ..CallOptions::new(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let request = &server.received_requests().await.unwrap()[0]; + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["tool_choice"], json!("auto")); + assert!(body.get("tools").is_some()); + } + + #[tokio::test] + async fn tool_choice_none() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_basic_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + tools: Some(vec![test_tool()]), + tool_choice: Some(LanguageModelV4ToolChoice::None), + ..CallOptions::new(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let request = &server.received_requests().await.unwrap()[0]; + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["tool_choice"], json!("none")); + } + + #[tokio::test] + async fn tool_choice_required() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_basic_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + tools: Some(vec![test_tool()]), + tool_choice: Some(LanguageModelV4ToolChoice::Required), + ..CallOptions::new(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let request = &server.received_requests().await.unwrap()[0]; + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["tool_choice"], json!("required")); + } + + #[tokio::test] + async fn tool_choice_specific_tool() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_basic_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let options = CallOptions { + tools: Some(vec![test_tool()]), + tool_choice: Some(LanguageModelV4ToolChoice::Tool { + tool_name: "get_weather".to_string(), + }), + ..CallOptions::new(test_prompt()) + }; + model.do_generate(&options).await.unwrap(); + + let request = &server.received_requests().await.unwrap()[0]; + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!( + body["tool_choice"], + json!({"type": "function", "name": "get_weather"}) + ); + } + // -- System messages tests -- #[tokio::test] @@ -1016,6 +1758,75 @@ mod do_generate_tests { // -- Multi-turn tool conversation -- + #[tokio::test] + async fn multi_turn_tool_conversation_request_body() { + let server = MockServer::start().await; + mock_json(&server, lmstudio_basic_json()).await; + + let model = make_model(&server, "gemma-7b-it"); + let prompt = vec![ + LanguageModelMessage::user_text("What is the weather in Tokyo?"), + LanguageModelMessage::Assistant { + content: vec![AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "call_weather_123".into(), + tool_name: "get_weather".into(), + input: json!({"location": "Tokyo"}), + provider_executed: None, + provider_options: None, + })], + provider_options: None, + }, + 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(), + + provider_options: None, + })], + provider_options: None, + }, + ]; + let options = CallOptions { + tools: Some(vec![Tool::from( + FunctionTool::new( + "get_weather", + json!({ + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"] + }), + ) + .with_description("Get the current weather for a location"), + )]), + ..CallOptions::new(prompt) + }; + model.do_generate(&options).await.unwrap(); + + let request = &server.received_requests().await.unwrap()[0]; + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!( + body, + json!({ + "model": "gemma-7b-it", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is the weather in Tokyo?"}]}, + {"type": "function_call", "call_id": "call_weather_123", "name": "get_weather", "arguments": "{\"location\":\"Tokyo\"}"}, + {"type": "function_call_output", "call_id": "call_weather_123", "output": "{\"temperature\":22,\"condition\":\"sunny\",\"humidity\":65}"} + ], + "tools": [{ + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, + "description": "Get the current weather for a location" + }] + }) + ); + } + // -- PDF input file -- #[tokio::test] @@ -1023,7 +1834,7 @@ mod do_generate_tests { let server = MockServer::start().await; mock_json(&server, openai_pdf_json()).await; - let model = make_model(&server, "gpt-4.1-nano"); + let model = make_model(&server, "test-pdf-model"); let prompt = vec![LanguageModelMessage::User { content: vec![ UserPart::Text(TextPart { @@ -1033,8 +1844,8 @@ 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, + }, media_type: "application/pdf".into(), filename: None, provider_options: None @@ -1049,7 +1860,7 @@ provider_options: None assert_eq!( body, json!({ - "model": "gpt-4.1-nano", + "model": "test-pdf-model", "input": [{ "type": "message", "role": "user", @@ -1067,7 +1878,7 @@ provider_options: None let server = MockServer::start().await; mock_json(&server, openai_pdf_json()).await; - let model = make_model(&server, "gpt-4.1-nano"); + let model = make_model(&server, "test-pdf-model"); let prompt = vec![LanguageModelMessage::User { content: vec![ UserPart::Text(TextPart { @@ -1077,8 +1888,8 @@ 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, + }, media_type: "application/pdf".into(), filename: None, provider_options: None @@ -1103,7 +1914,7 @@ provider_options: None let server = MockServer::start().await; mock_json(&server, openai_pdf_json()).await; - let model = make_model(&server, "gpt-4.1-nano"); + let model = make_model(&server, "test-pdf-model"); let prompt = vec![LanguageModelMessage::User { content: vec![ UserPart::Text(TextPart { @@ -1113,8 +1924,8 @@ 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, + }, media_type: "application/pdf".into(), filename: None, provider_options: None @@ -1328,7 +2139,7 @@ mod do_stream_tests { let (tc_id, tc_name, tc_input) = tool_call.unwrap(); assert_eq!(tc_id, "call_1"); assert_eq!(tc_name, "weather"); - assert_eq!(tc_input, r#"{"location":"SF"}"#); + assert_eq!(tc_input, &Value::String(r#"{"location":"SF"}"#.into())); // Finish with tool-calls reason let finish = parts.iter().find_map(|p| match p { @@ -1387,7 +2198,7 @@ mod do_stream_tests { ); mock_sse(&server, body).await; - let model = make_model(&server, "gpt-4.1-nano"); + let model = make_model(&server, "test-pdf-model"); let prompt = vec![LanguageModelMessage::User { content: vec![ UserPart::Text(TextPart { @@ -1397,8 +2208,8 @@ 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, + }, media_type: "application/pdf".into(), filename: None, provider_options: None @@ -1459,89 +2270,3 @@ provider_options: None let _parts = collect_stream(result).await; } } - -// ============================================================================ -// config_snapshot / api_key_source (M2b) -// ============================================================================ - -#[cfg(test)] -mod config_snapshot_tests { - use super::*; - use aimux_core::language_model::LanguageModel; - - /// No `headers` closure → no auth → `api_key_source == "none"` (e.g. a local - /// LM Studio server). - #[test] - fn no_headers_means_none_source() { - let provider = OpenResponsesProvider::new(make_config_dummy()); - let snap = provider.model("gemma-7b-it").config_snapshot(); - assert_eq!(snap.provider, "lmstudio"); - assert_eq!(snap.api_key_source, "none"); - } - - /// A `headers` closure that carries an `Authorization` header → - /// `api_key_source == "explicit"` (auth detected from the closure; the - /// secret value is never serialized). - #[test] - fn headers_with_auth_means_explicit_source() { - let config = - OpenResponsesConfig::new("lmstudio", "lmstudio", "https://example/v1/responses") - .with_headers(|| { - let mut h = HashMap::new(); - h.insert( - "Authorization".to_string(), - "Bearer super-secret".to_string(), - ); - h - }); - let snap = OpenResponsesProvider::new(config) - .model("gemma-7b-it") - .config_snapshot(); - assert_eq!(snap.api_key_source, "explicit"); - let json = serde_json::to_string(&snap).unwrap(); - assert!( - !json.contains("super-secret"), - "plaintext secret leaked: {json}" - ); - } - - /// An explicit `with_api_key_source` marker overrides closure inference - /// (lets a caller record an `env:VAR` source precisely). - #[test] - fn explicit_source_overrides_closure_inference() { - let config = - OpenResponsesConfig::new("lmstudio", "lmstudio", "https://example/v1/responses") - .with_headers(|| { - let mut h = HashMap::new(); - h.insert("Authorization".to_string(), "Bearer k".to_string()); - h - }) - .with_api_key_source(Some("env:OPEN_RESPONSES_API_KEY")); - let snap = OpenResponsesProvider::new(config) - .model("gemma-7b-it") - .config_snapshot(); - assert_eq!(snap.api_key_source, "env:OPEN_RESPONSES_API_KEY"); - } - - /// A `headers` closure carrying only non-auth custom headers → `none` - /// (auth detection is key-name based, not "closure present"). - #[test] - fn headers_without_auth_means_none_source() { - let config = - OpenResponsesConfig::new("lmstudio", "lmstudio", "https://example/v1/responses") - .with_headers(|| { - let mut h = HashMap::new(); - h.insert("X-Custom".to_string(), "value".to_string()); - h - }); - let snap = OpenResponsesProvider::new(config) - .model("gemma-7b-it") - .config_snapshot(); - assert_eq!(snap.api_key_source, "none"); - } - - /// Minimal config without a mock server (snapshot needs no HTTP). - fn make_config_dummy() -> OpenResponsesConfig { - OpenResponsesConfig::new("lmstudio", "lmstudio", "https://example/v1/responses") - } -} diff --git a/aimux-providers/tests/openai_compatible_factory_test.rs b/aimux-providers/tests/openai_compatible_factory_test.rs new file mode 100644 index 000000000..1dada4c3c --- /dev/null +++ b/aimux-providers/tests/openai_compatible_factory_test.rs @@ -0,0 +1,854 @@ +//! `create_openai_compatible` against the recorded AI SDK behavior. +//! +//! `fixtures/aisdk/openai-compatible/*.json` hold what +//! `@ai-sdk/openai-compatible`'s `createOpenAICompatible` sent over a mocked +//! fetch for a fixed input: URL, headers, body and the model's `provider` +//! string. Each fixture is replayed here through an injected +//! [`Fetch`] (the factory's `fetch` setting), so the test also proves the +//! transport is the one the settings name. The expected request comes from the +//! fixture; the inputs are rebuilt from the fixture's recorded `sdk.input`. +//! +//! The recorded `authorization` is redacted, so the key the case used is +//! substituted back. The recorded SDK and runtime user-agent identifiers are +//! replaced with the pinned provider package's user-agent suffix. + +use std::collections::{BTreeMap, HashMap, VecDeque}; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::StreamExt; +use reqwest::StatusCode; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use serde_json::{Value, json}; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, TextPart, +}; +use aimux_core::options::CallOptions; +use aimux_core::result::{GenerateContent, GenerateResult}; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{FunctionTool, LanguageModelV4ToolChoice, Tool}; +use aimux_core::types::FinishReasonUnified; +use aimux_provider_utils::{ + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, +}; +use aimux_providers::openai_compatible::{ + OpenAICompatibleProvider, OpenAICompatibleProviderSettings, create_openai_compatible, +}; + +// ── injected transport ─────────────────────────────────────────────────────── + +const KEY: &str = "sk-test-fixture"; + +const ENV_KEY: &str = "sk-env-should-not-be-used"; + +#[derive(Clone)] +struct Canned { + status: u16, + headers: Vec<(String, String)>, + body: Vec, +} + +impl Canned { + fn json(body: &Value) -> Self { + Self { + status: 200, + headers: vec![("content-type".into(), "application/json".into())], + body: serde_json::to_vec(body).unwrap(), + } + } +} + +/// One request as the transport saw it. +#[derive(Clone, Debug)] +struct Seen { + method: String, + url: String, + /// Lower-cased names (the `http` crate normalizes them). + headers: BTreeMap, + body: Vec, +} + +impl Seen { + fn json_body(&self) -> Value { + serde_json::from_slice(&self.body).expect("request body is JSON") + } +} + +#[derive(Default)] +struct MockFetch { + canned: Mutex>, + seen: Mutex>, +} + +impl MockFetch { + fn new(canned: Vec) -> Arc { + Arc::new(Self { + canned: Mutex::new(canned.into()), + seen: Mutex::default(), + }) + } + + fn seen(&self) -> Vec { + self.seen.lock().unwrap().clone() + } + + fn transport(self: &Arc) -> FetchFunction { + self.clone() + } +} + +#[async_trait] +impl Fetch for MockFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + self.seen.lock().unwrap().push(Seen { + method: request.method.to_string(), + url: request.url.to_string(), + headers: request + .headers + .iter() + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().expect("ASCII header").to_string(), + ) + }) + .collect(), + body: request.body.to_vec(), + }); + let canned = self + .canned + .lock() + .unwrap() + .pop_front() + .ok_or_else(|| FetchError::Other("no canned response left".into()))?; + let mut headers = HeaderMap::new(); + for (name, value) in &canned.headers { + headers.insert( + HeaderName::try_from(name.as_str()).unwrap(), + HeaderValue::try_from(value.as_str()).unwrap(), + ); + } + Ok(FetchResponse::from_bytes( + StatusCode::from_u16(canned.status).unwrap(), + headers, + request.url, + Bytes::from(canned.body), + )) + } +} + +/// The settings every Groq-named fixture was recorded with +/// (`scripts/aisdk-fixtures/cases/openai-compatible.mjs`). +fn groq_settings(mock: &Arc) -> OpenAICompatibleProviderSettings { + OpenAICompatibleProviderSettings { + name: "groq".to_string(), + base_url: "https://api.groq.com/openai/v1".to_string(), + api_key: Some(KEY.to_string()), + fetch: Some(mock.transport()), + ..Default::default() + } +} + +/// Sets (or removes) an environment variable and restores it on drop. Tests +/// that use it are `#[serial]`. +struct EnvVar { + name: &'static str, + saved: Option, +} + +impl EnvVar { + fn set(name: &'static str, value: Option<&str>) -> Self { + let saved = std::env::var(name).ok(); + unsafe { + match value { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + Self { name, saved } + } +} + +impl Drop for EnvVar { + fn drop(&mut self) { + unsafe { + match &self.saved { + Some(value) => std::env::set_var(self.name, value), + None => std::env::remove_var(self.name), + } + } + } +} + +// ── fixtures ───────────────────────────────────────────────────────────────── + +/// Every fixture under `fixtures/aisdk/openai/`; the directory listing is +/// compared to this so a new fixture cannot go unreplayed. +const FIXTURES: [&str; 7] = [ + "chat-basic", + "chat-no-api-key", + "chat-provider-options-generic-key-unknown", + "chat-provider-options-namespace", + "chat-query-params-and-headers", + "chat-stream-basic", + "embedding-basic", +]; + +fn fixtures_dir() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../fixtures/aisdk/openai-compatible") +} + +struct Fixture { + name: &'static str, + json: Value, +} + +impl Fixture { + fn load(name: &'static str) -> Self { + assert!(FIXTURES.contains(&name), "unlisted fixture {name}"); + let path = fixtures_dir().join(format!("{name}.json")); + let text = std::fs::read_to_string(&path) + .unwrap_or_else(|e| panic!("cannot read {}: {e}", path.display())); + Self { + name, + json: serde_json::from_str(&text).expect("fixture is JSON"), + } + } + + fn provider_string(&self) -> &str { + self.json["model"]["provider"].as_str().unwrap() + } + + fn model_id(&self) -> &str { + self.json["model"]["modelId"].as_str().unwrap() + } + + fn sdk_input(&self) -> &Value { + &self.json["sdk"]["input"] + } + + /// The transport answering with the recorded response (or a plain chat + /// completion for the case that never reached the wire). + fn mock(&self) -> Arc { + let response = &self.json["response"]; + if response.is_null() { + return MockFetch::new(vec![Canned::json(&json!({ + "id": "chatcmpl-unused", "object": "chat.completion", "created": 1, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))]); + } + let body = match &response["body"] { + Value::String(text) => text.clone().into_bytes(), + other => serde_json::to_vec(other).unwrap(), + }; + MockFetch::new(vec![Canned { + status: response["status"].as_u64().unwrap() as u16, + headers: response["headers"] + .as_object() + .unwrap() + .iter() + .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string())) + .collect(), + body, + }]) + } + + /// The recorded request must equal what went out: method, URL, headers + /// (names case-insensitively, with the pinned user-agent) and body. + /// `authorization` is the value the redacted recording stands for, or + /// `None` when the recording carries no `authorization` header at all. + fn assert_request(&self, seen: &Seen, authorization: Option<&str>) { + let recorded = &self.json["request"]; + assert_eq!( + seen.method, + recorded["method"].as_str().unwrap(), + "{}", + self.name + ); + assert_eq!(seen.url, recorded["url"].as_str().unwrap(), "{}", self.name); + + let mut expected: BTreeMap = recorded["headers"] + .as_object() + .unwrap() + .iter() + .map(|(k, v)| (k.to_ascii_lowercase(), v.as_str().unwrap().to_string())) + .collect(); + expected.insert( + "user-agent".to_string(), + "ai-sdk-openai-compatible/3.0.59 ai-sdk-provider-utils/5.0.51".to_string(), + ); + match authorization { + Some(value) => { + assert_eq!( + expected.get("authorization").map(String::as_str), + Some(""), + "{}: the recording redacts the key", + self.name + ); + expected.insert("authorization".to_string(), value.to_string()); + } + None => assert!( + !expected.contains_key("authorization"), + "{}: the recording carries no authorization", + self.name + ), + } + assert_eq!(seen.headers, expected, "{}: headers", self.name); + + assert_eq!(seen.json_body(), recorded["body"], "{}: body", self.name); + } + + /// The call options the recording was made with. + fn call_options(&self) -> CallOptions { + call_options_from(self.sdk_input()) + } +} + +fn message(role: &str, text: &str) -> LanguageModelMessage { + match role { + "system" => LanguageModelMessage::System { + content: text.to_string(), + provider_options: None, + }, + "user" => LanguageModelMessage::user_text(text), + "assistant" => LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: text.to_string(), + provider_options: None, + })], + provider_options: None, + }, + other => unreachable!("fixture has no {other} messages"), + } +} + +/// Rebuild aimux call options from a recorded `sdk.input` (the subset the +/// OpenAI fixtures use). +fn call_options_from(input: &Value) -> CallOptions { + let mut prompt: LanguageModelPrompt = Vec::new(); + if let Some(system) = input.get("system").and_then(Value::as_str) { + prompt.push(message("system", system)); + } + if let Some(text) = input.get("prompt").and_then(Value::as_str) { + prompt.push(message("user", text)); + } + for entry in input + .get("messages") + .and_then(Value::as_array) + .into_iter() + .flatten() + { + prompt.push(message( + entry["role"].as_str().unwrap(), + entry["content"].as_str().unwrap(), + )); + } + let mut options = CallOptions::new(prompt); + options.temperature = input.get("temperature").and_then(Value::as_f64); + options.top_p = input.get("topP").and_then(Value::as_f64); + options.max_output_tokens = input + .get("maxOutputTokens") + .and_then(Value::as_u64) + .map(|v| v as u32); + options.seed = input.get("seed").and_then(Value::as_u64); + options.stop_sequences = input + .get("stopSequences") + .and_then(Value::as_array) + .map(|a| a.iter().map(|v| v.as_str().unwrap().to_string()).collect()); + if let Some(tools) = input.get("tools").and_then(Value::as_object) { + options.tools = Some( + tools + .iter() + .map(|(name, tool)| { + let mut function = FunctionTool::new(name.clone(), tool["inputSchema"].clone()); + if let Some(description) = tool["description"].as_str() { + function = function.with_description(description); + } + Tool::Function(function) + }) + .collect(), + ); + } + options.tool_choice = match input.get("toolChoice").and_then(Value::as_str) { + Some("required") => Some(LanguageModelV4ToolChoice::Required), + Some("none") => Some(LanguageModelV4ToolChoice::None), + Some("auto") => Some(LanguageModelV4ToolChoice::Auto), + None => options + .tools + .as_ref() + .filter(|tools| !tools.is_empty()) + .map(|_| LanguageModelV4ToolChoice::Auto), + Some(other) => panic!("unmapped toolChoice {other}"), + }; + options.provider_options = input + .get("providerOptions") + .map(|o| serde_json::from_value(o.clone()).unwrap()); + options.headers = input.get("headers").and_then(Value::as_object).map(|h| { + h.iter() + .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string())) + .collect() + }); + options +} + +fn generated_text(result: &GenerateResult) -> String { + result + .content + .iter() + .filter_map(|part| match part { + GenerateContent::Text { text, .. } => Some(text.as_str()), + _ => None, + }) + .collect() +} + +/// Result fields every recorded `generateText` case pins. +fn assert_generate_result(fixture: &Fixture, result: &GenerateResult) { + let recorded = &fixture.json["result"]; + assert_eq!( + generated_text(result), + recorded["text"].as_str().unwrap(), + "{}: text", + fixture.name + ); + assert_eq!( + result.response.as_ref().and_then(|r| r.id.as_deref()), + recorded["response"]["id"].as_str(), + "{}: 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 _), + "{}: input tokens", + fixture.name + ); + assert_eq!( + result.usage.output_tokens.total, + recorded["usage"]["outputTokens"].as_u64().map(|v| v as _), + "{}: output tokens", + fixture.name + ); + let unified = match recorded["finishReason"].as_str().unwrap() { + "stop" => FinishReasonUnified::Stop, + "tool-calls" => FinishReasonUnified::ToolCalls, + other => panic!("unmapped finish reason {other}"), + }; + assert_eq!( + result.finish_reason.unified, unified, + "{}: finish", + fixture.name + ); + if let Some(raw) = recorded["rawFinishReason"].as_str() { + assert_eq!( + result.finish_reason.raw.as_deref(), + Some(raw), + "{}: raw finish", + fixture.name + ); + } +} + +fn bearer(key: &str) -> String { + format!("Bearer {key}") +} +fn opt_map(entries: &[(&str, Option<&str>)]) -> HeaderMapOpt { + entries + .iter() + .map(|(k, v)| ((*k).to_string(), v.map(str::to_string))) + .collect() +} + +// ── the fixtures ───────────────────────────────────────────────────────────── + +#[test] +fn every_fixture_is_replayed() { + let mut on_disk: Vec = std::fs::read_dir(fixtures_dir()) + .expect("fixtures/aisdk/openai-compatible exists") + .map(|entry| { + let name = entry.unwrap().file_name().into_string().unwrap(); + name.strip_suffix(".json") + .unwrap_or_else(|| panic!("unexpected file {name}")) + .to_string() + }) + .collect(); + on_disk.sort(); + assert_eq!( + on_disk, FIXTURES, + "add a test for the new fixture and list it" + ); +} + +/// Metadata is reported under the provider name, with the namespace entry +/// present even when empty. +fn assert_metadata_key(fixture: &Fixture, result: &GenerateResult, key: &str) { + assert_eq!( + fixture.json["result"]["providerMetadata"], + json!({ key: {} }), + "{}: recorded", + fixture.name + ); + assert_eq!( + serde_json::to_value(result.provider_metadata.as_ref().unwrap()).unwrap(), + json!({ key: {} }), + "{}: metadata key", + fixture.name + ); +} + +/// Replay a `generateText` fixture recorded with the Groq-named provider. +async fn replay_groq_chat(name: &'static str) -> GenerateResult { + let fixture = Fixture::load(name); + let mock = fixture.mock(); + let provider = create_openai_compatible(groq_settings(&mock)).unwrap(); + let model = provider.chat(fixture.model_id()); + assert_eq!(model.provider(), "groq.chat"); + assert_eq!( + model.provider(), + fixture.provider_string(), + "{name}: provider()" + ); + assert_eq!(model.model_id(), fixture.model_id()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1, "{name}: exactly one request"); + fixture.assert_request(&seen[0], Some(&bearer(KEY))); + assert_generate_result(&fixture, &result); + assert_metadata_key(&fixture, &result, "groq"); + result +} + +#[tokio::test] +async fn chat_basic() { + replay_groq_chat("chat-basic").await; +} + +#[tokio::test] +async fn chat_provider_options_namespace() { + // `groq.foo` passes through at the top level; `openaiCompatible.user` is a + // schema field. + replay_groq_chat("chat-provider-options-namespace").await; + let fixture = Fixture::load("chat-provider-options-namespace"); + assert_eq!( + fixture.json["observations"], + json!({"bodyHasFooAtTopLevel": true, "bodyUser": "u1", "providerIsGroqChat": true}) + ); +} + +#[tokio::test] +async fn unknown_fields_of_the_generic_namespace_are_not_forwarded() { + replay_groq_chat("chat-provider-options-generic-key-unknown").await; + let fixture = Fixture::load("chat-provider-options-generic-key-unknown"); + assert_eq!( + fixture.json["observations"]["bodyHasExtraField"], + json!(false) + ); +} + +#[tokio::test] +async fn chat_query_params_and_headers() { + let fixture = Fixture::load("chat-query-params-and-headers"); + let mock = fixture.mock(); + // The recording was made with headers { "x-a": "1", "X-B": "2" } and + // queryParams { "api-version": "1" }. + let provider = create_openai_compatible(OpenAICompatibleProviderSettings { + headers: Some(opt_map(&[("x-a", Some("1")), ("X-B", Some("2"))])), + query_params: Some(HashMap::from([( + "api-version".to_string(), + "1".to_string(), + )])), + ..groq_settings(&mock) + }) + .unwrap(); + let model = provider.chat(fixture.model_id()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + let seen = mock.seen(); + fixture.assert_request(&seen[0], Some(&bearer(KEY))); + assert_generate_result(&fixture, &result); + // Case-insensitive override: the call's `x-b` replaced the provider's `X-B`. + assert_eq!(seen[0].headers["x-b"], "3"); + let sent: Vec<&String> = seen[0] + .headers + .keys() + .filter(|n| n.starts_with("x-")) + .collect(); + assert_eq!(sent, ["x-a", "x-b", "x-c"]); + assert_eq!( + fixture.json["observations"]["headerNamesAsSent"], + json!(["x-a", "x-b", "x-c"]) + ); +} + +#[tokio::test] +async fn chat_without_an_api_key_sends_no_authorization() { + let fixture = Fixture::load("chat-no-api-key"); + let mock = fixture.mock(); + // A key in the environment must not leak in either. + let _env = EnvVar::set("OPENAI_API_KEY", Some(ENV_KEY)); + let provider = create_openai_compatible(OpenAICompatibleProviderSettings { + name: "local".to_string(), + base_url: "http://localhost:1234/v1".to_string(), + fetch: Some(mock.transport()), + ..Default::default() + }) + .expect("no key is needed to create the provider"); + let model = provider.chat(fixture.model_id()); + assert_eq!(model.provider(), "local.chat"); + assert_eq!(model.provider(), fixture.provider_string()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + fixture.assert_request(&seen[0], None); + assert!(!seen[0].headers.contains_key("authorization")); + assert_eq!( + fixture.json["observations"]["hasAuthorizationHeader"], + json!(false) + ); + assert_generate_result(&fixture, &result); + assert_metadata_key(&fixture, &result, "local"); +} + +/// `supportedUrls` is read through the settings function each time, and its +/// value may be async (openai-compatible@3.0.62 +/// `openai-compatible-chat-language-model.ts:84`, `:161-163`); completion +/// models get none (`openai-compatible-provider.ts:189-193`). +#[tokio::test] +async fn chat_supported_urls_come_from_the_async_settings_function() { + let mock = MockFetch::new(Vec::new()); + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let counter = calls.clone(); + let provider = create_openai_compatible(OpenAICompatibleProviderSettings { + supported_urls: Some(Arc::new(move || { + counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Box::pin(async { + tokio::task::yield_now().await; + aimux_core::language_model::SupportedUrls(HashMap::from([( + "image/*".to_string(), + vec![regex::Regex::new("^https://").unwrap()], + )])) + }) + })), + ..groq_settings(&mock) + }) + .unwrap(); + let chat = provider.chat("m"); + for _ in 0..2 { + let urls = chat.supported_urls().await; + assert_eq!(urls.0.keys().collect::>(), ["image/*"]); + } + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); + assert!(provider.completion("m").supported_urls().await.is_empty()); + assert!(mock.seen().is_empty()); +} + +#[tokio::test] +async fn chat_stream_basic() { + let fixture = Fixture::load("chat-stream-basic"); + let mock = fixture.mock(); + let model = create_openai_compatible(groq_settings(&mock)) + .unwrap() + .chat(fixture.model_id()); + assert_eq!(model.provider(), fixture.provider_string()); + + let result = model.do_stream(&fixture.call_options()).await.unwrap(); + let parts: Vec = result + .stream + .map(|part| part.expect("no stream error")) + .collect() + .await; + + // Unlike the native OpenAI package, the baseline asks for no + // `stream_options` unless `include_usage` is set: the recorded body has + // none. (The usage chunk is read anyway when the server sends it.) + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + fixture.assert_request(&seen[0], Some(&bearer(KEY))); + assert!(seen[0].json_body().get("stream_options").is_none()); + + let text: String = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextDelta { delta, .. } => Some(delta.as_str()), + _ => None, + }) + .collect(); + assert_eq!(text, "Hello world"); + let (reason, usage, metadata) = parts + .iter() + .find_map(|part| match part { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => Some((finish_reason, usage, provider_metadata)), + _ => None, + }) + .expect("a finish part"); + assert_eq!(reason.unified, FinishReasonUnified::Stop); + assert_eq!(usage.input_tokens.total, Some(12)); + assert_eq!(usage.output_tokens.total, Some(2)); + assert_eq!( + serde_json::to_value(metadata.as_ref().unwrap()).unwrap(), + json!({"groq": {}}) + ); +} + +#[tokio::test] +async fn embedding_basic() { + let fixture = Fixture::load("embedding-basic"); + let mock = fixture.mock(); + let model = create_openai_compatible(groq_settings(&mock)) + .unwrap() + .embedding(fixture.model_id()); + assert_eq!(model.provider(), "groq.embedding"); + assert_eq!(model.provider(), fixture.provider_string()); + + let values: Vec = fixture.sdk_input()["values"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap().to_string()) + .collect(); + let options = EmbeddingCallOptions { + values, + ..EmbeddingCallOptions::new("") + }; + let result = model.do_embed(&options).await.unwrap(); + + fixture.assert_request(&mock.seen()[0], Some(&bearer(KEY))); + let recorded: Vec> = + serde_json::from_value(fixture.json["result"]["embeddings"].clone()).unwrap(); + assert_eq!(result.embeddings, recorded); + assert_eq!( + result.usage.map(|u| u.tokens), + fixture.json["result"]["usage"]["tokens"] + .as_u64() + .map(|v| v as u32) + ); +} + +fn _assert_provider_is_shareable() { + fn shareable() {} + shareable::(); +} + +// ── compat request body key order ──────────────────────────────────────────── + +/// One chat request with `options`; returns the body bytes as sent. +async fn sent_chat_body(options: CallOptions) -> String { + let mock = MockFetch::new(vec![Canned::json(&json!({ + "id": "chatcmpl-order", "object": "chat.completion", "created": 1, + "model": "m", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))]); + let provider = create_openai_compatible(groq_settings(&mock)).unwrap(); + provider.chat("m").do_generate(&options).await.unwrap(); + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + String::from_utf8(seen[0].body.clone()).unwrap() +} + +/// Byte offset of top-level key `key` in a compact JSON object. +fn key_at(body: &str, key: &str) -> usize { + body.find(&format!("\"{key}\":")) + .unwrap_or_else(|| panic!("{key} missing in {body}")) +} + +/// A user "hi" prompt with a `weather` tool, `toolChoice: auto` and the given +/// `groq` provider options. +fn tool_call_options(groq: Value) -> CallOptions { + let mut options = CallOptions::new(vec![LanguageModelMessage::user_text("hi")]); + options.temperature = Some(0.5); + options.tools = Some(vec![Tool::Function(FunctionTool::new( + "weather".to_string(), + json!({"type": "object", "properties": {}}), + ))]); + options.tool_choice = Some(LanguageModelV4ToolChoice::Auto); + options.provider_options = Some(serde_json::from_value(json!({ "groq": groq })).unwrap()); + options +} + +/// The body bytes follow openai-compatible@3.0.62 +/// `openai-compatible-chat-language-model.ts:270-325` (the alignment doc's +/// "compat 请求体顺序"): model, user → sampling, response_format, stop, seed +/// → unknown provider-namespace fields → reasoning_effort, verbosity → +/// messages, tools, tool_choice. `serde_json::Value` ignores order, so this +/// asserts on the bytes. Ported from li's fork PR #35 (f9c9653b). +#[tokio::test] +async fn compat_body_keys_follow_the_documented_order() { + let mut options = tool_call_options(json!({ + "zzzUnknown": 1, + "user": "u1", + "reasoningEffort": "low", + "textVerbosity": "low" + })); + options.stop_sequences = Some(vec!["END".to_string()]); + options.seed = Some(7); + let body = sent_chat_body(options).await; + let order = [ + "model", + "user", + "temperature", + "stop", + "seed", + "zzzUnknown", + "reasoning_effort", + "verbosity", + "messages", + "tools", + "tool_choice", + ]; + let at: Vec = order.iter().map(|k| key_at(&body, k)).collect(); + for (i, pair) in at.windows(2).enumerate() { + assert!( + pair[0] < pair[1], + "{} must come before {}: {body}", + order[i], + order[i + 1] + ); + } +} + +/// An unknown field may override earlier fields (sampling), but not the +/// fields assigned after it: `messages` / `tools` / `tool_choice` from the +/// provider namespace are replaced by the real ones (the same spread order). +#[tokio::test] +async fn unknown_fields_override_sampling_but_not_messages_or_tools() { + let body = sent_chat_body(tool_call_options(json!({ + "temperature": 0.9, + "messages": "bogus", + "tools": "bogus", + "tool_choice": "bogus" + }))) + .await; + let sent: Value = serde_json::from_str(&body).unwrap(); + assert_eq!(sent["temperature"], json!(0.9), "{body}"); + assert_eq!(sent["messages"][0]["role"], json!("user"), "{body}"); + assert_eq!( + sent["tools"][0]["function"]["name"], + json!("weather"), + "{body}" + ); + assert_eq!(sent["tool_choice"], json!("auto"), "{body}"); + assert!(!body.contains("bogus"), "{body}"); +} diff --git a/aimux-providers/tests/openai_compatible_test.rs b/aimux-providers/tests/openai_compatible_test.rs index 96830dd28..575df1494 100644 --- a/aimux-providers/tests/openai_compatible_test.rs +++ b/aimux-providers/tests/openai_compatible_test.rs @@ -1,7 +1,7 @@ //! Wiremock tests for the OpenAI-compatible provider wrappers. //! -//! Each of the 8 wrappers (groq, deepseek, togetherai, fireworks, perplexity, -//! cerebras, xai, moonshotai) is a thin layer over [`OpenAIProvider`] that +//! Each of the wrappers (groq, deepseek, togetherai, fireworks, perplexity, +//! cerebras, moonshotai, ...) is a registry preset over the OpenAI-compatible package that //! only fixes the default base URL and the API-key environment variable. These //! tests verify, for every wrapper, the four behaviours that the wrapper is //! responsible for getting right: @@ -18,6 +18,7 @@ use aimux_core::tool::RawToolCall; use futures::StreamExt; use serde_json::{Value, json}; +use std::sync::Arc; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -29,30 +30,25 @@ use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; use aimux_core::types::FinishReasonUnified; -use aimux_providers::{ - ProviderOptions, bedrock_mantle::BedrockMantleConfig, bedrock_mantle::BedrockMantleProvider, - huggingface::HuggingFaceConfig, huggingface::HuggingFaceProvider, provider, - vertex_ai_ai21_models::VertexAiAi21ModelsConfig, - vertex_ai_ai21_models::VertexAiAi21ModelsProvider, - vertex_ai_anthropic_models::VertexAiAnthropicModelsConfig, - vertex_ai_anthropic_models::VertexAiAnthropicModelsProvider, - vertex_ai_deepseek_models::VertexAiDeepseekModelsConfig, - vertex_ai_deepseek_models::VertexAiDeepseekModelsProvider, - vertex_ai_llama_models::VertexAiLlamaModelsConfig, - vertex_ai_llama_models::VertexAiLlamaModelsProvider, - vertex_ai_minimax_models::VertexAiMinimaxModelsConfig, - vertex_ai_minimax_models::VertexAiMinimaxModelsProvider, - vertex_ai_mistral_models::VertexAiMistralModelsConfig, - vertex_ai_mistral_models::VertexAiMistralModelsProvider, - vertex_ai_moonshot_models::VertexAiMoonshotModelsConfig, - vertex_ai_moonshot_models::VertexAiMoonshotModelsProvider, - vertex_ai_openai_models::VertexAiOpenaiModelsConfig, - vertex_ai_openai_models::VertexAiOpenaiModelsProvider, - vertex_ai_qwen_models::VertexAiQwenModelsConfig, - vertex_ai_qwen_models::VertexAiQwenModelsProvider, - vertex_ai_zai_models::VertexAiZaiModelsConfig, vertex_ai_zai_models::VertexAiZaiModelsProvider, - xai::XAIConfig, xai::XAIProvider, -}; +use aimux_providers::{PresetSettings, create_provider}; + +/// A registry preset pointed at `base_url`, as a language model. +fn registry_model( + name: &str, + api_key: String, + model_id: &str, + base_url: String, +) -> Result, AiMuxError> { + create_provider( + name, + PresetSettings { + api_key: Some(api_key.into()), + base_url: Some(base_url), + ..Default::default() + }, + )? + .language_model(model_id) +} // ── shared helpers ─────────────────────────────────────────────────────────── @@ -134,16 +130,14 @@ fn text_deltas(parts: &[StreamPart]) -> Vec { macro_rules! openai_compatible_tests { ( $mod_name:ident, - $config:ty, - $provider:ty, + $ctor:path, $model_id:literal ) => { mod $mod_name { use super::*; - fn make_provider(server: &MockServer) -> Box { - let config = <$config>::new("test-api-key").with_base_url(server.uri()); - Box::new(<$provider>::new(config).model($model_id)) + fn make_provider(server: &MockServer) -> Arc { + $ctor(server.uri(), $model_id) } mod $mod_name { use super::*; @@ -310,16 +304,8 @@ macro_rules! openai_compatible_tests { mod $mod_name { use super::*; - fn make_provider(server: &MockServer) -> Box { - provider( - $provider_name, - Some("test-api-key".to_string()), - $model_id, - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), - ) + fn make_provider(server: &MockServer) -> Arc { + registry_model($provider_name, "test-api-key".to_string(), $model_id, server.uri()) .expect("provider construction") } mod $mod_name { @@ -491,7 +477,6 @@ openai_compatible_tests!( ); openai_compatible_tests!(perplexity, "perplexity", "sonar"); openai_compatible_tests!(cerebras, "cerebras", "llama-3.3-70b"); -openai_compatible_tests!(xai, XAIConfig, XAIProvider, "grok-2"); openai_compatible_tests!(moonshotai, "moonshotai", "moonshot-v1-8k"); // Second batch of OpenAI-compatible thin wrappers. @@ -501,12 +486,6 @@ openai_compatible_tests!( "meta-llama/Meta-Llama-3-70B-Instruct" ); openai_compatible_tests!(baseten, "baseten", "deepseek-ai/DeepSeek-V3-0324"); -openai_compatible_tests!( - huggingface, - HuggingFaceConfig, - HuggingFaceProvider, - "meta-llama/Llama-3.3-70B-Instruct" -); openai_compatible_tests!(alibaba, "alibaba", "qwen-max"); openai_compatible_tests!(bytedance, "bytedance", "doubao-pro-32k"); openai_compatible_tests!(vercel, "vercel", "v0-1.5-md"); @@ -551,16 +530,14 @@ fn tool_call_completion_body() -> Value { macro_rules! openai_compatible_tool_tests { ( $mod_name:ident, - $config:ty, - $provider:ty, + $ctor:path, $model_id:literal ) => { mod $mod_name { use super::*; - fn make_provider(server: &MockServer) -> Box { - let config = <$config>::new("test-api-key").with_base_url(server.uri()); - Box::new(<$provider>::new(config).model($model_id)) + fn make_provider(server: &MockServer) -> Arc { + $ctor(server.uri(), $model_id) } mod $mod_name { use super::*; @@ -590,7 +567,7 @@ macro_rules! openai_compatible_tool_tests { GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. }) => { assert_eq!(tool_call_id, "call_abc"); assert_eq!(tool_name, "get-weather"); - assert_eq!(input, r#"{"city":"SF"}"#); + assert_eq!(input, &Value::String(r#"{"city":"SF"}"#.into())); } other => panic!("expected ToolCall, got {:?}", other), } @@ -633,7 +610,7 @@ macro_rules! openai_compatible_tool_tests { let (id, name, input) = tool_call.expect("should have ToolCall"); assert_eq!(id, "call_abc"); assert_eq!(name, "get-weather"); - assert_eq!(input, r#"{"city":"SF"}"#); + assert_eq!(input, Value::String(r#"{"city":"SF"}"#.into())); } } } @@ -646,16 +623,8 @@ macro_rules! openai_compatible_tool_tests { mod $mod_name { use super::*; - fn make_provider(server: &MockServer) -> Box { - provider( - $provider_name, - Some("test-api-key".to_string()), - $model_id, - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), - ) + fn make_provider(server: &MockServer) -> Arc { + registry_model($provider_name, "test-api-key".to_string(), $model_id, server.uri()) .expect("provider construction") } mod $mod_name { @@ -686,7 +655,7 @@ macro_rules! openai_compatible_tool_tests { GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. }) => { assert_eq!(tool_call_id, "call_abc"); assert_eq!(tool_name, "get-weather"); - assert_eq!(input, r#"{"city":"SF"}"#); + assert_eq!(input, &Value::String(r#"{"city":"SF"}"#.into())); } other => panic!("expected ToolCall, got {:?}", other), } @@ -729,7 +698,7 @@ macro_rules! openai_compatible_tool_tests { let (id, name, input) = tool_call.expect("should have ToolCall"); assert_eq!(id, "call_abc"); assert_eq!(name, "get-weather"); - assert_eq!(input, r#"{"city":"SF"}"#); + assert_eq!(input, Value::String(r#"{"city":"SF"}"#.into())); } } } @@ -750,7 +719,6 @@ openai_compatible_tool_tests!( ); openai_compatible_tool_tests!(perplexity_tools, "perplexity", "sonar"); openai_compatible_tool_tests!(cerebras_tools, "cerebras", "llama-3.3-70b"); -openai_compatible_tool_tests!(xai_tools, XAIConfig, XAIProvider, "grok-2"); openai_compatible_tool_tests!(moonshotai_tools, "moonshotai", "moonshot-v1-8k"); openai_compatible_tool_tests!( deepinfra_tools, @@ -758,12 +726,6 @@ openai_compatible_tool_tests!( "meta-llama/Meta-Llama-3-70B-Instruct" ); openai_compatible_tool_tests!(baseten_tools, "baseten", "deepseek-ai/DeepSeek-V3-0324"); -openai_compatible_tool_tests!( - huggingface_tools, - HuggingFaceConfig, - HuggingFaceProvider, - "meta-llama/Llama-3.3-70B-Instruct" -); openai_compatible_tool_tests!(alibaba_tools, "alibaba", "qwen-max"); openai_compatible_tool_tests!(bytedance_tools, "bytedance", "doubao-pro-32k"); openai_compatible_tool_tests!(vercel_tools, "vercel", "v0-1.5-md"); @@ -790,17 +752,14 @@ mod default_base_urls { /// a smoke test that the wrapper wires `with_base_url` through correctly. #[test] fn base_url_override_is_applied() { - // provider() options must override the registry default. Spot-check + // the base URL setting must override the registry default. Spot-check // groq here; the per-provider request tests above cover the remaining // providers end-to-end against a mock server. - let model = provider( + let model = registry_model( "groq", - Some("k".to_string()), + "k".to_string(), "llama-3.3-70b", - Some(ProviderOptions { - base_url: Some("https://example.test/v1".to_string()), - ..Default::default() - }), + "https://example.test/v1".to_string(), ) .expect("provider construction should succeed"); let _ = model; @@ -823,14 +782,11 @@ mod default_base_urls { // important assertion is that WITHOUT the override the wrapper would // target api.groq.com (we can't mock that here), so this test merely // confirms the override mechanism the other tests depend on. - let model = provider( + let model = registry_model( "groq", - Some("test-api-key".to_string()), + "test-api-key".to_string(), "llama-3.3-70b-versatile", - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), + server.uri(), ) .expect("provider construction"); let _ = model @@ -866,14 +822,11 @@ async fn groq_extracts_usage() { .mount(&server) .await; - let model = provider( + let model = registry_model( "groq", - Some("test-api-key".to_string()), + "test-api-key".to_string(), "llama-3.3-70b-versatile", - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), + server.uri(), ) .expect("provider construction"); @@ -898,14 +851,11 @@ async fn deepseek_rate_limit_maps_to_rate_limited() { .mount(&server) .await; - let model = provider( + let model = registry_model( "deepseek", - Some("test-api-key".to_string()), + "test-api-key".to_string(), "deepseek-chat", - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), + server.uri(), ) .expect("provider construction"); @@ -916,6 +866,41 @@ async fn deepseek_rate_limit_maps_to_rate_limited() { ); } +/// TS: response headers are exposed on the generate result. +#[tokio::test] +async fn groq_exposes_response_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("test-header", "test-value") + .set_body_json(text_completion_body()), + ) + .mount(&server) + .await; + + let model = registry_model( + "groq", + "test-api-key".to_string(), + "llama-3.3-70b-versatile", + server.uri(), + ) + .expect("provider construction"); + + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("should succeed"); + + let headers = result + .response + .as_ref() + .and_then(|response| response.headers.as_ref()) + .expect("response_headers should be Some"); + assert_eq!(headers.get("test-header"), Some(&"test-value".to_string())); +} + // P1 thin-wrapper providers (provider-research batch). openai_compatible_tests!(chutes, "chutes", "default"); openai_compatible_tests!(meta, "meta", "meta/muse-spark-1.1"); @@ -1012,8 +997,7 @@ openai_compatible_tests!( ); openai_compatible_tests!( bedrock_mantle, - BedrockMantleConfig, - BedrockMantleProvider, + "bedrock_mantle", "bedrock_mantle/openai.gpt-oss-120b" ); openai_compatible_tests!(cherryin, "cherryin", "BAAI/bge-reranker-v2-m3(free)"); @@ -1114,64 +1098,55 @@ openai_compatible_tests!(zeldoc, "zeldoc", "z-code"); openai_compatible_tests!(privatemode_ai, "privatemode_ai", "gpt-oss-120b"); openai_compatible_tests!(snowflake, "snowflake", "claude-sonnet-4-5"); -// Vertex AI MaaS partner-model providers (OpenAI-compatible thin wrappers). +// Vertex AI MaaS partner-model providers (registry presets with template parameters; +// the tests override the whole base URL). openai_compatible_tests!( vertex_ai_ai21_models, - VertexAiAi21ModelsConfig, - VertexAiAi21ModelsProvider, + "vertex_ai_ai21_models", "ai21/jamba-1.5-large" ); openai_compatible_tests!( vertex_ai_anthropic_models, - VertexAiAnthropicModelsConfig, - VertexAiAnthropicModelsProvider, + "vertex_ai_anthropic_models", "anthropic/claude-sonnet-4" ); openai_compatible_tests!( vertex_ai_deepseek_models, - VertexAiDeepseekModelsConfig, - VertexAiDeepseekModelsProvider, + "vertex_ai_deepseek_models", "deepseek-ai/deepseek-v3.1-maas" ); openai_compatible_tests!( vertex_ai_llama_models, - VertexAiLlamaModelsConfig, - VertexAiLlamaModelsProvider, + "vertex_ai_llama_models", "meta/llama-4-scout-17b-16e-instruct-maas" ); openai_compatible_tests!( vertex_ai_minimax_models, - VertexAiMinimaxModelsConfig, - VertexAiMinimaxModelsProvider, + "vertex_ai_minimax_models", "minimax/minimax-m2-maas" ); openai_compatible_tests!( vertex_ai_mistral_models, - VertexAiMistralModelsConfig, - VertexAiMistralModelsProvider, + "vertex_ai_mistral_models", "mistralai/mistral-large-2411" ); openai_compatible_tests!( vertex_ai_moonshot_models, - VertexAiMoonshotModelsConfig, - VertexAiMoonshotModelsProvider, + "vertex_ai_moonshot_models", "moonshotai/kimi-k2-thinking-maas" ); openai_compatible_tests!( vertex_ai_openai_models, - VertexAiOpenaiModelsConfig, - VertexAiOpenaiModelsProvider, + "vertex_ai_openai_models", "openai/gpt-oss-120b-maas" ); openai_compatible_tests!( vertex_ai_qwen_models, - VertexAiQwenModelsConfig, - VertexAiQwenModelsProvider, + "vertex_ai_qwen_models", "qwen/qwen3-coder-480b-a35b-instruct-maas" ); openai_compatible_tests!( vertex_ai_zai_models, - VertexAiZaiModelsConfig, - VertexAiZaiModelsProvider, + "vertex_ai_zai_models", "zai-org/glm-4.7-maas" ); diff --git a/aimux-providers/tests/openai_decision_test.rs b/aimux-providers/tests/openai_decision_test.rs index 37866f974..f8d3203fb 100644 --- a/aimux-providers/tests/openai_decision_test.rs +++ b/aimux-providers/tests/openai_decision_test.rs @@ -4,13 +4,12 @@ use std::sync::{ Arc, atomic::{AtomicUsize, Ordering}, }; -use std::time::Duration; use aimux_core::decision_model::*; use aimux_core::recording::{self, RingRecorder}; use aimux_core::replay::{MockDecisionReplayModel, replay_decision_with_model}; use aimux_core::{AiMuxError, Provider}; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::{OpenAIProvider, OpenAIProviderSettings, PresetSettings, create_openai}; use serde_json::{Value, json}; use serial_test::serial; use wiremock::matchers::{body_json, header, method, path}; @@ -31,21 +30,40 @@ fn fixture() -> Value { fn options() -> DecisionCallOptions { serde_json::from_value(fixture()["options"].clone()).unwrap() } -fn config(server: &MockServer) -> OpenAIConfig { - OpenAIConfig::new("decision-secret") - .with_base_url(format!("{}/v1", server.uri())) - .with_org_id("org-test") - .with_project("proj-test") - .with_headers(HashMap::from([( +fn provider(server: &MockServer, api_key: &str) -> OpenAIProvider { + create_openai(OpenAIProviderSettings { + api_key: Some(api_key.into()), + base_url: Some(format!("{}/v1", server.uri())), + organization: Some("org-test".into()), + project: Some("proj-test".into()), + headers: Some(HashMap::from([( "X-Client".into(), - "decision-tests".into(), - )])) + Some("decision-tests".into()), + )])), + ..Default::default() + }) + .unwrap() } fn model(server: &MockServer) -> Box { - OpenAIProvider::new(config(server)) + provider(server, "decision-secret") .decision_model("gpt-6-luna") .unwrap() } +/// The caller's registry at replay time: the same OpenAI settings at the +/// replay server, with another key. +fn rebuild( + record: &aimux_core::recording::ProviderRecord, + server: &MockServer, +) -> Box { + let registry = aimux_core::create_provider_registry( + std::collections::BTreeMap::from([( + "openai".into(), + Arc::new(provider(server, "replacement-key")) as Arc, + )]), + Default::default(), + ); + aimux_providers::rebuild_decision_provider(record, ®istry).unwrap() +} async fn mount(server: &MockServer, response: Value) { Mock::given(method("POST")) .and(path("/v1/decisions")) @@ -257,11 +275,7 @@ async fn retry_and_standard_provider_errors_use_the_shared_transport() { }) .mount(&server) .await; - let mut config = config(&server); - config.retry_config.initial_delay = Duration::ZERO; - let model = OpenAIProvider::new(config) - .decision_model("gpt-6-luna") - .unwrap(); + let model = model(&server); let mut request = options(); request.max_retries = Some(1); decide(model.as_ref(), request).await.unwrap(); @@ -283,11 +297,25 @@ async fn retry_and_standard_provider_errors_use_the_shared_transport() { #[test] fn factory_is_explicit_and_preserves_the_provider_boundary() { - let provider = aimux_providers::provider_handle("openai", Some("key".into()), None).unwrap(); + let provider = aimux_providers::create_provider( + "openai", + PresetSettings { + api_key: Some("key".to_string().into()), + ..Default::default() + }, + ) + .unwrap(); let model = provider.decision_model("deployment-alias").unwrap(); assert_eq!(model.model_id(), "deployment-alias"); assert!(provider.decision_model(" ").is_err()); - let other = aimux_providers::provider_handle("deepseek", Some("key".into()), None).unwrap(); + let other = aimux_providers::create_provider( + "deepseek", + PresetSettings { + api_key: Some("key".to_string().into()), + ..Default::default() + }, + ) + .unwrap(); assert!(matches!( other.decision_model("any"), Err(AiMuxError::UnsupportedFunctionality(_)) @@ -304,7 +332,7 @@ impl Drop for StopRecording { #[tokio::test] #[serial] -async fn unified_recording_replays_refusals_and_rebuilds_the_same_provider_config() { +async fn unified_recording_replays_refusals_and_rebuilds_from_the_registered_provider() { let server = MockServer::start().await; mount(&server, fixture()["response"].clone()).await; let recorder = Arc::new(RingRecorder::new()); @@ -328,10 +356,7 @@ async fn unified_recording_replays_refusals_and_rebuilds_the_same_provider_confi ); let replay_server = MockServer::start().await; common::replay::mount_recording(&replay_server, &records[0]).await; - let mut snapshot = records[0].provider.clone(); - snapshot.base_url = Some(format!("{}/v1", replay_server.uri())); - let rebuilt = - aimux_providers::rebuild_decision_provider(&snapshot, Some("replacement-key")).unwrap(); + let rebuilt = rebuild(&records[0].provider, &replay_server); let again = replay_decision_with_model(&records[0], rebuilt.as_ref()) .await .unwrap(); @@ -383,10 +408,7 @@ async fn images_typed_choices_and_score_descriptions_roundtrip_and_key_replay() assert!(decide(&mock, changed).await.is_err()); let replay_server = MockServer::start().await; common::replay::mount_recording(&replay_server, &records[0]).await; - let mut config = records[0].provider.clone(); - config.base_url = Some(format!("{}/v1", replay_server.uri())); - let rebuilt = - aimux_providers::rebuild_decision_provider(&config, Some("replacement-key")).unwrap(); + let rebuilt = rebuild(&records[0].provider, &replay_server); let again = replay_decision_with_model(&records[0], rebuilt.as_ref()) .await .unwrap(); diff --git a/aimux-providers/tests/openai_embedding_test.rs b/aimux-providers/tests/openai_embedding_test.rs index 84324fedf..4f3f743b6 100644 --- a/aimux-providers/tests/openai_embedding_test.rs +++ b/aimux-providers/tests/openai_embedding_test.rs @@ -13,8 +13,9 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_core::embedding_model::EmbeddingCallOptions; +use aimux_core::provider::Provider; +use aimux_providers::{OpenAIProviderSettings, create_openai}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -76,9 +77,13 @@ async fn should_extract_embedding() { .mount(&server) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.embedding_model("text-embedding-3-large"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.embedding_model("text-embedding-3-large").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -122,9 +127,13 @@ async fn should_expose_raw_response_headers() { .mount(&server) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.embedding_model("text-embedding-3-large"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.embedding_model("text-embedding-3-large").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -152,9 +161,13 @@ async fn should_expose_raw_response_body() { .mount(&server) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.embedding_model("text-embedding-3-large"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.embedding_model("text-embedding-3-large").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -180,9 +193,13 @@ async fn should_extract_usage() { .mount(&server) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.embedding_model("text-embedding-3-large"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.embedding_model("text-embedding-3-large").unwrap(); let result = model .do_embed(&default_options(test_values())) @@ -203,9 +220,13 @@ async fn should_pass_model_and_values() { .mount(&server) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.embedding_model("text-embedding-3-large"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.embedding_model("text-embedding-3-large").unwrap(); let _ = model .do_embed(&default_options(test_values())) @@ -244,13 +265,21 @@ async fn should_pass_headers() { "provider-header-value".to_string(), ); - let config = OpenAIConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_org_id("test-organization") - .with_project("test-project") - .with_headers(config_headers); - let provider = OpenAIProvider::new(config); - let model = provider.embedding_model("text-embedding-3-large"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + organization: Some("test-organization".into()), + project: Some("test-project".into()), + headers: Some( + config_headers + .into_iter() + .map(|(k, v)| (k, Some(v))) + .collect(), + ), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.embedding_model("text-embedding-3-large").unwrap(); let mut request_headers = HashMap::new(); request_headers.insert( diff --git a/aimux-providers/tests/openai_factory_test.rs b/aimux-providers/tests/openai_factory_test.rs new file mode 100644 index 000000000..fe162a676 --- /dev/null +++ b/aimux-providers/tests/openai_factory_test.rs @@ -0,0 +1,888 @@ +//! `create_openai` against the recorded AI SDK behavior. +//! +//! `fixtures/aisdk/openai/*.json` hold what `@ai-sdk/openai`'s `createOpenAI` +//! sent over a mocked fetch for a fixed input: URL, headers, body and the +//! model's `provider` string. Each fixture is replayed here through an injected +//! [`Fetch`] (the factory's `fetch` setting), so the test also proves the +//! transport is the one the settings name. The expected request comes from the +//! fixture; the inputs are rebuilt from the fixture's recorded `sdk.input`. +//! +//! The recorded `authorization` is redacted, so the key the case used is +//! substituted back. The recorded SDK and runtime user-agent identifiers are +//! replaced with the pinned provider package's user-agent suffix. + +use aimux_core::tool::RawToolCall; +use std::collections::{BTreeMap, VecDeque}; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::StreamExt; +use reqwest::StatusCode; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use serde_json::{Value, json}; +use serial_test::serial; + +use aimux_core::AiMuxError; +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, TextPart, +}; +use aimux_core::options::CallOptions; +use aimux_core::result::{GenerateContent, GenerateResult}; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{FunctionTool, LanguageModelV4ToolChoice, Tool}; +use aimux_core::types::FinishReasonUnified; +use aimux_provider_utils::{ + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, +}; +use aimux_providers::openai::{OpenAIProvider, OpenAIProviderSettings, create_openai}; + +// ── injected transport ─────────────────────────────────────────────────────── + +const KEY: &str = "sk-test-fixture"; + +const ENV_KEY: &str = "sk-env-should-not-be-used"; + +#[derive(Clone)] +struct Canned { + status: u16, + headers: Vec<(String, String)>, + body: Vec, +} + +impl Canned { + fn json(body: &Value) -> Self { + Self { + status: 200, + headers: vec![("content-type".into(), "application/json".into())], + body: serde_json::to_vec(body).unwrap(), + } + } +} + +/// One request as the transport saw it. +#[derive(Clone, Debug)] +struct Seen { + method: String, + url: String, + /// Lower-cased names (the `http` crate normalizes them). + headers: BTreeMap, + body: Vec, +} + +impl Seen { + fn json_body(&self) -> Value { + serde_json::from_slice(&self.body).expect("request body is JSON") + } +} + +#[derive(Default)] +struct MockFetch { + canned: Mutex>, + seen: Mutex>, +} + +impl MockFetch { + fn new(canned: Vec) -> Arc { + Arc::new(Self { + canned: Mutex::new(canned.into()), + seen: Mutex::default(), + }) + } + + fn seen(&self) -> Vec { + self.seen.lock().unwrap().clone() + } + + fn transport(self: &Arc) -> FetchFunction { + self.clone() + } +} + +#[async_trait] +impl Fetch for MockFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + self.seen.lock().unwrap().push(Seen { + method: request.method.to_string(), + url: request.url.to_string(), + headers: request + .headers + .iter() + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().expect("ASCII header").to_string(), + ) + }) + .collect(), + body: request.body.to_vec(), + }); + let canned = self + .canned + .lock() + .unwrap() + .pop_front() + .ok_or_else(|| FetchError::Other("no canned response left".into()))?; + let mut headers = HeaderMap::new(); + for (name, value) in &canned.headers { + headers.insert( + HeaderName::try_from(name.as_str()).unwrap(), + HeaderValue::try_from(value.as_str()).unwrap(), + ); + } + Ok(FetchResponse::from_bytes( + StatusCode::from_u16(canned.status).unwrap(), + headers, + request.url, + Bytes::from(canned.body), + )) + } +} + +fn settings(mock: &Arc) -> OpenAIProviderSettings { + OpenAIProviderSettings { + api_key: Some(KEY.to_string()), + fetch: Some(mock.transport()), + ..Default::default() + } +} + +/// Sets (or removes) an environment variable and restores it on drop. Tests +/// that use it are `#[serial]`. +struct EnvVar { + name: &'static str, + saved: Option, +} + +impl EnvVar { + fn set(name: &'static str, value: Option<&str>) -> Self { + let saved = std::env::var(name).ok(); + unsafe { + match value { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + Self { name, saved } + } +} + +impl Drop for EnvVar { + fn drop(&mut self) { + unsafe { + match &self.saved { + Some(value) => std::env::set_var(self.name, value), + None => std::env::remove_var(self.name), + } + } + } +} + +// ── fixtures ───────────────────────────────────────────────────────────────── + +/// Every fixture under `fixtures/aisdk/openai/`; the directory listing is +/// compared to this so a new fixture cannot go unreplayed. +const FIXTURES: [&str; 11] = [ + "chat-basic", + "chat-custom-name", + "chat-headers-merge", + "chat-provider-options", + "chat-stream-basic", + "chat-system-and-settings", + "chat-tools", + "embedding-basic", + "empty-api-key", + "missing-api-key", + "responses-basic", +]; + +fn fixtures_dir() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../fixtures/aisdk/openai") +} + +struct Fixture { + name: &'static str, + json: Value, +} + +impl Fixture { + fn load(name: &'static str) -> Self { + assert!(FIXTURES.contains(&name), "unlisted fixture {name}"); + let path = fixtures_dir().join(format!("{name}.json")); + let text = std::fs::read_to_string(&path) + .unwrap_or_else(|e| panic!("cannot read {}: {e}", path.display())); + Self { + name, + json: serde_json::from_str(&text).expect("fixture is JSON"), + } + } + + fn provider_string(&self) -> &str { + self.json["model"]["provider"].as_str().unwrap() + } + + fn model_id(&self) -> &str { + self.json["model"]["modelId"].as_str().unwrap() + } + + fn sdk_input(&self) -> &Value { + &self.json["sdk"]["input"] + } + + /// The transport answering with the recorded response (or a plain chat + /// completion for the case that never reached the wire). + fn mock(&self) -> Arc { + let response = &self.json["response"]; + if response.is_null() { + return MockFetch::new(vec![Canned::json(&json!({ + "id": "chatcmpl-unused", "object": "chat.completion", "created": 1, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))]); + } + let body = match &response["body"] { + Value::String(text) => text.clone().into_bytes(), + other => serde_json::to_vec(other).unwrap(), + }; + MockFetch::new(vec![Canned { + status: response["status"].as_u64().unwrap() as u16, + headers: response["headers"] + .as_object() + .unwrap() + .iter() + .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string())) + .collect(), + body, + }]) + } + + /// The recorded request must equal what went out: method, URL, headers + /// (names case-insensitively, with the pinned user-agent) and body. + fn assert_request(&self, seen: &Seen, authorization: &str) { + let recorded = &self.json["request"]; + assert_eq!( + seen.method, + recorded["method"].as_str().unwrap(), + "{}", + self.name + ); + assert_eq!(seen.url, recorded["url"].as_str().unwrap(), "{}", self.name); + + let mut expected: BTreeMap = recorded["headers"] + .as_object() + .unwrap() + .iter() + .map(|(k, v)| (k.to_ascii_lowercase(), v.as_str().unwrap().to_string())) + .collect(); + expected.insert( + "user-agent".to_string(), + "ai-sdk-openai/4.0.80 ai-sdk-provider-utils/5.0.51".to_string(), + ); + assert_eq!( + expected.get("authorization").map(String::as_str), + Some(""), + "{}: the recording redacts the key", + self.name + ); + expected.insert("authorization".to_string(), authorization.to_string()); + assert_eq!(seen.headers, expected, "{}: headers", self.name); + + assert_eq!(seen.json_body(), recorded["body"], "{}: body", self.name); + } + + /// The call options the recording was made with. + fn call_options(&self) -> CallOptions { + call_options_from(self.sdk_input()) + } +} + +fn message(role: &str, text: &str) -> LanguageModelMessage { + match role { + "system" => LanguageModelMessage::System { + content: text.into(), + provider_options: None, + }, + "user" => LanguageModelMessage::user_text(text), + "assistant" => LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: text.into(), + provider_options: None, + })], + provider_options: None, + }, + other => panic!("unmapped role {other}"), + } +} + +/// Rebuild aimux call options from a recorded `sdk.input` (the subset the +/// OpenAI fixtures use). +fn call_options_from(input: &Value) -> CallOptions { + let mut prompt: LanguageModelPrompt = Vec::new(); + if let Some(system) = input.get("system").and_then(Value::as_str) { + prompt.push(message("system", system)); + } + if let Some(text) = input.get("prompt").and_then(Value::as_str) { + prompt.push(message("user", text)); + } + for entry in input + .get("messages") + .and_then(Value::as_array) + .into_iter() + .flatten() + { + prompt.push(message( + entry["role"].as_str().unwrap(), + entry["content"].as_str().unwrap(), + )); + } + let mut options = CallOptions::new(prompt); + options.temperature = input.get("temperature").and_then(Value::as_f64); + options.top_p = input.get("topP").and_then(Value::as_f64); + options.max_output_tokens = input + .get("maxOutputTokens") + .and_then(Value::as_u64) + .map(|v| v as u32); + options.seed = input.get("seed").and_then(Value::as_u64); + options.stop_sequences = input + .get("stopSequences") + .and_then(Value::as_array) + .map(|a| a.iter().map(|v| v.as_str().unwrap().to_string()).collect()); + if let Some(tools) = input.get("tools").and_then(Value::as_object) { + options.tools = Some( + tools + .iter() + .map(|(name, tool)| { + let mut function = FunctionTool::new(name.clone(), tool["inputSchema"].clone()); + if let Some(description) = tool["description"].as_str() { + function = function.with_description(description); + } + Tool::Function(function) + }) + .collect(), + ); + } + options.tool_choice = match input.get("toolChoice").and_then(Value::as_str) { + Some("required") => Some(LanguageModelV4ToolChoice::Required), + Some("none") => Some(LanguageModelV4ToolChoice::None), + Some("auto") => Some(LanguageModelV4ToolChoice::Auto), + None => options + .tools + .as_ref() + .filter(|tools| !tools.is_empty()) + .map(|_| LanguageModelV4ToolChoice::Auto), + Some(other) => panic!("unmapped toolChoice {other}"), + }; + options.provider_options = input + .get("providerOptions") + .map(|value| serde_json::from_value(value.clone()).unwrap()); + options.headers = input.get("headers").and_then(Value::as_object).map(|h| { + h.iter() + .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string())) + .collect() + }); + options +} + +fn generated_text(result: &GenerateResult) -> String { + result + .content + .iter() + .filter_map(|part| match part { + GenerateContent::Text { text, .. } => Some(text.as_str()), + _ => None, + }) + .collect() +} + +/// Result fields every recorded `generateText` case pins. +fn assert_generate_result(fixture: &Fixture, result: &GenerateResult) { + let recorded = &fixture.json["result"]; + assert_eq!( + generated_text(result), + recorded["text"].as_str().unwrap(), + "{}: text", + fixture.name + ); + assert_eq!( + result.response.as_ref().and_then(|r| r.id.as_deref()), + recorded["response"]["id"].as_str(), + "{}: 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 _), + "{}: input tokens", + fixture.name + ); + assert_eq!( + result.usage.output_tokens.total, + recorded["usage"]["outputTokens"].as_u64().map(|v| v as _), + "{}: output tokens", + fixture.name + ); + let unified = match recorded["finishReason"].as_str().unwrap() { + "stop" => FinishReasonUnified::Stop, + "tool-calls" => FinishReasonUnified::ToolCalls, + other => panic!("unmapped finish reason {other}"), + }; + assert_eq!( + result.finish_reason.unified, unified, + "{}: finish", + fixture.name + ); + if let Some(raw) = recorded["rawFinishReason"].as_str() { + assert_eq!( + result.finish_reason.raw.as_deref(), + Some(raw), + "{}: raw finish", + fixture.name + ); + } +} + +fn bearer(key: &str) -> String { + format!("Bearer {key}") +} + +// ── the fixtures ───────────────────────────────────────────────────────────── + +#[test] +fn every_fixture_is_replayed() { + let mut on_disk: Vec = std::fs::read_dir(fixtures_dir()) + .expect("fixtures/aisdk/openai exists") + .map(|entry| { + let name = entry.unwrap().file_name().into_string().unwrap(); + name.strip_suffix(".json") + .unwrap_or_else(|| panic!("unexpected file {name}")) + .to_string() + }) + .collect(); + on_disk.sort(); + assert_eq!( + on_disk, FIXTURES, + "add a test for the new fixture and list it" + ); +} + +/// Replay a chat fixture whose request is a plain `generateText`. +async fn replay_chat(name: &'static str) { + let fixture = Fixture::load(name); + let mock = fixture.mock(); + let provider = create_openai(settings(&mock)).unwrap(); + let model = provider.chat(fixture.model_id()); + assert_eq!( + model.provider(), + fixture.provider_string(), + "{name}: provider()" + ); + assert_eq!(model.model_id(), fixture.model_id()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1, "{name}: exactly one request"); + fixture.assert_request(&seen[0], &bearer(KEY)); + assert_generate_result(&fixture, &result); +} + +#[tokio::test] +async fn chat_basic() { + replay_chat("chat-basic").await; +} + +#[tokio::test] +async fn chat_system_and_settings() { + replay_chat("chat-system-and-settings").await; +} + +#[tokio::test] +async fn chat_provider_options_use_the_fixed_openai_namespace() { + replay_chat("chat-provider-options").await; +} + +#[tokio::test] +async fn chat_tools() { + let fixture = Fixture::load("chat-tools"); + let mock = fixture.mock(); + let model = create_openai(settings(&mock)).unwrap().chat("gpt-4o"); + assert_eq!(model.provider(), fixture.provider_string()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + fixture.assert_request(&mock.seen()[0], &bearer(KEY)); + assert_generate_result(&fixture, &result); + let recorded = &fixture.json["result"]["content"][0]; + let call = result + .content + .iter() + .find_map(|part| match part { + GenerateContent::ToolCall(RawToolCall { + tool_call_id, + tool_name, + input, + .. + }) => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), + _ => None, + }) + .expect("a tool call"); + assert_eq!(call.0, recorded["toolCallId"].as_str().unwrap()); + assert_eq!(call.1, recorded["toolName"].as_str().unwrap()); + assert_eq!( + serde_json::from_str::(&call.2).unwrap(), + recorded["input"] + ); +} + +#[tokio::test] +async fn chat_custom_name() { + let fixture = Fixture::load("chat-custom-name"); + let mock = fixture.mock(); + let provider = create_openai(OpenAIProviderSettings { + name: Some("proxy".to_string()), + ..settings(&mock) + }) + .unwrap(); + let model = provider.chat(fixture.model_id()); + assert_eq!(model.provider(), "proxy.chat"); + assert_eq!(model.provider(), fixture.provider_string()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + // The name changes identity only, never the URL. + fixture.assert_request(&mock.seen()[0], &bearer(KEY)); + assert_generate_result(&fixture, &result); + assert_eq!( + fixture.json["observations"]["providerIsProxyChat"], + json!(true) + ); +} + +#[tokio::test] +async fn chat_headers_merge() { + let fixture = Fixture::load("chat-headers-merge"); + let mock = fixture.mock(); + // The recording was made with provider headers { "x-a": "1", "X-B": "2" } + // (scripts/aisdk-fixtures/cases/openai.mjs). + let headers: HeaderMapOpt = [("x-a", "1"), ("X-B", "2")] + .into_iter() + .map(|(k, v)| (k.to_string(), Some(v.to_string()))) + .collect(); + let provider = create_openai(OpenAIProviderSettings { + headers: Some(headers), + ..settings(&mock) + }) + .unwrap(); + let model = provider.chat(fixture.model_id()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + let seen = mock.seen(); + fixture.assert_request(&seen[0], &bearer(KEY)); + assert_generate_result(&fixture, &result); + // Case-insensitive override: the call's `x-b` replaced the provider's + // `X-B`, and no second `x-b` went out. + assert_eq!(seen[0].headers["x-b"], "3"); + let sent: Vec<&String> = seen[0] + .headers + .keys() + .filter(|n| n.starts_with("x-")) + .collect(); + assert_eq!(sent, ["x-a", "x-b", "x-c"]); + assert_eq!( + fixture.json["observations"]["headerNamesAsSent"], + json!(["x-a", "x-b", "x-c"]) + ); +} + +#[tokio::test] +async fn chat_stream_basic() { + let fixture = Fixture::load("chat-stream-basic"); + let mock = fixture.mock(); + let model = create_openai(settings(&mock)).unwrap().chat("gpt-4o"); + assert_eq!(model.provider(), fixture.provider_string()); + + let result = model.do_stream(&fixture.call_options()).await.unwrap(); + let parts: Vec = result + .stream + .map(|part| part.expect("no stream error")) + .collect() + .await; + + // The recording's user-agent names the provider package; the stream case + // is checked like the others. + let mut seen = mock.seen(); + assert_eq!(seen.len(), 1); + let sent = seen.remove(0); + fixture.assert_request(&sent, &bearer(KEY)); + + // Compare the text parts one by one with the recorded AI SDK + // `result.parts`: only text-start/delta/end are kept (which also drops + // the ai-layer-only start, start-step and finish-step), compared by id + // and text in order; finish is compared by finishReason and usage. + let expected = &fixture.json["result"]["parts"]; + let expected = expected.as_array().expect("result.parts"); + let expected_text: Vec<(String, String, Option)> = expected + .iter() + .filter_map(|part| { + let kind = part["type"].as_str()?; + kind.starts_with("text-").then(|| { + ( + kind.to_string(), + part["id"].as_str().unwrap().to_string(), + part["text"].as_str().map(str::to_string), + ) + }) + }) + .collect(); + let actual_text: Vec<(String, String, Option)> = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextStart { id, .. } => Some(("text-start".into(), id.clone(), None)), + // The ai layer drops empty deltas without provider metadata + // (ai@7.0.127 stream-text.ts:2881-2891); the provider layer emits + // them (openai@4.0.83 openai-chat-language-model.ts:623-633), and + // so does aimux for the first `"content":""` chunk. + StreamPart::TextDelta { + id, + delta, + provider_metadata, + } if !delta.is_empty() || provider_metadata.is_some() => { + Some(("text-delta".into(), id.clone(), Some(delta.clone()))) + } + StreamPart::TextEnd { id, .. } => Some(("text-end".into(), id.clone(), None)), + _ => None, + }) + .collect(); + assert!( + parts + .iter() + .any(|part| matches!(part, StreamPart::TextDelta { delta, .. } if delta.is_empty())) + ); + assert_eq!(actual_text, expected_text); + + let expected_finish = expected + .iter() + .find(|part| part["type"] == "finish") + .expect("a recorded finish part"); + let finish = parts.iter().find_map(|part| match part { + StreamPart::Finish { + finish_reason, + usage, + .. + } => Some((finish_reason, usage)), + _ => None, + }); + let (reason, usage) = finish.expect("a finish part"); + assert_eq!(reason.unified, FinishReasonUnified::Stop); + assert_eq!(json!(reason.unified), expected_finish["finishReason"]); + assert_eq!( + reason.raw.as_deref(), + expected_finish["rawFinishReason"].as_str() + ); + let total = &expected_finish["totalUsage"]; + let n = |v: &Value| v.as_u64().map(|n| u32::try_from(n).unwrap()); + assert_eq!(usage.input_tokens.total, n(&total["inputTokens"])); + assert_eq!( + usage.input_tokens.no_cache, + n(&total["inputTokenDetails"]["noCacheTokens"]) + ); + assert_eq!( + usage.input_tokens.cache_read, + n(&total["inputTokenDetails"]["cacheReadTokens"]) + ); + assert_eq!(usage.output_tokens.total, n(&total["outputTokens"])); + assert_eq!( + usage.output_tokens.text, + n(&total["outputTokenDetails"]["textTokens"]) + ); + assert_eq!( + usage.output_tokens.reasoning, + n(&total["outputTokenDetails"]["reasoningTokens"]) + ); + assert_eq!( + usage.input_tokens.total.unwrap() + usage.output_tokens.total.unwrap(), + n(&total["totalTokens"]).unwrap() + ); +} + +async fn chat_error(status: u16, body: &Value) -> Box { + let mut canned = Canned::json(body); + canned.status = status; + let mock = MockFetch::new(vec![canned]); + let error = create_openai(settings(&mock)) + .unwrap() + .chat("test-model") + .do_generate(&CallOptions::new(vec![LanguageModelMessage::user_text( + "Hello", + )])) + .await + .unwrap_err(); + let AiMuxError::ApiCall(error) = error else { + panic!("expected an API call error, got {error:?}") + }; + error +} + +/// openai@4.0.83 `openai-error.test.ts` "should parse OpenRouter resource +/// exhausted error": the message is `error.message`, a numeric `code` is +/// accepted, and a 429 is retryable (provider@4.0.21 +/// `api-call-error.ts:28-32`: 408, 409, 429 and 5xx). +#[tokio::test] +async fn chat_error_429_is_retryable() { + 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 error = chat_error(429, &response).await; + assert_eq!(error.message, message); + assert_eq!(error.status_code, Some(429)); + assert_eq!(error.data, Some(response)); + assert!(error.is_retryable); +} + +/// A 400 is not retryable (provider@4.0.21 `api-call-error.ts:28-32`). +#[tokio::test] +async fn chat_error_400_is_not_retryable() { + let response = + json!({"error":{"message":"Invalid request","type":"invalid_request_error","code":null}}); + let error = chat_error(400, &response).await; + assert_eq!(error.message, "Invalid request"); + assert_eq!(error.status_code, Some(400)); + assert!(!error.is_retryable); +} + +#[tokio::test] +async fn responses_basic() { + let fixture = Fixture::load("responses-basic"); + let mock = fixture.mock(); + let model = create_openai(settings(&mock)).unwrap().responses("gpt-4o"); + assert_eq!(model.provider(), fixture.provider_string()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + fixture.assert_request(&mock.seen()[0], &bearer(KEY)); + assert_generate_result(&fixture, &result); +} + +#[tokio::test] +async fn embedding_basic() { + let fixture = Fixture::load("embedding-basic"); + let mock = fixture.mock(); + let model = create_openai(settings(&mock)) + .unwrap() + .embedding("text-embedding-3-small"); + assert_eq!(model.provider(), fixture.provider_string()); + + let values: Vec = fixture.sdk_input()["values"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap().to_string()) + .collect(); + let options = EmbeddingCallOptions { + values, + ..EmbeddingCallOptions::new("") + }; + let result = model.do_embed(&options).await.unwrap(); + + fixture.assert_request(&mock.seen()[0], &bearer(KEY)); + let recorded: Vec> = + serde_json::from_value(fixture.json["result"]["embeddings"].clone()).unwrap(); + assert_eq!(result.embeddings, recorded); + assert_eq!( + result.usage.map(|u| u.tokens), + fixture.json["result"]["usage"]["tokens"] + .as_u64() + .map(|v| v as u32) + ); +} + +#[serial] +#[tokio::test] +async fn missing_api_key_fails_the_call_not_the_factory() { + let fixture = Fixture::load("missing-api-key"); + let _env = EnvVar::set("OPENAI_API_KEY", None); + let mock = fixture.mock(); + + // Creating the provider and a model reads no key. + let provider = create_openai(OpenAIProviderSettings { + fetch: Some(mock.transport()), + ..Default::default() + }) + .expect("no key is read at creation"); + let model = provider.chat(fixture.model_id()); + assert_eq!(model.provider(), fixture.provider_string()); + + let error = model + .do_generate(&fixture.call_options()) + .await + .expect_err("the call has no key"); + match &error { + AiMuxError::LoadApiKey { + env_var, + description, + } => { + assert_eq!(env_var, "OPENAI_API_KEY"); + assert_eq!(description, "OpenAI"); + } + other => panic!("expected LoadApiKey, got {other:?}"), + } + // Same text as the SDK's: names the setting and the variable, no value. + let message = error.to_string(); + assert!( + message.contains("OpenAI") && message.contains("OPENAI_API_KEY"), + "{message}" + ); + assert_eq!( + fixture.json["error"]["name"].as_str(), + Some("AI_LoadAPIKeyError") + ); + assert!(mock.seen().is_empty(), "nothing was sent"); + assert_eq!(fixture.json["observations"]["requestsSent"], json!(0)); +} + +#[serial] +#[tokio::test] +async fn an_explicit_empty_key_is_sent_as_given() { + let fixture = Fixture::load("empty-api-key"); + // The environment holds a key; an explicit "" must not fall back to it. + let _env = EnvVar::set("OPENAI_API_KEY", Some(ENV_KEY)); + let mock = fixture.mock(); + let provider = create_openai(OpenAIProviderSettings { + api_key: Some(String::new()), + fetch: Some(mock.transport()), + ..Default::default() + }) + .unwrap(); + let model = provider.chat(fixture.model_id()); + + let result = model.do_generate(&fixture.call_options()).await.unwrap(); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + // The wire value is exactly `Bearer ` (the SDK's headers object trims the + // trailing space, which is what its observation records). + fixture.assert_request(&seen[0], "Bearer "); + assert_eq!(seen[0].headers["authorization"].trim(), "Bearer"); + assert!(!seen[0].headers["authorization"].contains(ENV_KEY)); + assert_eq!( + fixture.json["observations"]["authorizationIsEmptyBearer"], + json!(true) + ); + assert_eq!(fixture.json["observations"]["usedEnvKey"], json!(false)); + assert_generate_result(&fixture, &result); +} + +fn _assert_provider_is_shareable() { + fn shareable() {} + shareable::(); +} diff --git a/aimux-providers/tests/openai_image_test.rs b/aimux-providers/tests/openai_image_test.rs index c12a5a7af..70e57c43d 100644 --- a/aimux-providers/tests/openai_image_test.rs +++ b/aimux-providers/tests/openai_image_test.rs @@ -15,11 +15,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::image_model::{ - ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, -}; +use aimux_core::image_model::{ImageCallOptions, ImageFile, ImageFileData, ImageOutputs}; +use aimux_core::provider::Provider; use aimux_core::shared::{AspectRatio, Size}; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::{OpenAIProviderSettings, create_openai}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -132,9 +131,13 @@ async fn should_extract_the_generated_images() { let server = MockServer::start().await; mock_generations_response(&server, image_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("dall-e-3"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("dall-e-3").unwrap(); let result = model.do_generate(&options(PROMPT)).await.unwrap(); @@ -160,9 +163,13 @@ async fn should_return_warnings_for_unsupported_settings() { let server = MockServer::start().await; mock_generations_response(&server, image_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("dall-e-3"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("dall-e-3").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -196,16 +203,19 @@ async fn should_return_warnings_for_unsupported_settings() { /// TS: "should respect maxImagesPerCall setting" #[tokio::test] async fn should_respect_max_images_per_call_setting() { - let config = OpenAIConfig::new("test-api-key"); - let provider = OpenAIProvider::new(config); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); - let dall_e_2 = provider.image("dall-e-2"); + let dall_e_2 = provider.image_model("dall-e-2").unwrap(); assert_eq!(dall_e_2.max_images_per_call(), Some(10)); - let future_gpt = provider.image("gpt-image-99"); + let future_gpt = provider.image_model("gpt-image-99").unwrap(); assert_eq!(future_gpt.max_images_per_call(), Some(10)); - let unknown = provider.image("unknown-model"); + let unknown = provider.image_model("unknown-model").unwrap(); assert_eq!(unknown.max_images_per_call(), Some(1)); } @@ -223,9 +233,13 @@ async fn should_include_response_data_with_timestamp_model_id_and_headers() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("dall-e-3"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("dall-e-3").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -254,9 +268,13 @@ async fn should_use_real_date_when_no_custom_date_provider_is_specified() { let server = MockServer::start().await; mock_generations_response(&server, image_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("dall-e-3"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("dall-e-3").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -280,9 +298,13 @@ async fn should_not_include_response_format_for_gpt_image_1() { let server = MockServer::start().await; mock_generations_response(&server, image_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -305,9 +327,13 @@ async fn should_not_include_response_format_for_gpt_image_2() { let server = MockServer::start().await; mock_generations_response(&server, image_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-2"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-2").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -327,9 +353,13 @@ async fn should_not_include_response_format_for_future_gpt_image_models() { let server = MockServer::start().await; mock_generations_response(&server, image_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-99"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-99").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -349,9 +379,13 @@ async fn should_not_include_response_format_for_chatgpt_image_latest() { let server = MockServer::start().await; mock_generations_response(&server, image_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("chatgpt-image-latest"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("chatgpt-image-latest").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -371,9 +405,13 @@ async fn should_not_include_response_format_for_date_suffixed_gpt_image() { let server = MockServer::start().await; mock_generations_response(&server, image_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1.5-2025-12-16"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1.5-2025-12-16").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -405,9 +443,13 @@ async fn should_handle_null_revised_prompt_responses() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -443,9 +485,13 @@ async fn should_include_response_format_for_dall_e_3() { let server = MockServer::start().await; mock_generations_response(&server, image_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("dall-e-3"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("dall-e-3").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -482,9 +528,13 @@ async fn should_map_openai_usage_to_usage() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -530,9 +580,13 @@ async fn should_distribute_input_token_details_evenly_across_images() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 3; @@ -562,9 +616,13 @@ async fn should_call_edits_endpoint_when_files_are_provided() { let server = MockServer::start().await; mock_edits_response(&server, image_edit_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -584,9 +642,13 @@ async fn should_send_image_as_form_data_with_uint8array_input() { let server = MockServer::start().await; mock_edits_response(&server, image_edit_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -638,9 +700,13 @@ async fn should_send_image_as_form_data_with_base64_string_input() { let server = MockServer::start().await; mock_edits_response(&server, image_edit_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -671,9 +737,13 @@ async fn should_send_multiple_images_as_form_data_array() { let server = MockServer::start().await; mock_edits_response(&server, image_edit_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -699,9 +769,13 @@ async fn should_extract_the_edited_images_from_response() { let server = MockServer::start().await; mock_edits_response(&server, image_edit_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -732,9 +806,13 @@ async fn should_include_response_metadata_for_edited_images() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -757,9 +835,13 @@ async fn should_return_warnings_for_unsupported_settings_in_edit_mode() { let server = MockServer::start().await; mock_edits_response(&server, image_edit_response_body()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; @@ -805,9 +887,13 @@ async fn should_return_usage_information_for_edited_images() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.image("gpt-image-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.image_model("gpt-image-1").unwrap(); let mut opts = options(PROMPT); opts.n = 1; diff --git a/aimux-providers/tests/openai_output_round_trip_test.rs b/aimux-providers/tests/openai_output_round_trip_test.rs index cacdfe65d..79be16ecd 100644 --- a/aimux-providers/tests/openai_output_round_trip_test.rs +++ b/aimux-providers/tests/openai_output_round_trip_test.rs @@ -33,18 +33,20 @@ use aimux_core::generate::{ use aimux_core::openai_output::{ ChatCompletionChunk, OpenAiStreamOptions, encode_chunk_sse, to_chat_completion_stream, }; -use aimux_providers::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIProvider}; +use aimux_providers::deepseek::{DeepSeekProviderSettings, create_deepseek}; +use aimux_providers::openai::{OpenAIModel, OpenAIProviderSettings, create_openai}; // ── helpers ─────────────────────────────────────────────────────────────────── -/// Build an OpenAIProvider pointed at the replay server. -fn openai_provider(uri: &str) -> aimux_providers::openai::OpenAIModel { - let provider = OpenAIProvider::new( - OpenAIConfig::new("test-key") - .with_base_url(format!("{uri}/v1")) - .with_profile(OpenAICompatProfile::full()), - ); - provider.model("gpt-4o") +/// Build the native OpenAI chat model pointed at the replay server. +fn openai_provider(uri: &str) -> OpenAIModel { + create_openai(OpenAIProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{uri}/v1")), + ..Default::default() + }) + .unwrap() + .chat("gpt-4o") } /// Mount a single non-streaming OpenAI response on the mock server. @@ -359,15 +361,18 @@ async fn round_trip_multiple_tool_calls() { #[tokio::test] async fn cross_protocol_anthropic_to_openai_non_streaming() { - use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider}; + use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; let server = MockServer::start().await; replay::mount_cassettes(&server, "tests/cassettes/anthropic").await; - let provider = AnthropicProvider::new( - AnthropicConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())), - ); - let model = provider.model("claude-sonnet-4-6"); + let provider = create_anthropic(AnthropicProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap(); + let model = provider.messages("claude-sonnet-4-6"); let result = generate_text(&model, "Hello", GenerateTextOptions::default()) .await @@ -416,15 +421,18 @@ async fn cross_protocol_anthropic_to_openai_non_streaming() { #[tokio::test] async fn cross_protocol_anthropic_to_openai_streaming() { - use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider}; + use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; let server = MockServer::start().await; replay::mount_cassettes(&server, "tests/cassettes/anthropic").await; - let provider = AnthropicProvider::new( - AnthropicConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())), - ); - let model = provider.model("claude-sonnet-4-6"); + let provider = create_anthropic(AnthropicProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap(); + let model = provider.messages("claude-sonnet-4-6"); let result = stream_text(&model, "Hello", GenerateTextOptions::default()) .await @@ -838,12 +846,13 @@ async fn cross_protocol_deepseek_with_reasoning() { .mount(&server) .await; - let provider = OpenAIProvider::new( - OpenAIConfig::new("test-key") - .with_base_url(server.uri()) - .with_profile(OpenAICompatProfile::deepseek()), - ); - let model = provider.model("deepseek-chat"); + let provider = create_deepseek(DeepSeekProviderSettings { + api_key: Some("test-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.chat("deepseek-chat"); let result = generate_text( &model, diff --git a/aimux-providers/tests/openai_realtime_ws_replay_test.rs b/aimux-providers/tests/openai_realtime_ws_replay_test.rs new file mode 100644 index 000000000..734e54d2e --- /dev/null +++ b/aimux-providers/tests/openai_realtime_ws_replay_test.rs @@ -0,0 +1,203 @@ +//! Offline replay of the OpenAI realtime transcription WebSocket path +//! (RFC-0028 §3.2), driven through the connector that +//! `OpenAITranscriptionModel::with_web_socket` installs. +//! +//! `OpenAIProviderSettings::web_socket` carries a custom [`WsConnector`] that +//! never touches the network: it pairs the client with an in-process server +//! over an in-memory duplex pipe. The server records every frame the client +//! sends and replays the server events in +//! `fixtures/openai_realtime_transcription.json` once the audio is committed. +#![cfg(feature = "realtime")] + +use std::sync::{Arc, Mutex}; + +use aimux_core::error::AiMuxError; +use aimux_core::transcription_model::{ + AudioChunk, InputAudioFormat, TranscriptionModel, TranscriptionStreamOptions, + TranscriptionStreamPart, +}; +use aimux_provider_utils::ws::{WebSocketRequest, WsConnection, WsConnector}; +use aimux_providers::{OpenAIProviderSettings, create_openai}; +use futures::{SinkExt, StreamExt}; +use serde_json::Value; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::protocol::Role; + +/// What the in-process server saw. +#[derive(Default)] +struct Seen { + /// URL and headers of every connect request, in order. + connects: Vec<(String, Vec<(String, String)>)>, + /// Every text frame the client sent, parsed. + frames: Vec, + /// Whether the client closed the socket. + closed: bool, +} + +/// A connector that serves the fixture over an in-memory pipe. +struct ReplayConnector { + seen: Arc>, + events: Vec, +} + +#[async_trait::async_trait] +impl WsConnector for ReplayConnector { + async fn connect(&self, request: &WebSocketRequest) -> Result { + self.seen + .lock() + .unwrap() + .connects + .push((request.url.clone(), request.headers.clone())); + let (client_io, server_io) = tokio::io::duplex(64 * 1024); + let client = WebSocketStream::from_raw_socket(client_io, Role::Client, None).await; + let server = WebSocketStream::from_raw_socket(server_io, Role::Server, None).await; + tokio::spawn(serve(server, Arc::clone(&self.seen), self.events.clone())); + Ok(WsConnection::from_transport(Box::new(client), request)) + } +} + +/// Record client frames; after `input_audio_buffer.commit`, replay `events`. +async fn serve( + mut ws: WebSocketStream, + seen: Arc>, + events: Vec, +) { + while let Some(message) = ws.next().await { + match message { + Ok(Message::Text(text)) => { + let frame: Value = serde_json::from_str(&text).expect("client sends JSON"); + let commit = frame["type"] == "input_audio_buffer.commit"; + seen.lock().unwrap().frames.push(frame); + if commit { + for event in &events { + ws.send(Message::Text(event.to_string())) + .await + .expect("replay event"); + } + } + } + Ok(Message::Close(_)) | Err(_) => { + seen.lock().unwrap().closed = true; + break; + } + Ok(_) => {} + } + } +} + +fn fixture_events() -> Vec { + let fixture: Value = + serde_json::from_str(include_str!("fixtures/openai_realtime_transcription.json")) + .expect("fixture is JSON"); + fixture["events"].as_array().expect("events array").clone() +} + +#[tokio::test] +async fn realtime_transcription_replays_through_injected_web_socket() { + let seen = Arc::new(Mutex::new(Seen::default())); + let connector = ReplayConnector { + seen: Arc::clone(&seen), + events: fixture_events(), + }; + let provider = create_openai(OpenAIProviderSettings { + api_key: Some("sk-replay-test".to_string()), + base_url: Some("https://replay.invalid/v1".into()), + web_socket: Some(Arc::new(connector)), + ..Default::default() + }) + .expect("valid settings"); + let model = provider.transcription("gpt-realtime-whisper"); + + let options = TranscriptionStreamOptions { + audio: Box::pin(futures::stream::iter(vec![ + AudioChunk::Binary(vec![1, 2, 3]), + AudioChunk::Base64("BAUG".into()), + ])), + input_audio_format: InputAudioFormat { + format_type: "audio/pcm".to_string(), + rate: Some(24_000), + }, + provider_options: None, + abort_signal: None, + headers: None, + include_raw_chunks: false, + timeout: None, + }; + let result = model.do_stream(options).await.expect("stream opens"); + let parts: Vec = result + .stream + .map(|part| part.expect("no stream error")) + .collect() + .await; + + // Transcript events, mapped from the replayed server events. + let deltas: Vec<&str> = parts + .iter() + .filter_map(|part| match part { + TranscriptionStreamPart::TranscriptDelta { delta, id, .. } => { + assert_eq!(id.as_deref(), Some("item_1")); + Some(delta.as_str()) + } + _ => None, + }) + .collect(); + assert_eq!(deltas, ["Hello", " world"]); + assert!(parts.iter().any(|part| matches!( + part, + TranscriptionStreamPart::TranscriptFinal { text, id, .. } + if text == "Hello world" && id.as_deref() == Some("item_1") + ))); + assert!(matches!( + parts.last(), + Some(TranscriptionStreamPart::Finish { text, .. }) if text == "Hello world" + )); + assert!(matches!( + parts.first(), + Some(TranscriptionStreamPart::StreamStart { .. }) + )); + + // The injected connector carried the session: exactly one connect, to + // the realtime URL derived from the base URL, with the credential. + // The client closes the long-lived session after `completed` (B2). + for _ in 0..100 { + if seen.lock().unwrap().closed { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + let seen = seen.lock().unwrap(); + assert!(seen.closed, "client did not close the session"); + assert_eq!(seen.connects.len(), 1, "with_web_socket connector not used"); + let (url, headers) = &seen.connects[0]; + assert_eq!(url, "wss://replay.invalid/v1/realtime?intent=transcription"); + assert!(headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case("authorization") && value == "Bearer sk-replay-test" + })); + + // Frames the server saw: session.update, two appends, commit. + let types: Vec<&str> = seen + .frames + .iter() + .map(|frame| frame["type"].as_str().unwrap_or_default()) + .collect(); + assert_eq!( + types, + [ + "session.update", + "input_audio_buffer.append", + "input_audio_buffer.append", + "input_audio_buffer.commit", + ] + ); + let session = &seen.frames[0]["session"]; + assert_eq!(session["type"], "transcription"); + assert_eq!( + session["audio"]["input"]["transcription"]["model"], + "gpt-realtime-whisper" + ); + assert_eq!(session["audio"]["input"]["format"]["type"], "audio/pcm"); + assert_eq!(session["audio"]["input"]["format"]["rate"], 24_000); + assert_eq!(seen.frames[1]["audio"], "AQID"); + assert_eq!(seen.frames[2]["audio"], "BAUG"); +} diff --git a/aimux-providers/tests/openai_speech_test.rs b/aimux-providers/tests/openai_speech_test.rs index 6905d105c..dec09e1ee 100644 --- a/aimux-providers/tests/openai_speech_test.rs +++ b/aimux-providers/tests/openai_speech_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the OpenAI speech (TTS) model tests. +//! Rust translation of the OpenAI speech (TTS) model tests. //! //! Source: `reference/ai/packages/openai/src/speech/openai-speech-model.test.ts` //! @@ -17,8 +17,9 @@ use serde_json::Value; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::speech_model::{AudioData, SpeechCallOptions, SpeechModel}; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_core::provider::Provider; +use aimux_core::speech_model::{AudioData, SpeechCallOptions}; +use aimux_providers::{OpenAIProviderSettings, create_openai}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -75,9 +76,13 @@ async fn should_pass_the_model_and_text() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.speech("tts-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.speech_model("tts-1").unwrap().unwrap(); model .do_generate(&speech_options("Hello from the AI SDK!")) @@ -107,13 +112,21 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = OpenAIConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_org_id("test-organization") - .with_project("test-project") - .with_headers(provider_headers); - let provider = OpenAIProvider::new(config); - let model = provider.speech("tts-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + organization: Some("test-organization".into()), + project: Some("test-project".into()), + headers: Some( + provider_headers + .into_iter() + .map(|(k, v)| (k, Some(v))) + .collect(), + ), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.speech_model("tts-1").unwrap().unwrap(); let mut options = speech_options("Hello from the AI SDK!"); let mut request_headers = HashMap::new(); @@ -147,9 +160,13 @@ async fn should_pass_options() { let server = MockServer::start().await; mock_audio_response(&server, "opus").await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.speech("tts-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.speech_model("tts-1").unwrap().unwrap(); let mut options = speech_options("Hello from the AI SDK!"); options.voice = Some("nova".to_string()); @@ -181,9 +198,13 @@ async fn should_return_audio_data_with_correct_content_type() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.speech("tts-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.speech_model("tts-1").unwrap().unwrap(); let mut options = speech_options("Hello from the AI SDK!"); options.output_format = Some("opus".to_string()); @@ -214,9 +235,13 @@ async fn should_include_response_data_with_timestamp_modelid_and_headers() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.speech("tts-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.speech_model("tts-1").unwrap().unwrap(); let result = model .do_generate(&speech_options("Hello from the AI SDK!")) @@ -250,9 +275,13 @@ async fn should_use_real_date_when_no_custom_date_provider_is_specified() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.speech("tts-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.speech_model("tts-1").unwrap().unwrap(); let result = model .do_generate(&speech_options("Hello from the AI SDK!")) @@ -269,21 +298,30 @@ async fn should_use_real_date_when_no_custom_date_provider_is_specified() { /// TS: "should handle different audio formats" /// /// Iterates over the supported formats and verifies audio data is returned. -/// The TS test passes `providerOptions.openai.response_format`; the Rust model -/// does not parse openai-specific provider options (the TS schema does not -/// include `response_format` either), so we pass `output_format` directly. +/// Matches the TS input `providerOptions.openai.response_format`, which the +/// upstream schema does not include. #[tokio::test] async fn should_handle_different_audio_formats() { for format in &["mp3", "opus", "aac", "flac", "wav", "pcm"] { let server = MockServer::start().await; mock_audio_response(&server, format).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.speech("tts-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.speech_model("tts-1").unwrap().unwrap(); let mut options = speech_options("Hello from the AI SDK!"); - options.output_format = Some(format.to_string()); + options.provider_options = Some(aimux_core::shared::NamespaceMap::from([( + "openai".into(), + serde_json::Map::from_iter([( + "response_format".into(), + serde_json::Value::String((*format).into()), + )]), + )])); let result = model .do_generate(&options) @@ -305,9 +343,13 @@ async fn should_include_warnings_if_any_are_generated() { let server = MockServer::start().await; mock_audio_response(&server, "mp3").await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.speech("tts-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.speech_model("tts-1").unwrap().unwrap(); let result = model .do_generate(&speech_options("Hello from the AI SDK!")) @@ -320,36 +362,3 @@ async fn should_include_warnings_if_any_are_generated() { result.warnings ); } - -// ── Additional: unsupported language emits a warning ───────────────────────── - -/// The TS model emits an `unsupported` warning when `language` is set. The TS -/// suite does not have a dedicated test for this, but the behaviour is -/// exercised here to guard the Rust translation. -#[tokio::test] -async fn language_option_emits_unsupported_warning() { - let server = MockServer::start().await; - mock_audio_response(&server, "mp3").await; - - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.speech("tts-1"); - - let mut options = speech_options("Hello from the AI SDK!"); - options.language = Some("en".to_string()); - - let result = model.do_generate(&options).await.unwrap(); - - assert_eq!( - result.warnings.len(), - 1, - "expected one warning, got {:?}", - result.warnings - ); - match &result.warnings[0] { - aimux_core::types::Warning::Unsupported { feature, .. } => { - assert_eq!(feature, "language"); - } - other => panic!("expected Unsupported warning, got {other:?}"), - } -} diff --git a/aimux-providers/tests/openai_transcription_stream_test.rs b/aimux-providers/tests/openai_transcription_stream_test.rs index 7ed19a2f4..cdd8429c4 100644 --- a/aimux-providers/tests/openai_transcription_stream_test.rs +++ b/aimux-providers/tests/openai_transcription_stream_test.rs @@ -8,7 +8,11 @@ //! //! These replace the TS mock-WebSocket tests that //! `openai_transcription_test.rs` omitted ("not practical to translate" — -//! they became practical with RFC-0028's WebSocket infrastructure). +//! they became practical with RFC-0028's WebSocket infrastructure). The +//! provider is built with `create_openai`, whose base URL also names the +//! realtime endpoint (`ws://` for an `http://` base URL). + +#![cfg(feature = "realtime")] use std::time::Duration; @@ -22,7 +26,7 @@ use aimux_core::transcription_model::{ AudioChunk, InputAudioFormat, TranscriptionModel, TranscriptionStreamOptions, TranscriptionStreamPart, }; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::{OpenAIProvider, OpenAIProviderSettings, create_openai}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -135,10 +139,18 @@ async fn serve(ws: &mut tokio_tungstenite::WebSocketStream, api_key: &str) -> OpenAIProvider { + create_openai(OpenAIProviderSettings { + api_key: Some(api_key.to_string()), + base_url: Some(base_url.into()), + ..Default::default() + }) + .expect("valid settings") +} + fn realtime_model(base_url: &str) -> aimux_providers::OpenAITranscriptionModel { - let config = OpenAIConfig::new("test-api-key").with_base_url(base_url.to_string()); - let provider = OpenAIProvider::new(config); - provider.transcription("gpt-realtime-whisper") + openai_at(base_url, "test-api-key").transcription("gpt-realtime-whisper") } fn stream_options( @@ -354,8 +366,7 @@ async fn stream_abort_mid_session() { /// Non-realtime models are rejected from do_stream (inverse gating). #[tokio::test] async fn stream_rejects_non_realtime_models() { - let config = OpenAIConfig::new("k").with_base_url("http://127.0.0.1:1"); - let provider = OpenAIProvider::new(config); + let provider = openai_at("http://127.0.0.1:1", "k"); let model = provider.transcription("whisper-1"); let err = model .do_stream(stream_options(vec![], None)) @@ -374,8 +385,7 @@ async fn stream_rejects_non_realtime_models() { #[tokio::test] async fn stream_connect_failure_errors() { // Port 1 is never listening on loopback: TCP refuses immediately. - let config = OpenAIConfig::new("k").with_base_url("http://127.0.0.1:1".to_string()); - let provider = OpenAIProvider::new(config); + let provider = openai_at("http://127.0.0.1:1".to_string(), "k"); let model = provider.transcription("gpt-realtime-whisper"); let err = model @@ -448,8 +458,7 @@ async fn stream_connect_timeout_fires() { std::future::pending::<()>().await; }); - let config = OpenAIConfig::new("k").with_base_url(format!("http://127.0.0.1:{port}")); - let provider = OpenAIProvider::new(config); + let provider = openai_at(format!("http://127.0.0.1:{port}"), "k"); let model = provider.transcription("gpt-realtime-whisper"); let options = TranscriptionStreamOptions { @@ -501,8 +510,7 @@ async fn stream_first_chunk_timeout_fires() { while let Ok(Some(_)) = tokio::time::timeout(Duration::from_secs(2), ws.next()).await {} }); - let config = OpenAIConfig::new("k").with_base_url(format!("http://127.0.0.1:{port}")); - let provider = OpenAIProvider::new(config); + let provider = openai_at(format!("http://127.0.0.1:{port}"), "k"); let model = provider.transcription("gpt-realtime-whisper"); let options = TranscriptionStreamOptions { @@ -657,8 +665,7 @@ async fn stream_chunk_idle_timeout_fires() { while let Ok(Some(_)) = tokio::time::timeout(Duration::from_secs(2), ws.next()).await {} }); - let config = OpenAIConfig::new("k").with_base_url(format!("http://127.0.0.1:{port}")); - let provider = OpenAIProvider::new(config); + let provider = openai_at(format!("http://127.0.0.1:{port}"), "k"); let model = provider.transcription("gpt-realtime-whisper"); let options = TranscriptionStreamOptions { @@ -721,8 +728,7 @@ async fn stream_total_timeout_fires() { } }); - let config = OpenAIConfig::new("k").with_base_url(format!("http://127.0.0.1:{port}")); - let provider = OpenAIProvider::new(config); + let provider = openai_at(format!("http://127.0.0.1:{port}"), "k"); let model = provider.transcription("gpt-realtime-whisper"); let options = TranscriptionStreamOptions { diff --git a/aimux-providers/tests/openai_transcription_test.rs b/aimux-providers/tests/openai_transcription_test.rs index f533f960e..2f488fe0a 100644 --- a/aimux-providers/tests/openai_transcription_test.rs +++ b/aimux-providers/tests/openai_transcription_test.rs @@ -12,7 +12,9 @@ //! rather than asserting an exact timestamp value. //! //! The TS `doStream` tests are translated in -//! `openai_transcription_stream_test.rs` (RFC-0028: local WebSocket server). +//! `openai_transcription_stream_test.rs` (RFC-0028: local WebSocket server) +//! and `openai_realtime_ws_replay_test.rs` (offline, through the provider's +//! `web_socket` connector). use std::collections::HashMap; @@ -21,8 +23,9 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::error::AiMuxError; -use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_core::provider::Provider; +use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions}; +use aimux_providers::{OpenAIProviderSettings, create_openai}; // ── helpers ───────────────────────────────────────────────────────────────── @@ -141,9 +144,16 @@ fn fixture_response() -> Value { #[tokio::test] async fn should_reject_gpt_realtime_whisper_for_non_streaming() { let server = MockServer::start().await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("gpt-realtime-whisper"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider + .transcription_model("gpt-realtime-whisper") + .unwrap() + .unwrap(); let result = model .do_generate(&transcription_options(mock_audio_data(), "audio/wav")) @@ -164,9 +174,13 @@ async fn should_pass_the_model() { let server = MockServer::start().await; mock_json_response(&server, &fixture_response()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("whisper-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.transcription_model("whisper-1").unwrap().unwrap(); model .do_generate(&transcription_options(mock_audio_data(), "audio/wav")) @@ -185,9 +199,13 @@ async fn should_default_whisper_1_to_verbose_json() { let server = MockServer::start().await; mock_json_response(&server, &fixture_response()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("whisper-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.transcription_model("whisper-1").unwrap().unwrap(); let result = model .do_generate(&transcription_options(mock_audio_data(), "audio/wav")) @@ -213,13 +231,21 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = OpenAIConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_org_id("test-organization") - .with_project("test-project") - .with_headers(provider_headers); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("whisper-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + organization: Some("test-organization".into()), + project: Some("test-project".into()), + headers: Some( + provider_headers + .into_iter() + .map(|(k, v)| (k, Some(v))) + .collect(), + ), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.transcription_model("whisper-1").unwrap().unwrap(); let mut options = transcription_options(mock_audio_data(), "audio/wav"); let mut request_headers = HashMap::new(); @@ -253,9 +279,13 @@ async fn should_extract_the_transcription_text() { let server = MockServer::start().await; mock_json_response(&server, &fixture_response()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("whisper-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.transcription_model("whisper-1").unwrap().unwrap(); let result = model .do_generate(&transcription_options(mock_audio_data(), "audio/wav")) @@ -282,9 +312,13 @@ async fn should_include_response_data() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("whisper-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.transcription_model("whisper-1").unwrap().unwrap(); let result = model .do_generate(&transcription_options(mock_audio_data(), "audio/wav")) @@ -305,9 +339,13 @@ async fn should_use_real_date() { let server = MockServer::start().await; mock_json_response(&server, &fixture_response()).await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("whisper-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.transcription_model("whisper-1").unwrap().unwrap(); let result = model .do_generate(&transcription_options(mock_audio_data(), "audio/wav")) @@ -332,9 +370,13 @@ async fn should_work_without_optional_fields() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("whisper-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.transcription_model("whisper-1").unwrap().unwrap(); let result = model .do_generate(&transcription_options(mock_audio_data(), "audio/wav")) @@ -369,9 +411,13 @@ async fn should_handle_empty_segments() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("whisper-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.transcription_model("whisper-1").unwrap().unwrap(); let result = model .do_generate(&transcription_options(mock_audio_data(), "audio/wav")) @@ -404,9 +450,13 @@ async fn should_handle_segments_with_missing_optional_fields() { ) .await; - let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("whisper-1"); + let settings = OpenAIProviderSettings { + api_key: Some("test-api-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }; + let provider = create_openai(settings).unwrap(); + let model = provider.transcription_model("whisper-1").unwrap().unwrap(); let result = model .do_generate(&transcription_options(mock_audio_data(), "audio/wav")) diff --git a/aimux-providers/tests/openrouter_test.rs b/aimux-providers/tests/openrouter_test.rs index cca5ff453..73bfc3a7a 100644 --- a/aimux-providers/tests/openrouter_test.rs +++ b/aimux-providers/tests/openrouter_test.rs @@ -1,7 +1,7 @@ //! Provider-specific tests for the OpenRouter provider. //! -//! OpenRouter is a thin OpenAI-compatible wrapper over [`OpenAIProvider`]. The -//! behaviours verified here are the ones the wrapper is responsible for: +//! OpenRouter is a registry preset of the OpenAI-compatible package. The +//! behaviours verified here are the ones the preset is responsible for: //! //! - Provider configuration: name, `OPENROUTER_API_KEY` env var, custom API //! key, custom headers, base-URL override. @@ -34,7 +34,8 @@ use aimux_core::stream_part::StreamPart; use aimux_core::tool::FunctionTool; use aimux_core::types::FinishReasonUnified; -use aimux_providers::{OpenRouterConfig, OpenRouterProvider}; +use aimux_providers::openai_compatible::OpenAICompatibleProvider; +use aimux_providers::{PresetSettings, create_provider, preset}; // Shared cassette-replay infrastructure (same `mod common` used by // `conformance_test.rs`). @@ -113,23 +114,27 @@ async fn collect_stream(result: aimux_core::result::StreamResult) -> Vec OpenRouterProvider { - let config = OpenRouterConfig::new("test-api-key").with_base_url(server.uri()); - OpenRouterProvider::new(config) +/// The OpenRouter preset at `base_url` with an explicit key. +fn openrouter_at(base_url: String, key: &str) -> OpenAICompatibleProvider { + preset::create( + "openrouter", + PresetSettings { + api_key: Some(aimux_provider_utils::Resolvable::Value(key.to_string())), + base_url: Some(base_url), + ..Default::default() + }, + ) + .unwrap() +} + +fn make_provider(server: &MockServer) -> OpenAICompatibleProvider { + openrouter_at(server.uri(), "test-api-key") } // ════════════════════════════════════════════════════════════════════════════ // Provider configuration // ════════════════════════════════════════════════════════════════════════════ -/// `createOpenRouter()` produces a provider whose name is "openrouter". -#[test] -fn provider_name_is_openrouter() { - let config = OpenRouterConfig::new("test-key"); - let provider = OpenRouterProvider::new(config); - assert_eq!(provider.name(), "openrouter"); -} - /// Custom API key is sent in the `Authorization: Bearer` header. #[tokio::test] async fn custom_api_key_used_in_auth_header() { @@ -141,9 +146,8 @@ async fn custom_api_key_used_in_auth_header() { .mount(&server) .await; - let config = OpenRouterConfig::new("my-custom-key").with_base_url(server.uri()); - let provider = OpenRouterProvider::new(config); - let model = provider.model("openai/gpt-4o-mini"); + let provider = openrouter_at(server.uri(), "my-custom-key"); + let model = provider.chat("openai/gpt-4o-mini"); model .do_generate(&default_options(test_prompt())) @@ -162,9 +166,8 @@ async fn custom_headers_forwarded() { .mount(&server) .await; - let config = OpenRouterConfig::new("test-key").with_base_url(server.uri()); - let provider = OpenRouterProvider::new(config); - let model = provider.model("openai/gpt-4o-mini"); + let provider = openrouter_at(server.uri(), "test-key"); + let model = provider.chat("openai/gpt-4o-mini"); let mut options = default_options(test_prompt()); options.headers = Some( @@ -190,8 +193,7 @@ async fn language_model_via_trait() { .mount(&server) .await; - let config = OpenRouterConfig::new("test-key").with_base_url(server.uri()); - let provider = OpenRouterProvider::new(config); + let provider = openrouter_at(server.uri(), "test-key"); let model = provider .language_model("openai/gpt-4o-mini") .expect("language_model should succeed"); @@ -211,8 +213,9 @@ fn from_env_loads_openrouter_api_key() { std::env::set_var("OPENROUTER_API_KEY", "env-test-key"); } - let config = OpenRouterConfig::from_env(); - assert!(config.is_ok(), "from_env should succeed with env var set"); + let model = create_provider("openrouter", PresetSettings::default()) + .and_then(|p| p.language_model("openai/gpt-4o-mini")); + assert!(model.is_ok(), "the by-name entry point reads the variable"); unsafe { match saved { @@ -224,15 +227,21 @@ fn from_env_loads_openrouter_api_key() { /// Without the env var, `createOpenRouter()` fails. #[serial] -#[test] -fn from_env_fails_without_env_var() { +#[tokio::test] +async fn from_env_fails_without_env_var() { let saved = std::env::var("OPENROUTER_API_KEY").ok(); unsafe { std::env::remove_var("OPENROUTER_API_KEY"); } - let config = OpenRouterConfig::from_env(); - assert!(config.is_err(), "from_env should fail without env var"); + let model = create_provider("openrouter", PresetSettings::default()) + .and_then(|p| p.language_model("openai/gpt-4o-mini")) + .unwrap(); + let err = model + .do_generate(&default_options(test_prompt())) + .await + .expect_err("no key variable, no key"); + assert!(matches!(err, AiMuxError::LoadApiKey { .. }), "{err:?}"); unsafe { if let Some(v) = saved { @@ -256,7 +265,7 @@ async fn sends_correct_request_body() { .await; let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let _ = model .do_generate(&default_options(test_prompt())) .await @@ -293,7 +302,7 @@ async fn stream_request_body_has_stream_flag() { .await; let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let _ = model .do_stream(&default_options(test_prompt())) .await @@ -320,7 +329,7 @@ async fn extracts_usage() { .await; let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let result = model .do_generate(&default_options(test_prompt())) .await @@ -345,7 +354,7 @@ async fn do_generate_returns_text() { .await; let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let result = model .do_generate(&default_options(test_prompt())) @@ -371,7 +380,7 @@ async fn do_generate_extracts_tool_call() { .await; let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let tool = FunctionTool::new( "get-weather", @@ -402,7 +411,7 @@ async fn do_generate_extracts_tool_call() { }) => { assert_eq!(tool_call_id, "call_abc"); assert_eq!(tool_name, "get-weather"); - assert_eq!(input, r#"{"city":"SF"}"#); + assert_eq!(input, &Value::String(r#"{"city":"SF"}"#.into())); } other => panic!("expected ToolCall, got {other:?}"), } @@ -435,7 +444,7 @@ async fn do_stream_returns_text() { .await; let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let result = model .do_stream(&default_options(test_prompt())) @@ -479,7 +488,7 @@ async fn do_stream_emits_tool_call() { .await; let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let tool = FunctionTool::new( "get-weather", @@ -507,7 +516,7 @@ async fn do_stream_emits_tool_call() { let (id, name, input) = tool_call.expect("should have ToolCall"); assert_eq!(id, "call_abc"); assert_eq!(name, "get-weather"); - assert_eq!(input, r#"{"city":"SF"}"#); + assert_eq!(input, Value::String(r#"{"city":"SF"}"#.into())); } /// A 401 response maps to `AiMuxError::ApiCall` (401 in `status_code`). @@ -523,7 +532,7 @@ async fn status_401_maps_to_auth_error() { .await; let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let result = model.do_generate(&default_options(test_prompt())).await; assert!( @@ -545,7 +554,7 @@ async fn status_429_maps_to_rate_limited() { .await; let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let result = model.do_generate(&default_options(test_prompt())).await; assert!( @@ -570,10 +579,8 @@ mod conformance { /// Build a provider whose requests land at `/api/v1/chat/completions`, /// matching the path recorded in the OpenRouter cassettes. - fn make_cassette_provider(server: &MockServer) -> OpenRouterProvider { - let base = format!("{}/api/v1", server.uri()); - let config = OpenRouterConfig::new("test-key").with_base_url(base); - OpenRouterProvider::new(config) + fn make_cassette_provider(server: &MockServer) -> OpenAICompatibleProvider { + openrouter_at(format!("{}/api/v1", server.uri()), "test-key") } fn has_text(content: &[GenerateContent]) -> bool { @@ -607,7 +614,7 @@ mod conformance { let provider = make_cassette_provider(&server); // `openai/gpt-4o-mini` matches the `completion_smoke` cassette (model // field weight 10), so a real non-streaming JSON response is returned. - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let result = model.do_generate(&default_options(test_prompt())).await; @@ -636,7 +643,7 @@ mod conformance { mount_cassettes(&server, "tests/cassettes/openrouter").await; let provider = make_cassette_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); + let model = provider.chat("openai/gpt-4o-mini"); let result = model.do_stream(&default_options(test_prompt())).await; diff --git a/aimux-providers/tests/parallel_ai_test.rs b/aimux-providers/tests/parallel_ai_test.rs index 82ca6261a..36a44dcd2 100644 --- a/aimux-providers/tests/parallel_ai_test.rs +++ b/aimux-providers/tests/parallel_ai_test.rs @@ -10,10 +10,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{ParallelAiConfig, ParallelAiProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{ParallelAiProvider, ParallelAiProviderSettings, create_parallel_ai}; const API_KEY: &str = "test-parallel-key"; @@ -38,8 +38,12 @@ fn search_body() -> Value { } fn provider(server: &MockServer) -> ParallelAiProvider { - let config = ParallelAiConfig::new(API_KEY).with_base_url(server.uri()); - ParallelAiProvider::new(config) + let config = ParallelAiProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_parallel_ai(config).unwrap() } fn opts(query: &str) -> SearchCallOptions { @@ -192,33 +196,14 @@ async fn status_403_maps_to_provider_error() { // -- Provider trait ---------------------------------------------------------- -#[tokio::test] -async fn provider_name_is_parallel_ai() { - let config = ParallelAiConfig::new(API_KEY); - let provider = ParallelAiProvider::new(config); - assert_eq!(provider.name(), "parallel_ai"); -} - -#[test] -fn language_model_returns_unsupported_error() { - let config = ParallelAiConfig::new(API_KEY); - let provider = ParallelAiProvider::new(config); - match provider.language_model("parallel-search") { - Err(AiMuxError::UnsupportedFunctionality(msg)) => { - assert!( - msg.contains("provider 'parallel_ai' does not provide language models"), - "unexpected message: {msg}" - ); - } - _ => panic!("expected Unsupported error, got success or another error variant"), - } -} - #[test] fn model_id_is_parallel_search() { - let config = ParallelAiConfig::new(API_KEY); - let provider = ParallelAiProvider::new(config); + let config = ParallelAiProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + ..Default::default() + }; + let provider = create_parallel_ai(config).unwrap(); let model = provider.search_model(); assert_eq!(model.model_id(), "parallel-search"); - assert_eq!(model.provider(), "parallel_ai"); + assert_eq!(model.provider(), "parallel_ai.search"); } diff --git a/aimux-providers/tests/perplexity_test.rs b/aimux-providers/tests/perplexity_test.rs index e59e4ea6d..a5476ad7c 100644 --- a/aimux-providers/tests/perplexity_test.rs +++ b/aimux-providers/tests/perplexity_test.rs @@ -4,7 +4,7 @@ //! - `packages/perplexity/src/perplexity-language-model.test.ts` //! - `packages/perplexity/src/convert-to-perplexity-messages.test.ts` //! -//! Perplexity is a thin OpenAI-compatible wrapper over [`OpenAIProvider`]. The +//! Perplexity is a registry preset of the OpenAI-compatible package. The //! behaviours verified here are the ones the wrapper is responsible for: //! //! - Provider configuration: name, `PERPLEXITY_API_KEY` env var, custom API @@ -23,6 +23,7 @@ use futures::StreamExt; use serde_json::{Value, json}; use serial_test::serial; +use std::sync::Arc; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -34,7 +35,7 @@ use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; use aimux_core::types::FinishReasonUnified; -use aimux_providers::{ProviderOptions, provider, provider_from_env}; +use aimux_providers::{PresetSettings, create_provider}; // ── shared helpers ─────────────────────────────────────────────────────────── @@ -86,16 +87,17 @@ async fn collect_stream(result: aimux_core::result::StreamResult) -> Vec Box { - provider( +fn make_provider(server: &MockServer) -> Arc { + create_provider( "perplexity", - Some("test-api-key".to_string()), - "sonar", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-api-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("sonar") .expect("perplexity provider should build") } @@ -106,8 +108,16 @@ fn make_provider(server: &MockServer) -> Box { /// TS: `createPerplexity()` produces a provider whose name is "perplexity". #[test] fn provider_builds_perplexity_model() { - let model = provider("perplexity", Some("test-key".to_string()), "sonar", None) - .expect("perplexity provider should build"); + let model = create_provider( + "perplexity", + PresetSettings { + api_key: Some("test-key".to_string().into()), + ..Default::default() + }, + ) + .expect("provider should build") + .language_model("sonar") + .expect("perplexity provider should build"); assert_eq!(model.model_id(), "sonar"); } @@ -122,15 +132,16 @@ async fn custom_api_key_used_in_auth_header() { .mount(&server) .await; - let model = provider( + let model = create_provider( "perplexity", - Some("my-custom-key".to_string()), - "sonar", - Some(ProviderOptions { + PresetSettings { + api_key: Some("my-custom-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("sonar") .expect("perplexity provider should build"); model @@ -150,15 +161,16 @@ async fn custom_headers_forwarded() { .mount(&server) .await; - let model = provider( + let model = create_provider( "perplexity", - Some("test-key".to_string()), - "sonar", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("sonar") .expect("perplexity provider should build"); let mut options = default_options(test_prompt()); @@ -188,15 +200,16 @@ async fn language_model_via_trait() { .mount(&server) .await; - let model = provider( + let model = create_provider( "perplexity", - Some("test-key".to_string()), - "sonar", - Some(ProviderOptions { + PresetSettings { + api_key: Some("test-key".to_string().into()), base_url: Some(server.uri()), ..Default::default() - }), + }, ) + .expect("provider should build") + .language_model("sonar") .expect("provider should build a model"); model @@ -214,7 +227,9 @@ fn from_env_loads_perplexity_api_key() { std::env::set_var("PERPLEXITY_API_KEY", "env-test-key"); } - let model = provider_from_env("perplexity", "sonar", None); + let model = create_provider("perplexity", PresetSettings::default()) + .expect("provider should build") + .language_model("sonar"); assert!(model.is_ok(), "from_env should succeed with env var set"); unsafe { @@ -227,15 +242,19 @@ fn from_env_loads_perplexity_api_key() { /// TS: without the env var, `createPerplexity()` fails. #[serial] -#[test] -fn from_env_fails_without_env_var() { +#[tokio::test] +async fn from_env_fails_without_env_var() { let saved = std::env::var("PERPLEXITY_API_KEY").ok(); unsafe { std::env::remove_var("PERPLEXITY_API_KEY"); } - let model = provider_from_env("perplexity", "sonar", None); - assert!(model.is_err(), "from_env should fail without env var"); + let model = create_provider("perplexity", PresetSettings::default()) + .expect("provider should build") + .language_model("sonar") + .expect("provider should build a model"); + let result = model.do_generate(&default_options(test_prompt())).await; + assert!(result.is_err(), "from_env should fail without env var"); unsafe { if let Some(v) = saved { @@ -445,3 +464,32 @@ async fn status_429_maps_to_rate_limited() { "expected RateLimited, got {result:?}" ); } + +/// TS: "should expose the raw response headers". +#[tokio::test] +async fn exposes_response_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("test-header", "test-value") + .set_body_json(text_completion_body()), + ) + .mount(&server) + .await; + + let model = make_provider(&server); + + let result = model + .do_generate(&default_options(test_prompt())) + .await + .expect("should succeed"); + + let headers = result + .response + .as_ref() + .and_then(|response| response.headers.as_ref()) + .expect("response_headers should be Some"); + assert_eq!(headers.get("test-header"), Some(&"test-value".to_string())); +} diff --git a/aimux-providers/tests/presets_test.rs b/aimux-providers/tests/presets_test.rs new file mode 100644 index 000000000..cbbaad627 --- /dev/null +++ b/aimux-providers/tests/presets_test.rs @@ -0,0 +1,367 @@ +//! Runtime registry validation and by-name preset behavior. + +use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use bytes::Bytes; +use reqwest::StatusCode; +use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue}; +use serde_json::json; +use serial_test::serial; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use aimux_core::AiMuxError; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelMessage; +use aimux_core::options::CallOptions; +use aimux_provider_utils::{Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse}; +use aimux_providers::preset; +use aimux_providers::{PresetSettings, create_provider}; + +// ── injected transport ─────────────────────────────────────────────────────── + +/// One request as the transport saw it. +#[derive(Clone, Debug)] +struct Seen { + url: String, + /// Lower-cased names (the `http` crate normalizes them). + headers: BTreeMap, +} + +#[derive(Default)] +struct MockFetch { + canned: Mutex>>, + seen: Mutex>, +} + +impl MockFetch { + fn new(canned: Vec>) -> Arc { + Arc::new(Self { + canned: Mutex::new(canned.into()), + seen: Mutex::default(), + }) + } + + fn seen(&self) -> Vec { + self.seen.lock().unwrap().clone() + } + + fn transport(self: &Arc) -> FetchFunction { + self.clone() + } +} + +#[async_trait] +impl Fetch for MockFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + self.seen.lock().unwrap().push(Seen { + url: request.url.to_string(), + headers: request + .headers + .iter() + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().expect("ASCII header").to_string(), + ) + }) + .collect(), + }); + let canned = self + .canned + .lock() + .unwrap() + .pop_front() + .ok_or_else(|| FetchError::Other("no canned response left".into()))?; + let mut headers = HeaderMap::new(); + headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + Ok(FetchResponse::from_bytes( + StatusCode::OK, + headers, + request.url, + Bytes::from(canned), + )) + } +} + +/// Temporarily change an environment variable; used by serial tests. +struct EnvVar { + name: &'static str, + saved: Option, +} + +impl EnvVar { + fn set(name: &'static str, value: Option<&str>) -> Self { + let saved = std::env::var(name).ok(); + unsafe { + match value { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + Self { name, saved } + } +} + +impl Drop for EnvVar { + fn drop(&mut self) { + unsafe { + match &self.saved { + Some(value) => std::env::set_var(self.name, value), + None => std::env::remove_var(self.name), + } + } + } +} + +// ── helpers ────────────────────────────────────────────────────────────────── + +fn chat_ok() -> Vec { + serde_json::to_vec(&json!({ + "id": "chatcmpl-1", "object": "chat.completion", "created": 1, "model": "m", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + })).unwrap() +} + +fn hello() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::user_text("hi")]) +} + +#[serial] +#[test] +fn every_registry_row_loads_and_creates_without_environment() { + let entries: Vec<_> = preset::entries().collect(); + assert_eq!(entries.len(), 282); + assert!( + preset::names() + .collect::>() + .windows(2) + .all(|pair| pair[0] < pair[1]) + ); + let mut variables = BTreeSet::new(); + for entry in &entries { + let d = entry.descriptor; + variables.extend((!d.env_var.is_empty()).then_some(d.env_var.as_str())); + variables.extend(d.base_url_env.as_deref()); + variables.extend( + d.params + .iter() + .flat_map(|param| param.env.iter().map(String::as_str)), + ); + } + let _guards: Vec<_> = variables + .into_iter() + .map(|var| EnvVar::set(var, None)) + .collect(); + for entry in entries { + preset::create(&entry.descriptor.name, PresetSettings::default()) + .unwrap_or_else(|error| panic!("{}: {error}", entry.descriptor.name)); + } +} + +#[tokio::test] +#[serial] +async fn a_keyless_row_sends_no_authorization() { + // The row's optional key variable must be unset for this assertion. + let _key = EnvVar::set("OLLAMA_API_KEY", None); + let mock = MockFetch::new(vec![chat_ok()]); + let preset = preset::create( + "ollama", + PresetSettings { + base_url: Some("http://127.0.0.1:9/v1".into()), + fetch: Some(mock.transport()), + ..Default::default() + }, + ) + .unwrap(); + preset.chat("m").do_generate(&hello()).await.unwrap(); + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert!(!seen[0].headers.contains_key("authorization")); + assert_eq!(seen[0].url, "http://127.0.0.1:9/v1/chat/completions"); +} + +#[test] +fn an_unknown_name_has_no_fallback() { + assert!(preset::lookup("no-such-preset").is_none()); + assert!( + matches!(create_provider("no-such-preset", PresetSettings::default()), + Err(AiMuxError::NoSuchProvider { provider_id, .. }) if provider_id == "no-such-preset") + ); +} + +#[tokio::test] +async fn a_template_parameter_must_be_a_plain_host_segment() { + let mock = MockFetch::new(vec![chat_ok()]); + let result = preset::create( + "neon", + PresetSettings { + params: [("branch_host".into(), "user@host/path".into())].into(), + fetch: Some(mock.transport()), + ..Default::default() + }, + ); + assert!( + matches!(result, Err(AiMuxError::InvalidArgument(message)) if message.contains("branch_host")) + ); + assert!(mock.seen().is_empty()); +} + +/// The built-in table holds the vendor packages and every registry row, and a +/// registry over it resolves both the same way. +#[test] +fn default_providers_serve_packages_and_registry_rows_through_one_registry() { + let providers = aimux_providers::default_providers(); + assert!( + providers.len() > 282, + "packages and rows: {}", + providers.len() + ); + let registry = aimux_core::create_provider_registry(providers, Default::default()); + + let model = registry.language_model("anthropic:claude-x").unwrap(); + assert_eq!(model.provider(), "anthropic.messages"); + let model = registry.language_model("abacus:some/model:v1").unwrap(); + assert_eq!(model.provider(), "abacus.chat"); + assert_eq!(model.model_id(), "some/model:v1"); + // A vendor package wins over the registry row of the same name. + let model = registry.language_model("groq:m").unwrap(); + assert_eq!(model.provider(), "groq.chat"); + + let error = registry.language_model("nope:m").err().unwrap(); + assert!( + matches!(error, AiMuxError::NoSuchProvider { .. }), + "{error:?}" + ); +} + +/// A local server that answers one chat completion at `/v1/chat/completions`. +async fn chat_server() -> MockServer { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_raw(chat_ok(), "application/json")) + .mount(&server) + .await; + server +} + +/// The row's `base_url_env` (`OLLAMA_BASE_URL`, provider_registry.json +/// :1303-1310) is read when a request is made, not when the provider is +/// created (preset.rs:375-384 builds a lazy base URL; :468-477 reads the +/// variable): setting it after `create` sends the call there. +#[tokio::test] +#[serial] +async fn the_base_url_variable_is_read_per_request() { + let server = chat_server().await; + let _key = EnvVar::set("OLLAMA_API_KEY", None); + let _unset = EnvVar::set("OLLAMA_BASE_URL", None); + let preset = preset::create("ollama", PresetSettings::default()).unwrap(); + + let url = format!("{}/v1", server.uri()); + let _set = EnvVar::set("OLLAMA_BASE_URL", Some(url.as_str())); + preset.chat("m").do_generate(&hello()).await.unwrap(); + + let received = server.received_requests().await.unwrap(); + assert_eq!(received.len(), 1); + assert_eq!(received[0].url.path(), "/v1/chat/completions"); +} + +/// An explicit `base_url` wins over the variable (preset.rs:376). +#[tokio::test] +#[serial] +async fn an_explicit_base_url_wins_over_the_variable() { + let server = chat_server().await; + let _key = EnvVar::set("OLLAMA_API_KEY", None); + // Port 9 (discard) refuses connections: reaching it fails the call. + let _url = EnvVar::set("OLLAMA_BASE_URL", Some("http://127.0.0.1:9/v1")); + let preset = preset::create( + "ollama", + PresetSettings { + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }, + ) + .unwrap(); + + preset.chat("m").do_generate(&hello()).await.unwrap(); + + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +/// An unusable `OLLAMA_BASE_URL` does not fail the factory; the request fails +/// with an InvalidArgument naming the variable, and nothing is sent +/// (preset.rs:474-476, validate_base_url in provider-utils url.rs:33-56). +#[tokio::test] +#[serial] +async fn an_invalid_base_url_variable_fails_the_request_not_the_factory() { + let _key = EnvVar::set("OLLAMA_API_KEY", None); + let _url = EnvVar::set("OLLAMA_BASE_URL", Some("ftp://models.internal/v1")); + let mock = MockFetch::new(vec![chat_ok()]); + let preset = preset::create( + "ollama", + PresetSettings { + fetch: Some(mock.transport()), + ..Default::default() + }, + ) + .expect("the variable is not read when the provider is created"); + + let Err(error) = preset.chat("m").do_generate(&hello()).await else { + panic!("an ftp:// base URL must fail the request"); + }; + assert!( + matches!(&error, AiMuxError::InvalidArgument(message) + if message.contains("OLLAMA_BASE_URL") && message.contains("http or https")), + "{error:?}" + ); + assert!(mock.seen().is_empty()); +} + +/// A keyless row with an `env_var` reads it per request as an optional key +/// (`Credential::OptionalEnv`, preset.rs:393; shared/mod.rs:85, :113): +/// `Authorization: Bearer ` when set, nothing when empty or unset. +#[tokio::test] +#[serial] +async fn the_optional_key_variable_is_sent_only_when_set() { + let server = chat_server().await; + let _url = EnvVar::set("OLLAMA_BASE_URL", None); + let _key = EnvVar::set("OLLAMA_API_KEY", None); + let preset = preset::create( + "ollama", + PresetSettings { + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }, + ) + .unwrap(); + let model = preset.chat("m"); + + { + let _set = EnvVar::set("OLLAMA_API_KEY", Some("ollama-secret")); + model.do_generate(&hello()).await.unwrap(); + } + { + let _empty = EnvVar::set("OLLAMA_API_KEY", Some("")); + model.do_generate(&hello()).await.unwrap(); + } + model.do_generate(&hello()).await.unwrap(); + + let received = server.received_requests().await.unwrap(); + assert_eq!(received.len(), 3); + assert_eq!( + received[0].headers["authorization"].to_str().unwrap(), + "Bearer ollama-secret" + ); + assert!( + !received[1].headers.contains_key("authorization"), + "empty key" + ); + assert!( + !received[2].headers.contains_key("authorization"), + "unset key" + ); +} diff --git a/aimux-providers/tests/prodia_image_test.rs b/aimux-providers/tests/prodia_image_test.rs index 38f6b624f..2dbffe600 100644 --- a/aimux-providers/tests/prodia_image_test.rs +++ b/aimux-providers/tests/prodia_image_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Prodia image model tests. +//! Rust translation of the Prodia image model tests. //! Source: `reference/ai/packages/prodia/src/prodia-image-model.test.ts` use serde_json::{Value, json}; @@ -6,7 +6,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; -use aimux_providers::{ProdiaConfig, ProdiaProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{ProdiaProviderSettings, create_prodia}; const PROMPT: &str = "A cute baby sea otter"; @@ -48,8 +49,12 @@ async fn should_extract_generated_image() { let job = json!({ "job": { "status": "succeeded" } }).to_string(); let body = multipart_response("testboundary", &job, b"fake-prodia-image"); mock_prodia(&server, body).await; - let config = ProdiaConfig::new("test-key").with_base_url(server.uri()); - let model = ProdiaProvider::new(config).image("sdxl"); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_prodia(config).unwrap().image("sdxl"); let result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { ImageOutputs::Binary(imgs) => { @@ -66,8 +71,12 @@ async fn should_pass_prompt() { let job = json!({ "job": { "status": "succeeded" } }).to_string(); let body = multipart_response("testboundary", &job, b"image"); mock_prodia(&server, body).await; - let config = ProdiaConfig::new("test-key").with_base_url(server.uri()); - let model = ProdiaProvider::new(config).image("sdxl"); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_prodia(config).unwrap().image("sdxl"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); let req_body: Value = serde_json::from_slice(&reqs[0].body).unwrap(); @@ -81,8 +90,12 @@ async fn should_pass_size() { let job = json!({ "job": { "status": "succeeded" } }).to_string(); let body = multipart_response("testboundary", &job, b"image"); mock_prodia(&server, body).await; - let config = ProdiaConfig::new("test-key").with_base_url(server.uri()); - let model = ProdiaProvider::new(config).image("sdxl"); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_prodia(config).unwrap().image("sdxl"); let mut opts = options(PROMPT); opts.size = Some(aimux_core::shared::Size::new(1024, 1024)); model.do_generate(&opts).await.unwrap(); @@ -98,8 +111,12 @@ async fn should_pass_seed() { let job = json!({ "job": { "status": "succeeded" } }).to_string(); let body = multipart_response("testboundary", &job, b"image"); mock_prodia(&server, body).await; - let config = ProdiaConfig::new("test-key").with_base_url(server.uri()); - let model = ProdiaProvider::new(config).image("sdxl"); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_prodia(config).unwrap().image("sdxl"); let mut opts = options(PROMPT); opts.seed = Some(42); model.do_generate(&opts).await.unwrap(); @@ -114,8 +131,12 @@ async fn should_pass_auth_headers() { let job = json!({ "job": { "status": "succeeded" } }).to_string(); let body = multipart_response("testboundary", &job, b"image"); mock_prodia(&server, body).await; - let config = ProdiaConfig::new("test-key").with_base_url(server.uri()); - let model = ProdiaProvider::new(config).image("sdxl"); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_prodia(config).unwrap().image("sdxl"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); assert_eq!( @@ -135,8 +156,12 @@ async fn should_return_provider_metadata() { let job = json!({ "status": "succeeded", "image": "url" }).to_string(); let body = multipart_response("testboundary", &job, b"image"); mock_prodia(&server, body).await; - let config = ProdiaConfig::new("test-key").with_base_url(server.uri()); - let model = ProdiaProvider::new(config).image("sdxl"); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_prodia(config).unwrap().image("sdxl"); let result = model.do_generate(&options(PROMPT)).await.unwrap(); let meta = result.provider_metadata.unwrap(); let prodia = meta.get("prodia").unwrap(); diff --git a/aimux-providers/tests/prodia_video_test.rs b/aimux-providers/tests/prodia_video_test.rs index e485ae7e8..2005bc7cc 100644 --- a/aimux-providers/tests/prodia_video_test.rs +++ b/aimux-providers/tests/prodia_video_test.rs @@ -1,10 +1,13 @@ -//! Rust translation of the Prodia video model tests. +//! Rust translation of the Prodia video model tests. //! Source: `reference/ai/packages/prodia/src/prodia-video-model.test.ts` +use aimux_core::download::{Download, DownloadRequest, DownloadedFile}; use aimux_core::video_model::{VideoCallOptions, generate_video}; -use aimux_providers::{ProdiaConfig, ProdiaProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{ProdiaProviderSettings, create_prodia}; use serde_json::{Value, json}; use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -18,9 +21,27 @@ fn fast_poll() -> Option { fn options(p: &str) -> VideoCallOptions { let mut o = VideoCallOptions::new(p); o.poll = fast_poll(); + o.download = Some(recording_download(Arc::default())); o } +/// The bytes the test download returns for every video URL. +const VIDEO_BYTES: &[u8] = b"video bytes"; + +/// A download that fetches nothing: it records each URL `generate_video` +/// asks for and returns [`VIDEO_BYTES`]. +fn recording_download(requests: Arc>>) -> Download { + Download::new(move |request: DownloadRequest| { + requests.lock().unwrap().push(request.url); + async { + Ok(DownloadedFile { + data: VIDEO_BYTES.to_vec(), + media_type: None, + }) + } + }) +} + async fn mock_job_and_result(server: &MockServer, result: &Value) { Mock::given(method("POST")) .and(path("/job")) @@ -42,8 +63,12 @@ async fn should_generate_video() { &json!({"status": "done", "videoUrl": "https://cdn.prodia.com/video.mp4"}), ) .await; - let config = ProdiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ProdiaProvider::new(config); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_prodia(config).unwrap(); let model = provider.video("svd-xt"); let r = generate_video(&model, options("A cat")).await.unwrap(); assert_eq!(r.videos.len(), 1); @@ -57,8 +82,12 @@ async fn should_pass_model_type_and_prompt() { &json!({"status": "done", "videoUrl": "https://cdn.prodia.com/video.mp4"}), ) .await; - let config = ProdiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ProdiaProvider::new(config); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_prodia(config).unwrap(); let model = provider.video("svd-xt"); generate_video(&model, options("A cat")).await.unwrap(); let requests = server.received_requests().await.unwrap(); @@ -77,10 +106,13 @@ async fn should_pass_headers() { .await; let mut ph = HashMap::new(); ph.insert("Custom-Header".to_string(), "val".to_string()); - let config = ProdiaConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = ProdiaProvider::new(config); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_prodia(config).unwrap(); let model = provider.video("svd-xt"); let mut opts = options("test"); let mut rh = HashMap::new(); @@ -101,8 +133,12 @@ async fn should_include_response_data() { &json!({"status": "done", "videoUrl": "https://cdn.prodia.com/video.mp4"}), ) .await; - let config = ProdiaConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ProdiaProvider::new(config); + let config = ProdiaProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_prodia(config).unwrap(); let model = provider.video("svd-xt"); let r = generate_video(&model, options("test")).await.unwrap(); assert!(r.response.timestamp.is_some()); diff --git a/aimux-providers/tests/provider_error_test.rs b/aimux-providers/tests/provider_error_test.rs index bfbe8fd40..e667afa90 100644 --- a/aimux-providers/tests/provider_error_test.rs +++ b/aimux-providers/tests/provider_error_test.rs @@ -28,13 +28,23 @@ use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::stream_part::StreamPart; -use aimux_provider_utils::RetryConfig; -use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider}; -use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; +use aimux_providers::openai::{OpenAIProvider, OpenAIProviderSettings, create_openai}; use futures::StreamExt; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +/// A native OpenAI provider against `base_url`. The key is an explicit value, +/// so the environment is never consulted. +fn provider_with(api_key: &str, base_url: impl Into) -> OpenAIProvider { + create_openai(OpenAIProviderSettings { + api_key: Some(api_key.to_string()), + base_url: Some(base_url.into()), + ..Default::default() + }) + .expect("settings are valid") +} + // ── helpers ───────────────────────────────────────────────────────────────── /// The TS `TEST_PROMPT`: a single user text message "Hello". @@ -73,13 +83,7 @@ mod openai_generate_errors { use super::*; fn model(server: &MockServer) -> impl LanguageModel { - let config = OpenAIConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_retry_config(RetryConfig { - max_retries: 0, - ..Default::default() - }); - OpenAIProvider::new(config).model("gpt-4o") + provider_with("test-api-key", server.uri()).chat("gpt-4o") } /// TS response-handler: 401 → AuthenticationError. @@ -216,13 +220,7 @@ mod openai_stream_errors { use super::*; fn model(server: &MockServer) -> impl LanguageModel { - let config = OpenAIConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_retry_config(RetryConfig { - max_retries: 0, - ..Default::default() - }); - OpenAIProvider::new(config).model("gpt-4o") + provider_with("test-api-key", server.uri()).chat("gpt-4o") } /// A non-success HTTP status on the stream endpoint makes `do_stream` @@ -412,13 +410,13 @@ mod anthropic_generate_errors { use super::*; fn model(server: &MockServer) -> impl LanguageModel { - let config = AnthropicConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_retry_config(RetryConfig { - max_retries: 0, - ..Default::default() - }); - AnthropicProvider::new(config).model("claude-3-haiku-20240307") + create_anthropic(AnthropicProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap() + .messages("claude-3-haiku-20240307") } /// Anthropic 401 → `ApiCall` (401). Anthropic body shape: @@ -546,13 +544,13 @@ mod anthropic_stream_errors { use super::*; fn model(server: &MockServer) -> impl LanguageModel { - let config = AnthropicConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_retry_config(RetryConfig { - max_retries: 0, - ..Default::default() - }); - AnthropicProvider::new(config).model("claude-3-haiku-20240307") + create_anthropic(AnthropicProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(format!("{}/v1", server.uri())), + ..Default::default() + }) + .unwrap() + .messages("claude-3-haiku-20240307") } /// TS: "should throw an api error when the server is returning a 529 diff --git a/aimux-providers/tests/reasoning_map_test.rs b/aimux-providers/tests/reasoning_map_test.rs deleted file mode 100644 index 0f9bdc6d1..000000000 --- a/aimux-providers/tests/reasoning_map_test.rs +++ /dev/null @@ -1,361 +0,0 @@ -//! stage2-002（RFC-0017 阶段 2）：退役回归 + max_tokens_key 矩阵 + reasoning_effort 直传。 -//! -//! 设计来源：[stage2-reasoning-map.md](../../docs/plan/analysis/stage2-reasoning-map.md) §4 I4-I5、§5。 -//! -//! 覆盖： -//! - **I4**：退役后 DeepSeek 请求体不含 `thinking` 注入（除非用户 bodyOverrides 注入） -//! - **I5**：用户 `body_overrides: { thinking: { type: 'disabled' } }` → 请求体含之 -//! （阶段 1 能力回归） -//! - **max_tokens_key 矩阵**：8 家接线（stepfun/siliconflow/sarvam/reka_ai/publicai/ -//! perplexity → `"max_tokens"`；groq/heroku → `"max_completion_tokens"`）× -//! 推理/非推理两分支。profile 取自注册表 `provider_registry_entry(name)` -//! （锁注册表接线，防死代码）。 -//! - **无 warning 断言**：直传语义下 reasoning 不再产生"未翻译"warning（防未来误加）。 -//! - **reasoning_effort 直传**：7 档无归一化（none/minimal/low/medium/high/xhigh 原样 -//! 透传；provider-default 不发字段）。 - -use aimux_core::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; -use aimux_core::options::CallOptions; -use aimux_core::types::{ReasoningEffort, Warning}; -use aimux_providers::openai::OpenAICompatProfile; -use aimux_providers::openai::convert::build_request_body_with_warnings; -use aimux_providers::provider_registry_entry; -use serde_json::json; - -fn user_prompt() -> LanguageModelPrompt { - vec![LanguageModelMessage::user_text("Hello")] -} - -fn opts_with_max_tokens(n: u32) -> CallOptions { - CallOptions { - prompt: user_prompt(), - max_output_tokens: Some(n), - ..CallOptions::default() - } -} - -fn opts_with_reasoning(r: ReasoningEffort) -> CallOptions { - CallOptions { - prompt: user_prompt(), - reasoning: Some(r), - ..CallOptions::default() - } -} - -fn has_reasoning_warning(warnings: &[Warning]) -> bool { - warnings - .iter() - .any(|w| matches!(w, Warning::Compatibility { feature, .. } if feature == "reasoning")) -} - -// ════════════════════════════════════════════════════════════════════════════ -// I4: 退役后 DeepSeek 请求体不含 thinking（除非用户 bodyOverrides） -// ════════════════════════════════════════════════════════════════════════════ - -/// I4: `reasoning:'none'` 透传为 `reasoning_effort:"none"`，请求体**不含** -/// `thinking` 注入（退役语义——thinking 注入不再由内置特化产生）。 -#[test] -fn i4_deepseek_retired_no_thinking_injection() { - let result = build_request_body_with_warnings( - "deepseek-reasoner", - &opts_with_reasoning(ReasoningEffort::None), - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - assert_eq!(result.body["reasoning_effort"], json!("none")); - assert!( - result.body.get("thinking").is_none(), - "I4: 退役后 DeepSeek 请求体不应含 thinking 注入: {:?}", - result.body - ); -} - -/// I4 补充: 完全不设 reasoning 时同样不含 thinking（依赖 API 默认）。 -#[test] -fn i4_deepseek_no_thinking_when_reasoning_unset() { - let result = build_request_body_with_warnings( - "deepseek-reasoner", - &CallOptions::new(user_prompt()), - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - assert!(result.body.get("thinking").is_none()); - assert!(result.body.get("reasoning_effort").is_none()); -} - -// ════════════════════════════════════════════════════════════════════════════ -// I5: 用户 body_overrides 注入 thinking → 请求体含之（阶段 1 能力回归） -// ════════════════════════════════════════════════════════════════════════════ - -/// I5: 用户 `body_overrides: { thinking: { type: 'disabled' } }` 原样进入请求体。 -/// 关思考的语义完全由用户定义（退役后不再由 reasoning:'none' 自动注入）。 -#[test] -fn i5_body_overrides_injects_thinking_disabled() { - let opts = CallOptions { - prompt: user_prompt(), - reasoning: Some(ReasoningEffort::None), - body_overrides: Some(json!({ "thinking": { "type": "disabled" } })), - ..CallOptions::default() - }; - let result = build_request_body_with_warnings( - "deepseek-reasoner", - &opts, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - assert_eq!(result.body["thinking"], json!({ "type": "disabled" })); - assert_eq!(result.body["reasoning_effort"], json!("none")); -} - -/// I5 补充: 用户注入 `thinking: { type: 'enabled' }` 同样原样进入请求体 -/// （开思考回归——退役前由特化注入，现由用户定义）。 -#[test] -fn i5_body_overrides_injects_thinking_enabled() { - let opts = CallOptions { - prompt: user_prompt(), - body_overrides: Some(json!({ "thinking": { "type": "enabled" } })), - ..CallOptions::default() - }; - let result = build_request_body_with_warnings( - "deepseek-reasoner", - &opts, - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - assert_eq!(result.body["thinking"], json!({ "type": "enabled" })); -} - -// ════════════════════════════════════════════════════════════════════════════ -// max_tokens_key 矩阵：8 家接线 × 推理/非推理两分支 -// -// profile 从注册表 `provider_registry_entry(name)` 实际构造取出——直接锁注册表接线 -// （若注册表行被改回 full()，此处 profile.max_tokens_key 断言即失败，防死代码）。 -// ════════════════════════════════════════════════════════════════════════════ - -/// 8 家接线清单：(provider 名, 实际 profile, 期望 max_tokens_key)。 -fn wired_vendors() -> Vec<(&'static str, OpenAICompatProfile, &'static str)> { - vec![ - ( - "stepfun", - provider_registry_entry("stepfun").unwrap(), - "max_tokens", - ), - ( - "siliconflow", - provider_registry_entry("siliconflow").unwrap(), - "max_tokens", - ), - ( - "sarvam", - provider_registry_entry("sarvam").unwrap(), - "max_tokens", - ), - ( - "reka_ai", - provider_registry_entry("reka_ai").unwrap(), - "max_tokens", - ), - ( - "publicai", - provider_registry_entry("publicai").unwrap(), - "max_tokens", - ), - ( - "perplexity", - provider_registry_entry("perplexity").unwrap(), - "max_tokens", - ), - ( - "groq", - provider_registry_entry("groq").unwrap(), - "max_completion_tokens", - ), - ( - "heroku", - provider_registry_entry("heroku").unwrap(), - "max_completion_tokens", - ), - ] -} - -/// 断言单个厂商在指定模型（推理/非推理）下请求体 key 名正确、另一 key 缺席。 -fn assert_vendor_key( - provider: &str, - profile: &OpenAICompatProfile, - expected_key: &str, - model_id: &str, - branch: &str, -) { - let result = build_request_body_with_warnings( - model_id, - &opts_with_max_tokens(100), - false, - provider, - profile, - ) - .unwrap(); - assert_eq!( - result.body[expected_key], - json!(100), - "[{provider}] {branch} 分支: 请求体应含 {expected_key}" - ); - let other_key = if expected_key == "max_tokens" { - "max_completion_tokens" - } else { - "max_tokens" - }; - assert!( - result.body.get(other_key).is_none(), - "[{}] {} 分支: 不应含 {}（只认 {}）: {:?}", - provider, - branch, - other_key, - expected_key, - result.body - ); -} - -/// 接线本身：每个厂商 profile.max_tokens_key 与清单一致（防注册表行被改回）。 -#[test] -fn max_tokens_key_wiring_registry() { - for (provider, profile, expected) in wired_vendors() { - assert_eq!( - profile.max_tokens_key, - Some(expected), - "[{provider}] 注册表接线错误" - ); - } -} - -/// 矩阵 × 推理分支（o3-mini 是推理模型 → 若未接线会推断发 mct，接线后按厂商 key 发）。 -#[test] -fn max_tokens_key_matrix_reasoning_branch() { - for (provider, profile, expected) in wired_vendors() { - assert_vendor_key(provider, &profile, expected, "o3-mini", "推理"); - } -} - -/// 矩阵 × 非推理分支（gpt-4o 非推理 → 若未接线会推断发 max_tokens，接线后按厂商 key 发）。 -#[test] -fn max_tokens_key_matrix_non_reasoning_branch() { - for (provider, profile, expected) in wired_vendors() { - assert_vendor_key(provider, &profile, expected, "gpt-4o", "非推理"); - } -} - -/// 未接线厂商（None）保持现状推断：推理模型发 mct、非推理发 max_tokens -/// （回归护栏——full() 默认不受接线影响）。 -#[test] -fn max_tokens_key_none_keeps_inference() { - let profile = OpenAICompatProfile::full(); - assert_eq!(profile.max_tokens_key, None); - - let reasoning = build_request_body_with_warnings( - "o3-mini", - &opts_with_max_tokens(100), - false, - "openai", - &profile, - ) - .unwrap(); - assert_eq!(reasoning.body["max_completion_tokens"], json!(100)); - assert!(reasoning.body.get("max_tokens").is_none()); - - let non_reasoning = build_request_body_with_warnings( - "gpt-4o", - &opts_with_max_tokens(100), - false, - "openai", - &profile, - ) - .unwrap(); - assert_eq!(non_reasoning.body["max_tokens"], json!(100)); - assert!(non_reasoning.body.get("max_completion_tokens").is_none()); -} - -// ════════════════════════════════════════════════════════════════════════════ -// 无 warning 断言（直传语义下无"未翻译"warning——防未来误加） -// ════════════════════════════════════════════════════════════════════════════ - -/// 直传语义下 `reasoning` 各档一律映射为 `reasoning_effort`，不应产生 -/// "reasoning 未翻译/无映射"兼容性 warning（stage2 已删除该死代码块）。 -#[test] -fn no_reasoning_warning_on_direct_passthrough() { - for effort in [ - ReasoningEffort::None, - ReasoningEffort::Minimal, - ReasoningEffort::Low, - ReasoningEffort::Medium, - ReasoningEffort::High, - ReasoningEffort::Xhigh, - ] { - let result = build_request_body_with_warnings( - "deepseek-reasoner", - &opts_with_reasoning(effort), - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - assert!( - !has_reasoning_warning(&result.warnings), - "effort={} 直传不应产生 reasoning warning: {:?}", - effort, - result.warnings - ); - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// reasoning_effort 直传：7 档无归一化 -// ════════════════════════════════════════════════════════════════════════════ - -/// 7 档（provider-default 不发字段；其余 6 档原样透传，含 none/minimal/xhigh， -/// 无 xhigh→max / minimal→low 归一化）。 -#[test] -fn reasoning_effort_passthrough_all_seven_levels() { - // provider-default: 不发 reasoning_effort（非自定义）。 - let default_body = build_request_body_with_warnings( - "deepseek-reasoner", - &opts_with_reasoning(ReasoningEffort::ProviderDefault), - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - assert!( - default_body.body.get("reasoning_effort").is_none(), - "provider-default 不应发 reasoning_effort" - ); - - for (effort, expected) in [ - (ReasoningEffort::None, "none"), - (ReasoningEffort::Minimal, "minimal"), - (ReasoningEffort::Low, "low"), - (ReasoningEffort::Medium, "medium"), - (ReasoningEffort::High, "high"), - (ReasoningEffort::Xhigh, "xhigh"), - ] { - let result = build_request_body_with_warnings( - "deepseek-reasoner", - &opts_with_reasoning(effort), - false, - "deepseek", - &OpenAICompatProfile::deepseek(), - ) - .unwrap(); - assert_eq!( - result.body["reasoning_effort"], - json!(expected), - "reasoning 档位应原样透传（无归一化）" - ); - } -} diff --git a/aimux-providers/tests/recraft_test.rs b/aimux-providers/tests/recraft_test.rs index 3a7df642c..dfb81dfbf 100644 --- a/aimux-providers/tests/recraft_test.rs +++ b/aimux-providers/tests/recraft_test.rs @@ -10,11 +10,11 @@ use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; use aimux_core::shared::Size; -use aimux_providers::{RecraftConfig, RecraftProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{RecraftProviderSettings, create_recraft}; const PROMPT: &str = "A cute baby sea otter"; const API_KEY: &str = "test-recraft-key"; @@ -43,8 +43,12 @@ fn options(prompt: &str) -> ImageCallOptions { } fn make_model(server: &MockServer) -> impl ImageModel { - let config = RecraftConfig::new(API_KEY).with_base_url(server.uri()); - RecraftProvider::new(config).image(MODEL) + let config = RecraftProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_recraft(config).unwrap().image(MODEL) } // ════════════════════════════════════════════════════════════════════════════ @@ -209,25 +213,6 @@ async fn status_500_maps_to_provider_error() { // Provider trait // ════════════════════════════════════════════════════════════════════════════ -#[tokio::test] -async fn provider_name_is_recraft() { - let server = MockServer::start().await; - let config = RecraftConfig::new(API_KEY).with_base_url(server.uri()); - let provider = RecraftProvider::new(config); - assert_eq!(provider.name(), "recraft"); -} - -#[test] -fn language_model_returns_unsupported_error() { - let config = RecraftConfig::new(API_KEY); - let provider = RecraftProvider::new(config); - let result = provider.language_model("recraftv3"); - assert!( - matches!(result, Err(AiMuxError::UnsupportedFunctionality(_))), - "expected Unsupported error" - ); -} - // ════════════════════════════════════════════════════════════════════════════ // 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 4ee5d239d..385cf110a 100644 --- a/aimux-providers/tests/replicate_image_test.rs +++ b/aimux-providers/tests/replicate_image_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Replicate image model tests. +//! Rust translation of the Replicate image model tests. //! Source: `reference/ai/packages/replicate/src/replicate-image-model.test.ts` use serde_json::{Value, json}; @@ -6,7 +6,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; -use aimux_providers::{ReplicateConfig, ReplicateProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{ReplicateProviderSettings, create_replicate}; const PROMPT: &str = "A cute baby sea otter"; @@ -32,8 +33,14 @@ fn options(prompt: &str) -> ImageCallOptions { async fn should_extract_generated_images() { let server = MockServer::start().await; mock_replicate(&server, json!(format!("{}/image.png", server.uri()))).await; - let config = ReplicateConfig::new("test-token").with_base_url(server.uri()); - let model = ReplicateProvider::new(config).image("black-forest-labs/flux-schnell"); + 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 result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { ImageOutputs::Binary(imgs) => { @@ -48,8 +55,14 @@ async fn should_extract_generated_images() { async fn should_pass_prompt_and_n() { let server = MockServer::start().await; mock_replicate(&server, json!(format!("{}/image.png", server.uri()))).await; - let config = ReplicateConfig::new("test-token").with_base_url(server.uri()); - let model = ReplicateProvider::new(config).image("black-forest-labs/flux-schnell"); + 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; model.do_generate(&opts).await.unwrap(); @@ -63,8 +76,14 @@ async fn should_pass_prompt_and_n() { async fn should_pass_aspect_ratio() { let server = MockServer::start().await; mock_replicate(&server, json!(format!("{}/image.png", server.uri()))).await; - let config = ReplicateConfig::new("test-token").with_base_url(server.uri()); - let model = ReplicateProvider::new(config).image("black-forest-labs/flux-schnell"); + 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.aspect_ratio = Some(aimux_core::shared::AspectRatio::new(16, 9)); @@ -78,8 +97,14 @@ async fn should_pass_aspect_ratio() { async fn should_pass_seed() { let server = MockServer::start().await; mock_replicate(&server, json!(format!("{}/image.png", server.uri()))).await; - let config = ReplicateConfig::new("test-token").with_base_url(server.uri()); - let model = ReplicateProvider::new(config).image("black-forest-labs/flux-schnell"); + 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.seed = Some(42); @@ -93,8 +118,14 @@ async fn should_pass_seed() { async fn should_pass_auth_headers() { let server = MockServer::start().await; mock_replicate(&server, json!(format!("{}/image.png", server.uri()))).await; - let config = ReplicateConfig::new("test-token").with_base_url(server.uri()); - let model = ReplicateProvider::new(config).image("black-forest-labs/flux-schnell"); + 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"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); assert_eq!( @@ -117,8 +148,14 @@ async fn should_handle_array_output() { let server = MockServer::start().await; let img_url = format!("{}/image.png", server.uri()); mock_replicate(&server, json!([img_url, img_url])).await; - let config = ReplicateConfig::new("test-token").with_base_url(server.uri()); - let model = ReplicateProvider::new(config).image("black-forest-labs/flux-schnell"); + 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 result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { ImageOutputs::Binary(imgs) => assert_eq!(imgs.len(), 2), @@ -130,8 +167,14 @@ async fn should_handle_array_output() { async fn should_use_versioned_endpoint() { let server = MockServer::start().await; mock_replicate(&server, json!(format!("{}/image.png", server.uri()))).await; - let config = ReplicateConfig::new("test-token").with_base_url(server.uri()); - let model = ReplicateProvider::new(config).image("black-forest-labs/flux-schnell:abc123"); + 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:abc123"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&reqs[0].body).unwrap(); diff --git a/aimux-providers/tests/replicate_video_test.rs b/aimux-providers/tests/replicate_video_test.rs index 41eb81624..20ea6d96a 100644 --- a/aimux-providers/tests/replicate_video_test.rs +++ b/aimux-providers/tests/replicate_video_test.rs @@ -1,10 +1,13 @@ -//! Rust translation of the Replicate video model tests. +//! Rust translation of the Replicate video model tests. //! Source: `reference/ai/packages/replicate/src/replicate-video-model.test.ts` +use aimux_core::download::{Download, DownloadRequest, DownloadedFile}; use aimux_core::video_model::{VideoCallOptions, generate_video}; -use aimux_providers::{ReplicateConfig, ReplicateProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{ReplicateProviderSettings, create_replicate}; use serde_json::{Value, json}; use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -18,9 +21,27 @@ fn fast_poll() -> Option { fn options(p: &str) -> VideoCallOptions { let mut o = VideoCallOptions::new(p); o.poll = fast_poll(); + o.download = Some(recording_download(Arc::default())); o } +/// The bytes the test download returns for every video URL. +const VIDEO_BYTES: &[u8] = b"video bytes"; + +/// A download that fetches nothing: it records each URL `generate_video` +/// asks for and returns [`VIDEO_BYTES`]. +fn recording_download(requests: Arc>>) -> Download { + Download::new(move |request: DownloadRequest| { + requests.lock().unwrap().push(request.url); + async { + Ok(DownloadedFile { + data: VIDEO_BYTES.to_vec(), + media_type: None, + }) + } + }) +} + async fn mock_predict_and_result(server: &MockServer, output: &Value) { Mock::given(method("POST")) .and(path("/predictions")) @@ -41,8 +62,12 @@ async fn mock_predict_and_result(server: &MockServer, output: &Value) { async fn should_generate_video() { let server = MockServer::start().await; mock_predict_and_result(&server, &json!("https://cdn.replicate.com/video.mp4")).await; - let config = ReplicateConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ReplicateProvider::new(config); + let config = ReplicateProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_replicate(config).unwrap(); let model = provider.video("wan-lab/wan-2.1-t2v-14b"); let r = generate_video(&model, options("A cat")).await.unwrap(); assert_eq!(r.videos.len(), 1); @@ -52,8 +77,12 @@ async fn should_generate_video() { async fn should_pass_model_and_prompt() { let server = MockServer::start().await; mock_predict_and_result(&server, &json!("https://cdn.replicate.com/video.mp4")).await; - let config = ReplicateConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ReplicateProvider::new(config); + let config = ReplicateProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_replicate(config).unwrap(); let model = provider.video("wan-lab/wan-2.1-t2v-14b"); generate_video(&model, options("A cat")).await.unwrap(); let requests = server.received_requests().await.unwrap(); @@ -68,10 +97,13 @@ async fn should_pass_headers() { mock_predict_and_result(&server, &json!("https://cdn.replicate.com/video.mp4")).await; let mut ph = HashMap::new(); ph.insert("Custom-Header".to_string(), "val".to_string()); - let config = ReplicateConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = ReplicateProvider::new(config); + let config = ReplicateProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_replicate(config).unwrap(); let model = provider.video("test-model"); let mut opts = options("test"); let mut rh = HashMap::new(); @@ -88,8 +120,12 @@ async fn should_pass_headers() { async fn should_include_response_data() { let server = MockServer::start().await; mock_predict_and_result(&server, &json!("https://cdn.replicate.com/video.mp4")).await; - let config = ReplicateConfig::new("test-api-key").with_base_url(server.uri()); - let provider = ReplicateProvider::new(config); + let config = ReplicateProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_replicate(config).unwrap(); let model = provider.video("test-model"); let r = generate_video(&model, options("test")).await.unwrap(); assert!(r.response.timestamp.is_some()); diff --git a/aimux-providers/tests/revai_transcription_test.rs b/aimux-providers/tests/revai_transcription_test.rs index c6fdefc9c..e206d6131 100644 --- a/aimux-providers/tests/revai_transcription_test.rs +++ b/aimux-providers/tests/revai_transcription_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Rev.ai transcription model tests. +//! Rust translation of the Rev.ai transcription model tests. //! Source: `reference/ai/packages/revai/src/revai-transcription-model.test.ts` use std::collections::HashMap; @@ -8,7 +8,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; -use aimux_providers::{RevaiConfig, RevaiProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{RevaiProviderSettings, create_revai}; fn mock_audio() -> Vec { vec![1u8, 2, 3] @@ -95,8 +96,12 @@ async fn should_pass_model_in_config() { let server = MockServer::start().await; mock_all_endpoints(&server, &fixture_transcript(), &[]).await; - let config = RevaiConfig::new("test-api-key").with_base_url(server.uri()); - let provider = RevaiProvider::new(config); + let config = RevaiProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_revai(config).unwrap(); let model = provider.transcription("machine"); model @@ -122,10 +127,13 @@ async fn should_pass_headers() { "Custom-Provider-Header".to_string(), "provider-header-value".to_string(), ); - let config = RevaiConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(ph); - let provider = RevaiProvider::new(config); + let config = RevaiProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some((ph).into_iter().map(|(k, v)| (k, Some(v))).collect()), + ..Default::default() + }; + let provider = create_revai(config).unwrap(); let model = provider.transcription("machine"); let mut opts = options(mock_audio(), "audio/wav"); @@ -156,8 +164,12 @@ async fn should_extract_text() { let server = MockServer::start().await; mock_all_endpoints(&server, &fixture_transcript(), &[]).await; - let config = RevaiConfig::new("test-api-key").with_base_url(server.uri()); - let provider = RevaiProvider::new(config); + let config = RevaiProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_revai(config).unwrap(); let model = provider.transcription("machine"); let result = model @@ -181,8 +193,12 @@ async fn should_include_response_data() { ) .await; - let config = RevaiConfig::new("test-api-key").with_base_url(server.uri()); - let provider = RevaiProvider::new(config); + let config = RevaiProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_revai(config).unwrap(); let model = provider.transcription("machine"); let result = model @@ -199,8 +215,12 @@ async fn should_use_real_date() { let server = MockServer::start().await; mock_all_endpoints(&server, &fixture_transcript(), &[]).await; - let config = RevaiConfig::new("test-api-key").with_base_url(server.uri()); - let provider = RevaiProvider::new(config); + let config = RevaiProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_revai(config).unwrap(); let model = provider.transcription("machine"); let result = model diff --git a/aimux-providers/tests/runtime_decision_test.rs b/aimux-providers/tests/runtime_decision_test.rs index 6599a9717..b92894c31 100644 --- a/aimux-providers/tests/runtime_decision_test.rs +++ b/aimux-providers/tests/runtime_decision_test.rs @@ -10,7 +10,7 @@ use aimux_core::decision_model::*; use aimux_core::recording::{self, RingRecorder}; use aimux_core::replay::{MockDecisionReplayModel, replay_decision_with_model}; use aimux_core::{AbortSignal, AiMuxError, Provider}; -use aimux_providers::{JevConfig, JevProvider, ProviderOptions, provider_handle}; +use aimux_providers::{JevConfig, JevProvider, PresetSettings, create_provider}; use serde_json::{Value, json}; use serial_test::serial; use wiremock::matchers::{body_json, header, method, path}; @@ -31,16 +31,23 @@ fn fixture() -> Value { fn options() -> DecisionCallOptions { serde_json::from_value(fixture()["options"].clone()).unwrap() } -fn model(provider: &str, server: &MockServer) -> Box { - provider_handle( +fn runtime(provider: &str, api_key: Option<&str>, base_url: String) -> Arc { + create_provider( provider, - Some("runtime-secret".into()), - Some(ProviderOptions { - base_url: Some(format!("{}/v1", server.uri())), + PresetSettings { + api_key: api_key.map(|key| key.to_string().into()), + base_url: Some(base_url), ..Default::default() - }), + }, ) .unwrap() +} +fn model(provider: &str, server: &MockServer) -> Box { + runtime( + provider, + Some("runtime-secret"), + format!("{}/v1", server.uri()), + ) .decision_model(if provider.starts_with("cloudflare") { "@cf/cloudflare/clef" } else { @@ -48,6 +55,17 @@ fn model(provider: &str, server: &MockServer) -> Box { }) .unwrap() } +/// Rebuild from the caller's registry, which names the provider as recorded. +fn rebuild( + record: &aimux_core::recording::ProviderRecord, + provider: Arc, +) -> Result, AiMuxError> { + let registry = aimux_core::create_provider_registry( + std::collections::BTreeMap::from([(record.provider_id.clone(), provider)]), + Default::default(), + ); + aimux_providers::rebuild_decision_provider(record, ®istry) +} struct Stop; impl Drop for Stop { fn drop(&mut self) { @@ -86,9 +104,11 @@ async fn replay( ); let server = MockServer::start().await; common::replay::mount_recording(&server, &recordings[0]).await; - let mut snapshot = recordings[0].provider.clone(); - snapshot.base_url = Some(format!("{}/v1", server.uri())); - let rebuilt = aimux_providers::rebuild_decision_provider(&snapshot, Some("new-key")).unwrap(); + let rebuilt = rebuild( + &recordings[0].provider, + runtime(provider, Some("new-key"), format!("{}/v1", server.uri())), + ) + .unwrap(); let again = replay_decision_with_model(&recordings[0], rebuilt.as_ref()) .await .unwrap(); @@ -578,29 +598,56 @@ async fn retry_and_cancellation_use_shared_transport() { } #[test] -fn local_factories_and_wrapper_factories_preserve_credentials_for_replay() { - use aimux_core::Provider; +fn local_presets_record_identity_and_rebuild_from_the_registry() { for provider in ["ollama", "llamacpp", "localai", "laya", "sglang", "vllm"] { - let handle = provider_handle( - provider, - None, - Some(ProviderOptions { - base_url: Some("http://127.0.0.1:1/v1".into()), - ..Default::default() - }), - ) - .unwrap(); + let handle = runtime(provider, None, "http://127.0.0.1:1/v1".into()); let model = handle.decision_model("served-alias").unwrap(); assert_eq!(model.provider(), provider); - let record = model.config_snapshot(); - let rebuilt = aimux_providers::rebuild_decision_provider(&record, None).unwrap(); + let record = + aimux_core::recording::ProviderRecord::from_model(model.provider(), model.model_id()); + assert_eq!(record.provider_id, provider); + let rebuilt = rebuild(&record, handle.clone()).unwrap(); assert_eq!(rebuilt.provider(), provider); + // Another runtime registered under the recorded name is refused. + let other = if provider == "vllm" { "sglang" } else { "vllm" }; + assert!( + rebuild( + &record, + runtime(other, None, "http://127.0.0.1:1/v1".into()) + ) + .is_err() + ); } - let provider = - aimux_providers::VllmProvider::new(aimux_providers::VllmConfig::new("explicit-key")); - let model = provider.decision_model("served-alias").unwrap(); - assert_eq!(model.config_snapshot().api_key_source, "explicit"); - assert!(aimux_providers::rebuild_decision_provider(&model.config_snapshot(), None).is_err()); + assert!(matches!( + runtime("deepseek", Some("key"), "http://127.0.0.1:1/v1".into()).decision_model("any"), + Err(AiMuxError::UnsupportedFunctionality(_)) + )); +} + +#[tokio::test] +#[serial] +async fn local_runtime_key_is_read_from_its_optional_variable() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/systemone")) + .respond_with(ResponseTemplate::new(200).set_body_json(&fixture()["response"])) + .mount(&server) + .await; + let base_url = format!("{}/v1", server.uri()); + // SAFETY: serial test; no other thread reads the environment meanwhile. + unsafe { std::env::remove_var("VLLM_API_KEY") }; + let model = runtime("vllm", None, base_url.clone()) + .decision_model("served-decision") + .unwrap(); + decide(model.as_ref(), options()).await.unwrap(); + // SAFETY: as above. + unsafe { std::env::set_var("VLLM_API_KEY", "from-env") }; + decide(model.as_ref(), options()).await.unwrap(); + // SAFETY: as above. + unsafe { std::env::remove_var("VLLM_API_KEY") }; + let requests = server.received_requests().await.unwrap(); + assert!(requests[0].headers.get("authorization").is_none()); + assert_eq!(requests[1].headers["authorization"], "Bearer from-env"); } #[tokio::test] diff --git a/aimux-providers/tests/runwayml_test.rs b/aimux-providers/tests/runwayml_test.rs index 3ed02fed9..fbd7afa87 100644 --- a/aimux-providers/tests/runwayml_test.rs +++ b/aimux-providers/tests/runwayml_test.rs @@ -6,13 +6,16 @@ //! async polling loop without a real delay. use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +use aimux_core::download::{Download, DownloadRequest, DownloadedFile}; use aimux_core::video_model::{VideoCallOptions, VideoModel, generate_video}; -use aimux_providers::{RunwaymlConfig, RunwaymlProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{RunwaymlProviderSettings, create_runwayml}; const MODEL_ID: &str = "gen3a_turbo"; const VIDEO_URL: &str = "https://cdn.runwayml.com/video.mp4"; @@ -27,14 +30,36 @@ fn fast_poll() -> Option { fn options(prompt: &str) -> VideoCallOptions { let mut o = VideoCallOptions::new(prompt); o.poll = fast_poll(); + o.download = Some(recording_download(Arc::default())); o } +/// The bytes the test download returns for every video URL. +const VIDEO_BYTES: &[u8] = b"video bytes"; + +/// A download that fetches nothing: it records each URL `generate_video` +/// asks for and returns [`VIDEO_BYTES`]. +fn recording_download(requests: Arc>>) -> Download { + Download::new(move |request: DownloadRequest| { + requests.lock().unwrap().push(request.url); + async { + Ok(DownloadedFile { + data: VIDEO_BYTES.to_vec(), + media_type: None, + }) + } + }) +} + /// A config pointed at the mock server. Poll pacing is a per-call option /// ([`fast_poll`]), as the AI SDK's `generateVideo` `poll`. Uses a dummy, /// non-secret API key. -fn config(server_uri: String) -> RunwaymlConfig { - RunwaymlConfig::new("test-api-key").with_base_url(server_uri) +fn config(server_uri: String) -> RunwaymlProviderSettings { + RunwaymlProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server_uri.to_string()), + ..Default::default() + } } /// Mount a submit mock (returning task id `task-123`) and a polling sequence: @@ -83,19 +108,23 @@ async fn should_generate_video_from_prompt() { let server = MockServer::start().await; mock_task_sequence(&server).await; - let provider = RunwaymlProvider::new(config(server.uri())); + let provider = create_runwayml(config(server.uri())).unwrap(); let model = provider.video(MODEL_ID); - let result = generate_video(&model, options("A cat playing")) - .await - .unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let mut opts = options("A cat playing"); + opts.download = Some(recording_download(requests.clone())); + let result = generate_video(&model, opts).await.unwrap(); + // `generate_video` downloads the returned URL, as the AI SDK does. + assert_eq!(*requests.lock().unwrap(), [VIDEO_URL]); assert_eq!(result.videos.len(), 1); match &result.videos[0] { - aimux_core::video_model::VideoData::Url { url, .. } => { - assert_eq!(url, VIDEO_URL); + aimux_core::video_model::VideoData::Binary { data, media_type } => { + assert_eq!(data, VIDEO_BYTES); + assert_eq!(media_type, "video/mp4"); } - v => panic!("expected Url, got: {v:?}"), + v => panic!("expected the downloaded video, got: {v:?}"), } assert!(result.response.timestamp.is_some()); assert_eq!(result.response.model_id, Some(MODEL_ID.to_string())); @@ -106,7 +135,7 @@ async fn should_post_to_text_to_video_endpoint() { let server = MockServer::start().await; mock_task_sequence(&server).await; - let provider = RunwaymlProvider::new(config(server.uri())); + let provider = create_runwayml(config(server.uri())).unwrap(); let model = provider.video(MODEL_ID); generate_video(&model, options("A cat playing")) @@ -127,7 +156,7 @@ async fn should_send_auth_and_version_headers() { let server = MockServer::start().await; mock_task_sequence(&server).await; - let provider = RunwaymlProvider::new(config(server.uri())); + let provider = create_runwayml(config(server.uri())).unwrap(); let model = provider.video(MODEL_ID); generate_video(&model, options("test")).await.unwrap(); @@ -161,10 +190,18 @@ async fn should_pass_custom_headers() { let mut provider_headers = HashMap::new(); provider_headers.insert("Custom-Provider-Header".to_string(), "val".to_string()); - let cfg = RunwaymlConfig::new("test-api-key") - .with_base_url(server.uri()) - .with_headers(provider_headers); - let provider = RunwaymlProvider::new(cfg); + let cfg = RunwaymlProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + headers: Some( + (provider_headers) + .into_iter() + .map(|(k, v)| (k, Some(v))) + .collect(), + ), + ..Default::default() + }; + let provider = create_runwayml(cfg).unwrap(); let model = provider.video(MODEL_ID); let mut opts = options("test"); @@ -198,7 +235,7 @@ async fn should_return_error_when_task_failed() { .mount(&server) .await; - let provider = RunwaymlProvider::new(config(server.uri())); + let provider = create_runwayml(config(server.uri())).unwrap(); let model = provider.video(MODEL_ID); let result = generate_video(&model, options("test")).await; @@ -219,7 +256,7 @@ async fn should_return_error_on_submit_failure() { .mount(&server) .await; - let provider = RunwaymlProvider::new(config(server.uri())); + let provider = create_runwayml(config(server.uri())).unwrap(); let model = provider.video(MODEL_ID); let result = generate_video(&model, options("test")).await; @@ -228,9 +265,13 @@ async fn should_return_error_on_submit_failure() { #[tokio::test] async fn should_return_max_videos_per_call() { - let provider = RunwaymlProvider::new(RunwaymlConfig::new("test-api-key")); + let provider = create_runwayml(RunwaymlProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + ..Default::default() + }) + .unwrap(); let model = provider.video(MODEL_ID); assert_eq!(model.max_videos_per_call(), Some(1)); - assert_eq!(model.provider(), "runwayml"); + assert_eq!(model.provider(), "runwayml.video"); assert_eq!(model.model_id(), MODEL_ID); } diff --git a/aimux-providers/tests/searxng_test.rs b/aimux-providers/tests/searxng_test.rs index fac774f35..9f8150b9c 100644 --- a/aimux-providers/tests/searxng_test.rs +++ b/aimux-providers/tests/searxng_test.rs @@ -11,10 +11,9 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path, query_param}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{SearxngConfig, SearxngProvider}; +use aimux_providers::{SearxngProvider, SearxngProviderSettings, create_searxng}; // -- helpers ----------------------------------------------------------------- @@ -41,8 +40,11 @@ fn search_body() -> Value { } fn provider(server: &MockServer) -> SearxngProvider { - let config = SearxngConfig::new(server.uri()); - SearxngProvider::new(config) + let config = SearxngProviderSettings { + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_searxng(config).unwrap() } fn opts(query: &str) -> SearchCallOptions { @@ -174,50 +176,16 @@ async fn status_401_maps_to_auth_error() { // -- Provider trait ---------------------------------------------------------- -#[tokio::test] -async fn provider_name_is_searxng() { - let config = SearxngConfig::new("http://localhost:8080"); - let provider = SearxngProvider::new(config); - assert_eq!(provider.name(), "searxng"); -} - -#[test] -fn language_model_returns_unsupported_error() { - let config = SearxngConfig::new("http://localhost:8080"); - let provider = SearxngProvider::new(config); - match provider.language_model("searxng-search") { - Err(AiMuxError::UnsupportedFunctionality(msg)) => { - assert!( - msg.contains("provider 'searxng' does not provide language models"), - "unexpected message: {msg}" - ); - } - _ => panic!("expected Unsupported error, got success or another error variant"), - } -} - #[test] fn model_id_is_searxng_search() { - let config = SearxngConfig::new("http://localhost:8080"); - let provider = SearxngProvider::new(config); + let config = SearxngProviderSettings { + base_url: Some("http://localhost:8080".to_string()), + ..Default::default() + }; + let provider = create_searxng(config).unwrap(); let model = provider.search_model(); assert_eq!(model.model_id(), "searxng-search"); - assert_eq!(model.provider(), "searxng"); + assert_eq!(model.provider(), "searxng.search"); } -// -- from_env ---------------------------------------------------------------- - -#[test] -fn from_env_requires_searxng_url() { - // Ensure the variable is unset for this test. - // SAFETY: this is the only test in this binary that touches `SEARXNG_URL`, - // so there is no concurrent access from other tests. - unsafe { - std::env::remove_var("SEARXNG_URL"); - } - let result = SearxngConfig::from_env(); - assert!( - matches!(result, Err(AiMuxError::InvalidArgument(ref m)) if m.contains("SEARXNG_URL")), - "expected InvalidArgument error, got {result:?}" - ); -} +// -- instance URL ----------------------------------------------------------- diff --git a/aimux-providers/tests/serper_test.rs b/aimux-providers/tests/serper_test.rs index b673194a9..d7050343f 100644 --- a/aimux-providers/tests/serper_test.rs +++ b/aimux-providers/tests/serper_test.rs @@ -6,22 +6,25 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::provider::Provider; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{SerperConfig, SerperProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{SerperProvider, SerperProviderSettings, create_serper}; fn make_provider(server: &MockServer) -> SerperProvider { - let config = SerperConfig::new("test-api-key").with_base_url(server.uri()); - SerperProvider::new(config) -} - -#[test] -fn provider_name_is_serper() { - let provider = SerperProvider::new(SerperConfig::new("test-key")); - assert_eq!(provider.name(), "serper"); + let config = SerperProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_serper(config).unwrap() } #[test] fn language_model_returns_unsupported() { - let provider = SerperProvider::new(SerperConfig::new("test-key")); + let provider = create_serper(SerperProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + ..Default::default() + }) + .unwrap(); assert!(provider.language_model("any").is_err()); } @@ -64,8 +67,12 @@ async fn uses_x_api_key_header() { .mount(&server) .await; - let config = SerperConfig::new("my-key").with_base_url(server.uri()); - let provider = SerperProvider::new(config); + let config = SerperProviderSettings { + api_key: Some(Resolvable::Value("my-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_serper(config).unwrap(); let model = provider.search_model(); model .do_search(&SearchCallOptions::new("test")) diff --git a/aimux-providers/tests/sse_terminal_error_test.rs b/aimux-providers/tests/sse_terminal_error_test.rs index 2d110c389..8fa09556d 100644 --- a/aimux-providers/tests/sse_terminal_error_test.rs +++ b/aimux-providers/tests/sse_terminal_error_test.rs @@ -6,7 +6,7 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::stream_part::StreamPart; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::{OpenAIProviderSettings, create_openai}; use futures::StreamExt; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; @@ -29,9 +29,14 @@ async fn corrupted_stream(prefix: Value) -> Vec> ) .mount(&server) .await; - let provider = OpenAIProvider::new(OpenAIConfig::new("test-key").with_base_url(server.uri())); + let provider = create_openai(OpenAIProviderSettings { + api_key: Some("test-key".into()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); provider - .model("gpt-4o") + .chat("gpt-4o") .do_stream(&CallOptions::new(vec![])) .await .expect("the first event is valid") diff --git a/aimux-providers/tests/stability_test.rs b/aimux-providers/tests/stability_test.rs index 04f4e06f0..84fc5090c 100644 --- a/aimux-providers/tests/stability_test.rs +++ b/aimux-providers/tests/stability_test.rs @@ -9,7 +9,8 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::error::AiMuxError; use aimux_core::image_model::{ImageCallOptions, ImageModel, ImageOutputs}; -use aimux_providers::{StabilityConfig, StabilityProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{StabilityProviderSettings, create_stability}; const PROMPT: &str = "A cute baby sea otter"; const API_KEY: &str = "test-key"; @@ -36,8 +37,14 @@ async fn should_extract_generated_image() { let server = MockServer::start().await; mock_ultra_image(&server, b"fake-stability-image".to_vec()).await; - let config = StabilityConfig::new(API_KEY).with_base_url(server.uri()); - let model = StabilityProvider::new(config).image("stable-image-ultra"); + let config = StabilityProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_stability(config) + .unwrap() + .image("stable-image-ultra"); let result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { @@ -54,8 +61,14 @@ async fn should_use_correct_url_for_ultra() { let server = MockServer::start().await; mock_ultra_image(&server, b"image".to_vec()).await; - let config = StabilityConfig::new(API_KEY).with_base_url(server.uri()); - let model = StabilityProvider::new(config).image("stable-image-ultra"); + let config = StabilityProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_stability(config) + .unwrap() + .image("stable-image-ultra"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); @@ -76,8 +89,12 @@ async fn should_use_correct_url_for_core() { .mount(&server) .await; - let config = StabilityConfig::new(API_KEY).with_base_url(server.uri()); - let model = StabilityProvider::new(config).image("stable-image-core"); + let config = StabilityProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_stability(config).unwrap().image("stable-image-core"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); @@ -98,8 +115,12 @@ async fn should_use_correct_url_for_sd3() { .mount(&server) .await; - let config = StabilityConfig::new(API_KEY).with_base_url(server.uri()); - let model = StabilityProvider::new(config).image("sd3"); + let config = StabilityProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_stability(config).unwrap().image("sd3"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); @@ -112,8 +133,14 @@ async fn should_pass_auth_header() { let server = MockServer::start().await; mock_ultra_image(&server, b"image".to_vec()).await; - let config = StabilityConfig::new(API_KEY).with_base_url(server.uri()); - let model = StabilityProvider::new(config).image("stable-image-ultra"); + let config = StabilityProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_stability(config) + .unwrap() + .image("stable-image-ultra"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); @@ -133,8 +160,14 @@ async fn should_pass_prompt_and_output_format_in_multipart_body() { let server = MockServer::start().await; mock_ultra_image(&server, b"image".to_vec()).await; - let config = StabilityConfig::new(API_KEY).with_base_url(server.uri()); - let model = StabilityProvider::new(config).image("stable-image-ultra"); + let config = StabilityProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_stability(config) + .unwrap() + .image("stable-image-ultra"); model.do_generate(&options(PROMPT)).await.unwrap(); let reqs = server.received_requests().await.unwrap(); @@ -156,8 +189,14 @@ async fn should_pass_seed_and_aspect_ratio() { let server = MockServer::start().await; mock_ultra_image(&server, b"image".to_vec()).await; - let config = StabilityConfig::new(API_KEY).with_base_url(server.uri()); - let model = StabilityProvider::new(config).image("stable-image-ultra"); + let config = StabilityProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_stability(config) + .unwrap() + .image("stable-image-ultra"); let mut opts = options(PROMPT); opts.seed = Some(42); opts.aspect_ratio = Some(aimux_core::shared::AspectRatio::new(16, 9)); @@ -194,8 +233,14 @@ async fn should_map_401_to_auth_error() { .mount(&server) .await; - let config = StabilityConfig::new(API_KEY).with_base_url(server.uri()); - let model = StabilityProvider::new(config).image("stable-image-ultra"); + let config = StabilityProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let model = create_stability(config) + .unwrap() + .image("stable-image-ultra"); let result = model.do_generate(&options(PROMPT)).await; assert!(result.is_err()); @@ -223,12 +268,23 @@ async fn should_decode_base64_json_response() { .mount(&server) .await; - let config = StabilityConfig::new(API_KEY).with_base_url(server.uri()); + let config = StabilityProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; // Override Accept to request the JSON/base64 response shape. - let mut headers = std::collections::HashMap::new(); - headers.insert("Accept".to_string(), "application/json".to_string()); - let config = config.with_headers(headers); - let model = StabilityProvider::new(config).image("stable-image-ultra"); + let config = StabilityProviderSettings { + headers: Some( + [("Accept".to_string(), Some("application/json".to_string()))] + .into_iter() + .collect(), + ), + ..config + }; + let model = create_stability(config) + .unwrap() + .image("stable-image-ultra"); let result = model.do_generate(&options(PROMPT)).await.unwrap(); match result.images { diff --git a/aimux-providers/tests/tavily_test.rs b/aimux-providers/tests/tavily_test.rs index 8572b1600..bd454a3a0 100644 --- a/aimux-providers/tests/tavily_test.rs +++ b/aimux-providers/tests/tavily_test.rs @@ -6,22 +6,25 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::provider::Provider; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{TavilyConfig, TavilyProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{TavilyProvider, TavilyProviderSettings, create_tavily}; fn make_provider(server: &MockServer) -> TavilyProvider { - let config = TavilyConfig::new("test-api-key").with_base_url(server.uri()); - TavilyProvider::new(config) -} - -#[test] -fn provider_name_is_tavily() { - let provider = TavilyProvider::new(TavilyConfig::new("test-key")); - assert_eq!(provider.name(), "tavily"); + let config = TavilyProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_tavily(config).unwrap() } #[test] fn language_model_returns_unsupported() { - let provider = TavilyProvider::new(TavilyConfig::new("test-key")); + let provider = create_tavily(TavilyProviderSettings { + api_key: Some(Resolvable::Value("test-key".to_string())), + ..Default::default() + }) + .unwrap(); assert!(provider.language_model("any").is_err()); } @@ -70,8 +73,12 @@ async fn uses_bearer_auth_header() { .mount(&server) .await; - let config = TavilyConfig::new("my-key").with_base_url(server.uri()); - let provider = TavilyProvider::new(config); + let config = TavilyProviderSettings { + api_key: Some(Resolvable::Value("my-key".to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + let provider = create_tavily(config).unwrap(); let model = provider.search_model(); model .do_search(&SearchCallOptions::new("test")) diff --git a/aimux-providers/tests/thin_wrapper_config_test.rs b/aimux-providers/tests/thin_wrapper_config_test.rs index 5eb2a5c71..085931416 100644 --- a/aimux-providers/tests/thin_wrapper_config_test.rs +++ b/aimux-providers/tests/thin_wrapper_config_test.rs @@ -7,7 +7,7 @@ //! suites verify: //! //! - `from_env()` loads the correct environment variable. -//! - `Provider::name()` returns the documented provider string. +//! - the provider's models report the documented provider string. //! - Custom headers are forwarded to the HTTP request. //! - The default base URL constant is wired correctly (verified via a mock //! round-trip that relies on `with_base_url`). @@ -21,13 +21,15 @@ use serde_json::{Value, json}; use wiremock::matchers::{header, 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, LanguageModelPrompt}; use aimux_core::options::CallOptions; use aimux_core::provider::Provider; use aimux_providers::{ - HuggingFaceConfig, HuggingFaceProvider, ProviderOptions, provider, provider_from_env, + HuggingFaceProvider, HuggingFaceProviderSettings, PresetSettings, create_huggingface, + create_provider, }; // ── shared helpers ─────────────────────────────────────────────────────────── @@ -40,6 +42,42 @@ fn default_options(prompt: LanguageModelPrompt) -> CallOptions { CallOptions::new(prompt) } +/// A minimal Responses API body returning one text message. +/// A registry preset as a language model. +fn preset_model( + name: &str, + api_key: Option<&str>, + model_id: &str, + settings: PresetSettings, +) -> Result, AiMuxError> { + create_provider( + name, + PresetSettings { + api_key: api_key.map(|key| key.to_string().into()), + ..settings + }, + )? + .language_model(model_id) +} + +fn responses_body() -> Value { + json!({ + "id": "resp_test", + "model": "test-model", + "object": "response", + "created_at": 1741257730, + "status": "completed", + "usage": { "input_tokens": 12, "output_tokens": 25, "total_tokens": 37 }, + "output": [{ + "id": "msg_test", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{ "type": "output_text", "text": "Hello, World!" }] + }] + }) +} + fn text_completion_body() -> Value { json!({ "id": "chatcmpl-test", @@ -62,14 +100,23 @@ fn text_completion_body() -> Value { mod huggingface_config { use super::*; + fn hf_provider(api_key: &str, base_url: Option) -> HuggingFaceProvider { + create_huggingface(HuggingFaceProviderSettings { + api_key: Some(api_key.to_string().into()), + base_url, + ..Default::default() + }) + .expect("valid settings") + } + /// TS: `createHuggingFace()` should create a provider with default - /// configuration. In Rust we verify the provider name and that a model - /// can be created. + /// configuration. In Rust we verify that a model can be created. #[test] - fn provider_name_is_huggingface() { - let config = HuggingFaceConfig::new("test-key"); - let provider = HuggingFaceProvider::new(config); - assert_eq!(provider.name(), "huggingface"); + fn model_is_created_with_default_configuration() { + let provider = hf_provider("test-key", None); + let model = provider.responses("any"); + assert_eq!(model.model_id(), "any"); + assert_eq!(model.provider(), "huggingface.responses"); } /// TS: `createHuggingFace({ apiKey: 'custom-key' })` �?custom API key @@ -78,15 +125,14 @@ mod huggingface_config { async fn custom_api_key_used_in_auth_header() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/chat/completions")) + .and(path("/responses")) .and(header("authorization", "Bearer my-custom-key")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) + .respond_with(ResponseTemplate::new(200).set_body_json(responses_body())) .mount(&server) .await; - let config = HuggingFaceConfig::new("my-custom-key").with_base_url(server.uri()); - let provider = HuggingFaceProvider::new(config); - let model = provider.model("meta-llama/Llama-3.3-70B-Instruct"); + let provider = hf_provider("my-custom-key", Some(server.uri())); + let model = provider.responses("meta-llama/Llama-3.3-70B-Instruct"); model .do_generate(&default_options(test_prompt())) @@ -100,15 +146,14 @@ mod huggingface_config { async fn custom_headers_forwarded() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/chat/completions")) + .and(path("/responses")) .and(header("x-custom-header", "test-value")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) + .respond_with(ResponseTemplate::new(200).set_body_json(responses_body())) .mount(&server) .await; - let config = HuggingFaceConfig::new("test-key").with_base_url(server.uri()); - let provider = HuggingFaceProvider::new(config); - let model = provider.model("meta-llama/Llama-3.3-70B-Instruct"); + let provider = hf_provider("test-key", Some(server.uri())); + let model = provider.responses("meta-llama/Llama-3.3-70B-Instruct"); let mut options = default_options(test_prompt()); options.headers = Some( @@ -123,64 +168,56 @@ mod huggingface_config { .expect("should succeed with custom headers"); } - /// TS: `provider.languageModel(modelId)` �?the provider should create a - /// language model via the `Provider` trait. + /// TS: `provider.languageModel(modelId)` — the provider creates the + /// Responses model (the AI SDK package has no chat model) via the + /// `Provider` trait. + #[test] + fn language_model_via_trait_is_the_responses_model() { + let provider = hf_provider("test-key", Some("http://127.0.0.1:1".to_string())); + let model = provider + .language_model("meta-llama/Llama-3.3-70B-Instruct") + .expect("language_model should succeed"); + assert_eq!(model.provider(), "huggingface.responses"); + } + + /// TS: the key is loaded from `HUGGINGFACE_API_KEY` when a request is + /// made, not when the provider is created. #[tokio::test] - async fn language_model_via_trait() { + #[serial_test::serial] + async fn key_is_loaded_from_env_per_request() { + let saved = std::env::var("HUGGINGFACE_API_KEY").ok(); + unsafe { std::env::remove_var("HUGGINGFACE_API_KEY") }; + let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) + .and(path("/responses")) + .and(header("authorization", "Bearer env-test-key")) + .respond_with(ResponseTemplate::new(200).set_body_json(responses_body())) .mount(&server) .await; - - let config = HuggingFaceConfig::new("test-key").with_base_url(server.uri()); - let provider = HuggingFaceProvider::new(config); - let model = provider - .language_model("meta-llama/Llama-3.3-70B-Instruct") - .expect("language_model should succeed"); - - model + let provider = create_huggingface(HuggingFaceProviderSettings { + base_url: Some(server.uri()), + ..Default::default() + }) + .expect("created without a key"); + let model = provider.responses("any"); + + let err = model .do_generate(&default_options(test_prompt())) .await - .expect("do_generate should succeed"); - } + .expect_err("no key yet"); + assert!(matches!(err, AiMuxError::LoadApiKey { .. }), "{err:?}"); - /// TS: `from_env` should load `HUGGINGFACE_API_KEY`. - #[test] - fn from_env_loads_correct_env_var() { - // Save and restore env var. - let saved = std::env::var("HUGGINGFACE_API_KEY").ok(); unsafe { std::env::set_var("HUGGINGFACE_API_KEY", "env-test-key") }; + let result = model.do_generate(&default_options(test_prompt())).await; - let config = HuggingFaceConfig::from_env(); - assert!(config.is_ok(), "from_env should succeed with env var set"); - - // Restore. unsafe { match saved { Some(v) => std::env::set_var("HUGGINGFACE_API_KEY", v), None => std::env::remove_var("HUGGINGFACE_API_KEY"), } } - } - - /// TS: `from_env` should fail without the env var. - #[test] - #[ignore = "flaky: parallel test env var race"] - fn from_env_fails_without_env_var() { - let saved = std::env::var("HUGGINGFACE_API_KEY").ok(); - unsafe { std::env::remove_var("HUGGINGFACE_API_KEY") }; - - let config = HuggingFaceConfig::from_env(); - assert!(config.is_err(), "from_env should fail without env var"); - - // Restore. - unsafe { - if let Some(v) = saved { - std::env::set_var("HUGGINGFACE_API_KEY", v); - } - } + result.expect("the key set after creation is used"); } } @@ -196,12 +233,12 @@ mod togetherai_config { #[test] fn provider_name_is_togetherai() { // Phase 4: the shell TogetherAIConfig/TogetherAIProvider pair is - // retired — registry-backed provider() replaces it. - let model = provider( + // retired — registry-backed create_provider() replaces it. + let model = preset_model( "togetherai", - Some("test-key".to_string()), + Some("test-key"), "meta-llama/Llama-3-70b-chat-hf", - None, + PresetSettings::default(), ) .expect("togetherai should construct from registry"); assert_eq!(model.model_id(), "meta-llama/Llama-3-70b-chat-hf"); @@ -219,14 +256,14 @@ mod togetherai_config { .mount(&server) .await; - let model = provider( + let model = preset_model( "togetherai", - Some("my-custom-key".to_string()), + Some("my-custom-key"), "meta-llama/Llama-3-70b-chat-hf", - Some(ProviderOptions { + PresetSettings { base_url: Some(server.uri()), ..Default::default() - }), + }, ) .expect("togetherai should construct from registry"); @@ -247,19 +284,22 @@ mod togetherai_config { .mount(&server) .await; - let model = provider( + let model = preset_model( "togetherai", - Some("test-key".to_string()), + Some("test-key"), "meta-llama/Llama-3-70b-chat-hf", - Some(ProviderOptions { + PresetSettings { base_url: Some(server.uri()), headers: Some( - vec![("x-custom-header".to_string(), "test-value".to_string())] - .into_iter() - .collect(), + vec![( + "x-custom-header".to_string(), + Some("test-value".to_string()), + )] + .into_iter() + .collect(), ), ..Default::default() - }), + }, ) .expect("togetherai should construct from registry"); @@ -275,10 +315,15 @@ mod togetherai_config { let saved = std::env::var("TOGETHER_API_KEY").ok(); unsafe { std::env::set_var("TOGETHER_API_KEY", "env-test-key") }; - let result = provider_from_env("togetherai", "meta-llama/Llama-3-70b-chat-hf", None); + let result = preset_model( + "togetherai", + None, + "meta-llama/Llama-3-70b-chat-hf", + PresetSettings::default(), + ); assert!( result.is_ok(), - "provider_from_env should succeed with env var set" + "the registry preset should succeed with env var set" ); unsafe { @@ -288,26 +333,6 @@ mod togetherai_config { } } } - - /// TS: `from_env` should fail without the env var. - #[test] - #[ignore = "flaky: parallel test env var race"] - fn from_env_fails_without_env_var() { - let saved = std::env::var("TOGETHER_API_KEY").ok(); - unsafe { std::env::remove_var("TOGETHER_API_KEY") }; - - let result = provider_from_env("togetherai", "meta-llama/Llama-3-70b-chat-hf", None); - assert!( - result.is_err(), - "provider_from_env should fail without env var" - ); - - unsafe { - if let Some(v) = saved { - std::env::set_var("TOGETHER_API_KEY", v); - } - } - } } // ════════════════════════════════════════════════════════════════════════════ @@ -322,9 +347,14 @@ mod vercel_config { #[test] fn provider_name_is_vercel() { // Phase 4: the shell VercelConfig/VercelProvider pair is retired — - // registry-backed provider() replaces it. - let model = provider("vercel", Some("test-key".to_string()), "v0-1.5-md", None) - .expect("vercel should construct from registry"); + // registry-backed create_provider() replaces it. + let model = preset_model( + "vercel", + Some("test-key"), + "v0-1.5-md", + PresetSettings::default(), + ) + .expect("vercel should construct from registry"); assert_eq!(model.model_id(), "v0-1.5-md"); } @@ -340,14 +370,14 @@ mod vercel_config { .mount(&server) .await; - let model = provider( + let model = preset_model( "vercel", - Some("my-custom-key".to_string()), + Some("my-custom-key"), "v0-1.5-md", - Some(ProviderOptions { + PresetSettings { base_url: Some(server.uri()), ..Default::default() - }), + }, ) .expect("vercel should construct from registry"); @@ -363,10 +393,10 @@ mod vercel_config { let saved = std::env::var("VERCEL_API_KEY").ok(); unsafe { std::env::set_var("VERCEL_API_KEY", "env-test-key") }; - let result = provider_from_env("vercel", "v0-1.5-md", None); + let result = preset_model("vercel", None, "v0-1.5-md", PresetSettings::default()); assert!( result.is_ok(), - "provider_from_env should succeed with env var set" + "the registry preset should succeed with env var set" ); unsafe { @@ -376,24 +406,4 @@ mod vercel_config { } } } - - /// TS: `from_env` should fail without the env var. - #[test] - #[ignore = "flaky: parallel test env var race"] - fn from_env_fails_without_env_var() { - let saved = std::env::var("VERCEL_API_KEY").ok(); - unsafe { std::env::remove_var("VERCEL_API_KEY") }; - - let result = provider_from_env("vercel", "v0-1.5-md", None); - assert!( - result.is_err(), - "provider_from_env should fail without env var" - ); - - unsafe { - if let Some(v) = saved { - std::env::set_var("VERCEL_API_KEY", v); - } - } - } } diff --git a/aimux-providers/tests/tinyfish_test.rs b/aimux-providers/tests/tinyfish_test.rs index 5258634b3..e011f0b52 100644 --- a/aimux-providers/tests/tinyfish_test.rs +++ b/aimux-providers/tests/tinyfish_test.rs @@ -10,10 +10,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{TinyfishConfig, TinyfishProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{TinyfishProvider, TinyfishProviderSettings, create_tinyfish}; const API_KEY: &str = "test-tinyfish-key"; @@ -48,8 +48,12 @@ fn search_response_body() -> Value { } fn provider(server: &MockServer) -> TinyfishProvider { - let config = TinyfishConfig::new(API_KEY).with_base_url(server.uri()); - TinyfishProvider::new(config) + let config = TinyfishProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_tinyfish(config).unwrap() } fn opts(query: &str, max_results: Option) -> SearchCallOptions { @@ -196,25 +200,3 @@ async fn status_401_maps_to_auth_error() { } // -- Provider trait ---------------------------------------------------------- - -#[tokio::test] -async fn provider_name_is_tinyfish() { - let config = TinyfishConfig::new(API_KEY); - let provider = TinyfishProvider::new(config); - assert_eq!(provider.name(), "tinyfish"); -} - -#[test] -fn language_model_returns_unsupported_error() { - let config = TinyfishConfig::new(API_KEY); - let provider = TinyfishProvider::new(config); - match provider.language_model("tinyfish-search") { - Err(AiMuxError::UnsupportedFunctionality(msg)) => { - assert!( - msg.contains("provider 'tinyfish' does not provide language models"), - "unexpected message: {msg}" - ); - } - _ => panic!("expected Unsupported error, got success or another error variant"), - } -} diff --git a/aimux-providers/tests/vertex_transcription_test.rs b/aimux-providers/tests/vertex_transcription_test.rs index 15accfcb0..23212dd7d 100644 --- a/aimux-providers/tests/vertex_transcription_test.rs +++ b/aimux-providers/tests/vertex_transcription_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Google Vertex transcription model tests. +//! Rust translation of the Google Vertex transcription model tests. //! Source: `reference/ai/packages/google-vertex/src/google-vertex-transcription-model.test.ts` use std::collections::HashMap; @@ -7,8 +7,9 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel}; -use aimux_providers::{VertexProvider, VertexProviderConfig}; +use aimux_core::Provider; +use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions}; +use aimux_providers::{VertexProvider, VertexProviderSettings, create_google_vertex}; fn mock_audio() -> Vec { vec![1u8, 2, 3, 4, 5, 6, 7, 8] @@ -34,9 +35,14 @@ fn default_body() -> Value { } fn make_provider(server_uri: &str) -> VertexProvider { - let config = VertexProviderConfig::new("test-token", "test-project", "us-central1") - .with_base_url(server_uri); - VertexProvider::new(config) + create_google_vertex(VertexProviderSettings { + access_token: Some("test-token".to_string().into()), + project: Some("test-project".to_string()), + location: Some("us-central1".to_string()), + base_url: Some(server_uri.to_string()), + ..Default::default() + }) + .unwrap() } async fn mock_response(server: &MockServer, body: &Value, headers: &[(&str, &str)]) { @@ -60,7 +66,7 @@ async fn should_send_model_language_codes_features_and_base64() { mock_response(&server, &default_body(), &[]).await; let provider = make_provider(&server.uri()); - let model = provider.transcription("chirp_2").unwrap(); + let model = provider.transcription_model("chirp_2").unwrap().unwrap(); model .do_generate(&options(mock_audio(), "audio/wav")) @@ -86,7 +92,7 @@ async fn should_extract_text_segments_language_and_duration() { mock_response(&server, &default_body(), &[]).await; let provider = make_provider(&server.uri()); - let model = provider.transcription("chirp_2").unwrap(); + let model = provider.transcription_model("chirp_2").unwrap().unwrap(); let result = model .do_generate(&options(mock_audio(), "audio/wav")) @@ -110,10 +116,24 @@ async fn should_pass_headers() { let server = MockServer::start().await; mock_response(&server, &default_body(), &[]).await; - let config = VertexProviderConfig::new("test-token", "test-project", "us-central1") - .with_base_url(server.uri()); - let provider = VertexProvider::new(config); - let model = provider.transcription("chirp_2").unwrap(); + let provider = create_google_vertex(VertexProviderSettings { + access_token: Some("test-token".to_string().into()), + project: Some("test-project".to_string()), + location: Some("us-central1".to_string()), + base_url: Some(server.uri()), + headers: Some( + [( + "Custom-Provider-Header".to_string(), + Some("provider-header-value".to_string()), + )] + .into_iter() + .collect::>() + .into(), + ), + ..Default::default() + }) + .unwrap(); + let model = provider.transcription_model("chirp_2").unwrap().unwrap(); let mut opts = options(mock_audio(), "audio/wav"); let mut rh = HashMap::new(); @@ -128,6 +148,10 @@ async fn should_pass_headers() { let requests = server.received_requests().await.unwrap(); let h = &requests[0].headers; assert_eq!(h.get("authorization").unwrap(), "Bearer test-token"); + assert_eq!( + h.get("custom-provider-header").unwrap(), + "provider-header-value" + ); assert_eq!( h.get("custom-request-header").unwrap(), "request-header-value" @@ -140,24 +164,7 @@ async fn should_include_response_data() { mock_response(&server, &default_body(), &[("x-request-id", "test-req")]).await; let provider = make_provider(&server.uri()); - let model = provider.transcription("chirp_2").unwrap(); - - let result = model - .do_generate(&options(mock_audio(), "audio/wav")) - .await - .unwrap(); - - assert!(result.response.timestamp.is_some()); - assert_eq!(result.response.model_id, Some("chirp_2".to_string())); -} - -#[tokio::test] -async fn should_use_real_date() { - let server = MockServer::start().await; - mock_response(&server, &default_body(), &[]).await; - - let provider = make_provider(&server.uri()); - let model = provider.transcription("chirp_2").unwrap(); + let model = provider.transcription_model("chirp_2").unwrap().unwrap(); let result = model .do_generate(&options(mock_audio(), "audio/wav")) @@ -166,6 +173,15 @@ async fn should_use_real_date() { assert!(result.response.timestamp.is_some()); assert_eq!(result.response.model_id, Some("chirp_2".to_string())); + assert_eq!( + result + .response + .headers + .unwrap() + .get("x-request-id") + .map(String::as_str), + Some("test-req") + ); } #[tokio::test] @@ -174,7 +190,7 @@ async fn should_handle_no_results() { mock_response(&server, &json!({}), &[]).await; let provider = make_provider(&server.uri()); - let model = provider.transcription("chirp_2").unwrap(); + let model = provider.transcription_model("chirp_2").unwrap().unwrap(); let result = model .do_generate(&options(mock_audio(), "audio/wav")) diff --git a/aimux-providers/tests/vertex_video_test.rs b/aimux-providers/tests/vertex_video_test.rs index b1abe8909..42a556d62 100644 --- a/aimux-providers/tests/vertex_video_test.rs +++ b/aimux-providers/tests/vertex_video_test.rs @@ -1,10 +1,13 @@ -//! Rust translation of the Google Vertex video model tests. +//! Rust translation of the Google Vertex video model tests. //! Source: `reference/ai/packages/google-vertex/src/google-vertex-video-model.test.ts` +use std::sync::{Arc, Mutex}; + +use aimux_core::Provider; +use aimux_core::download::{Download, DownloadRequest, DownloadedFile}; use aimux_core::video_model::{VideoCallOptions, generate_video}; -use aimux_providers::{VertexProvider, VertexProviderConfig}; +use aimux_providers::{VertexProvider, VertexProviderSettings, create_google_vertex}; use serde_json::{Value, json}; -use std::collections::HashMap; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -18,13 +21,36 @@ fn fast_poll() -> Option { fn options(p: &str) -> VideoCallOptions { let mut o = VideoCallOptions::new(p); o.poll = fast_poll(); + o.download = Some(recording_download(Arc::default())); o } +/// The bytes the test download returns for every video URL. +const VIDEO_BYTES: &[u8] = b"video bytes"; + +/// A download that fetches nothing: it records each URL `generate_video` +/// asks for and returns [`VIDEO_BYTES`]. +fn recording_download(requests: Arc>>) -> Download { + Download::new(move |request: DownloadRequest| { + requests.lock().unwrap().push(request.url); + async { + Ok(DownloadedFile { + data: VIDEO_BYTES.to_vec(), + media_type: None, + }) + } + }) +} + fn make_provider(server_uri: &str) -> VertexProvider { - let config = VertexProviderConfig::new("test-token", "test-project", "us-central1") - .with_base_url(server_uri); - VertexProvider::new(config) + create_google_vertex(VertexProviderSettings { + access_token: Some("test-token".to_string().into()), + project: Some("test-project".to_string()), + location: Some("us-central1".to_string()), + base_url: Some(server_uri.to_string()), + ..Default::default() + }) + .unwrap() } async fn mock_predict_and_poll(server: &MockServer, result: &Value) { @@ -35,8 +61,8 @@ async fn mock_predict_and_poll(server: &MockServer, result: &Value) { ) .mount(server) .await; - Mock::given(method("GET")) - .and(path("/operations/test-op")) + Mock::given(method("POST")) + .and(path("/models/veo-3.0-generate-001:fetchPredictOperation")) .respond_with(ResponseTemplate::new(200).set_body_json(result.clone())) .mount(server) .await; @@ -49,9 +75,25 @@ async fn should_generate_video() { json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); mock_predict_and_poll(&server, &result).await; let provider = make_provider(&server.uri()); - let model = provider.video("veo-3.0-generate-001").unwrap(); - let r = generate_video(&model, options("A cat")).await.unwrap(); + let model = provider + .video_model("veo-3.0-generate-001") + .unwrap() + .unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let mut opts = options("A cat"); + opts.download = Some(recording_download(requests.clone())); + let r = generate_video(model.as_ref(), opts).await.unwrap(); + // `generate_video` downloads the returned URL with the injected download + // (the guarded default fetches only HTTP(S), as upstream's `fetch`). + assert_eq!(*requests.lock().unwrap(), ["gs://bucket/video.mp4"]); assert_eq!(r.videos.len(), 1); + match &r.videos[0] { + aimux_core::video_model::VideoData::Binary { data, media_type } => { + assert_eq!(data, VIDEO_BYTES); + assert_eq!(media_type, "video/mp4"); + } + _ => panic!("expected the downloaded video"), + } } #[tokio::test] @@ -61,30 +103,24 @@ async fn should_pass_prompt() { json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); mock_predict_and_poll(&server, &result).await; let provider = make_provider(&server.uri()); - let model = provider.video("veo-3.0-generate-001").unwrap(); - generate_video(&model, options("A cat")).await.unwrap(); + let model = provider + .video_model("veo-3.0-generate-001") + .unwrap() + .unwrap(); + generate_video(model.as_ref(), options("A cat")) + .await + .unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!(body["instances"][0]["prompt"], "A cat"); -} - -#[tokio::test] -async fn should_pass_headers() { - let server = MockServer::start().await; - let result = - json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); - mock_predict_and_poll(&server, &result).await; - let provider = make_provider(&server.uri()); - let model = provider.video("veo-3.0-generate-001").unwrap(); - let mut opts = options("test"); - let mut rh = HashMap::new(); - rh.insert("Custom-Header".to_string(), "val".to_string()); - opts.headers = Some(rh); - generate_video(&model, opts).await.unwrap(); - let requests = server.received_requests().await.unwrap(); - let h = &requests[0].headers; - assert_eq!(h.get("authorization").unwrap(), "Bearer test-token"); - assert_eq!(h.get("custom-header").unwrap(), "val"); + assert_eq!( + body, + json!({ + "instances": [{"prompt": "A cat"}], + "parameters": {"sampleCount": 1}, + }) + ); + let poll_body: Value = serde_json::from_slice(&requests[1].body).unwrap(); + assert_eq!(poll_body, json!({"operationName": "operations/test-op"})); } #[tokio::test] @@ -94,8 +130,13 @@ async fn should_include_response_data() { json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); mock_predict_and_poll(&server, &result).await; let provider = make_provider(&server.uri()); - let model = provider.video("veo-3.0-generate-001").unwrap(); - let r = generate_video(&model, options("test")).await.unwrap(); + let model = provider + .video_model("veo-3.0-generate-001") + .unwrap() + .unwrap(); + let r = generate_video(model.as_ref(), options("test")) + .await + .unwrap(); assert!(r.response.timestamp.is_some()); assert_eq!( r.response.model_id, diff --git a/aimux-providers/tests/voyage_embedding_test.rs b/aimux-providers/tests/voyage_embedding_test.rs index fede560e0..332268881 100644 --- a/aimux-providers/tests/voyage_embedding_test.rs +++ b/aimux-providers/tests/voyage_embedding_test.rs @@ -10,7 +10,16 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; -use aimux_providers::{VoyageConfig, VoyageProvider}; +use aimux_providers::{VoyageProvider, VoyageProviderSettings, create_voyage}; + +fn test_provider(api_key: &str, base_url: impl Into) -> VoyageProvider { + create_voyage(VoyageProviderSettings { + api_key: Some(api_key.to_string().into()), + base_url: Some(base_url.into()), + ..Default::default() + }) + .expect("valid settings") +} const TEST_VALUES: &[&str] = &["sunny day at the beach", "rainy day in the city"]; @@ -56,9 +65,8 @@ async fn should_extract_embedding() { .mount(&server) .await; - let config = VoyageConfig::new("test-api-key").with_base_url(server.uri()); - let provider = VoyageProvider::new(config); - let model = provider.embedding_model("voyage-3.5"); + let provider = test_provider("test-api-key", server.uri()); + let model = provider.embedding("voyage-3.5"); let result = model .do_embed(&default_options(test_values())) @@ -101,9 +109,8 @@ async fn should_expose_raw_response() { .mount(&server) .await; - let config = VoyageConfig::new("test-api-key").with_base_url(server.uri()); - let provider = VoyageProvider::new(config); - let model = provider.embedding_model("voyage-3.5"); + let provider = test_provider("test-api-key", server.uri()); + let model = provider.embedding("voyage-3.5"); let result = model .do_embed(&default_options(test_values())) @@ -130,9 +137,8 @@ async fn should_extract_usage() { .mount(&server) .await; - let config = VoyageConfig::new("test-api-key").with_base_url(server.uri()); - let provider = VoyageProvider::new(config); - let model = provider.embedding_model("voyage-3.5"); + let provider = test_provider("test-api-key", server.uri()); + let model = provider.embedding("voyage-3.5"); let result = model .do_embed(&default_options(test_values())) @@ -152,9 +158,8 @@ async fn should_pass_model_and_values() { .mount(&server) .await; - let config = VoyageConfig::new("test-api-key").with_base_url(server.uri()); - let provider = VoyageProvider::new(config); - let model = provider.embedding_model("voyage-3.5"); + let provider = test_provider("test-api-key", server.uri()); + let model = provider.embedding("voyage-3.5"); let _ = model .do_embed(&default_options(test_values())) @@ -183,9 +188,8 @@ async fn should_pass_headers() { .mount(&server) .await; - let config = VoyageConfig::new("test-api-key").with_base_url(server.uri()); - let provider = VoyageProvider::new(config); - let model = provider.embedding_model("voyage-3.5"); + let provider = test_provider("test-api-key", server.uri()); + let model = provider.embedding("voyage-3.5"); let mut request_headers = HashMap::new(); request_headers.insert( diff --git a/aimux-providers/tests/voyage_reranking_test.rs b/aimux-providers/tests/voyage_reranking_test.rs index 43432c769..1546d872b 100644 --- a/aimux-providers/tests/voyage_reranking_test.rs +++ b/aimux-providers/tests/voyage_reranking_test.rs @@ -8,7 +8,16 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::reranking_model::{RerankingCallOptions, RerankingDocuments, RerankingModel}; -use aimux_providers::{VoyageConfig, VoyageProvider}; +use aimux_providers::{VoyageProvider, VoyageProviderSettings, create_voyage}; + +fn test_provider(api_key: &str, base_url: impl Into) -> VoyageProvider { + create_voyage(VoyageProviderSettings { + api_key: Some(api_key.to_string().into()), + base_url: Some(base_url.into()), + ..Default::default() + }) + .expect("valid settings") +} // -- helpers ----------------------------------------------------------------- @@ -25,8 +34,7 @@ fn rerank_response_body() -> Value { } fn provider(server: &MockServer) -> VoyageProvider { - let config = VoyageConfig::new("test-api-key").with_base_url(server.uri()); - VoyageProvider::new(config) + test_provider("test-api-key", server.uri()) } fn text_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { @@ -81,7 +89,7 @@ async fn object_docs_should_send_request_with_stringified_json_documents() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); model .do_rerank(&json_docs_opts("rainy day", 2)) @@ -109,7 +117,7 @@ async fn object_docs_should_send_request_with_correct_headers() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); model .do_rerank(&json_docs_opts("rainy day", 2)) @@ -134,7 +142,7 @@ async fn object_docs_should_return_result_with_warnings() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); let result = model .do_rerank(&json_docs_opts("rainy day", 2)) @@ -161,7 +169,7 @@ async fn object_docs_should_return_result_with_correct_ranking() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); let result = model .do_rerank(&json_docs_opts("rainy day", 2)) @@ -181,7 +189,7 @@ async fn object_docs_should_return_result_with_correct_response_body() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); let result = model .do_rerank(&json_docs_opts("rainy day", 2)) @@ -205,7 +213,7 @@ async fn text_docs_should_send_request_with_text_documents() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); model .do_rerank(&text_docs_opts("rainy day", 2)) @@ -227,7 +235,7 @@ async fn text_docs_should_send_request_with_correct_headers() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); model .do_rerank(&text_docs_opts("rainy day", 2)) @@ -252,7 +260,7 @@ async fn text_docs_should_return_result_without_warnings() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); let result = model .do_rerank(&text_docs_opts("rainy day", 2)) @@ -269,7 +277,7 @@ async fn text_docs_should_return_result_with_correct_ranking() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); let result = model .do_rerank(&text_docs_opts("rainy day", 2)) @@ -287,7 +295,7 @@ async fn text_docs_should_return_result_with_correct_response_body() { mount_rerank_mock(&server).await; let provider = provider(&server); - let model = provider.reranking_model("rerank-2.5"); + let model = provider.reranking("rerank-2.5"); let result = model .do_rerank(&text_docs_opts("rainy day", 2)) diff --git a/aimux-providers/tests/you_com_test.rs b/aimux-providers/tests/you_com_test.rs index 9c7b11fed..9d0accd97 100644 --- a/aimux-providers/tests/you_com_test.rs +++ b/aimux-providers/tests/you_com_test.rs @@ -10,10 +10,10 @@ use serde_json::{Value, json}; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::Provider; use aimux_core::error::AiMuxError; use aimux_core::search_model::{SearchCallOptions, SearchModel}; -use aimux_providers::{YouComConfig, YouComProvider}; +use aimux_provider_utils::Resolvable; +use aimux_providers::{YouComProvider, YouComProviderSettings, create_you_com}; const API_KEY: &str = "test-ydc-key"; @@ -44,8 +44,12 @@ fn search_response_body() -> Value { } fn provider(server: &MockServer) -> YouComProvider { - let config = YouComConfig::new(API_KEY).with_base_url(server.uri()); - YouComProvider::new(config) + let config = YouComProviderSettings { + api_key: Some(Resolvable::Value(API_KEY.to_string())), + base_url: Some(server.uri().to_string()), + ..Default::default() + }; + create_you_com(config).unwrap() } fn opts(query: &str, max_results: Option) -> SearchCallOptions { @@ -177,25 +181,3 @@ async fn status_401_maps_to_auth_error() { } // -- Provider trait ---------------------------------------------------------- - -#[tokio::test] -async fn provider_name_is_you_com() { - let config = YouComConfig::new(API_KEY); - let provider = YouComProvider::new(config); - assert_eq!(provider.name(), "you_com"); -} - -#[test] -fn language_model_returns_unsupported_error() { - let config = YouComConfig::new(API_KEY); - let provider = YouComProvider::new(config); - match provider.language_model("youcom-search") { - Err(AiMuxError::UnsupportedFunctionality(msg)) => { - assert!( - msg.contains("provider 'you_com' does not provide language models"), - "unexpected message: {msg}" - ); - } - _ => panic!("expected Unsupported error, got success or another error variant"), - } -} diff --git a/bindings/flutter/lib/aimux.dart b/bindings/flutter/lib/aimux.dart index bfe4d06ed..4a5c8bc8c 100644 --- a/bindings/flutter/lib/aimux.dart +++ b/bindings/flutter/lib/aimux.dart @@ -503,7 +503,9 @@ class Model implements Finalizable { /// /// [apiKey] may be null to read the provider's env var from the registry /// entry. [configJson] is an optional JSON object of ProviderOptions - /// (`{"base_url": "...", "headers": {...}, "max_retries": 0, "body_overrides": {...}}`). + /// (`{"base_url": "...", "headers": {...}, "organization": "...", + /// "project": "...", "params": {...}}`). `max_retries` (a per-call option) + /// and `body_overrides` (removed) are rejected as [InvalidArgumentError]. factory Model.provider(String name, String modelId, {String? apiKey, String? configJson}) { checkJson(configJson, 'config_json', emptyIsDefault: true); @@ -746,20 +748,21 @@ class ProviderConfig { /// OpenAI project ID (`OpenAI-Project` header). final String? project; - /// Retry count override; `0` disables retries. - final int? maxRetries; - - /// Request-body overrides (deep-merged into every request). Any - /// JSON-serializable value (typically a `Map`). - final Object? bodyOverrides; + /// Values of a preset's template parameters (`account_id`, `region`, ...); + /// only the ones the preset declares are accepted. + /// + /// Retry is a per-call setting ([GenerateTextOptions.maxRetries]) and + /// request-body overrides no longer exist: a `max_retries` or + /// `body_overrides` key in the config is rejected as [InvalidArgumentError], + /// so this class has neither. + final Map? params; const ProviderConfig({ this.baseUrl, this.headers, this.organization, this.project, - this.maxRetries, - this.bodyOverrides, + this.params, }); /// Serialize to the `ProviderOptions` JSON wire format (snake_case keys, @@ -771,8 +774,7 @@ class ProviderConfig { if (headers != null) map['headers'] = headers; if (organization != null) map['organization'] = organization; if (project != null) map['project'] = project; - if (maxRetries != null) map['max_retries'] = maxRetries; - if (bodyOverrides != null) map['body_overrides'] = bodyOverrides; + if (params != null) map['params'] = params; if (map.isEmpty) return null; return encodeJson(map, 'config_json'); } @@ -1018,9 +1020,9 @@ void registerProviders(String configJSON) { (p) => expectAimuxError(_ffi.registerProviders(p), 'register_providers')); } -/// Set the global proxy configuration (M6, RFC-0016). Must be called before -/// the first `generateText` / `streamText` call; a no-op if the shared HTTP -/// client is already initialised. +/// Replace the default transport with one routed by this proxy configuration +/// (M6, RFC-0016). Requests started afterwards use it; calling it again +/// replaces the previous configuration. /// /// `configJSON` shape: `{ "http_url", "https_url", "all_url", "no_proxy" }` /// (all fields optional). Wrong shape throws [InvalidArgumentError]. diff --git a/bindings/flutter/lib/errors.dart b/bindings/flutter/lib/errors.dart index f9f6cab81..5e3407630 100644 --- a/bindings/flutter/lib/errors.dart +++ b/bindings/flutter/lib/errors.dart @@ -4,7 +4,7 @@ // Transport (aimux-error.h): every fallible C call returns // `aimux_error_t *` — NULL on success (the result is in the trailing // out-param), non-NULL on failure (out-param at its sentinel: 0 / NULL). The -// unified code selects AiMuxError (1..18), RecordingError (100..105), or a +// unified code selects AiMuxError (1..20), RecordingError (100..105), or a // C ABI failure (200..206). The last range maps to // StateError('aimux ffi: …'); Dart does not expose seven additional classes. // Every field is copied before the error is released with `aimux_error_free` @@ -22,7 +22,8 @@ import 'package:ffi/ffi.dart'; // ───────────────────────────────────────────────────────────────────────────── /// Machine-readable codes. Values match C `aimux_error_code_t` / Go `Code`. -/// 17 variant codes: 1–18 (1 is the catch-all; 4 is retired). A code outside +/// 19 variant codes: 1–20 (1 is the catch-all; 4 is retired; 19 / 20 are a +/// missing API key / provider setting). A code outside /// that set is a header/library mismatch and fails with [StateError], not an /// error type. Every HTTP-shaped failure /// arrives as [apiCall], classified @@ -46,6 +47,8 @@ abstract final class AimuxErrorCode { static const int invalidToolInput = 16; static const int toolCallRepair = 17; static const int noOutputGenerated = 18; + static const int loadApiKey = 19; + static const int loadSetting = 20; static const int other = 1; static const int retry = 14; @@ -65,6 +68,8 @@ abstract final class AimuxErrorCode { noSuchTool: 'NoSuchTool', invalidToolInput: 'InvalidToolInput', toolCallRepair: 'ToolCallRepair', + loadApiKey: 'LoadApiKey', + loadSetting: 'LoadSetting', noOutputGenerated: 'NoOutputGenerated', other: 'Other', retry: 'Retry', @@ -80,7 +85,7 @@ abstract final class AimuxErrorCode { /// Decode the `aimux_error_t *` [e] returned by a call that can fail in /// `AiMuxError` (`[AiMuxError]` in aimux-ffi.h). NULL → returns (success). -/// Codes 1..18 become [AimuxException]; 200..206 become [StateError]. +/// Codes 1..20 become [AimuxException]; 200..206 become [StateError]. /// The returned error is always freed. void expectAimuxError(Pointer e, String context) { if (e == nullptr) return; @@ -515,7 +520,7 @@ class AimuxException implements Exception { /// Build the typed subclass from a returned `const aimux_error_t *` [error] /// via the `aimux_error_*` getters (payload getters only under the owning /// code; getter strings freed here). The caller ([expectAimuxError]) - /// frees it. A code outside 1..18 is a + /// frees it. A code outside 1..20 is a /// contract violation and throws [StateError]. factory AimuxException._decode(Pointer error, String context) { final code = _errorCode(error); @@ -570,6 +575,15 @@ class AimuxException implements Exception { return ToolCallRepairError(message, retryable: retryable, originalError: original == null ? null : jsonDecode(original)); + case AimuxErrorCode.loadApiKey: + // The consulted environment variable rides the provider-code getter. + return LoadAPIKeyError(message, + retryable: retryable, + envVar: _errStr(_errorProviderCode, error) ?? ''); + case AimuxErrorCode.loadSetting: + return LoadSettingError(message, + retryable: retryable, + envVar: _errStr(_errorProviderCode, error) ?? ''); default: try { return AimuxException.fromCode(code, message, retryable: retryable); @@ -631,6 +645,10 @@ class AimuxException implements Exception { return ToolCallRepairError(message, status: status, retryMs: retryMs, retryable: retryable); case AimuxErrorCode.noOutputGenerated: return NoOutputGeneratedError(message, status: status, retryMs: retryMs, retryable: retryable); + case AimuxErrorCode.loadApiKey: + return LoadAPIKeyError(message, status: status, retryMs: retryMs, retryable: retryable); + case AimuxErrorCode.loadSetting: + return LoadSettingError(message, status: status, retryMs: retryMs, retryable: retryable); case AimuxErrorCode.other: return OtherError(message, status: status, retryMs: retryMs, retryable: retryable); default: @@ -723,6 +741,28 @@ class InvalidPromptError extends AimuxException { : super(code: AimuxErrorCode.invalidPrompt); } +/// No API key was passed and the fallback environment variable is unset (AI SDK +/// `LoadAPIKeyError`). No request was made. +class LoadAPIKeyError extends AimuxException { + /// The environment variable that was consulted, e.g. `OPENAI_API_KEY`. + final String envVar; + + LoadAPIKeyError(super.message, + {super.status, super.retryMs, super.retryable, this.envVar = ''}) + : super(code: AimuxErrorCode.loadApiKey); +} + +/// A required provider setting was not passed and its fallback environment +/// variable is unset (AI SDK `LoadSettingError`). +class LoadSettingError extends AimuxException { + /// The environment variable that was consulted. + final String envVar; + + LoadSettingError(super.message, + {super.status, super.retryMs, super.retryable, this.envVar = ''}) + : super(code: AimuxErrorCode.loadSetting); +} + /// Access token expired. class TokenExpiredError extends AimuxException { TokenExpiredError(super.message, {super.status = 401, super.retryMs, super.retryable}) diff --git a/bindings/flutter/lib/types.dart b/bindings/flutter/lib/types.dart index 87657a1ae..3daf20341 100644 --- a/bindings/flutter/lib/types.dart +++ b/bindings/flutter/lib/types.dart @@ -474,13 +474,6 @@ class ResponseMetadata { Map toJson() => _$ResponseMetadataToJson(this); } -class GenerateRequest { - final dynamic body; - GenerateRequest({this.body}); - factory GenerateRequest.fromJson(Map json) => GenerateRequest(body: json['body']); - Map toJson() => {'body': body}; -} - class GenerateResponse { final String? id; final String? timestamp; @@ -488,8 +481,15 @@ class GenerateResponse { final Map? headers; final dynamic body; GenerateResponse({this.id, this.timestamp, this.modelId, this.headers, this.body}); - factory GenerateResponse.fromJson(Map json) => GenerateResponse(id: json['id'] as String?, timestamp: json['timestamp'] as String?, modelId: json['model_id'] as String?, headers: (json['headers'] as Map?)?.cast(), body: json['body']); - Map toJson() => {if (id != null) 'id': id, 'timestamp': timestamp, 'model_id': modelId, 'headers': headers, 'body': body}; + factory GenerateResponse.fromJson(Map json) => GenerateResponse(id: json['id'] as String?, timestamp: json['timestamp'] as String?, modelId: json['model_id'] as String?, headers: (json['headers'] as Map?)?.map((key, value) => MapEntry(key, value as String)), body: json['body']); + Map toJson() => {'id': id, 'timestamp': timestamp, 'model_id': modelId, if (headers != null) 'headers': headers, if (body != null) 'body': body}; +} + +class GenerateRequest { + final dynamic body; + GenerateRequest({this.body}); + factory GenerateRequest.fromJson(Map json) => GenerateRequest(body: json['body']); + Map toJson() => {if (body != null) 'body': body}; } /// A content item in a `GenerateResult`. Mirrors `GenerateContent.ts` @@ -853,7 +853,7 @@ class GenerateResult { response: json['response'] == null ? null : GenerateResponse.fromJson(json['response'] as Map), request: json['request'] == null ? null : GenerateRequest.fromJson(json['request'] as Map), ); - Map toJson() => {'content': content.map((c) => c.toJson()).toList(), 'finish_reason': finishReason.toJson(), 'usage': usage.toJson(), 'warnings': warnings, 'response': response?.toJson(), if (providerMetadata != null) 'provider_metadata': providerMetadata, 'request': request?.toJson()}; + Map toJson() => {'content': content.map((c) => c.toJson()).toList(), 'finish_reason': finishReason.toJson(), 'usage': usage.toJson(), 'warnings': warnings, 'response': response?.toJson(), if (providerMetadata != null) 'provider_metadata': providerMetadata, if (request != null) 'request': request?.toJson()}; } // ───────────────────────────────────────────────────────────────────────────── @@ -1036,8 +1036,6 @@ class GenerateTextOptions { final Map? providerOptions; final ReasoningEffort? reasoning; final String? instructions; - @JsonKey(name: 'body_overrides') - final Map? bodyOverrides; @JsonKey(name: 'max_retries') final int? maxRetries; final TimeoutConfiguration? timeout; @@ -1080,7 +1078,6 @@ class GenerateTextOptions { this.providerOptions, this.reasoning, this.instructions, - this.bodyOverrides, this.maxRetries, this.timeout, this.includeRawChunks, @@ -1111,7 +1108,6 @@ class GenerateTextOptions { ? ReasoningEffort.fromJson(json['reasoning'] as String) : null, instructions: json['instructions'] as String?, - bodyOverrides: json['body_overrides'] as Map?, maxRetries: json['max_retries'] as int?, timeout: json['timeout'] != null ? TimeoutConfiguration.fromJson(json['timeout'] as Map) @@ -1136,7 +1132,6 @@ class GenerateTextOptions { if (providerOptions != null) 'provider_options': providerOptions, if (reasoning != null) 'reasoning': reasoning!.toJson(), if (instructions != null) 'instructions': instructions, - if (bodyOverrides != null) 'body_overrides': bodyOverrides, if (maxRetries != null) 'max_retries': maxRetries, if (timeout != null) 'timeout': timeout!.toJson(), if (includeRawChunks != null) 'include_raw_chunks': includeRawChunks, diff --git a/bindings/flutter/lib/types.g.dart b/bindings/flutter/lib/types.g.dart index 70f61cdcf..30fd4efb9 100644 --- a/bindings/flutter/lib/types.g.dart +++ b/bindings/flutter/lib/types.g.dart @@ -273,7 +273,6 @@ GenerateTextOptions _$GenerateTextOptionsFromJson(Map json) => reasoning: $enumDecodeNullable(_$ReasoningEffortEnumMap, json['reasoning']), instructions: json['instructions'] as String?, - bodyOverrides: json['body_overrides'] as Map?, maxRetries: (json['max_retries'] as num?)?.toInt(), timeout: json['timeout'] == null ? null @@ -299,7 +298,6 @@ Map _$GenerateTextOptionsToJson( 'provider_options': instance.providerOptions, 'reasoning': instance.reasoning, 'instructions': instance.instructions, - 'body_overrides': instance.bodyOverrides, 'max_retries': instance.maxRetries, 'timeout': instance.timeout, }; diff --git a/bindings/flutter/test/errors_test.dart b/bindings/flutter/test/errors_test.dart index 49f498be1..be832d69c 100644 --- a/bindings/flutter/test/errors_test.dart +++ b/bindings/flutter/test/errors_test.dart @@ -75,6 +75,9 @@ void main() { AimuxErrorCode.noSuchTool: NoSuchToolError, AimuxErrorCode.invalidToolInput: InvalidToolInputError, AimuxErrorCode.toolCallRepair: ToolCallRepairError, + AimuxErrorCode.noOutputGenerated: NoOutputGeneratedError, + AimuxErrorCode.loadApiKey: LoadAPIKeyError, + AimuxErrorCode.loadSetting: LoadSettingError, AimuxErrorCode.other: OtherError, AimuxErrorCode.retry: RetryError, }; @@ -88,12 +91,12 @@ void main() { test('unknown code is rejected with StateError', () { // A code outside the published table is an ABI mismatch, not an error // kind. 1 is AIMUX_E_OTHER now because Other inherited the old UNKNOWN - // slot, so it resolves; 4 is retired, 18 is NoOutputGenerated and 19 is + // slot, so it resolves; 4 is retired, 18 is NoOutputGenerated and 21 is // the first unassigned value. expect(() => AimuxException.fromCode(999, 'future'), throwsStateError); expect(() => AimuxException.fromCode(4, 'retired'), throwsStateError); expect(AimuxException.fromCode(18, 'no output'), isA()); - expect(() => AimuxException.fromCode(19, 'unassigned'), throwsStateError); + expect(() => AimuxException.fromCode(21, 'unassigned'), throwsStateError); }); test('bare retry code synthesizes a single-attempt RetryError', () { @@ -183,7 +186,7 @@ void main() { expect(AimuxErrorCode.retry, 14); expect(AimuxErrorCode.name(AimuxErrorCode.retry), 'Retry'); // The AIMUX_E_UNKNOWN catch-all is gone and Other took its slot; the - // engine codes are 1–17 (4 retired). + // engine codes are 1–20 (4 retired). expect(AimuxErrorCode.other, 1); expect(AimuxErrorCode.aborted, 13); expect(AimuxErrorCode.noSuchTool, 15); @@ -194,6 +197,14 @@ void main() { expect(AimuxErrorCode.toolCallRepair, 17); expect(AimuxErrorCode.name(AimuxErrorCode.toolCallRepair), 'ToolCallRepair'); + // A missing API key / provider setting have their own codes. + expect(AimuxErrorCode.noOutputGenerated, 18); + expect(AimuxErrorCode.name(AimuxErrorCode.noOutputGenerated), + 'NoOutputGenerated'); + expect(AimuxErrorCode.loadApiKey, 19); + expect(AimuxErrorCode.name(AimuxErrorCode.loadApiKey), 'LoadApiKey'); + expect(AimuxErrorCode.loadSetting, 20); + expect(AimuxErrorCode.name(AimuxErrorCode.loadSetting), 'LoadSetting'); }); }); diff --git a/bindings/flutter/test/structured_e2e_test.dart b/bindings/flutter/test/structured_e2e_test.dart index c680acff8..0ae5887d5 100644 --- a/bindings/flutter/test/structured_e2e_test.dart +++ b/bindings/flutter/test/structured_e2e_test.dart @@ -28,9 +28,10 @@ import 'dart:convert'; import 'dart:io'; import 'dart:isolate'; -import 'package:aimux/aimux.dart'; import 'package:test/test.dart'; +import 'typed_model_test.dart' show mockChatModel; + // ───────────────────────────────────────────────────────────────────────────── // Mock server // ───────────────────────────────────────────────────────────────────────────── @@ -115,7 +116,7 @@ class _GenerateArgs { /// Run a non-streaming `generateText` in a worker isolate. Future> runGenerateText(_GenerateArgs args) { return Isolate.run(() { - final model = Model.openai(args.apiKey, args.modelId, baseUrl: args.baseUrl); + final model = mockChatModel(args.baseUrl, args.apiKey, args.modelId); try { return model.generateText(args.prompt, args.options); } finally { @@ -127,7 +128,7 @@ Future> runGenerateText(_GenerateArgs args) { /// Run a streaming `streamText` in a worker isolate, returning all parts. Future>> runStreamText(_GenerateArgs args) { return Isolate.run(() async { - final model = Model.openai(args.apiKey, args.modelId, baseUrl: args.baseUrl); + final model = mockChatModel(args.baseUrl, args.apiKey, args.modelId); try { // streamText blocks synchronously until the stream completes, then // returns an already-closed stream containing the buffered parts. diff --git a/bindings/flutter/test/tool_call_repair_typed_test.dart b/bindings/flutter/test/tool_call_repair_typed_test.dart index 181cb5c5b..116bb770d 100644 --- a/bindings/flutter/test/tool_call_repair_typed_test.dart +++ b/bindings/flutter/test/tool_call_repair_typed_test.dart @@ -10,13 +10,12 @@ import 'dart:convert'; import 'dart:isolate'; -import 'package:aimux/aimux.dart'; import 'package:aimux/typed_model.dart'; import 'package:aimux/types.dart'; import 'package:test/test.dart'; import 'typed_model_test.dart' - show buildToolCallSse, startMockServer, toolCallOpenAIResponse; + show buildToolCallSse, mockChatModel, startMockServer, toolCallOpenAIResponse; /// Accepts `city` only, so the canned `{"location":"Tokyo"}` is invalid. GenerateTextOptions _strictOptions() => GenerateTextOptions( @@ -48,7 +47,7 @@ void main() { final baseUrl = server.baseUrl; final result = await Isolate.run(() { - final typed = TypedModel(Model.openai('sk-test', 'gpt-4o', baseUrl: baseUrl)); + final typed = TypedModel(mockChatModel(baseUrl, 'sk-test', 'gpt-4o')); try { return jsonDecode(jsonEncode(typed.generateText('weather?', _strictOptions()))) as Map; @@ -75,7 +74,7 @@ void main() { final baseUrl = server.baseUrl; final arguments = await Isolate.run(() { - final typed = TypedModel(Model.openai('sk-test', 'gpt-4o', baseUrl: baseUrl)); + final typed = TypedModel(mockChatModel(baseUrl, 'sk-test', 'gpt-4o')); try { final completion = typed.generateTextAsOpenAI('weather?', _strictOptions()); return completion.choices.first.message.toolCalls!.single.function.arguments; @@ -95,7 +94,7 @@ void main() { final baseUrl = server.baseUrl; final input = await Isolate.run(() async { - final typed = TypedModel(Model.openai('sk-test', 'gpt-4o', baseUrl: baseUrl)); + final typed = TypedModel(mockChatModel(baseUrl, 'sk-test', 'gpt-4o')); try { final parts = await typed.streamText('weather?', _strictOptions()).toList(); final call = parts.whereType().single; diff --git a/bindings/flutter/test/typed_model_test.dart b/bindings/flutter/test/typed_model_test.dart index 3c94acb16..d7773cd0d 100644 --- a/bindings/flutter/test/typed_model_test.dart +++ b/bindings/flutter/test/typed_model_test.dart @@ -37,6 +37,12 @@ import 'package:test/test.dart'; // Mock server (mirrors structured_e2e_test.dart) // ───────────────────────────────────────────────────────────────────────────── +/// A chat-completions model against [baseUrl]: the registry's `deepseek` entry +/// with a `base_url` setting (`Model.openai` speaks the Responses API). +Model mockChatModel(String baseUrl, String apiKey, String modelId) => + Model.provider('deepseek', modelId, + apiKey: apiKey, configJson: jsonEncode({'base_url': baseUrl})); + class MockOpenAIServer { final HttpServer _server; final List> _responses; @@ -130,7 +136,7 @@ class _TypedArgs { /// instances (not plain maps), which cannot be re-parsed by `fromJson`. Future> runTypedGenerateText(_TypedArgs args) { return Isolate.run(() { - final model = Model.openai(args.apiKey, args.modelId, baseUrl: args.baseUrl); + final model = mockChatModel(args.baseUrl, args.apiKey, args.modelId); final typed = TypedModel(model); try { final hasOptions = args.tools != null || args.toolChoice != null; @@ -151,7 +157,7 @@ Future> runTypedGenerateText(_TypedArgs args) { /// Run a typed `generateTextMessages` (multi-turn) in a worker isolate. Future> runTypedGenerateTextMessages(_TypedArgs args) { return Isolate.run(() { - final model = Model.openai(args.apiKey, args.modelId, baseUrl: args.baseUrl); + final model = mockChatModel(args.baseUrl, args.apiKey, args.modelId); final typed = TypedModel(model); try { final messages = args.messages! @@ -168,7 +174,7 @@ Future> runTypedGenerateTextMessages(_TypedArgs args) { /// Run a typed `streamText` in a worker isolate, returning the parsed parts. Future> runTypedStreamText(_TypedArgs args) { return Isolate.run(() async { - final model = Model.openai(args.apiKey, args.modelId, baseUrl: args.baseUrl); + final model = mockChatModel(args.baseUrl, args.apiKey, args.modelId); final typed = TypedModel(model); try { final options = (args.tools != null || args.includeRawChunks != null) diff --git a/bindings/go/aimux.go b/bindings/go/aimux.go index 8d9689921..288919f9e 100644 --- a/bindings/go/aimux.go +++ b/bindings/go/aimux.go @@ -364,10 +364,10 @@ func RegisterProviders(configJSON string) error { return expectAimuxError(C.aimux_register_providers(cJSON)) } -// InitProxy sets the global proxy configuration (M6, RFC-0016). Must be called -// before the first GenerateText / StreamText call; a no-op if the shared HTTP -// client is already initialised. configJSON is required (pass "{}" for -// defaults); shape: +// InitProxy replaces the default transport with one routed by this proxy +// configuration (M6, RFC-0016). Requests started afterwards use it; calling it +// again replaces the previous configuration. configJSON is required (pass "{}" +// for defaults); shape: // { "http_url": "...", "https_url": "...", "all_url": "...", "no_proxy": "..." } // (all fields optional). func InitProxy(configJSON string) error { @@ -663,14 +663,17 @@ func ProviderWithBase(name, apiKey, modelID, baseURL string) (*Model, error) { // ProviderConfig mirrors the Rust ProviderOptions accepted by // aimux_provider_new's config_json (RFC-0017 §3.4). Zero-value fields are -// omitted; MaxRetries is a pointer so 0 (disable retries) is expressible. +// omitted. Retry is a per-call setting (CallOptions.MaxRetries), and request-body +// overrides no longer exist: a config_json carrying max_retries or +// body_overrides is rejected as CodeInvalidArgument, so this struct has neither. type ProviderConfig struct { - BaseURL string `json:"base_url,omitempty"` - Headers map[string]string `json:"headers,omitempty"` - Organization string `json:"organization,omitempty"` - Project string `json:"project,omitempty"` - MaxRetries *uint32 `json:"max_retries,omitempty"` - BodyOverrides map[string]any `json:"body_overrides,omitempty"` + BaseURL string `json:"base_url,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + Organization string `json:"organization,omitempty"` + Project string `json:"project,omitempty"` + // Params are the values of a preset's template parameters (account_id, + // region, ...); only the ones the preset declares are accepted. + Params map[string]string `json:"params,omitempty"` } // ProviderWithConfig is Provider with the full ProviderOptions config. @@ -1149,7 +1152,7 @@ func cstr(p *C.char) string { // ── Error decoding ─────────────────────────────────────────────────────────── // // Every fallible C call returns *C.aimux_error_t: nil = success, non-nil -// = failure. The unified code space distinguishes AiMuxError (1..17), +// = failure. The unified code space distinguishes AiMuxError (1..20), // RecordingError (100..105), and failures detected by the C ABI (200..206). // The latter collapse to a plain error in Go; no public Go error type is added // for those implementation failures. Every helper frees the pointer once. @@ -1174,7 +1177,7 @@ func expectFfiError(e *C.aimux_error_t) error { return ffiError(e) } -// expectAimuxError decodes an [AiMuxError] call: nil → nil; 1..17 → +// expectAimuxError decodes an [AiMuxError] call: nil → nil; 1..20 → // *Error; 200..206 → plain C ABI error. Any other code is an ABI contract // violation. func expectAimuxError(e *C.aimux_error_t) error { @@ -1264,6 +1267,9 @@ func aimuxErrorFromC(e *C.aimux_error_t) *Error { if orig := str(C.aimux_error_original_error(e)); orig != "" { err.OriginalError = json.RawMessage(orig) } + case CodeLoadAPIKey, CodeLoadSetting: + // The consulted environment variable travels on the provider-code getter. + err.EnvVar = str(C.aimux_error_provider_code(e)) } // TokenExpired carries a 401 by contract even if C reports -1; every // other status is the observed one (ApiCall without a status = no HTTP diff --git a/bindings/go/aimux_test.go b/bindings/go/aimux_test.go index 027cae103..540de63be 100644 --- a/bindings/go/aimux_test.go +++ b/bindings/go/aimux_test.go @@ -13,6 +13,7 @@ import ( "fmt" "net/http" "net/http/httptest" + "os" "strings" "sync" "sync/atomic" @@ -205,14 +206,11 @@ func TestStreamTextContextDoesNotLoseRacingCancellation(t *testing.T) { } func TestProviderWithConfigFullOptions(t *testing.T) { - retries := uint32(0) m, err := ProviderWithConfig("groq", "sk-test-fake-key", "llama-3.3-70b", &ProviderConfig{ - BaseURL: "https://example.com/v1", - Headers: map[string]string{"X-Custom": "1"}, - Organization: "org-1", - Project: "proj-1", - MaxRetries: &retries, - BodyOverrides: map[string]any{"temperature": 0.1}, + BaseURL: "https://example.com/v1", + Headers: map[string]string{"X-Custom": "1"}, + Organization: "org-1", + Project: "proj-1", }) if err != nil { t.Fatalf("expected success, got %v", err) @@ -223,6 +221,48 @@ func TestProviderWithConfigFullOptions(t *testing.T) { } } +// config_json is ProviderOptions: params fill a preset's template parameters. +func TestProviderWithConfigParams(t *testing.T) { + m, err := ProviderWithConfig("cloudflare_workers_ai", "sk-test-fake-key", "m", &ProviderConfig{ + Params: map[string]string{"account_id": "acct123"}, + }) + if err != nil { + t.Fatalf("expected success, got %v", err) + } + defer m.Close() + + // A parameter the preset does not declare is an invalid argument. + _, err = ProviderWithConfig("groq", "sk-test-fake-key", "m", &ProviderConfig{ + Params: map[string]string{"account_id": "acct123"}, + }) + var e *Error + if !errors.As(err, &e) || e.Code != CodeInvalidArgument { + t.Fatalf("want CodeInvalidArgument, got %v", err) + } +} + +// Keys are resolved at call time: construction succeeds, then the missing key +// fails with CodeLoadAPIKey and names the variable it consulted. +func TestMissingAPIKeyIsLoadAPIKey(t *testing.T) { + if old, ok := os.LookupEnv("GROQ_API_KEY"); ok { + os.Unsetenv("GROQ_API_KEY") + defer os.Setenv("GROQ_API_KEY", old) + } + m, err := Provider("groq", "", "llama-3.3-70b") + if err != nil { + t.Fatalf("construction should defer key loading: %v", err) + } + defer m.Close() + _, err = m.GenerateText(`"hello"`, `{"max_retries":0}`) + var e *Error + if !errors.As(err, &e) || e.Code != CodeLoadAPIKey { + t.Fatalf("want CodeLoadAPIKey, got %v", err) + } + if e.EnvVar != "GROQ_API_KEY" { + t.Fatalf("EnvVar = %q, want GROQ_API_KEY", e.EnvVar) + } +} + func TestProviderWithConfigNil(t *testing.T) { m, err := ProviderWithConfig("groq", "sk-test-fake-key", "llama-3.3-70b", nil) if err != nil { diff --git a/bindings/go/batch1_e2e_test.go b/bindings/go/batch1_e2e_test.go index 9fe7cdb64..9875a5a14 100644 --- a/bindings/go/batch1_e2e_test.go +++ b/bindings/go/batch1_e2e_test.go @@ -5,7 +5,7 @@ package aimux // M10 Usage.raw preservation // M9 StartStep non-empty warnings // -// Mirrors the Rust provider tests (openai_model_test.rs / groq_test.rs) at +// Mirrors the Rust chat-provider tests at // the Go binding level, driving the real FFI path against a mock SSE server. import ( @@ -23,9 +23,12 @@ func TestE2E_StreamTextEmitsRawWhenEnabled(t *testing.T) { srv := newMockServer() defer srv.Close() srv.SetContentType("text/event-stream") - srv.SetResponse(buildTextDeltaSSE()) + srv.SetResponse(buildChatTextDeltaSSE()) - m := OpenAIWithBase("sk-test-fake-key", "gpt-4o", srv.URL) + m, err := ProviderWithBase("deepseek", "sk-test-fake-key", "deepseek-chat", srv.URL) + if err != nil { + t.Fatalf("ProviderWithBase: %v", err) + } defer m.Close() stream := m.StreamText(`"Say hello"`, mustMarshalOptions(t, &GenerateTextOptions{ @@ -34,6 +37,7 @@ func TestE2E_StreamTextEmitsRawWhenEnabled(t *testing.T) { var raws []json.RawMessage rawBeforeFirstDelta := false + sawDelta := false var textBuilder strings.Builder for part := range stream.Parts() { sp, err := ParseStreamPart(part) @@ -43,10 +47,14 @@ func TestE2E_StreamTextEmitsRawWhenEnabled(t *testing.T) { switch sp.Tag { case "Raw": raws = append(raws, sp.Payload) - if !rawBeforeFirstDelta { + if !sawDelta { rawBeforeFirstDelta = true } case "TextDelta": + sawDelta = true + if !rawBeforeFirstDelta { + t.Error("Raw must precede the first parsed TextDelta") + } var td TextDeltaPayload json.Unmarshal(sp.Payload, &td) textBuilder.WriteString(td.Text) @@ -56,7 +64,7 @@ func TestE2E_StreamTextEmitsRawWhenEnabled(t *testing.T) { t.Fatalf("stream error: %v", err) } - // buildTextDeltaSSE has 3 JSON events (2 content + 1 usage chunk); + // buildChatTextDeltaSSE has 3 JSON events (2 content + 1 usage chunk); // [DONE] emits no Raw. if len(raws) != 3 { t.Fatalf("expected 3 Raw parts, got %d", len(raws)) @@ -79,22 +87,20 @@ func TestE2E_StreamTextEmitsRawWhenEnabled(t *testing.T) { } } -// RFC-0016 M10: the streaming FinishStep usage carries the provider's raw usage object -// (buildTextDeltaSSE usage has prompt_tokens=3). - func TestE2E_StartStepCarriesNonEmptyWarnings(t *testing.T) { srv := newMockServer() defer srv.Close() srv.SetContentType("text/event-stream") - srv.SetResponse(buildTextDeltaSSE()) + srv.SetResponse(buildChatTextDeltaSSE()) - m := OpenAIWithBase("sk-test-fake-key", "gpt-4o", srv.URL) + m, err := ProviderWithBase("deepseek", "sk-test-fake-key", "deepseek-chat", srv.URL) + if err != nil { + t.Fatalf("ProviderWithBase: %v", err) + } defer m.Close() - // top_k on a provider that supports it produces no warning on the openai - // full profile; the Go assertion here is that a non-empty warnings array - // in StartStep decodes without breaking the stream (the warning itself - // is produced by the groq profile — covered in groq_test.rs). + // DeepSeek does not support top_k, so it must surface the warning through + // the Go FFI stream while continuing to parse text normally. stream := m.StreamText(`"Say hello"`, mustMarshalOptions(t, &GenerateTextOptions{ TopK: f64Ptr(0.5), })) @@ -113,9 +119,9 @@ func TestE2E_StartStepCarriesNonEmptyWarnings(t *testing.T) { if err := json.Unmarshal(sp.Payload, &start); err != nil { t.Fatalf("failed to decode StartStep payload: %v", err) } - // warnings may be empty on the full profile; the point is that a - // non-empty array would decode here without error. - t.Logf("StartStep warnings: %d", len(start.Warnings)) + if len(start.Warnings) == 0 { + t.Error("expected an unsupported topK warning") + } } } if err := stream.Err(); err != nil { diff --git a/bindings/go/e2e_test.go b/bindings/go/e2e_test.go index a5e44041e..1e0f0d355 100644 --- a/bindings/go/e2e_test.go +++ b/bindings/go/e2e_test.go @@ -1,6 +1,6 @@ // Structured end-to-end tests for the Go binding. // -// These spin up a local mock HTTP server speaking the OpenAI chat-completions +// These spin up a local mock HTTP server speaking the OpenAI Responses API // wire format, point an OpenAI Model at it via the base_url constructor, // and assert that the binding correctly: // 1. parses text from a provider response, @@ -39,6 +39,7 @@ type mockProviderServer struct { mu sync.Mutex lastRequestBody string lastMethod string + lastPath string status int contentType string @@ -64,6 +65,7 @@ func (s *mockProviderServer) handle(w http.ResponseWriter, r *http.Request) { s.mu.Lock() s.lastRequestBody = string(body) s.lastMethod = r.Method + s.lastPath = r.URL.Path status := s.status ct := s.contentType var resp string @@ -103,6 +105,12 @@ func (s *mockProviderServer) LastRequestBody() string { return s.lastRequestBody } +func (s *mockProviderServer) LastRequestPath() string { + s.mu.Lock() + defer s.mu.Unlock() + return s.lastPath +} + // SetResponse sets the response body and content type (thread-safe). func (s *mockProviderServer) SetResponse(body string) { s.mu.Lock() @@ -128,8 +136,9 @@ func newBackpressureSSEServer(chunks int) (*httptest.Server, <-chan struct{}, <- w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) flusher, _ := w.(http.Flusher) + _, _ = fmt.Fprint(w, "data: "+responsesTextStart+"\n\n") for i := 0; i < chunks; i++ { - chunk := `{"id":"1","model":"gpt-4o","choices":[{"delta":{"content":"x"}}]}` + chunk := `{"type":"response.output_text.delta","item_id":"msg-test","output_index":0,"content_index":0,"delta":"x"}` _, _ = fmt.Fprintf(w, "data: %s\n\n", chunk) if flusher != nil { flusher.Flush() @@ -159,32 +168,28 @@ func waitForChannelFull[T any](t *testing.T, channel chan T) { // plainOpenAIResponse is a plain OpenAI text response (no tool calls). const plainOpenAIResponse = `{ - "id": "chatcmpl-test", + "id": "resp-test", "model": "gpt-4o", - "choices": [{ - "message": {"role": "assistant", "content": "Rust is a systems programming language."}, - "finish_reason": "stop" + "created_at": 1, + "status": "completed", + "output": [{ + "type": "message", "id": "msg-test", "role": "assistant", "status": "completed", + "content": [{"type":"output_text","text":"Rust is a systems programming language.","annotations":[]}] }], - "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18} + "usage": {"input_tokens": 10, "output_tokens": 8, "total_tokens": 18} }` // toolCallOpenAIResponse is an OpenAI response that requests a tool call. const toolCallOpenAIResponse = `{ - "id": "chatcmpl-tc", + "id": "resp-tc", "model": "gpt-4o", - "choices": [{ - "message": { - "role": "assistant", - "content": null, - "tool_calls": [{ - "id": "call_abc", - "type": "function", - "function": {"name": "get_weather", "arguments": "{\"location\":\"Tokyo\"}"} - }] - }, - "finish_reason": "tool_calls" + "created_at": 1, + "status": "completed", + "output": [{ + "type": "function_call", "id": "fc-abc", "call_id": "call_abc", + "name": "get_weather", "arguments": "{\"location\":\"Tokyo\"}", "status": "completed" }], - "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30} + "usage": {"input_tokens": 20, "output_tokens": 10, "total_tokens": 30} }` // weatherToolsOpts builds a JSON options string with a get_weather tool. @@ -236,6 +241,9 @@ func TestE2E_GenerateTextParsesText(t *testing.T) { if len(parsed.ToolCalls) != 0 { t.Errorf("expected no tool calls, got %d", len(parsed.ToolCalls)) } + if srv.LastRequestPath() != "/responses" { + t.Errorf("expected /responses endpoint, got %q", srv.LastRequestPath()) + } } func TestE2E_GenerateTextParsesToolCalls(t *testing.T) { @@ -319,7 +327,7 @@ func TestE2E_MultiRoleMessagesReachProvider(t *testing.T) { if reqBody["model"] != "gpt-4o" { t.Errorf("expected model gpt-4o, got %v", reqBody["model"]) } - messages := reqBody["messages"].([]any) + messages := reqBody["input"].([]any) if len(messages) != 2 { t.Fatalf("expected 2 messages, got %d", len(messages)) } @@ -364,7 +372,7 @@ func TestE2E_StreamTextParsesToolCallStreamParts(t *testing.T) { defer srv.Close() srv.SetContentType("text/event-stream") - // SSE stream: tool_calls delta (name) -> arguments delta -> finish_reason. + // Responses SSE: output item -> arguments delta -> completed tool -> finish. sseResponse := buildToolCallSSE() srv.SetResponse(sseResponse) @@ -385,17 +393,36 @@ func TestE2E_StreamTextParsesToolCallStreamParts(t *testing.T) { } // The stream opens with Start and StartStep parts. - // We just verify we got JSON objects with known tags. + var gotToolCall, gotFinish bool for _, p := range parts { sp, err := ParseStreamPart(p) if err != nil { t.Errorf("failed to parse stream part: %v", err) continue } - // Tags we expect: Start, StartStep, TextDelta, ToolCall, Finish, etc. if sp.Tag == "" { t.Error("expected non-empty stream part tag") } + switch sp.Tag { + case "ToolCall": + gotToolCall = true + var call ToolCall + if err := json.Unmarshal(sp.Payload, &call); err != nil { + t.Fatalf("decode tool call: %v", err) + } + if call.ToolName != "get_weather" || call.ToolCallID != "call_xyz" { + t.Errorf("unexpected tool call: %#v", call) + } + var input map[string]any + if err := json.Unmarshal(call.Input, &input); err != nil || input["location"] != "Tokyo" { + t.Errorf("unexpected tool arguments: %q (%v)", call.Input, err) + } + case "Finish": + gotFinish = true + } + } + if !gotToolCall || !gotFinish { + t.Errorf("expected completed tool call and Finish; gotToolCall=%t gotFinish=%t", gotToolCall, gotFinish) } } @@ -539,7 +566,29 @@ func TestStreamCancelUnblocksFullPartsChannel(t *testing.T) { // ── SSE builders ─────────────────────────────────────────────────────────────── +const responsesTextStart = `{"type":"response.output_item.added","output_index":0,"item":{"type":"message","id":"msg-test","role":"assistant","content":[],"status":"in_progress"}}` + +func responsesSSE(events ...string) string { + var sb strings.Builder + for _, event := range events { + sb.WriteString("data: " + event + "\n\n") + } + return sb.String() +} + func buildTextDeltaSSE() string { + return responsesSSE( + `{"type":"response.created","response":{"id":"resp-test","model":"gpt-4o","created_at":1,"status":"in_progress","output":[]}}`, + responsesTextStart, + `{"type":"response.output_text.delta","item_id":"msg-test","output_index":0,"content_index":0,"delta":"Hello"}`, + `{"type":"response.output_text.delta","item_id":"msg-test","output_index":0,"content_index":0,"delta":" world"}`, + `{"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg-test","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Hello world","annotations":[]}]}}`, + `{"type":"response.completed","response":{"id":"resp-test","model":"gpt-4o","created_at":1,"status":"completed","output":[],"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}}`, + ) +} + +// Chat providers emit raw chunks; the Responses provider does not support them. +func buildChatTextDeltaSSE() string { chunk1 := `{"id":"1","model":"gpt-4o","choices":[{"delta":{"role":"assistant","content":"Hello"}}]}` chunk2 := `{"id":"1","model":"gpt-4o","choices":[{"delta":{"content":" world"}}]}` chunk3 := `{"id":"1","model":"gpt-4o","choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}` @@ -552,15 +601,13 @@ func buildTextDeltaSSE() string { } func buildToolCallSSE() string { - chunk1 := `{"id":"1","model":"gpt-4o","choices":[{"delta":{"role":"assistant","tool_calls":[{"index":0,"id":"call_xyz","type":"function","function":{"name":"get_weather","arguments":""}}]}}]}` - chunk2 := `{"id":"1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"location\":\"Tokyo\"}"}}]}}]}` - chunk3 := `{"id":"1","model":"gpt-4o","choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}` - var sb strings.Builder - sb.WriteString("data: " + chunk1 + "\n\n") - sb.WriteString("data: " + chunk2 + "\n\n") - sb.WriteString("data: " + chunk3 + "\n\n") - sb.WriteString("data: [DONE]\n\n") - return sb.String() + return responsesSSE( + `{"type":"response.created","response":{"id":"resp-tc","model":"gpt-4o","created_at":1,"status":"in_progress","output":[]}}`, + `{"type":"response.output_item.added","output_index":0,"item":{"type":"function_call","id":"fc-xyz","call_id":"call_xyz","name":"get_weather","arguments":"","status":"in_progress"}}`, + `{"type":"response.function_call_arguments.delta","output_index":0,"item_id":"fc-xyz","delta":"{\"location\":\"Tokyo\"}"}`, + `{"type":"response.output_item.done","output_index":0,"item":{"type":"function_call","id":"fc-xyz","call_id":"call_xyz","name":"get_weather","arguments":"{\"location\":\"Tokyo\"}","status":"completed"}}`, + `{"type":"response.completed","response":{"id":"resp-tc","model":"gpt-4o","created_at":1,"status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7}}}`, + ) } // ── Tool-call full round-trip ───────────────────────────────────────────────── @@ -608,8 +655,10 @@ func TestE2E_ToolCallFullRoundTrip(t *testing.T) { } conversation := []ModelMessage{ NewTextMessage(RoleUser, "What is the weather in Tokyo?"), - NewTextMessage(RoleAssistant, ""), // assistant's prior turn - toolResultMsg, // tool result fed back + {Role: RoleAssistant, Content: []map[string]any{ + {"type": "tool_call", "tool_call_id": call.ToolCallID, "tool_name": call.ToolName, "input": map[string]any{"location": "Tokyo"}}, + }}, + toolResultMsg, } prompt, _ := MarshalMessages(conversation) @@ -624,31 +673,42 @@ func TestE2E_ToolCallFullRoundTrip(t *testing.T) { } // Verify the second request carried the tool-result message correctly. - // The engine converts multi-part tool_result content into OpenAI's wire - // format: {"role":"tool","tool_call_id":"...","content":""}. + // Responses preserves the assistant's function_call and sends the tool's + // JSON result as a function_call_output with the same call_id. reqBody := srv.LastRequestBody() var body struct { - Messages []struct { - Role string `json:"role"` - Content json.RawMessage `json:"content"` - ToolCallID string `json:"tool_call_id,omitempty"` - } `json:"messages"` + Input []struct { + Role string `json:"role"` + Type string `json:"type"` + Name string `json:"name"` + Arguments string `json:"arguments"` + Output string `json:"output"` + ToolCallID string `json:"call_id"` + } `json:"input"` + } + if err := json.Unmarshal([]byte(reqBody), &body); err != nil { + t.Fatalf("decode request: %v", err) + } + if len(body.Input) != 3 { + t.Fatalf("expected 3 input items, got %d", len(body.Input)) + } + second := body.Input[1] + if second.Type != "function_call" || second.ToolCallID != call.ToolCallID || second.Name != call.ToolName { + t.Errorf("assistant tool call did not round-trip: %#v", second) } - json.Unmarshal([]byte(reqBody), &body) - if len(body.Messages) != 3 { - t.Fatalf("expected 3 messages, got %d", len(body.Messages)) + var arguments map[string]any + if err := json.Unmarshal([]byte(second.Arguments), &arguments); err != nil || arguments["location"] != "Tokyo" { + t.Errorf("assistant arguments did not round-trip: %q (%v)", second.Arguments, err) } - // The third message (tool result) in OpenAI format. - third := body.Messages[2] - if third.Role != "tool" { - t.Errorf("expected third role 'tool', got %q", third.Role) + third := body.Input[2] + if third.Type != "function_call_output" { + t.Errorf("expected function_call_output, got %q", third.Type) } if third.ToolCallID != call.ToolCallID { t.Errorf("expected tool_call_id %q, got %q", call.ToolCallID, third.ToolCallID) } - // The content should contain the result data (stringified by the engine). - contentStr := string(third.Content) - if !strings.Contains(contentStr, "temperature") { - t.Errorf("expected content to contain result data, got %s", contentStr) + var output map[string]any + if err := json.Unmarshal([]byte(third.Output), &output); err != nil || output["temperature"] != "20C" { + t.Errorf("tool output did not round-trip: %q (%v)", third.Output, err) } } diff --git a/bindings/go/error.go b/bindings/go/error.go index a92904e03..09eaa4ff4 100644 --- a/bindings/go/error.go +++ b/bindings/go/error.go @@ -10,8 +10,9 @@ import ( ) // Code is the machine-readable Aimux error code. Values match -// aimux-ffi aimux_error_code_t (1..3, 5..18; 4 is retired). Retry is 14; -// tool-call errors occupy 15..17. A code outside that set is a header/library +// aimux-ffi aimux_error_code_t (1..3, 5..20; 4 is retired). Retry is 14; +// tool-call errors occupy 15..17; a missing API key / provider setting are +// 19 / 20; NoOutputGenerated is 18. A code outside that set is a header/library // mismatch and expectAimuxError // panics rather than inventing an "unknown" variant. Recording failures are // a different type: see RecordingError. @@ -42,6 +43,12 @@ const ( CodeInvalidToolInput Code = 16 CodeToolCallRepair Code = 17 CodeNoOutputGenerated Code = 18 + // CodeLoadAPIKey: no API key was passed and the fallback environment + // variable is unset (the AI SDK's LoadAPIKeyError); no request was made. + CodeLoadAPIKey Code = 19 + // CodeLoadSetting: a required provider setting was not passed and its + // fallback environment variable is unset (LoadSettingError). + CodeLoadSetting Code = 20 ) // String returns the core error_type name (e.g. "ApiCall", "TokenExpired"). @@ -76,10 +83,14 @@ func (c Code) String() string { return "NoSuchTool" case CodeInvalidToolInput: return "InvalidToolInput" - case CodeToolCallRepair: - return "ToolCallRepair" case CodeNoOutputGenerated: return "NoOutputGenerated" + case CodeToolCallRepair: + return "ToolCallRepair" + case CodeLoadAPIKey: + return "LoadApiKey" + case CodeLoadSetting: + return "LoadSetting" case CodeOther: return "Other" case CodeRetry: @@ -89,10 +100,10 @@ func (c Code) String() string { } } -// codeFromC maps a C aimux_error_code_t (1..3, 5..18); false for any other +// codeFromC maps a C aimux_error_code_t (1..3, 5..20); false for any other // value, including retired code 4. func codeFromC(code int) (Code, bool) { - if code >= int(CodeOther) && code <= int(CodeNoOutputGenerated) && code != 4 { + if code >= int(CodeOther) && code <= int(CodeLoadSetting) && code != 4 { return Code(code), true } return 0, false @@ -128,6 +139,8 @@ func codeFromC(code int) (Code, bool) { // with CodeInvalidToolInput) // - ToolInput: CodeInvalidToolInput payload // - OriginalError: CodeToolCallRepair payload +// - EnvVar: CodeLoadAPIKey / CodeLoadSetting payload — the environment +// variable that was consulted type Error struct { Code Code Message string @@ -161,12 +174,13 @@ type Error struct { // Reason says why retrying stopped; Errors preserves each attempt as a // concrete *Error (oldest first; an attempt can itself be any Code, // with the full per-code payload). CodeRetry only. - Reason RetryErrorReason - Errors []*Error - ToolName string // CodeNoSuchTool / CodeInvalidToolInput: tool name - AvailableTools []string - ToolInput string // CodeInvalidToolInput: raw argument text - OriginalError json.RawMessage // CodeToolCallRepair: unrepaired error + Reason RetryErrorReason + Errors []*Error + ToolName string // CodeNoSuchTool / CodeInvalidToolInput: tool name + AvailableTools []string + ToolInput string // CodeInvalidToolInput: raw argument text + OriginalError json.RawMessage // CodeToolCallRepair: unrepaired error + EnvVar string // CodeLoadAPIKey / CodeLoadSetting: the environment variable consulted, e.g. "OPENAI_API_KEY" } // RetryErrorReason explains why operation retry stopped. Values are the diff --git a/bindings/go/error_test.go b/bindings/go/error_test.go index 1cfda5487..4fb904348 100644 --- a/bindings/go/error_test.go +++ b/bindings/go/error_test.go @@ -406,12 +406,6 @@ func TestMarshalledOptionsRejectRawPassThrough(t *testing.T) { _, e := requiredOptsJSON("ImageModel.Generate", &ImageCallOptions{ProviderOptions: loneSurrogate}) return e }, "opts: invalid JSON"}, - {"provider config lone surrogate", func() error { - _, e := ProviderWithConfig("groq", "k", "m", &ProviderConfig{ - BodyOverrides: map[string]any{"x": json.RawMessage(`"\udc00"`)}, - }) - return e - }, "config_json: invalid JSON: lone low surrogate \\uDC00"}, {"transcription session opts lone surrogate", func() error { b, e := marshalJSON("opts", &TranscriptionSessionOpts{ ProviderOptions: map[string]json.RawMessage{"x": json.RawMessage(`"\ud800"`)}, @@ -423,10 +417,19 @@ func TestMarshalledOptionsRejectRawPassThrough(t *testing.T) { // invalid UTF-8 becomes U+FFFD, NUL becomes a \u0000 escape. No error here. {"plain string field is sanitized, not rejected", func() error { s, e := marshalJSON("opts", &SpeechCallOptions{Text: "hi\xffthere\x00"}) - if e == nil && !strings.Contains(s, `"hi\ufffdthere\u0000"`) { - t.Errorf("expected U+FFFD / \\u0000 coercion, got %s", s) + if e != nil { + return e } - return e + var decoded struct { + Text string `json:"text"` + } + if e := json.Unmarshal([]byte(s), &decoded); e != nil { + return e + } + if decoded.Text != "hi\ufffdthere\x00" || strings.ContainsRune(s, '\x00') { + t.Errorf("expected decoded U+FFFD / NUL without a raw NUL byte, got text %q from %s", decoded.Text, s) + } + return nil }, ""}, } { t.Run(tc.name, func(t *testing.T) { @@ -472,8 +475,8 @@ func TestEmbedRejectsRawPassThroughOpts(t *testing.T) { func TestCodeFromCRejectsOutOfRange(t *testing.T) { // 4 is retired; 14 is Retry; 15..17 are tool-call errors; 18 is - // NoOutputGenerated. - for _, bad := range []int{0, 4, 19, 999} { + // NoOutputGenerated; 19 / 20 are the missing-key / missing-setting codes. + for _, bad := range []int{0, 4, 21, 999} { if _, ok := codeFromC(bad); ok { t.Fatalf("%d is not an AiMuxError variant", bad) } @@ -491,6 +494,12 @@ func TestCodeFromCRejectsOutOfRange(t *testing.T) { if c, ok := codeFromC(18); !ok || c != CodeNoOutputGenerated || c.String() != "NoOutputGenerated" { t.Fatalf("18 → %v, %v; want CodeNoOutputGenerated", c, ok) } + if c, ok := codeFromC(19); !ok || c != CodeLoadAPIKey || c.String() != "LoadApiKey" { + t.Fatalf("19 → %v, %v; want CodeLoadAPIKey", c, ok) + } + if c, ok := codeFromC(20); !ok || c != CodeLoadSetting || c.String() != "LoadSetting" { + t.Fatalf("20 → %v, %v; want CodeLoadSetting", c, ok) + } } // A CodeRetry error keeps the per-attempt history in order; LastError is the diff --git a/bindings/go/typed_test.go b/bindings/go/typed_test.go index 35dfc2b0d..cccf459ca 100644 --- a/bindings/go/typed_test.go +++ b/bindings/go/typed_test.go @@ -54,7 +54,7 @@ func TestTypedGenerateMessages(t *testing.T) { // Verify the provider received both messages. var reqBody map[string]any json.Unmarshal([]byte(srv.LastRequestBody()), &reqBody) - messages := reqBody["messages"].([]any) + messages := reqBody["input"].([]any) if len(messages) != 2 { t.Fatalf("expected 2 messages, got %d", len(messages)) } diff --git a/bindings/go/types.go b/bindings/go/types.go index 0cb79f032..e1dbefa7c 100644 --- a/bindings/go/types.go +++ b/bindings/go/types.go @@ -211,18 +211,18 @@ type ResponseMetadata struct { ModelID *string `json:"model_id,omitempty"` } -// GenerateRequest contains the provider HTTP request body. -type GenerateRequest struct { - Body json.RawMessage `json:"body,omitempty"` -} - -// GenerateResponseMetadata is response metadata with optional headers and body. +// GenerateResponseMetadata is response metadata with optional HTTP headers and body. type GenerateResponseMetadata struct { ResponseMetadata Headers map[string]string `json:"headers,omitempty"` Body json.RawMessage `json:"body,omitempty"` } +// GenerateRequest is provider request information with an optional HTTP body. +type GenerateRequest struct { + Body json.RawMessage `json:"body,omitempty"` +} + // GenerateResult is the raw provider result. // Mirrors Kotlin GenerateResult (Types.kt:853-871). type GenerateResult struct { @@ -434,7 +434,6 @@ type GenerateTextOptions struct { ProviderOptions json.RawMessage `json:"provider_options,omitempty"` Reasoning *ReasoningEffort `json:"reasoning,omitempty"` Instructions *string `json:"instructions,omitempty"` - BodyOverrides json.RawMessage `json:"body_overrides,omitempty"` MaxRetries *uint32 `json:"max_retries,omitempty"` Timeout *TimeoutConfiguration `json:"timeout,omitempty"` IncludeRawChunks *bool `json:"include_raw_chunks,omitempty"` diff --git a/bindings/java/src/main/java/ai/arcships/aimux/Aimux.java b/bindings/java/src/main/java/ai/arcships/aimux/Aimux.java index 6525e1c34..71947081c 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/Aimux.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/Aimux.java @@ -151,9 +151,9 @@ public static void registerProviders(String configJson) { } /** - * Set the global proxy configuration (M6, RFC-0016). Must be called before - * the first {@code generateText} / {@code streamText} call; a no-op if the - * shared HTTP client is already initialised. + * Replace the default transport with one routed by this proxy configuration + * (M6, RFC-0016). Requests started afterwards use it; calling it again + * replaces the previous configuration. * * @param configJson ProxyConfig JSON ({@code "http_url"}, {@code "https_url"}, * {@code "all_url"}, {@code "no_proxy"} — all optional). diff --git a/bindings/java/src/main/java/ai/arcships/aimux/AimuxException.java b/bindings/java/src/main/java/ai/arcships/aimux/AimuxException.java index c73ca92ff..79b256515 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/AimuxException.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/AimuxException.java @@ -11,7 +11,7 @@ /** * AiMuxError hierarchy (OpenAI Java / Vercel AI SDK style). * - *

Raised when a fallible C ABI call returns an AiMuxError code (1–18). + *

Raised when a fallible C ABI call returns an AiMuxError code (1–20). * Recording failures use the * independent {@link RecordingException} type; C ABI failures (bad raw * wire JSON, use-after-close, re-entrant call) surface as plain @@ -29,7 +29,7 @@ * } * } * - *

Every instance carries {@link #getCode()} (C {@code aimux_error_code_t} 1–18), + *

Every instance carries {@link #getCode()} (C {@code aimux_error_code_t} 1–20), * {@link #getStatusCode()} (HTTP or {@code -1}), {@link #getRetryMs()} (hint * or {@code -1}; {@code 0} = retry now) and {@link #isRetryable()}. Message * text comes from the C layer. @@ -42,8 +42,9 @@ public class AimuxException extends RuntimeException { private static final long serialVersionUID = 1L; // ── aimux_error_code_t (aimux-error.h) ────────────────────────────────── - // 17 variant codes (1–18; 1 is the catch-all OTHER; 4 is retired — the - // legacy Tool variant; 14 = RETRY); every HTTP-shaped failure + // 19 variant codes (1–20; 1 is the catch-all OTHER; 4 is retired — the + // legacy Tool variant; 14 = RETRY; 19 / 20 = missing API key / setting); + // every HTTP-shaped failure // arrives as AIMUX_E_API_CALL. A code outside that set is a header/library // mismatch and fails with IllegalStateException, never an AimuxException. // Recording failures are a different type: see RecordingException. @@ -64,6 +65,8 @@ public class AimuxException extends RuntimeException { public static final int AIMUX_E_INVALID_TOOL_INPUT = 16; public static final int AIMUX_E_TOOL_CALL_REPAIR = 17; public static final int AIMUX_E_NO_OUTPUT_GENERATED = 18; + public static final int AIMUX_E_LOAD_API_KEY = 19; + public static final int AIMUX_E_LOAD_SETTING = 20; public static final int AIMUX_E_OTHER = 1; public static final int AIMUX_E_RETRY = 14; @@ -104,7 +107,7 @@ public AimuxException(String message, int code, int status, long retryMs, Throwa // ── Accessors ─────────────────────────────────────────────────────────── - /** C {@code aimux_error_code_t} value (1–18). */ + /** C {@code aimux_error_code_t} value (1–20). */ public int getCode() { return code; } @@ -141,7 +144,7 @@ public boolean isRetryable() { * every returned string. Does not own the pointer: the caller * ({@link AimuxResult#expectAimuxError}) frees the returned error afterwards * (retry attempt errors are new owned copies and are freed here). - * A code outside 1–18 is a header/library mismatch → + * A code outside 1–20 is a header/library mismatch → * {@link IllegalStateException}. */ static AimuxException fromC(Pointer error, String prefix) { @@ -195,6 +198,14 @@ static AimuxException fromC(Pointer error, String prefix) { ex = new ToolCallRepairError(msg, -1, -1L, parseWireJson(AimuxResult.takeString(ffi.aimux_error_original_error(error)))); break; + case AIMUX_E_LOAD_API_KEY: + ex = new LoadAPIKeyError(msg, -1, -1L, + AimuxResult.takeString(ffi.aimux_error_provider_code(error))); + break; + case AIMUX_E_LOAD_SETTING: + ex = new LoadSettingError(msg, -1, -1L, + AimuxResult.takeString(ffi.aimux_error_provider_code(error))); + break; default: ex = createByCode(code, msg, -1, -1L); } @@ -303,6 +314,10 @@ private static AimuxException createByCode(int code, String message, int status, return new InvalidArgumentError(message, status, retryMs); case AIMUX_E_INVALID_PROMPT: return new InvalidPromptError(message, status, retryMs); + case AIMUX_E_LOAD_API_KEY: + return new LoadAPIKeyError(message, status, retryMs); + case AIMUX_E_LOAD_SETTING: + return new LoadSettingError(message, status, retryMs); case AIMUX_E_TOKEN_EXPIRED: return new TokenExpiredError(message, status, retryMs); case AIMUX_E_UNSUPPORTED_FUNCTIONALITY: @@ -348,6 +363,10 @@ public static String codeName(int code) { return "InvalidArgument"; case AIMUX_E_INVALID_PROMPT: return "InvalidPrompt"; + case AIMUX_E_LOAD_API_KEY: + return "LoadApiKey"; + case AIMUX_E_LOAD_SETTING: + return "LoadSetting"; case AIMUX_E_TOKEN_EXPIRED: return "TokenExpired"; case AIMUX_E_UNSUPPORTED_FUNCTIONALITY: @@ -411,6 +430,50 @@ public InvalidPromptError(String message, int status, long retryMs) { } } + /** + * No API key was passed and the fallback environment variable is unset + * (AI SDK {@code LoadAPIKeyError}). No request was made. + */ + public static class LoadAPIKeyError extends AimuxException { + private final String envVar; + + public LoadAPIKeyError(String message, int status, long retryMs) { + this(message, status, retryMs, null); + } + + public LoadAPIKeyError(String message, int status, long retryMs, String envVar) { + super(message, AIMUX_E_LOAD_API_KEY, status, retryMs); + this.envVar = envVar; + } + + /** The environment variable that was consulted, e.g. {@code "OPENAI_API_KEY"}, or {@code null}. */ + public String getEnvVar() { + return envVar; + } + } + + /** + * A required provider setting was not passed and its fallback environment + * variable is unset (AI SDK {@code LoadSettingError}). + */ + public static class LoadSettingError extends AimuxException { + private final String envVar; + + public LoadSettingError(String message, int status, long retryMs) { + this(message, status, retryMs, null); + } + + public LoadSettingError(String message, int status, long retryMs, String envVar) { + super(message, AIMUX_E_LOAD_SETTING, status, retryMs); + this.envVar = envVar; + } + + /** The environment variable that was consulted, or {@code null}. */ + public String getEnvVar() { + return envVar; + } + } + public static class TokenExpiredError extends AimuxException { public TokenExpiredError(String message, int status, long retryMs) { super(message, AIMUX_E_TOKEN_EXPIRED, status == -1 ? 401 : status, retryMs); diff --git a/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java b/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java index e5a1d6340..6416736bc 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java @@ -223,7 +223,10 @@ Pointer aimux_apply_tool_call_repair_to_result(String resultJson, String optsJso /** AIMUX_E_API_CALL: rate-limit hint in ms, or -1. */ long aimux_error_retry_ms(Pointer err); - /** AIMUX_E_API_CALL: owned string or NULL. */ + /** + * AIMUX_E_API_CALL: owned string or NULL. AIMUX_E_LOAD_API_KEY / + * AIMUX_E_LOAD_SETTING: the environment variable that was consulted. + */ Pointer aimux_error_provider_code(Pointer err); /** AIMUX_E_API_CALL: owned string or NULL. */ @@ -379,7 +382,7 @@ Pointer aimux_file_upload(long handle, String dataBase64, /** Register external OpenAI-compatible providers from JSON config (RFC-0020). Returns the error or null. */ Pointer aimux_register_providers(String configJson); - /** Set the global proxy configuration (M6, RFC-0016). Returns the error or null. */ + /** Replace the default transport's proxy configuration (M6, RFC-0016). Returns the error or null. */ Pointer aimux_init_proxy(String configJson); /** Create a RouterModel (RFC-0021) over child handles. {@code handles} is pinned by diff --git a/bindings/java/src/main/java/ai/arcships/aimux/AimuxResult.java b/bindings/java/src/main/java/ai/arcships/aimux/AimuxResult.java index 5ff3b786f..269773f42 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/AimuxResult.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/AimuxResult.java @@ -9,7 +9,7 @@ * *

Every fallible C call returns an {@code aimux_error_t *} ({@code null} * = success, result in the out-parameter). Its code identifies an AiMuxError - * (1–18), RecordingError (100–105), or a failure detected by the C ABI + * (1–20), RecordingError (100–105), or a failure detected by the C ABI * (200–206). The last range collapses to {@link IllegalStateException} * ({@code "aimux ffi: "} + message); Java does not expose seven additional * exception types. Each helper frees the pointer exactly once. User-triggerable @@ -64,7 +64,7 @@ private static String prefix(String context) { } /** - * Decode an error from a call that may return {@code AiMuxError}: 1–18 → + * Decode an error from a call that may return {@code AiMuxError}: 1–20 → * {@link AimuxException}; 200–206 → {@link IllegalStateException}. * Frees {@code e}. */ @@ -79,9 +79,10 @@ static RuntimeException expectAimuxError(Pointer e, String context) { if (isFfiCode(code)) { return ffiError(e, prefix); } - // Code 4 is retired; Retry is 14 and tool errors are 15..17. + // Code 4 is retired; Retry is 14, tool errors are 15..17, and a missing + // API key / setting are 19 / 20. if (code < AimuxException.AIMUX_E_OTHER - || code > AimuxException.AIMUX_E_NO_OUTPUT_GENERATED + || code > AimuxException.AIMUX_E_LOAD_SETTING || code == 4) { return codeMismatch(code, prefix, "AiMuxError"); } diff --git a/bindings/java/src/main/java/ai/arcships/aimux/Model.java b/bindings/java/src/main/java/ai/arcships/aimux/Model.java index a6c5c185c..28c4b297d 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/Model.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/Model.java @@ -368,8 +368,10 @@ public static Model azureWithBase(String apiKey, String baseUrl, String deployme * from the registry entry. * @param modelId Model id. * @param configJson Optional JSON object of ProviderOptions - * ({@code {"base_url": "...", "headers": {...}, "max_retries": 0, - * "body_overrides": {...}}}); {@code null} for defaults. + * ({@code {"base_url": "...", "headers": {...}, "organization": "...", + * "project": "...", "params": {...}}}); {@code null} for defaults. + * {@code max_retries} (a per-call option) and {@code body_overrides} + * (removed) are rejected as {@link AimuxException.InvalidArgumentError}. * @return A new {@link Model}. * @throws AimuxException if the provider could not be constructed * (unknown provider, bad config, missing env key). diff --git a/bindings/java/src/main/java/ai/arcships/aimux/Types.java b/bindings/java/src/main/java/ai/arcships/aimux/Types.java index 2d13adc07..edaaa157e 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/Types.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/Types.java @@ -2108,7 +2108,6 @@ public static class GenerateTextOptions { @JsonProperty("provider_options") private Map providerOptions; @JsonProperty("reasoning") private ReasoningEffort reasoning; @JsonProperty("instructions") private String instructions; - @JsonProperty("body_overrides") private JsonNode bodyOverrides; @JsonProperty("max_retries") private Long maxRetries; @JsonProperty("include_raw_chunks") private Boolean includeRawChunks; @JsonProperty("timeout") private TimeoutConfiguration timeout; @@ -2125,7 +2124,7 @@ private GenerateTextOptions(Long maxOutputTokens, Double temperature, List tools, ToolChoice toolChoice, Map headers, Map providerOptions, ReasoningEffort reasoning, String instructions, - JsonNode bodyOverrides, Long maxRetries, Boolean includeRawChunks, + Long maxRetries, Boolean includeRawChunks, TimeoutConfiguration timeout, String sessionId) { this.maxOutputTokens = maxOutputTokens; this.temperature = temperature; @@ -2142,7 +2141,6 @@ private GenerateTextOptions(Long maxOutputTokens, Double temperature, List getProviderOptions() { return providerOptions; } public ReasoningEffort getReasoning() { return reasoning; } public String getInstructions() { return instructions; } - public JsonNode getBodyOverrides() { return bodyOverrides; } public Long getMaxRetries() { return maxRetries; } public Boolean getIncludeRawChunks() { return includeRawChunks; } public TimeoutConfiguration getTimeout() { return timeout; } @@ -2211,7 +2208,6 @@ public static class Builder { private Map providerOptions; private ReasoningEffort reasoning; private String instructions; - private JsonNode bodyOverrides; private Long maxRetries; private Boolean includeRawChunks; private TimeoutConfiguration timeout; @@ -2233,7 +2229,6 @@ public static class Builder { public Builder providerOptions(Map v) { this.providerOptions = v; return this; } public Builder reasoning(ReasoningEffort v) { this.reasoning = v; return this; } public Builder instructions(String v) { this.instructions = v; return this; } - public Builder bodyOverrides(JsonNode v) { this.bodyOverrides = v; return this; } public Builder maxRetries(Long v) { this.maxRetries = v; return this; } public Builder includeRawChunks(Boolean v) { this.includeRawChunks = v; return this; } public Builder timeout(TimeoutConfiguration v) { this.timeout = v; return this; } @@ -2244,7 +2239,7 @@ public static class Builder { public GenerateTextOptions build() { GenerateTextOptions options = new GenerateTextOptions(maxOutputTokens, temperature, stopSequences, topP, topK, presencePenalty, frequencyPenalty, output, seed, - tools, toolChoice, headers, providerOptions, reasoning, instructions, bodyOverrides, + tools, toolChoice, headers, providerOptions, reasoning, instructions, maxRetries, includeRawChunks, timeout, sessionId); // Set outside the constructor: the hook is host-side state, not // part of the serialized option set. @@ -2273,7 +2268,6 @@ public boolean equals(Object o) { && Objects.equals(providerOptions, that.providerOptions) && Objects.equals(reasoning, that.reasoning) && Objects.equals(instructions, that.instructions) - && Objects.equals(bodyOverrides, that.bodyOverrides) && Objects.equals(maxRetries, that.maxRetries) && Objects.equals(includeRawChunks, that.includeRawChunks) && Objects.equals(timeout, that.timeout) @@ -2284,7 +2278,7 @@ public boolean equals(Object o) { public int hashCode() { return Objects.hash(maxOutputTokens, temperature, stopSequences, topP, topK, presencePenalty, frequencyPenalty, output, seed, tools, toolChoice, headers, providerOptions, reasoning, - instructions, bodyOverrides, maxRetries, includeRawChunks, timeout, sessionId); + instructions, maxRetries, includeRawChunks, timeout, sessionId); } } @@ -3206,34 +3200,8 @@ public GenerateContent deserialize(JsonParser p, DeserializationContext ctxt) th } } - /** Provider HTTP request information. */ - public static class GenerateRequest { - @JsonProperty("body") private JsonNode body; - - @JsonCreator - GenerateRequest() {} - - private GenerateRequest(JsonNode body) { this.body = body; } - public JsonNode getBody() { return body; } - public static Builder builder() { return new Builder(); } - - public static class Builder { - private JsonNode body; - public Builder body(JsonNode v) { this.body = v; return this; } - public GenerateRequest build() { return new GenerateRequest(body); } - } - - @Override - public boolean equals(Object o) { - return o instanceof GenerateRequest && Objects.equals(body, ((GenerateRequest) o).body); - } - - @Override - public int hashCode() { return Objects.hash(body); } - } - - /** Provider response metadata, headers, and body. */ - public static class GenerateResponse { + /** Response metadata, headers, and body returned by the provider. */ + public static class ResponseInfo { @JsonProperty("id") private String id; @JsonProperty("timestamp") private String timestamp; @JsonProperty("model_id") private String modelId; @@ -3241,10 +3209,9 @@ public static class GenerateResponse { @JsonProperty("body") private JsonNode body; @JsonCreator - GenerateResponse() {} + ResponseInfo() {} - private GenerateResponse(String id, String timestamp, String modelId, - Map headers, JsonNode body) { + private ResponseInfo(String id, String timestamp, String modelId, Map headers, JsonNode body) { this.id = id; this.timestamp = timestamp; this.modelId = modelId; @@ -3257,6 +3224,7 @@ private GenerateResponse(String id, String timestamp, String modelId, public String getModelId() { return modelId; } public Map getHeaders() { return headers; } public JsonNode getBody() { return body; } + public static Builder builder() { return new Builder(); } public static class Builder { @@ -3265,20 +3233,25 @@ public static class Builder { private String modelId; private Map headers; private JsonNode body; + public Builder id(String v) { this.id = v; return this; } public Builder timestamp(String v) { this.timestamp = v; return this; } public Builder modelId(String v) { this.modelId = v; return this; } public Builder headers(Map v) { this.headers = v; return this; } public Builder body(JsonNode v) { this.body = v; return this; } - public GenerateResponse build() { return new GenerateResponse(id, timestamp, modelId, headers, body); } + + public ResponseInfo build() { return new ResponseInfo(id, timestamp, modelId, headers, body); } } @Override public boolean equals(Object o) { - if (!(o instanceof GenerateResponse)) return false; - GenerateResponse that = (GenerateResponse) o; - return Objects.equals(id, that.id) && Objects.equals(timestamp, that.timestamp) - && Objects.equals(modelId, that.modelId) && Objects.equals(headers, that.headers) + if (this == o) return true; + if (!(o instanceof ResponseInfo)) return false; + ResponseInfo that = (ResponseInfo) o; + return Objects.equals(id, that.id) + && Objects.equals(timestamp, that.timestamp) + && Objects.equals(modelId, that.modelId) + && Objects.equals(headers, that.headers) && Objects.equals(body, that.body); } @@ -3286,6 +3259,40 @@ public boolean equals(Object o) { public int hashCode() { return Objects.hash(id, timestamp, modelId, headers, body); } } + public static class RequestInfo { + @JsonProperty("body") private JsonNode body; + + @JsonCreator + RequestInfo() {} + + private RequestInfo(JsonNode body) { + this.body = body; + } + + public JsonNode getBody() { return body; } + + public static Builder builder() { return new Builder(); } + + public static class Builder { + private JsonNode body; + + public Builder body(JsonNode v) { this.body = v; return this; } + + public RequestInfo build() { return new RequestInfo(body); } + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof RequestInfo)) return false; + RequestInfo that = (RequestInfo) o; + return Objects.equals(body, that.body); + } + + @Override + public int hashCode() { return Objects.hash(body); } + } + /** * Result of `LanguageModel::do_generate` (non-streaming) — the raw provider * result surfaced via {@link GenerateTextResult#getRaw()}. @@ -3299,15 +3306,15 @@ public static class GenerateResult { @JsonProperty("usage") private Usage usage = new Usage(); @JsonProperty("warnings") private List warnings = new ArrayList<>(); @JsonProperty("provider_metadata") private JsonNode providerMetadata; - @JsonProperty("response") private GenerateResponse response; - @JsonProperty("request") private GenerateRequest request; + @JsonProperty("response") private ResponseInfo response; + @JsonProperty("request") private RequestInfo request; @JsonCreator GenerateResult() {} private GenerateResult(List content, FinishReason finishReason, Usage usage, - List warnings, JsonNode providerMetadata, GenerateResponse response, - GenerateRequest request) { + List warnings, JsonNode providerMetadata, ResponseInfo response, + RequestInfo request) { this.content = content; this.finishReason = finishReason; this.usage = usage; @@ -3322,8 +3329,8 @@ private GenerateResult(List content, FinishReason finishReason, public Usage getUsage() { return usage; } public List getWarnings() { return warnings; } public JsonNode getProviderMetadata() { return providerMetadata; } - public GenerateResponse getResponse() { return response; } - public GenerateRequest getRequest() { return request; } + public ResponseInfo getResponse() { return response; } + public RequestInfo getRequest() { return request; } /** Names of the variant tags present in {@link #getContent()} (e.g. "Text", "ToolCall"). */ public List getContentVariantTags() { @@ -3350,16 +3357,16 @@ public static class Builder { private Usage usage = new Usage(); private List warnings = new ArrayList<>(); private JsonNode providerMetadata; - private GenerateResponse response; - private GenerateRequest request; + private ResponseInfo response; + private RequestInfo request; public Builder content(List v) { this.content = v; return this; } public Builder finishReason(FinishReason v) { this.finishReason = v; return this; } public Builder usage(Usage v) { this.usage = v; return this; } public Builder warnings(List v) { this.warnings = v; return this; } public Builder providerMetadata(JsonNode v) { this.providerMetadata = v; return this; } - public Builder response(GenerateResponse v) { this.response = v; return this; } - public Builder request(GenerateRequest v) { this.request = v; return this; } + public Builder response(ResponseInfo v) { this.response = v; return this; } + public Builder request(RequestInfo v) { this.request = v; return this; } public GenerateResult build() { return new GenerateResult(content, finishReason, usage, warnings, providerMetadata, response, diff --git a/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java b/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java index c91690972..49e74e4c2 100644 --- a/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java +++ b/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java @@ -226,17 +226,37 @@ void ofMapsEveryVariantToExpectedClass() { .isInstanceOf(AimuxException.OtherError.class); assertThat(AimuxException.of(AimuxException.AIMUX_E_RETRY, "m")) .isInstanceOf(AimuxException.RetryError.class); + assertThat(AimuxException.of(AimuxException.AIMUX_E_LOAD_API_KEY, "m")) + .isInstanceOf(AimuxException.LoadAPIKeyError.class); + assertThat(AimuxException.of(AimuxException.AIMUX_E_LOAD_SETTING, "m")) + .isInstanceOf(AimuxException.LoadSettingError.class); + } + + @Test + void loadErrorsHaveTheirOwnCodesNamesAndNoEnvVarLocally() { + assertThat(AimuxException.AIMUX_E_LOAD_API_KEY).isEqualTo(19); + assertThat(AimuxException.AIMUX_E_LOAD_SETTING).isEqualTo(20); + assertThat(AimuxException.codeName(19)).isEqualTo("LoadApiKey"); + assertThat(AimuxException.codeName(20)).isEqualTo("LoadSetting"); + AimuxException key = AimuxException.of(AimuxException.AIMUX_E_LOAD_API_KEY, "no key"); + assertThat(key.getCode()).isEqualTo(19); + assertThat(((AimuxException.LoadAPIKeyError) key).getEnvVar()).isNull(); + AimuxException setting = AimuxException.of(AimuxException.AIMUX_E_LOAD_SETTING, "no region"); + assertThat(setting.getCode()).isEqualTo(20); + assertThat(((AimuxException.LoadSettingError) setting).getEnvVar()).isNull(); } @Test void codesOutsideTheRustEnumAreRejected() { - // Tool errors occupy 15..17 and 18 is NoOutputGenerated; retired code 4 - // and values beyond 18 are header/library mismatches. + // Tool errors occupy 15..17, 18 is NoOutputGenerated and 19 / 20 + // are the missing key / setting; 4 is retired and 21 is unassigned. assertThat(AimuxException.of(15, "")).isInstanceOf(AimuxException.NoSuchToolError.class); + assertThat(AimuxException.AIMUX_E_NO_OUTPUT_GENERATED).isEqualTo(18); + assertThat(AimuxException.codeName(18)).isEqualTo("NoOutputGenerated"); assertThat(AimuxException.of(18, "")).isInstanceOf(AimuxException.NoOutputGeneratedError.class); assertThatThrownBy(() -> AimuxException.of(4, "")) .isInstanceOf(IllegalStateException.class); - assertThatThrownBy(() -> AimuxException.of(19, "")) + assertThatThrownBy(() -> AimuxException.of(21, "")) .isInstanceOf(IllegalStateException.class); assertThatThrownBy(() -> AimuxException.of(999, "m")) .isInstanceOf(IllegalStateException.class); diff --git a/bindings/java/src/test/java/ai/arcships/aimux/MockProviderServer.java b/bindings/java/src/test/java/ai/arcships/aimux/MockProviderServer.java index 599e35663..0e8afa7db 100644 --- a/bindings/java/src/test/java/ai/arcships/aimux/MockProviderServer.java +++ b/bindings/java/src/test/java/ai/arcships/aimux/MockProviderServer.java @@ -28,6 +28,9 @@ public class MockProviderServer implements Closeable { /** Last received request method (e.g. "POST"). */ private volatile String lastRequestMethod = ""; + /** Last received request path, used to verify the selected API. */ + private volatile String lastRequestPath = ""; + /** Status code to return. */ private volatile int status = 200; @@ -55,6 +58,7 @@ private void handle(HttpExchange exchange) throws IOException { byte[] reqBytes = readAll(exchange.getRequestBody()); lastRequestBody = new String(reqBytes, StandardCharsets.UTF_8); lastRequestMethod = exchange.getRequestMethod(); + lastRequestPath = exchange.getRequestURI().getPath(); // Pop the next queued response if any, otherwise fall back to the // single-response field (backward compatible with existing tests). @@ -113,6 +117,10 @@ public String lastRequestMethod() { return lastRequestMethod; } + public String lastRequestPath() { + return lastRequestPath; + } + public void setStatus(int status) { this.status = status; } diff --git a/bindings/java/src/test/java/ai/arcships/aimux/ModelTest.java b/bindings/java/src/test/java/ai/arcships/aimux/ModelTest.java index 0edddeec5..72ee89256 100644 --- a/bindings/java/src/test/java/ai/arcships/aimux/ModelTest.java +++ b/bindings/java/src/test/java/ai/arcships/aimux/ModelTest.java @@ -1,6 +1,5 @@ package ai.arcships.aimux; -import org.json.JSONArray; import org.json.JSONObject; import org.junit.jupiter.api.Test; @@ -74,6 +73,7 @@ void generateTextReturnsJsonString() { assertThat(result.getString("text")) .isEqualTo("Rust is a systems programming language."); assertThat(server.lastRequestMethod()).isEqualTo("POST"); + assertThat(server.lastRequestPath()).isEqualTo("/responses"); assertThat(server.lastRequestBody()).contains("gpt-4o-mini"); } } @@ -239,54 +239,13 @@ void concurrentCloseAndGenerateDoesNotCrash() throws InterruptedException { // ── canned OpenAI responses ──────────────────────────────────────────── - /** Plain OpenAI chat-completions response (no tool calls). */ + /** Plain OpenAI Responses response (no tool calls). */ private static String plainOpenAiResponse() { - return new JSONObject() - .put("id", "chatcmpl-test") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject() - .put("message", new JSONObject() - .put("role", "assistant") - .put("content", "Rust is a systems programming language.")) - .put("finish_reason", "stop"))) - .put("usage", new JSONObject() - .put("prompt_tokens", 10) - .put("completion_tokens", 8) - .put("total_tokens", 18)) - .toString(); + return OpenAiResponsesFixtures.text("Rust is a systems programming language.", 10, 8); } - /** OpenAI SSE stream: two content deltas + finish + [DONE]. */ + /** OpenAI Responses SSE: two text deltas, item done and completed. */ private static String openAiSseResponse() { - StringBuilder sb = new StringBuilder(); - sb.append("data: ").append(new JSONObject() - .put("id", "1") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject().put("delta", new JSONObject() - .put("role", "assistant") - .put("content", "Hello"))))) - .append("\n\n"); - sb.append("data: ").append(new JSONObject() - .put("id", "1") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject().put("delta", new JSONObject() - .put("content", " world"))))) - .append("\n\n"); - sb.append("data: ").append(new JSONObject() - .put("id", "1") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject().put("delta", new JSONObject()) - .put("finish_reason", "stop"))) - .put("usage", new JSONObject() - .put("prompt_tokens", 5) - .put("completion_tokens", 2) - .put("total_tokens", 7))) - .append("\n\n"); - sb.append("data: [DONE]\n\n"); - return sb.toString(); + return OpenAiResponsesFixtures.textStream(); } } diff --git a/bindings/java/src/test/java/ai/arcships/aimux/MultimodalCoverageTest.java b/bindings/java/src/test/java/ai/arcships/aimux/MultimodalCoverageTest.java index b5e01c8c3..57d936bd5 100644 --- a/bindings/java/src/test/java/ai/arcships/aimux/MultimodalCoverageTest.java +++ b/bindings/java/src/test/java/ai/arcships/aimux/MultimodalCoverageTest.java @@ -176,9 +176,11 @@ void deepseekGenerateTextViaMock() { + "\"choices\":[{\"index\":0,\"message\":{\"role\":\"assistant\",\"content\":\"Hi\"},\"finish_reason\":\"stop\"}]," + "\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1,\"total_tokens\":2}}"); - try (Model model = Model.openaiWithBase("sk-test", "deepseek-chat", server.baseUrl())) { + String config = new JSONObject().put("base_url", server.baseUrl()).toString(); + try (Model model = Model.provider("deepseek", "sk-test", "deepseek-chat", config)) { String result = model.generateText("\"Say hi\""); assertThat(new JSONObject(result).getString("text")).isEqualTo("Hi"); + assertThat(server.lastRequestPath()).isEqualTo("/chat/completions"); } } } diff --git a/bindings/java/src/test/java/ai/arcships/aimux/OpenAiResponsesFixtures.java b/bindings/java/src/test/java/ai/arcships/aimux/OpenAiResponsesFixtures.java new file mode 100644 index 000000000..f2d58f43d --- /dev/null +++ b/bindings/java/src/test/java/ai/arcships/aimux/OpenAiResponsesFixtures.java @@ -0,0 +1,94 @@ +package ai.arcships.aimux; + +import org.json.JSONArray; +import org.json.JSONObject; + +/** Wire fixtures for the Responses model selected by the OpenAI FFI factory. */ +final class OpenAiResponsesFixtures { + private OpenAiResponsesFixtures() { } + + private static JSONObject usage(int inputTokens, int outputTokens) { + return new JSONObject() + .put("input_tokens", inputTokens) + .put("output_tokens", outputTokens) + .put("total_tokens", inputTokens + outputTokens) + .put("input_tokens_details", new JSONObject().put("cached_tokens", 0)) + .put("output_tokens_details", new JSONObject().put("reasoning_tokens", 0)); + } + + private static JSONObject response(JSONArray output, int inputTokens, int outputTokens) { + return new JSONObject() + .put("id", "resp-test") + .put("object", "response") + .put("created_at", 1) + .put("model", "gpt-4o") + .put("status", "completed") + .put("output", output) + .put("usage", usage(inputTokens, outputTokens)); + } + + private static JSONObject message(String text) { + return new JSONObject() + .put("id", "msg-test") + .put("type", "message") + .put("role", "assistant") + .put("status", "completed") + .put("content", new JSONArray().put(new JSONObject() + .put("type", "output_text") + .put("text", text) + .put("annotations", new JSONArray()))); + } + + private static JSONObject toolCall(String callId, String arguments) { + return new JSONObject() + .put("id", "fc-test") + .put("type", "function_call") + .put("call_id", callId) + .put("name", "get_weather") + .put("arguments", arguments) + .put("status", "completed"); + } + + static String text(String text, int inputTokens, int outputTokens) { + return response(new JSONArray().put(message(text)), inputTokens, outputTokens).toString(); + } + + static String toolCall() { + return response(new JSONArray().put(toolCall("call_abc", "{\"location\":\"Tokyo\"}")), + 20, 10).toString(); + } + + private static String event(String type, JSONObject payload) { + return "event: " + type + "\ndata: " + payload.put("type", type) + "\n\n"; + } + + static String textStream() { + JSONObject item = message("Hello world"); + return event("response.created", new JSONObject().put("response", + response(new JSONArray(), 0, 0).put("status", "in_progress"))) + + event("response.output_item.added", new JSONObject().put("output_index", 0) + .put("item", message("").put("status", "in_progress"))) + + event("response.output_text.delta", new JSONObject().put("item_id", "msg-test") + .put("output_index", 0).put("content_index", 0).put("delta", "Hello")) + + event("response.output_text.delta", new JSONObject().put("item_id", "msg-test") + .put("output_index", 0).put("content_index", 0).put("delta", " world")) + + event("response.output_item.done", new JSONObject().put("output_index", 0).put("item", item)) + + event("response.completed", new JSONObject().put("response", + response(new JSONArray().put(item), 5, 2))); + } + + static String toolCallStream() { + JSONObject item = toolCall("call_xyz", "{\"location\":\"Tokyo\"}"); + return event("response.created", new JSONObject().put("response", + response(new JSONArray(), 0, 0).put("status", "in_progress"))) + + event("response.output_item.added", new JSONObject().put("output_index", 0) + .put("item", toolCall("call_xyz", "").put("status", "in_progress"))) + + event("response.function_call_arguments.delta", new JSONObject() + .put("item_id", "fc-test").put("output_index", 0).put("delta", "{\"location\":")) + + event("response.function_call_arguments.delta", new JSONObject() + .put("item_id", "fc-test").put("output_index", 0).put("delta", "\"Tokyo\"}")) + + event("response.output_item.done", new JSONObject().put("output_index", 0).put("item", item)) + + event("response.completed", new JSONObject().put("response", + response(new JSONArray().put(item), 5, 2))); + } +} diff --git a/bindings/java/src/test/java/ai/arcships/aimux/StructuredE2ETest.java b/bindings/java/src/test/java/ai/arcships/aimux/StructuredE2ETest.java index aee2db261..7ab343616 100644 --- a/bindings/java/src/test/java/ai/arcships/aimux/StructuredE2ETest.java +++ b/bindings/java/src/test/java/ai/arcships/aimux/StructuredE2ETest.java @@ -20,7 +20,7 @@ * the Kotlin binding's StructuredE2ETest.kt). * *

These spin up a local mock HTTP server speaking the OpenAI - * chat-completions wire format, point an OpenAI {@code Model} at it via the + * Responses wire format, point an OpenAI {@code Model} at it via the * base-url constructor, and assert that the binding correctly: *

    *
  1. parses tool_calls out of a provider response,
    2. @@ -50,107 +50,22 @@ void tearDown() { /** Plain OpenAI text response (no tool calls). */ private static String plainOpenAiResponse() { - return new JSONObject() - .put("id", "chatcmpl-test") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject() - .put("message", new JSONObject() - .put("role", "assistant") - .put("content", "Rust is a systems programming language.")) - .put("finish_reason", "stop"))) - .put("usage", new JSONObject() - .put("prompt_tokens", 10) - .put("completion_tokens", 8) - .put("total_tokens", 18)) - .toString(); + return OpenAiResponsesFixtures.text("Rust is a systems programming language.", 10, 8); } /** OpenAI response that requests a tool call (get_weather). */ private static String toolCallOpenAiResponse() { - return new JSONObject() - .put("id", "chatcmpl-tc") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject() - .put("message", new JSONObject() - .put("role", "assistant") - .put("content", JSONObject.NULL) - .put("tool_calls", new JSONArray().put( - new JSONObject() - .put("id", "call_abc") - .put("type", "function") - .put("function", new JSONObject() - .put("name", "get_weather") - .put("arguments", "{\"location\":\"Tokyo\"}"))))) - .put("finish_reason", "tool_calls"))) - .put("usage", new JSONObject() - .put("prompt_tokens", 20) - .put("completion_tokens", 10) - .put("total_tokens", 30)) - .toString(); + return OpenAiResponsesFixtures.toolCall(); } /** A plain text response used as the second leg of the tool-call round-trip. */ private static String finalTextResponse() { - return new JSONObject() - .put("id", "chatcmpl-2") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject() - .put("message", new JSONObject() - .put("role", "assistant") - .put("content", "The weather in Tokyo is sunny.")) - .put("finish_reason", "stop"))) - .put("usage", new JSONObject() - .put("prompt_tokens", 30) - .put("completion_tokens", 8) - .put("total_tokens", 38)) - .toString(); + return OpenAiResponsesFixtures.text("The weather in Tokyo is sunny.", 30, 8); } - /** OpenAI SSE stream: tool_calls name delta -> arguments delta -> finish -> [DONE]. */ + /** OpenAI Responses SSE: item added -> argument deltas -> item done -> completed. */ private static String toolCallSseResponse() { - StringBuilder sb = new StringBuilder(); - sb.append("data: ").append(new JSONObject() - .put("id", "1") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject().put("delta", new JSONObject() - .put("role", "assistant") - .put("tool_calls", new JSONArray().put( - new JSONObject() - .put("index", 0) - .put("id", "call_xyz") - .put("type", "function") - .put("function", new JSONObject() - .put("name", "get_weather") - .put("arguments", "")))))))) - .append("\n\n"); - sb.append("data: ").append(new JSONObject() - .put("id", "1") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject().put("delta", new JSONObject() - .put("tool_calls", new JSONArray().put( - new JSONObject() - .put("index", 0) - .put("function", new JSONObject() - .put("arguments", "{\"location\":\"Tokyo\"}")))))))) - .append("\n\n"); - sb.append("data: ").append(new JSONObject() - .put("id", "1") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject().put("delta", new JSONObject()) - .put("finish_reason", "tool_calls"))) - .put("usage", new JSONObject() - .put("prompt_tokens", 5) - .put("completion_tokens", 2) - .put("total_tokens", 7))) - .append("\n\n"); - sb.append("data: [DONE]\n\n"); - return sb.toString(); + return OpenAiResponsesFixtures.toolCallStream(); } /** A {@code tools} option describing a single get_weather function tool. */ @@ -224,7 +139,7 @@ void multiRoleMessagesReachProvider() { // The provider received both messages (system + user). JSONObject reqBody = new JSONObject(server.lastRequestBody()); assertThat(reqBody.getString("model")).isEqualTo("gpt-4o"); - JSONArray messages = reqBody.getJSONArray("messages"); + JSONArray messages = reqBody.getJSONArray("input"); assertThat(messages.length()).isEqualTo(2); assertThat(messages.getJSONObject(0).getString("role")).isEqualTo("system"); assertThat(messages.getJSONObject(1).getString("role")).isEqualTo("user"); @@ -247,7 +162,7 @@ void toolChoiceReachesProvider() { @Test void streamTextParsesToolCallStreamParts() { - // SSE: tool_calls delta (name) -> arguments delta -> finish_reason. + // Responses SSE: function-call item -> argument deltas -> completed. server.setContentType("text/event-stream"); server.setResponseBody(toolCallSseResponse()); @@ -363,14 +278,17 @@ void toolCallFullRoundTrip() { assertThat(secondResult.getString("text")) .isEqualTo("The weather in Tokyo is sunny."); - // The provider received 3 messages; the last is the tool result, - // carrying the matching tool_call_id. + // The provider received 3 input items; the last is the function-call output, + // carrying the matching call_id. JSONObject reqBody = new JSONObject(server.lastRequestBody()); - JSONArray reqMessages = reqBody.getJSONArray("messages"); + JSONArray reqMessages = reqBody.getJSONArray("input"); assertThat(reqMessages.length()).isEqualTo(3); JSONObject lastMsg = reqMessages.getJSONObject(2); - assertThat(lastMsg.getString("role")).isEqualTo("tool"); - assertThat(lastMsg.getString("tool_call_id")).isEqualTo("call_abc"); + assertThat(lastMsg.getString("type")).isEqualTo("function_call_output"); + assertThat(lastMsg.getString("call_id")).isEqualTo("call_abc"); + JSONObject output = new JSONObject(lastMsg.getString("output")); + assertThat(output.getInt("temperature")).isEqualTo(22); + assertThat(output.getString("condition")).isEqualTo("sunny"); } } } diff --git a/bindings/java/src/test/java/ai/arcships/aimux/TypedModelTest.java b/bindings/java/src/test/java/ai/arcships/aimux/TypedModelTest.java index 883fba548..e3da4beae 100644 --- a/bindings/java/src/test/java/ai/arcships/aimux/TypedModelTest.java +++ b/bindings/java/src/test/java/ai/arcships/aimux/TypedModelTest.java @@ -23,7 +23,7 @@ * {@link Model} (mirror of the Kotlin binding's TypedModelTest.kt). * *

    Like {@link StructuredE2ETest}, these spin up the shared - * {@link MockProviderServer} (OpenAI chat-completions wire format) and assert + * {@link MockProviderServer} (OpenAI Responses wire format) and assert * that the typed wrapper: *

      *
    1. returns a {@link Types.GenerateTextResult} object (no manual JSON parsing),
      2. @@ -53,89 +53,17 @@ void tearDown() { /** Plain OpenAI text response (no tool calls). */ private static String plainOpenAiResponse() { - return new JSONObject() - .put("id", "chatcmpl-test") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject() - .put("message", new JSONObject() - .put("role", "assistant") - .put("content", "Rust is a systems programming language.")) - .put("finish_reason", "stop"))) - .put("usage", new JSONObject() - .put("prompt_tokens", 10) - .put("completion_tokens", 8) - .put("total_tokens", 18)) - .toString(); + return OpenAiResponsesFixtures.text("Rust is a systems programming language.", 10, 8); } /** OpenAI response that requests a tool call (get_weather). */ private static String toolCallOpenAiResponse() { - return new JSONObject() - .put("id", "chatcmpl-tc") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject() - .put("message", new JSONObject() - .put("role", "assistant") - .put("content", JSONObject.NULL) - .put("tool_calls", new JSONArray().put( - new JSONObject() - .put("id", "call_abc") - .put("type", "function") - .put("function", new JSONObject() - .put("name", "get_weather") - .put("arguments", "{\"location\":\"Tokyo\"}"))))) - .put("finish_reason", "tool_calls"))) - .put("usage", new JSONObject() - .put("prompt_tokens", 20) - .put("completion_tokens", 10) - .put("total_tokens", 30)) - .toString(); + return OpenAiResponsesFixtures.toolCall(); } - /** OpenAI SSE stream: tool_calls name delta -> arguments delta -> finish -> [DONE]. */ + /** OpenAI Responses SSE: item added -> argument deltas -> item done -> completed. */ private static String toolCallSseResponse() { - StringBuilder sb = new StringBuilder(); - sb.append("data: ").append(new JSONObject() - .put("id", "1") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject().put("delta", new JSONObject() - .put("role", "assistant") - .put("tool_calls", new JSONArray().put( - new JSONObject() - .put("index", 0) - .put("id", "call_xyz") - .put("type", "function") - .put("function", new JSONObject() - .put("name", "get_weather") - .put("arguments", "")))))))) - .append("\n\n"); - sb.append("data: ").append(new JSONObject() - .put("id", "1") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject().put("delta", new JSONObject() - .put("tool_calls", new JSONArray().put( - new JSONObject() - .put("index", 0) - .put("function", new JSONObject() - .put("arguments", "{\"location\":\"Tokyo\"}")))))))) - .append("\n\n"); - sb.append("data: ").append(new JSONObject() - .put("id", "1") - .put("model", "gpt-4o") - .put("choices", new JSONArray().put( - new JSONObject().put("delta", new JSONObject()) - .put("finish_reason", "tool_calls"))) - .put("usage", new JSONObject() - .put("prompt_tokens", 5) - .put("completion_tokens", 2) - .put("total_tokens", 7))) - .append("\n\n"); - sb.append("data: [DONE]\n\n"); - return sb.toString(); + return OpenAiResponsesFixtures.toolCallStream(); } // ── typed builders ────────────────────────────────────────────────── @@ -304,10 +232,10 @@ void typedToolsAndToolChoiceReachProvider() { assertThat(tools.length()).isEqualTo(1); JSONObject tool = tools.getJSONObject(0); // The engine forwards the typed tool to the provider in OpenAI's - // wire format: {type:"function", function:{name, parameters}}. + // Responses wire format: {type:"function", name, parameters}. assertThat(tool.getString("type")).isEqualTo("function"); - assertThat(tool.getJSONObject("function").getString("name")).isEqualTo("get_weather"); - assertThat(tool.getJSONObject("function").getJSONObject("parameters").getString("type")) + assertThat(tool.getString("name")).isEqualTo("get_weather"); + assertThat(tool.getJSONObject("parameters").getString("type")) .isEqualTo("object"); } } @@ -329,19 +257,21 @@ void multiRoleModelMessageListReachesProvider() { // The provider received both messages (system + user), in order. JSONObject reqBody = new JSONObject(server.lastRequestBody()); assertThat(reqBody.getString("model")).isEqualTo("gpt-4o"); - JSONArray reqMessages = reqBody.getJSONArray("messages"); + JSONArray reqMessages = reqBody.getJSONArray("input"); assertThat(reqMessages.length()).isEqualTo(2); assertThat(reqMessages.getJSONObject(0).getString("role")).isEqualTo("system"); assertThat(reqMessages.getJSONObject(0).getString("content")) .isEqualTo("You are a helpful assistant."); assertThat(reqMessages.getJSONObject(1).getString("role")).isEqualTo("user"); - assertThat(reqMessages.getJSONObject(1).getString("content")).isEqualTo("What is Rust?"); + assertThat(reqMessages.getJSONObject(1).getJSONArray("content").getJSONObject(0).getString("type")) + .isEqualTo("input_text"); + assertThat(reqMessages.getJSONObject(1).getJSONArray("content").getJSONObject(0).getString("text")).isEqualTo("What is Rust?"); } } @Test void streamTextStreamDecodesTypedParts() { - // SSE: tool_calls delta (name) -> arguments delta -> finish_reason. + // Responses SSE: function-call item -> argument deltas -> completed. server.setContentType("text/event-stream"); server.setResponseBody(toolCallSseResponse()); diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Errors.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Errors.kt index 353be46dd..e3e8eadaa 100644 --- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Errors.kt +++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Errors.kt @@ -10,8 +10,9 @@ import kotlinx.serialization.json.JsonPrimitive /** * Machine-readable codes matching aimux-ffi `aimux_error_code_t` (aimux-error.h). - * 1..18 mirror the 17 core variants (1 is the catch-all `Other`; - * 4 is retired — the legacy `Tool` variant; 14 = `Retry`). The + * 1..20 mirror the 19 core variants (1 is the catch-all `Other`; + * 4 is retired — the legacy `Tool` variant; 14 = `Retry`; 19 / 20 = a + * missing API key / provider setting). The * per-status codes (Provider, Http, RateLimited, Auth, ModelNotFound) are * gone, every HTTP-shaped failure arrives as [AIMUX_E_API_CALL]. */ @@ -31,6 +32,8 @@ const val AIMUX_E_NO_SUCH_TOOL: Int = 15 const val AIMUX_E_INVALID_TOOL_INPUT: Int = 16 const val AIMUX_E_TOOL_CALL_REPAIR: Int = 17 const val AIMUX_E_NO_OUTPUT_GENERATED: Int = 18 +const val AIMUX_E_LOAD_API_KEY: Int = 19 +const val AIMUX_E_LOAD_SETTING: Int = 20 const val AIMUX_E_OTHER: Int = 1 const val AIMUX_E_RETRY: Int = 14 @@ -56,7 +59,7 @@ const val AIMUX_E_RETRY: Int = 14 * } * ``` * - * Transport: Rust → C `aimux_error_t *` with code 1..18 → [fromC]. + * Transport: Rust → C `aimux_error_t *` with code 1..20 → [fromC]. * Primary path is not a JSON * error envelope. */ @@ -87,7 +90,7 @@ sealed class AimuxException( * string. Does not own the pointer: the caller ([expectAimuxError]) frees * the returned error afterwards (retry attempt errors are new owned * copies and are freed here). Code [AIMUX_OK] or a code outside - * 1..18 is a header/library mismatch and throws [IllegalStateException]. + * 1..20 is a header/library mismatch and throws [IllegalStateException]. */ @JvmStatic internal fun fromC(error: Pointer, prefix: String = ""): AimuxException { @@ -153,6 +156,13 @@ sealed class AimuxException( originalError = takeString(lib.aimux_error_original_error(error)) ?.let(AimuxJson::parseToJsonElement), ) + AIMUX_E_LOAD_API_KEY, AIMUX_E_LOAD_SETTING -> createByCode( + code, + msg, + retryable = retryable, + // The consulted environment variable rides the provider-code getter. + envVar = takeString(lib.aimux_error_provider_code(error)) ?: "", + ) else -> createByCode(code, msg, retryable = retryable) } } @@ -190,7 +200,7 @@ sealed class AimuxException( } /** - * Build the subclass for a core / C error code (1..18). + * Build the subclass for a core / C error code (1..20). * * Any other code — [AIMUX_OK] on a failure path or a code this binding * does not know — is a header/library mismatch and throws @@ -220,11 +230,14 @@ sealed class AimuxException( availableTools: List? = null, toolInput: String = "", originalError: JsonElement? = null, + envVar: String = "", ): AimuxException = when (code) { AIMUX_E_JSON_PARSE -> JSONParseError(message, status, retryMs, cause, retryable) AIMUX_E_INVALID_RESPONSE_DATA -> InvalidResponseDataError(message, status, retryMs, cause, retryable) AIMUX_E_INVALID_ARGUMENT -> InvalidArgumentError(message, status, retryMs, cause, retryable) AIMUX_E_INVALID_PROMPT -> InvalidPromptError(message, status, retryMs, cause, retryable) + AIMUX_E_LOAD_API_KEY -> LoadAPIKeyError(message, status, retryMs, cause, retryable, envVar) + AIMUX_E_LOAD_SETTING -> LoadSettingError(message, status, retryMs, cause, retryable, envVar) AIMUX_E_TOKEN_EXPIRED -> TokenExpiredError( message, status = if (status == -1) 401 else status, @@ -265,6 +278,8 @@ sealed class AimuxException( AIMUX_E_INVALID_RESPONSE_DATA -> "InvalidResponseData" AIMUX_E_INVALID_ARGUMENT -> "InvalidArgument" AIMUX_E_INVALID_PROMPT -> "InvalidPrompt" + AIMUX_E_LOAD_API_KEY -> "LoadApiKey" + AIMUX_E_LOAD_SETTING -> "LoadSetting" AIMUX_E_TOKEN_EXPIRED -> "TokenExpired" AIMUX_E_UNSUPPORTED_FUNCTIONALITY -> "UnsupportedFunctionality" AIMUX_E_NO_SUCH_MODEL -> "NoSuchModel" @@ -323,6 +338,34 @@ class InvalidPromptError( retryable: Boolean = false, ) : AimuxException(message, AIMUX_E_INVALID_PROMPT, status, retryMs, cause, retryable) +/** + * No API key was passed and the fallback environment variable is unset (AI SDK + * `LoadAPIKeyError`). No request was made. + */ +class LoadAPIKeyError( + message: String, + status: Int = -1, + retryMs: Long = -1, + cause: Throwable? = null, + retryable: Boolean = false, + /** The environment variable that was consulted, e.g. "OPENAI_API_KEY" ("" when synthesized locally). */ + val envVar: String = "", +) : AimuxException(message, AIMUX_E_LOAD_API_KEY, status, retryMs, cause, retryable) + +/** + * A required provider setting was not passed and its fallback environment + * variable is unset (AI SDK `LoadSettingError`). + */ +class LoadSettingError( + message: String, + status: Int = -1, + retryMs: Long = -1, + cause: Throwable? = null, + retryable: Boolean = false, + /** The environment variable that was consulted ("" when synthesized locally). */ + val envVar: String = "", +) : AimuxException(message, AIMUX_E_LOAD_SETTING, status, retryMs, cause, retryable) + /** Access token expired and must be refreshed (HTTP 401, but retry-after-refresh). */ class TokenExpiredError( message: String, diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Model.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Model.kt index f435819f8..520a88945 100644 --- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Model.kt +++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Model.kt @@ -8,7 +8,7 @@ * Errors: every fallible C call returns an `aimux_error_t *` ([Pointer]?): * null = success, result written to the trailing out-parameter * ([LongByReference] for handles, [PointerByReference] for JSON strings); - * non-null = failure. Its unified code identifies [AimuxException] (1..18), + * non-null = failure. Its unified code identifies [AimuxException] (1..20), * [RecordingException] (100..105), or a C ABI failure (200..206). The last * range maps to `IllegalStateException("aimux ffi: …")`. A decoder releases * the returned pointer with `aimux_error_free`. @@ -224,7 +224,7 @@ internal interface AimuxFFI : Library { // [AiMuxError] Register external providers from JSON config (RFC-0020). fun aimux_register_providers(configJson: String): Pointer? - // [AiMuxError] Set the global proxy configuration (M6, RFC-0016). + // [AiMuxError] Replace the default transport's proxy configuration (M6, RFC-0016). fun aimux_init_proxy(configJson: String): Pointer? // [AiMuxError] Create a RouterModel (RFC-0021) over child handles (LongArray is @@ -290,7 +290,7 @@ private fun ffiError(e: Pointer, prefix: String): IllegalStateException { } /** - * Decode an error from a call that may return `AiMuxError`: 1..18 → + * Decode an error from a call that may return `AiMuxError`: 1..20 → * [AimuxException]; 200..206 → [IllegalStateException]. Frees [e]. */ internal fun expectAimuxError(e: Pointer, context: String = ""): RuntimeException { @@ -298,8 +298,9 @@ internal fun expectAimuxError(e: Pointer, context: String = ""): RuntimeExceptio try { val code = FFI.lib.aimux_error_code(e) if (isFfiCode(code)) return ffiError(e, prefix) - // Code 4 is retired; Retry is 14 and tool errors are 15..17. - check(code in AIMUX_E_OTHER..AIMUX_E_NO_OUTPUT_GENERATED && code != 4) { + // Code 4 is retired; Retry is 14, tool errors are 15..17 and a missing + // API key / setting are 19 / 20. + check(code in AIMUX_E_OTHER..AIMUX_E_LOAD_SETTING && code != 4) { "${prefix}aimux ffi: expected AiMuxError code, got $code" } return AimuxException.fromC(e, prefix) @@ -998,9 +999,9 @@ class Model internal constructor(handle: Long) : Closeable { } /** - * Set the global proxy configuration (M6, RFC-0016). Must be called - * before the first `generateText` / `streamText` call; a no-op if the - * shared HTTP client is already initialised. + * Replace the default transport with one routed by this proxy + * configuration (M6, RFC-0016). Requests started afterwards use it; + * calling it again replaces the previous configuration. * * @param configJson ProxyConfig JSON (`http_url`, `https_url`, * `all_url`, `no_proxy` — all optional). diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Multimodal.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Multimodal.kt index 993011c34..dd0fc2f62 100644 --- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Multimodal.kt +++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Multimodal.kt @@ -8,7 +8,7 @@ * JSON strings (base64 for binary), matching the C ABI wire format. * * Every fallible C call returns an `aimux_error_t *` (null = success) that - * [expectAimuxError] decodes codes 1..17 as [AimuxException] and 200..206 as + * [expectAimuxError] decodes codes 1..20 as [AimuxException] and 200..206 as * [IllegalStateException] `"aimux ffi: …"` (malformed raw JSON is * caught before the C call by [requireJson] as [IllegalArgumentException]). * No JSON envelope on the primary path. diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt index cee944943..8822027b6 100644 --- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt +++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt @@ -178,12 +178,7 @@ data class ResponseMetadata( ) @Serializable -data class GenerateRequest( - val body: JsonElement? = null, -) - -@Serializable -data class GenerateResponse( +data class ResponseInfo( val id: String? = null, val timestamp: String? = null, @SerialName("model_id") val modelId: String? = null, @@ -191,6 +186,9 @@ data class GenerateResponse( val body: JsonElement? = null, ) +@Serializable +data class RequestInfo(val body: JsonElement? = null) + /** * A tool call requested by the model. * @@ -678,8 +676,6 @@ data class GenerateTextOptions( @SerialName("provider_options") val providerOptions: Map? = null, val reasoning: ReasoningEffort? = null, val instructions: String? = null, - /** Per-call JSON deep-merge overrides for the request body. Untyped ([JsonElement]). */ - @SerialName("body_overrides") val bodyOverrides: JsonElement? = null, /** Per-call retry count (0 = disable retries). */ @SerialName("max_retries") val maxRetries: Long? = null, /** Emit raw provider stream chunks as `StreamPart.Raw` (debugging; OpenAI-compatible family only). */ @@ -1098,8 +1094,8 @@ data class GenerateResult( val usage: Usage = Usage(), val warnings: List = emptyList(), @SerialName("provider_metadata") val providerMetadata: JsonElement? = null, - val response: GenerateResponse? = null, - val request: GenerateRequest? = null, + val response: ResponseInfo? = null, + val request: RequestInfo? = null, ) { /** Names of the variant tags present in [content] (e.g. "Text", "ToolCall"). */ val contentVariantTags: List diff --git a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt index 8e228669c..56faa55aa 100644 --- a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt +++ b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt @@ -215,15 +215,34 @@ class ErrorsTest { assertThat(noKey.retryable).isFalse() } + /** 19 / 20: a missing API key / provider setting carry the consulted environment variable. */ + @Test + fun `createByCode maps the missing key and setting codes`() { + assertThat(AIMUX_E_LOAD_API_KEY).isEqualTo(19) + assertThat(AIMUX_E_LOAD_SETTING).isEqualTo(20) + assertThat(AimuxException.codeName(19)).isEqualTo("LoadApiKey") + assertThat(AimuxException.codeName(20)).isEqualTo("LoadSetting") + val key = AimuxException.createByCode(AIMUX_E_LOAD_API_KEY, "no key", envVar = "OPENAI_API_KEY") + assertThat(key).isInstanceOf(LoadAPIKeyError::class.java) + assertThat(key.code).isEqualTo(19) + assertThat((key as LoadAPIKeyError).envVar).isEqualTo("OPENAI_API_KEY") + val setting = AimuxException.createByCode(AIMUX_E_LOAD_SETTING, "no region", envVar = "AWS_REGION") + assertThat(setting).isInstanceOf(LoadSettingError::class.java) + assertThat(setting.code).isEqualTo(20) + assertThat((setting as LoadSettingError).envVar).isEqualTo("AWS_REGION") + } + /** A code outside the enum is a header/library mismatch, not an error type. */ @Test fun `createByCode rejects codes outside the enum`() { assertThatThrownBy { AimuxException.createByCode(999, "?") } .isInstanceOf(IllegalStateException::class.java) - // 18 is NoOutputGenerated; 19 is the first unassigned value and is rejected. - assertThat(AimuxException.createByCode(AIMUX_E_NO_OUTPUT_GENERATED, "?")) + // 18 is NoOutputGenerated; 21 is the first unassigned value. + assertThat(AIMUX_E_NO_OUTPUT_GENERATED).isEqualTo(18) + assertThat(AimuxException.codeName(18)).isEqualTo("NoOutputGenerated") + assertThat(AimuxException.createByCode(18, "?")) .isInstanceOf(NoOutputGeneratedError::class.java) - assertThatThrownBy { AimuxException.createByCode(19, "?") } + assertThatThrownBy { AimuxException.createByCode(21, "?") } .isInstanceOf(IllegalStateException::class.java) assertThatThrownBy { AimuxException.createByCode(AIMUX_OK, "?") } .isInstanceOf(IllegalStateException::class.java) diff --git a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ModelTest.kt b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ModelTest.kt index 61aa7d1d3..a1c116664 100644 --- a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ModelTest.kt +++ b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ModelTest.kt @@ -4,8 +4,6 @@ import com.sun.jna.ptr.PointerByReference import org.assertj.core.api.Assertions.assertThat import org.assertj.core.api.Assertions.assertThatCode import org.assertj.core.api.Assertions.assertThatThrownBy -import org.json.JSONArray -import org.json.JSONObject import org.junit.jupiter.api.Test import java.util.concurrent.atomic.AtomicBoolean import java.util.concurrent.atomic.AtomicReference @@ -214,29 +212,7 @@ class ModelTest { // ── canned OpenAI responses ──────────────────────────────────────────── - /** Plain OpenAI chat-completions response (no tool calls). */ + /** Plain OpenAI Responses response (no tool calls). */ private fun plainOpenAiResponse(): String = - JSONObject().apply { - put("id", "chatcmpl-test") - put("model", "gpt-4o") - put( - "choices", - JSONArray().put( - JSONObject() - .put( - "message", - JSONObject() - .put("role", "assistant") - .put("content", "Rust is a systems programming language."), - ).put("finish_reason", "stop"), - ), - ) - put( - "usage", - JSONObject() - .put("prompt_tokens", 10) - .put("completion_tokens", 8) - .put("total_tokens", 18), - ) - }.toString() + OpenAiResponsesFixtures.text("Rust is a systems programming language.") } diff --git a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/OpenAiResponsesFixtures.kt b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/OpenAiResponsesFixtures.kt new file mode 100644 index 000000000..fa24ac522 --- /dev/null +++ b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/OpenAiResponsesFixtures.kt @@ -0,0 +1,65 @@ +package ai.arcships.aimux + +import org.json.JSONArray +import org.json.JSONObject + +/** Wire fixtures for the Responses model selected by the OpenAI FFI factory. */ +internal object OpenAiResponsesFixtures { + private fun usage(inputTokens: Int, outputTokens: Int) = JSONObject() + .put("input_tokens", inputTokens) + .put("output_tokens", outputTokens) + .put("total_tokens", inputTokens + outputTokens) + .put("input_tokens_details", JSONObject().put("cached_tokens", 0)) + .put("output_tokens_details", JSONObject().put("reasoning_tokens", 0)) + + private fun response(output: JSONArray, inputTokens: Int, outputTokens: Int) = JSONObject() + .put("id", "resp-test") + .put("object", "response") + .put("created_at", 1) + .put("model", "gpt-4o") + .put("status", "completed") + .put("output", output) + .put("usage", usage(inputTokens, outputTokens)) + + private fun message(text: String) = JSONObject() + .put("id", "msg-test") + .put("type", "message") + .put("role", "assistant") + .put("status", "completed") + .put("content", JSONArray().put(JSONObject() + .put("type", "output_text") + .put("text", text) + .put("annotations", JSONArray()))) + + private fun toolCall(callId: String, arguments: String) = JSONObject() + .put("id", "fc-test") + .put("type", "function_call") + .put("call_id", callId) + .put("name", "get_weather") + .put("arguments", arguments) + .put("status", "completed") + + fun text(text: String, inputTokens: Int = 10, outputTokens: Int = 8): String = + response(JSONArray().put(message(text)), inputTokens, outputTokens).toString() + + fun toolCall(): String = response( + JSONArray().put(toolCall("call_abc", "{\"location\":\"Tokyo\"}")), 20, 10, + ).toString() + + private fun event(type: String, payload: JSONObject): String = + "event: $type\ndata: ${payload.put("type", type)}\n\n" + + fun toolCallStream(): String { + val item = toolCall("call_xyz", "{\"location\":\"Tokyo\"}") + return event("response.created", JSONObject().put("response", + response(JSONArray(), 0, 0).put("status", "in_progress"))) + + event("response.output_item.added", JSONObject().put("output_index", 0) + .put("item", toolCall("call_xyz", "").put("status", "in_progress"))) + + event("response.function_call_arguments.delta", JSONObject() + .put("item_id", "fc-test").put("output_index", 0).put("delta", "{\"location\":")) + + event("response.function_call_arguments.delta", JSONObject() + .put("item_id", "fc-test").put("output_index", 0).put("delta", "\"Tokyo\"}")) + + event("response.output_item.done", JSONObject().put("output_index", 0).put("item", item)) + + event("response.completed", JSONObject().put("response", response(JSONArray().put(item), 5, 2))) + } +} diff --git a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/StructuredE2ETest.kt b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/StructuredE2ETest.kt index 6bf18c1e9..e7107819c 100644 --- a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/StructuredE2ETest.kt +++ b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/StructuredE2ETest.kt @@ -14,7 +14,7 @@ import java.nio.charset.StandardCharsets // ───────────────────────────────────────────────────────────────────────────── // Structured end-to-end tests for the Kotlin/JVM binding. // -// These spin up a local mock HTTP server speaking the OpenAI chat-completions +// These spin up a local mock HTTP server speaking the OpenAI Responses // wire format, point an OpenAI `Model` at it via the `base_url` constructor, // and assert that the binding correctly: // 1. parses tool_calls out of a provider response, @@ -41,6 +41,11 @@ internal class MockProviderServer { var lastRequestMethod: String = "" private set + /** Last received request path, used to verify the selected API. */ + @Volatile + var lastRequestPath: String = "" + private set + /** Status code to return. */ var status: Int = 200 @@ -79,6 +84,7 @@ internal class MockProviderServer { val reqBytes = exchange.requestBody.readBytes() lastRequestBody = String(reqBytes, StandardCharsets.UTF_8) lastRequestMethod = exchange.requestMethod + lastRequestPath = exchange.requestURI.path // Pop the next queued response if any, otherwise fall back to the // single-response field (backward compatible with the existing tests). @@ -122,79 +128,10 @@ class StructuredE2ETest { /** Plain OpenAI text response (no tool calls). */ private val plainOpenAiResponse: String = - JSONObject().apply { - put("id", "chatcmpl-test") - put("model", "gpt-4o") - put( - "choices", - JSONArray().put( - JSONObject().apply { - put( - "message", - JSONObject().apply { - put("role", "assistant") - put("content", "Rust is a systems programming language.") - }, - ) - put("finish_reason", "stop") - }, - ), - ) - put( - "usage", - JSONObject().apply { - put("prompt_tokens", 10) - put("completion_tokens", 8) - put("total_tokens", 18) - }, - ) - }.toString() + OpenAiResponsesFixtures.text("Rust is a systems programming language.") /** OpenAI response that requests a tool call (get_weather). */ - private val toolCallOpenAiResponse: String = - JSONObject().apply { - put("id", "chatcmpl-tc") - put("model", "gpt-4o") - put( - "choices", - JSONArray().put( - JSONObject().apply { - put( - "message", - JSONObject().apply { - put("role", "assistant") - put("content", JSONObject.NULL) - put( - "tool_calls", - JSONArray().put( - JSONObject().apply { - put("id", "call_abc") - put("type", "function") - put( - "function", - JSONObject().apply { - put("name", "get_weather") - put("arguments", "{\"location\":\"Tokyo\"}") - }, - ) - }, - ), - ) - }, - ) - put("finish_reason", "tool_calls") - }, - ), - ) - put( - "usage", - JSONObject().apply { - put("prompt_tokens", 20) - put("completion_tokens", 10) - put("total_tokens", 30) - }, - ) - }.toString() + private val toolCallOpenAiResponse: String = OpenAiResponsesFixtures.toolCall() /** A `tools` option describing a single get_weather function tool. */ private fun weatherToolsOpts(includeToolChoice: String? = null): String = @@ -278,7 +215,7 @@ class StructuredE2ETest { // The provider received both messages (system + user). val reqBody = JSONObject(server.lastRequestBody) assertThat(reqBody.getString("model")).isEqualTo("gpt-4o") - val messages = reqBody.getJSONArray("messages") + val messages = reqBody.getJSONArray("input") assertThat(messages.length()).isEqualTo(2) assertThat(messages.getJSONObject(0).getString("role")).isEqualTo("system") assertThat(messages.getJSONObject(1).getString("role")).isEqualTo("user") @@ -301,120 +238,21 @@ class StructuredE2ETest { @Test fun `streamText parses tool-call stream parts`() { - // SSE: tool_calls delta (name) -> arguments delta -> finish_reason. + // Responses SSE: function-call item -> argument deltas -> completed. server.contentType = "text/event-stream" - server.responseBody = buildString { - append("data: ") - append( - JSONObject().apply { - put("id", "1") - put("model", "gpt-4o") - put( - "choices", - JSONArray().put( - JSONObject().apply { - put( - "delta", - JSONObject().apply { - put("role", "assistant") - put( - "tool_calls", - JSONArray().put( - JSONObject().apply { - put("index", 0) - put("id", "call_xyz") - put("type", "function") - put( - "function", - JSONObject().apply { - put("name", "get_weather") - put("arguments", "") - }, - ) - }, - ), - ) - }, - ) - }, - ), - ) - }.toString(), - ) - append("\n\n") - append("data: ") - append( - JSONObject().apply { - put("id", "1") - put("model", "gpt-4o") - put( - "choices", - JSONArray().put( - JSONObject().apply { - put( - "delta", - JSONObject().apply { - put( - "tool_calls", - JSONArray().put( - JSONObject().apply { - put("index", 0) - put( - "function", - JSONObject().apply { - put("arguments", "{\"location\":\"Tokyo\"}") - }, - ) - }, - ), - ) - }, - ) - }, - ), - ) - }.toString(), - ) - append("\n\n") - append("data: ") - append( - JSONObject().apply { - put("id", "1") - put("model", "gpt-4o") - put( - "choices", - JSONArray().put( - JSONObject().apply { - put("delta", JSONObject()) - put("finish_reason", "tool_calls") - }, - ), - ) - put( - "usage", - JSONObject().apply { - put("prompt_tokens", 5) - put("completion_tokens", 2) - put("total_tokens", 7) - }, - ) - }.toString(), - ) - append("\n\n") - append("data: [DONE]\n\n") - } + server.responseBody = OpenAiResponsesFixtures.toolCallStream() Model.openai("sk-test-fake-key", "gpt-4o", server.baseUrl).use { model -> // Run the blocking FFI stream call on a separate thread so the JNA // callbacks (which attach to the calling thread) don't deadlock with // tokio's block_on. The main thread collects parts via a queue. - val parts = java.util.concurrent.LinkedBlockingQueue() + val parts = java.util.concurrent.LinkedBlockingQueue() val streamThread = Thread { model.streamText( "\"What is the weather in Tokyo?\"", weatherToolsOpts(), onPart = { parts.put(it) }, - onDone = { parts.put(null) }, + onDone = { parts.put("__done__") }, ) } streamThread.isDaemon = true @@ -423,6 +261,7 @@ class StructuredE2ETest { val collected = mutableListOf() while (true) { val part = parts.poll(30, java.util.concurrent.TimeUnit.SECONDS) ?: break + if (part == "__done__") break collected.add(part) } streamThread.join(5000) @@ -450,33 +289,8 @@ class StructuredE2ETest { // back-filled between them. The mock server returns a tool-call // response for the first request, then a plain text response for the // second (FIFO queue). - val finalTextResponse = JSONObject().apply { - put("id", "chatcmpl-2") - put("model", "gpt-4o") - put( - "choices", - JSONArray().put( - JSONObject().apply { - put( - "message", - JSONObject().apply { - put("role", "assistant") - put("content", "The weather in Tokyo is sunny.") - }, - ) - put("finish_reason", "stop") - }, - ), - ) - put( - "usage", - JSONObject().apply { - put("prompt_tokens", 30) - put("completion_tokens", 8) - put("total_tokens", 38) - }, - ) - }.toString() + val finalTextResponse = + OpenAiResponsesFixtures.text("The weather in Tokyo is sunny.", 30, 8) server.setResponses(toolCallOpenAiResponse, finalTextResponse) @@ -526,7 +340,7 @@ class StructuredE2ETest { put("type", "tool_result") put("tool_call_id", "call_abc") put( - "output", + "result", JSONObject().put("temperature", 22).put("condition", "sunny"), ) }, @@ -540,14 +354,17 @@ class StructuredE2ETest { assertThat(secondResult.getString("text")) .isEqualTo("The weather in Tokyo is sunny.") - // The provider received 3 messages; the last is the tool result, - // carrying the matching tool_call_id. + // The provider received 3 input items; the last is the function-call output, + // carrying the matching call_id. val reqBody = JSONObject(server.lastRequestBody) - val reqMessages = reqBody.getJSONArray("messages") + val reqMessages = reqBody.getJSONArray("input") assertThat(reqMessages.length()).isEqualTo(3) val lastMsg = reqMessages.getJSONObject(2) - assertThat(lastMsg.getString("role")).isEqualTo("tool") - assertThat(lastMsg.getString("tool_call_id")).isEqualTo("call_abc") + assertThat(lastMsg.getString("type")).isEqualTo("function_call_output") + assertThat(lastMsg.getString("call_id")).isEqualTo("call_abc") + val output = JSONObject(lastMsg.getString("output")) + assertThat(output.getInt("temperature")).isEqualTo(22) + assertThat(output.getString("condition")).isEqualTo("sunny") } } } diff --git a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/TypedModelTest.kt b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/TypedModelTest.kt index 25688c30b..192b2627e 100644 --- a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/TypedModelTest.kt +++ b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/TypedModelTest.kt @@ -15,7 +15,7 @@ import org.junit.jupiter.api.Test // Tests for the typed [TypedModel] wrapper over the raw JSON-string [Model]. // // Like StructuredE2ETest, these spin up the shared [MockProviderServer] -// (OpenAI chat-completions wire format) and assert that the typed wrapper: +// (OpenAI Responses wire format) and assert that the typed wrapper: // 1. returns a [GenerateTextResult] object (no manual JSON parsing), // 2. surfaces `.text`, `.toolCalls[0].toolName`, and `.content`, // 3. forwards typed `tools` / `tool_choice` onto the provider request, @@ -42,79 +42,10 @@ class TypedModelTest { /** Plain OpenAI text response (no tool calls). */ private val plainOpenAiResponse: String = - JSONObject().apply { - put("id", "chatcmpl-test") - put("model", "gpt-4o") - put( - "choices", - JSONArray().put( - JSONObject().apply { - put( - "message", - JSONObject().apply { - put("role", "assistant") - put("content", "Rust is a systems programming language.") - }, - ) - put("finish_reason", "stop") - }, - ), - ) - put( - "usage", - JSONObject().apply { - put("prompt_tokens", 10) - put("completion_tokens", 8) - put("total_tokens", 18) - }, - ) - }.toString() + OpenAiResponsesFixtures.text("Rust is a systems programming language.") /** OpenAI response that requests a tool call (get_weather). */ - private val toolCallOpenAiResponse: String = - JSONObject().apply { - put("id", "chatcmpl-tc") - put("model", "gpt-4o") - put( - "choices", - JSONArray().put( - JSONObject().apply { - put( - "message", - JSONObject().apply { - put("role", "assistant") - put("content", JSONObject.NULL) - put( - "tool_calls", - JSONArray().put( - JSONObject().apply { - put("id", "call_abc") - put("type", "function") - put( - "function", - JSONObject().apply { - put("name", "get_weather") - put("arguments", "{\"location\":\"Tokyo\"}") - }, - ) - }, - ), - ) - }, - ) - put("finish_reason", "tool_calls") - }, - ), - ) - put( - "usage", - JSONObject().apply { - put("prompt_tokens", 20) - put("completion_tokens", 10) - put("total_tokens", 30) - }, - ) - }.toString() + private val toolCallOpenAiResponse: String = OpenAiResponsesFixtures.toolCall() // ── typed builders ────────────────────────────────────────────────── @@ -148,6 +79,7 @@ class TypedModelTest { // .text is a plain Kotlin String. assertThat(result.text).isEqualTo("Rust is a systems programming language.") + assertThat(server.lastRequestPath).isEqualTo("/responses") // A plain-text response carries no tool calls. assertThat(result.toolCalls).isEmpty() @@ -208,10 +140,10 @@ class TypedModelTest { assertThat(tools.length()).isEqualTo(1) val tool = tools.getJSONObject(0) // The engine forwards the typed tool to the provider in OpenAI's - // wire format: {type:"function", function:{name, parameters}}. + // Responses wire format: {type:"function", name, parameters}. assertThat(tool.getString("type")).isEqualTo("function") - assertThat(tool.getJSONObject("function").getString("name")).isEqualTo("get_weather") - assertThat(tool.getJSONObject("function").getJSONObject("parameters").getString("type")) + assertThat(tool.getString("name")).isEqualTo("get_weather") + assertThat(tool.getJSONObject("parameters").getString("type")) .isEqualTo("object") } } @@ -233,13 +165,15 @@ class TypedModelTest { // The provider received both messages (system + user), in order. val reqBody = JSONObject(server.lastRequestBody) assertThat(reqBody.getString("model")).isEqualTo("gpt-4o") - val reqMessages = reqBody.getJSONArray("messages") + val reqMessages = reqBody.getJSONArray("input") assertThat(reqMessages.length()).isEqualTo(2) assertThat(reqMessages.getJSONObject(0).getString("role")).isEqualTo("system") assertThat(reqMessages.getJSONObject(0).getString("content")) .isEqualTo("You are a helpful assistant.") assertThat(reqMessages.getJSONObject(1).getString("role")).isEqualTo("user") - assertThat(reqMessages.getJSONObject(1).getString("content")).isEqualTo("What is Rust?") + assertThat(reqMessages.getJSONObject(1).getJSONArray("content").getJSONObject(0).getString("type")) + .isEqualTo("input_text") + assertThat(reqMessages.getJSONObject(1).getJSONArray("content").getJSONObject(0).getString("text")).isEqualTo("What is Rust?") } } } diff --git a/bindings/node/__test__/body_overrides.test.ts b/bindings/node/__test__/body_overrides.test.ts deleted file mode 100644 index 67c9e2372..000000000 --- a/bindings/node/__test__/body_overrides.test.ts +++ /dev/null @@ -1,254 +0,0 @@ -// body_overrides.test.ts — Node e2e tests for RFC-0017 phase 1: bodyOverrides -// (JSON deep-merge) and provider factory config (headers/maxRetries/bodyOverrides). -// -// These verify the full chain: JS → napi → Rust → HTTP mock, asserting the -// outbound request body carries the user's overrides. - -import test from 'ava' -import { createServer, type Server } from 'node:http' -import { openai, deepseek } from '../src/native.ts' - -function startMockServer(handler: (req: any, res: any) => void): Promise<{ server: Server; url: string }> { - return new Promise((resolve) => { - const server = createServer(handler) - server.listen(0, '127.0.0.1', () => { - const addr = server.address() as any - resolve({ server, url: `http://127.0.0.1:${addr.port}` }) - }) - }) -} - -function closeServer(server: Server): Promise { - return new Promise((resolve) => server.close(() => resolve())) -} - -function readBody(req: any): Promise { - return new Promise((resolve) => { - let body = '' - req.on('data', (chunk: string) => { body += chunk }) - req.on('end', () => resolve(JSON.parse(body))) - }) -} - -const CHAT_RESPONSE = JSON.stringify({ - id: 'chatcmpl-test', - model: 'gpt-4o', - choices: [{ - message: { role: 'assistant', content: 'Done.' }, - finish_reason: 'stop', - }], - usage: { prompt_tokens: 10, completion_tokens: 1, total_tokens: 11 }, -}) - -// ── per-call bodyOverrides ─────────────────────────────────────────────────── - -test('per-call bodyOverrides injects a field into the request body', async (t) => { - let requestBody: any = null - const { server, url } = await startMockServer(async (req, res) => { - requestBody = await readBody(req) - res.writeHead(200, { 'content-type': 'application/json' }) - res.end(CHAT_RESPONSE) - }) - - try { - const model = await openai('test-key', 'gpt-4o', url) - const opts = JSON.stringify({ - body_overrides: { enable_thinking: false }, - }) - await model.generateText(JSON.stringify('Hello'), opts) - - t.is(requestBody.enable_thinking, false) - t.is(requestBody.model, 'gpt-4o') - } finally { - await closeServer(server) - } -}) - -test('per-call bodyOverrides overwrites a standard field', async (t) => { - let requestBody: any = null - const { server, url } = await startMockServer(async (req, res) => { - requestBody = await readBody(req) - res.writeHead(200, { 'content-type': 'application/json' }) - res.end(CHAT_RESPONSE) - }) - - try { - const model = await openai('test-key', 'gpt-4o', url) - // temperature set via standard option, overridden by body_overrides - const opts = JSON.stringify({ - temperature: 0.9, - body_overrides: { temperature: 0.1 }, - }) - await model.generateText(JSON.stringify('Hello'), opts) - - t.is(requestBody.temperature, 0.1) - } finally { - await closeServer(server) - } -}) - -test('per-call bodyOverrides null deletes a field', async (t) => { - let requestBody: any = null - const { server, url } = await startMockServer(async (req, res) => { - requestBody = await readBody(req) - res.writeHead(200, { 'content-type': 'application/json' }) - res.end(CHAT_RESPONSE) - }) - - try { - const model = await openai('test-key', 'gpt-4o', url) - const opts = JSON.stringify({ - temperature: 0.5, - body_overrides: { temperature: null }, - }) - await model.generateText(JSON.stringify('Hello'), opts) - - t.true(requestBody.temperature === undefined, 'temperature should be deleted by null') - } finally { - await closeServer(server) - } -}) - -// ── provider-level bodyOverrides (factory config) ──────────────────────────── - -test('factory config bodyOverrides are merged into every request', async (t) => { - let requestBody: any = null - const { server, url } = await startMockServer(async (req, res) => { - requestBody = await readBody(req) - res.writeHead(200, { 'content-type': 'application/json' }) - res.end(CHAT_RESPONSE) - }) - - try { - const model = await openai('test-key', 'gpt-4o', { - baseUrl: url, - bodyOverrides: '{"X-Relay-Tag":"my-team"}', - }) - await model.generateText(JSON.stringify('Hello')) - - t.is(requestBody['X-Relay-Tag'], 'my-team') - } finally { - await closeServer(server) - } -}) - -test('factory config headers are sent on every request', async (t) => { - let headers: any = null - const { server, url } = await startMockServer(async (req, res) => { - headers = req.headers - // consume body - await readBody(req) - res.writeHead(200, { 'content-type': 'application/json' }) - res.end(CHAT_RESPONSE) - }) - - try { - const model = await openai('test-key', 'gpt-4o', { - baseUrl: url, - headers: JSON.stringify({ 'X-Custom-Header': 'custom-value' }), - }) - await model.generateText(JSON.stringify('Hello')) - - t.is(headers['x-custom-header'], 'custom-value') - } finally { - await closeServer(server) - } -}) - -// ── backward compatibility ─────────────────────────────────────────────────── - -test('factory accepts bare string baseUrl (backward compatible)', async (t) => { - const { server, url } = await startMockServer(async (req, res) => { - await readBody(req) - res.writeHead(200, { 'content-type': 'application/json' }) - res.end(CHAT_RESPONSE) - }) - - try { - // 3rd param is a plain string (old API) - const model = await openai('test-key', 'gpt-4o', url) - const result = JSON.parse(await model.generateText(JSON.stringify('Hello'))) - t.is(result.text, 'Done.') - } finally { - await closeServer(server) - } -}) - -// ── per-call overrides provider-level ──────────────────────────────────────── - -test('per-call bodyOverrides take precedence over provider-level', async (t) => { - let requestBody: any = null - const { server, url } = await startMockServer(async (req, res) => { - requestBody = await readBody(req) - res.writeHead(200, { 'content-type': 'application/json' }) - res.end(CHAT_RESPONSE) - }) - - try { - const model = await openai('test-key', 'gpt-4o', { - baseUrl: url, - bodyOverrides: '{"custom_field":"provider"}', - }) - const opts = JSON.stringify({ - body_overrides: { custom_field: 'call' }, - }) - await model.generateText(JSON.stringify('Hello'), opts) - - t.is(requestBody.custom_field, 'call', 'per-call should override provider-level') - } finally { - await closeServer(server) - } -}) - -// ── maxRetries ──────────────────────────────────────────────────────────────── - -test('maxRetries: 0 disables retries (single request on 500)', async (t) => { - let requestCount = 0 - const { server, url } = await startMockServer(async (req, res) => { - requestCount++ - await readBody(req) - res.writeHead(500, { 'content-type': 'application/json' }) - res.end(JSON.stringify({ error: { message: 'Internal server error', type: 'server_error' } })) - }) - - try { - const model = await openai('test-key', 'gpt-4o', url) - const opts = JSON.stringify({ max_retries: 0 }) - - await t.throwsAsync( - async () => model.generateText(JSON.stringify('Hello'), opts), - { message: /500|Internal|server_error/i }, - ) - - // With retries disabled, exactly 1 request should have been made. - t.is(requestCount, 1, 'should make exactly 1 request when maxRetries=0') - } finally { - await closeServer(server) - } -}) - -test('maxRetries: factory-level disables retries', async (t) => { - let requestCount = 0 - const { server, url } = await startMockServer(async (req, res) => { - requestCount++ - await readBody(req) - res.writeHead(500, { 'content-type': 'application/json' }) - res.end(JSON.stringify({ error: { message: 'Internal server error', type: 'server_error' } })) - }) - - try { - const model = await openai('test-key', 'gpt-4o', { - baseUrl: url, - maxRetries: 0, - }) - - await t.throwsAsync( - async () => model.generateText(JSON.stringify('Hello')), - { message: /500|Internal|server_error/i }, - ) - - t.is(requestCount, 1, 'factory maxRetries:0 should make exactly 1 request') - } finally { - await closeServer(server) - } -}) diff --git a/bindings/node/__test__/cassette-exhaustive.test.ts b/bindings/node/__test__/cassette-exhaustive.test.ts index f37a899c6..cac953c59 100644 --- a/bindings/node/__test__/cassette-exhaustive.test.ts +++ b/bindings/node/__test__/cassette-exhaustive.test.ts @@ -5,7 +5,7 @@ // Node.js → napi → Rust engine → single cassette → parse → result import test from 'ava' -import { openai } from '../src/native.ts' +import { deepseek } from '../src/native.ts' import { readFileSync, readdirSync, statSync } from 'node:fs' import { join, dirname } from 'node:path' import { fileURLToPath } from 'node:url' @@ -133,7 +133,7 @@ test(`exhaustive: ${cassettes.length} chat/completions cassettes replayed`, asyn const prompt = extractPrompt(cass) try { - const model = await openai('test-key', modelId, baseUrl) + const model = await deepseek('test-key', modelId, baseUrl) if (cass.isStream) { const gen = await model.streamText(JSON.stringify(prompt)) let parts = 0 diff --git a/bindings/node/__test__/cassette.test.ts b/bindings/node/__test__/cassette.test.ts index 4e3552632..f726f1c92 100644 --- a/bindings/node/__test__/cassette.test.ts +++ b/bindings/node/__test__/cassette.test.ts @@ -8,7 +8,7 @@ // are excluded (not tested here), not faked. import test from 'ava' -import { openai, anthropic } from '../src/native.ts' +import { openai, deepseek, anthropic } from '../src/native.ts' import { CassetteServer } from './cassette-replay.ts' // ── OpenAI ────────────────────────────────────────────────────────────────── @@ -63,7 +63,7 @@ test('cassette: DeepSeek generate', async (t) => { t.true(srv.count > 0) try { - const model = await openai('test-key', 'deepseek-chat', srv.url) + const model = await deepseek('test-key', 'deepseek-chat', srv.url) const resultJson = await model.generateText(JSON.stringify('Hello')) const result = JSON.parse(resultJson) t.false(!!result.error, `unexpected error: ${result.error}`) @@ -83,7 +83,7 @@ test('cassette: DeepSeek step usage.raw carries vendor fields', async (t) => { t.true(srv.count > 0) try { - const model = await openai('test-key', 'deepseek-chat', srv.url) + const model = await deepseek('test-key', 'deepseek-chat', srv.url) const resultJson = await model.generateText(JSON.stringify('Hello')) const result = JSON.parse(resultJson) const raw = result.steps?.[0]?.usage?.raw @@ -141,7 +141,7 @@ test('cassette: Groq generate', async (t) => { t.true(srv.count > 0) try { - const model = await openai('test-key', 'llama-3.3-70b-versatile', `${srv.url}/openai/v1`) + const model = await deepseek('test-key', 'llama-3.3-70b-versatile', `${srv.url}/openai/v1`) const resultJson = await model.generateText(JSON.stringify('Hello')) const result = JSON.parse(resultJson) t.false(!!result.error, `unexpected error: ${result.error}`) @@ -159,7 +159,7 @@ test('cassette: Mistral generate', async (t) => { t.true(srv.count > 0) try { - const model = await openai('test-key', 'ministral-8b-latest', `${srv.url}/v1`) + const model = await deepseek('test-key', 'ministral-8b-latest', `${srv.url}/v1`) const resultJson = await model.generateText(JSON.stringify('Hello')) const result = JSON.parse(resultJson) t.false(!!result.error, `unexpected error: ${result.error}`) @@ -177,7 +177,7 @@ test('cassette: Ollama generate', async (t) => { t.true(srv.count > 0) try { - const model = await openai('test-key', 'qwen3:4b', `${srv.url}/v1`) + const model = await deepseek('test-key', 'qwen3:4b', `${srv.url}/v1`) const resultJson = await model.generateText(JSON.stringify('Hello')) const result = JSON.parse(resultJson) t.false(!!result.error, `unexpected error: ${result.error}`) @@ -195,7 +195,7 @@ test('cassette: Perplexity generate', async (t) => { t.true(srv.count > 0) try { - const model = await openai('test-key', 'sonar', srv.url) + const model = await deepseek('test-key', 'sonar', srv.url) const resultJson = await model.generateText(JSON.stringify('Hello')) const result = JSON.parse(resultJson) t.false(!!result.error, `unexpected error: ${result.error}`) diff --git a/bindings/node/__test__/e2e.test.ts b/bindings/node/__test__/e2e.test.ts index f0fde9319..bec073eda 100644 --- a/bindings/node/__test__/e2e.test.ts +++ b/bindings/node/__test__/e2e.test.ts @@ -8,7 +8,7 @@ import test from 'ava' import { createServer, type Server } from 'node:http' -import { openai, anthropic } from '../src/native.ts' +import { deepseek, anthropic } from '../src/native.ts' // ── Mock server helpers ───────────────────────────────────────────────────── @@ -76,7 +76,7 @@ test('e2e: OpenAI generateText via mock server', async (t) => { }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const resultJson = await model.generateText(JSON.stringify('What is Rust?')) const r = JSON.parse(resultJson) @@ -95,7 +95,7 @@ test('e2e: OpenAI streamText via mock server', async (t) => { }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const parts: any[] = [] for await (const json of await model.streamText(JSON.stringify('Say hello'))) { @@ -172,7 +172,7 @@ test('e2e: OpenAI generateText with options (max_tokens, temperature)', async (t }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const opts = JSON.stringify({ max_output_tokens: 100, temperature: 0.5 }) await model.generateText(JSON.stringify('Hello'), opts) @@ -211,7 +211,7 @@ test('e2e: OpenAI generateText parses tool_calls (structured content)', async (t }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const opts = JSON.stringify({ tools: [{ type: 'function', @@ -259,7 +259,7 @@ test('e2e: OpenAI generateText with multi-role messages (system + user)', async }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const prompt = JSON.stringify([ { role: 'system', content: 'You are a helpful assistant.' }, { role: 'user', content: 'What is Rust?' }, @@ -294,7 +294,7 @@ test('e2e: OpenAI generateText with tool_choice: required', async (t) => { }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const opts = JSON.stringify({ tools: [{ type: 'function', @@ -328,7 +328,7 @@ test('e2e: OpenAI streamText parses tool-call stream parts', async (t) => { }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const opts = JSON.stringify({ tools: [{ type: 'function', @@ -376,7 +376,7 @@ test('e2e: OpenAI tool-call full round-trip (ToolCall → ToolResult → final t }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const opts = JSON.stringify({ tools: [{ type: 'function', diff --git a/bindings/node/__test__/error.test.ts b/bindings/node/__test__/error.test.ts index 5540c45e5..12eab5d7a 100644 --- a/bindings/node/__test__/error.test.ts +++ b/bindings/node/__test__/error.test.ts @@ -14,6 +14,8 @@ import { APICallError, RetryError, InvalidArgumentError, + LoadAPIKeyError, + LoadSettingError, NoSuchProviderError, RecordingError, initRecording, @@ -62,10 +64,9 @@ test('stream error path yields a typed AimuxError', async (t) => { // Unreachable port: connection is refused immediately; no retries. const model = await openai('test-key', 'gpt-4o', { baseUrl: 'http://127.0.0.1:1', - maxRetries: 0, }) const err = (await t.throwsAsync(async () => { - for await (const _part of streamText(model, 'hello')) { + for await (const _part of streamText(model, 'hello', { max_retries: 0 })) { // never yields } })) as AimuxError @@ -116,9 +117,9 @@ test('APICallError carries every field the response produced', async (t) => { res.end(body) }) try { - const model = await native.openai('test-key', 'gpt-4o', { baseUrl: url, maxRetries: 0 }) + const model = await native.openai('test-key', 'gpt-4o', { baseUrl: url }) const err = (await t.throwsAsync(() => - model.generateText(JSON.stringify('hi')), + model.generateText(JSON.stringify('hi'), JSON.stringify({ max_retries: 0 })), )) as APICallError t.true(err instanceof APICallError) t.true(err instanceof native.APICallError) @@ -153,9 +154,9 @@ test('exhausted retries throw RetryError with the per-attempt history', async (t res.end(body) }) try { - const model = await native.openai('test-key', 'gpt-4o', { baseUrl: url, maxRetries: 1 }) + const model = await native.openai('test-key', 'gpt-4o', { baseUrl: url }) const err = (await t.throwsAsync(() => - model.generateText(JSON.stringify('hi')), + model.generateText(JSON.stringify('hi'), JSON.stringify({ max_retries: 1 })), )) as RetryError t.true(err instanceof RetryError) t.true(err instanceof AimuxError) @@ -209,7 +210,7 @@ test('typed initRecordingRing keeps core errors in the AimuxError hierarchy', (t }) test('malformed prompt JSON is a plain napi InvalidArg error, not a core error', async (t) => { - const model = await openai('test-key', 'gpt-4o', { baseUrl: 'http://127.0.0.1:1', maxRetries: 0 }) + const model = await openai('test-key', 'gpt-4o', { baseUrl: 'http://127.0.0.1:1' }) // The typed generateText() always stringifies; hit the raw JSON-text path. const err = (await t.throwsAsync(() => model.generateText('{not json'))) as Error & { code?: string @@ -229,3 +230,54 @@ test('sync wire-JSON failure is napi InvalidArg; business validation stays core- // Well-formed JSON that fails the schema is what the core would reject. t.true(t.throws(() => mockReplay('{}')) instanceof InvalidArgumentError) }) + +// ── LoadAPIKeyError / LoadSettingError ─────────────────────────────────────── + +test.serial('a missing API key is a LoadAPIKeyError at call time carrying the environment variable', async (t) => { + const saved = process.env.GROQ_API_KEY + delete process.env.GROQ_API_KEY + try { + const model = await provider('groq', null, 'some-model') + const err = (await t.throwsAsync(() => + model.generateText(JSON.stringify('hi'), JSON.stringify({ max_retries: 0 })), + )) as LoadAPIKeyError + t.true(err instanceof LoadAPIKeyError) + t.true(err instanceof AimuxError) + t.false(err instanceof InvalidArgumentError) + t.is(err.constructor, LoadAPIKeyError) + t.is(err.name, 'LoadAPIKeyError') + t.is(err.envVar, 'GROQ_API_KEY') + t.is(err.description, 'Groq') + t.regex(err.message, /GROQ_API_KEY/) + // Pre-HTTP: none of the ApiCall-only fields exist. + t.false('status' in err) + t.false('retryable' in err) + } finally { + if (saved !== undefined) process.env.GROQ_API_KEY = saved + } +}) + +test.serial('an explicit empty API key is not a LoadAPIKeyError', async (t) => { + // The AI SDK sends an explicit empty key verbatim; only `null` consults the env. + const model = await provider('groq', '', 'some-model') + t.truthy(model) +}) + +test.serial('a missing provider setting is a LoadSettingError at call time', async (t) => { + const saved = process.env.AZURE_RESOURCE_NAME + delete process.env.AZURE_RESOURCE_NAME + try { + const model = await native.azure('test-key', '', 'some-deployment') + const err = (await t.throwsAsync(() => + model.generateText(JSON.stringify('hi'), JSON.stringify({ max_retries: 0 })), + )) as LoadSettingError + t.true(err instanceof LoadSettingError) + t.true(err instanceof AimuxError) + t.is(err.constructor, LoadSettingError) + t.is(err.envVar, 'AZURE_RESOURCE_NAME') + t.is(err.settingName, 'resourceName') + t.false('status' in err) + } finally { + if (saved !== undefined) process.env.AZURE_RESOURCE_NAME = saved + } +}) diff --git a/bindings/node/__test__/index.test.ts b/bindings/node/__test__/index.test.ts index fedf568d0..029db1701 100644 --- a/bindings/node/__test__/index.test.ts +++ b/bindings/node/__test__/index.test.ts @@ -42,10 +42,13 @@ test('provider() creates a registry model instance (RFC-0017 phase 4)', async (t }) test('provider() with undefined apiKey reads the registry env var', async (t) => { - // GROQ_API_KEY is unset in CI — the construction must fail with a clear + // ABACUS_API_KEY is unset in CI. As in the AI SDK, the key is loaded when a + // request is made, so construction succeeds and the call fails with a clear // error (missing env key), proving the env-var path is wired. + delete process.env.ABACUS_API_KEY + const model = await provider('abacus', undefined, 'm') await t.throwsAsync( - () => provider('abacus', undefined, 'm'), + () => model.generateText(JSON.stringify('hi')), { message: /api key/i }, ) }) diff --git a/bindings/node/__test__/provider_config.test.ts b/bindings/node/__test__/provider_config.test.ts new file mode 100644 index 000000000..27d5baff3 --- /dev/null +++ b/bindings/node/__test__/provider_config.test.ts @@ -0,0 +1,248 @@ +// provider_config.test.ts — Node e2e tests for the provider factory config +// (headers) and the call-level maxRetries, plus the removed bodyOverrides and +// the provider-level maxRetries, which are rejected instead of ignored. +// +// These verify the full chain: JS → napi → Rust → HTTP mock. + +import test from 'ava' +import { createServer, type Server } from 'node:http' +import { + anthropic, + anthropicAws, + azure, + bedrock, + cohere, + createProvider, + deepseek, + google, + mistral, + openai, + provider, + vertex, + xai, + InvalidArgumentError, +} from '../src/native.ts' + +function startMockServer(handler: (req: any, res: any) => void): Promise<{ server: Server; url: string }> { + return new Promise((resolve) => { + const server = createServer(handler) + server.listen(0, '127.0.0.1', () => { + const addr = server.address() as any + resolve({ server, url: `http://127.0.0.1:${addr.port}` }) + }) + }) +} + +function closeServer(server: Server): Promise { + return new Promise((resolve) => server.close(() => resolve())) +} + +function readBody(req: any): Promise { + return new Promise((resolve) => { + let body = '' + req.on('data', (chunk: string) => { body += chunk }) + req.on('end', () => resolve(JSON.parse(body))) + }) +} + +const RESPONSES_RESPONSE = JSON.stringify({ + id: 'resp-test', + model: 'm', + created_at: 1, + status: 'completed', + output: [{ + type: 'message', + id: 'msg-test', + role: 'assistant', + status: 'completed', + content: [{ type: 'output_text', text: 'Done.', annotations: [] }], + }], + usage: { input_tokens: 10, output_tokens: 1, total_tokens: 11 }, +}) + +// ── removed: bodyOverrides ─────────────────────────────────────────────────── + +test('per-call body_overrides no longer reaches the request body', async (t) => { + let requestBody: any = null + const { server, url } = await startMockServer(async (req, res) => { + requestBody = await readBody(req) + res.writeHead(200, { 'content-type': 'application/json' }) + res.end(RESPONSES_RESPONSE) + }) + + try { + const model = await openai('test-key', 'm', url) + const opts = JSON.stringify({ + temperature: 0.9, + body_overrides: { temperature: 0.1, enable_thinking: false }, + }) + await model.generateText(JSON.stringify('Hello'), opts) + + t.is(requestBody.temperature, 0.9, 'the standard option stands') + t.true(requestBody.enable_thinking === undefined, 'the override is not applied') + } finally { + await closeServer(server) + } +}) + +test('factory config bodyOverrides is rejected, not ignored', async (t) => { + const err = await t.throwsAsync(() => + openai('test-key', 'm', { + baseUrl: 'http://127.0.0.1:1', + bodyOverrides: '{"X-Relay-Tag":"my-team"}', + }), + ) + t.true(err instanceof Error) + t.regex(err!.message, /body_overrides/) +}) + +test('factory config maxRetries is rejected, not ignored', async (t) => { + const err = await t.throwsAsync(() => + openai('test-key', 'm', { baseUrl: 'http://127.0.0.1:1', maxRetries: 0 }), + ) + t.true(err instanceof Error) + t.regex(err!.message, /max_retries/) +}) + +// Every native constructor takes the same config and rejects the two removed +// keys; none of them drops a key silently. +const NATIVE_CONSTRUCTORS: Array<[string, (config: object) => Promise]> = [ + ['openai', (c) => openai('k', 'm', c)], + ['anthropic', (c) => anthropic('k', 'm', c)], + ['google', (c) => google('k', 'm', c)], + ['cohere', (c) => cohere('k', 'm', c)], + ['mistral', (c) => mistral('k', 'm', c)], + ['xai', (c) => xai('k', 'm', c)], + ['deepseek', (c) => deepseek('k', 'm', c)], + ['bedrock', (c) => bedrock('AKIA', 'secret', 'us-east-1', 'm', c)], + ['vertex', (c) => vertex('token', 'proj', 'us-central1', 'm', c)], + ['anthropicAws', (c) => anthropicAws('k', 'us-east-1', 'm', c)], + ['azure', (c) => azure('k', 'res', 'dep', null, c)], + ['provider', (c) => provider('groq', 'k', 'm', c)], + ['createProvider', (c) => createProvider('groq', 'k', c)], +] + +for (const [name, construct] of NATIVE_CONSTRUCTORS) { + test(`${name}: maxRetries / bodyOverrides in the config are InvalidArgumentError`, async (t) => { + const maxRetries = await t.throwsAsync(() => construct({ maxRetries: 1 })) + t.true(maxRetries instanceof InvalidArgumentError) + t.regex(maxRetries!.message, /max_retries/) + const overrides = await t.throwsAsync(() => construct({ bodyOverrides: '{"a":1}' })) + t.true(overrides instanceof InvalidArgumentError) + t.regex(overrides!.message, /body_overrides/) + }) + + test(`${name}: a plain config is accepted`, async (t) => { + await t.notThrowsAsync(() => construct({ baseUrl: 'http://127.0.0.1:1' })) + }) +} + +// ── params: preset template parameters ─────────────────────────────────────── + +test('params fills a preset template parameter', async (t) => { + const saved = process.env.CLOUDFLARE_ACCOUNT_ID + delete process.env.CLOUDFLARE_ACCOUNT_ID + try { + // Without the parameter the preset cannot build its base URL (reported at call time). + const unfilled = await provider('cloudflare_workers_ai', 'k', 'm') + const missing = await t.throwsAsync(() => + unfilled.generateText(JSON.stringify('hi'), JSON.stringify({ max_retries: 0 })), + ) + t.true(missing instanceof InvalidArgumentError) + t.regex(missing!.message, /account_id/) + + await t.notThrowsAsync(() => + provider('cloudflare_workers_ai', 'k', 'm', { params: { account_id: 'acct123' } }), + ) + await t.notThrowsAsync(() => + createProvider('cloudflare_workers_ai', 'k', { params: { account_id: 'acct123' } }), + ) + } finally { + if (saved !== undefined) process.env.CLOUDFLARE_ACCOUNT_ID = saved + } +}) + +test('params the preset does not declare are rejected', async (t) => { + const err = await t.throwsAsync(() => + provider('abacus', 'k', 'm', { params: { account_id: 'acct123' } }), + ) + t.true(err instanceof InvalidArgumentError) + t.regex(err!.message, /no template parameter `account_id`/) +}) + +test('params values that could change the URL are rejected', async (t) => { + const err = await t.throwsAsync(() => + provider('cloudflare_workers_ai', 'k', 'm', { params: { account_id: 'a/b' } }), + ) + t.true(err instanceof InvalidArgumentError) +}) + +test('factory config headers are sent on every request', async (t) => { + let headers: any = null + const { server, url } = await startMockServer(async (req, res) => { + headers = req.headers + // consume body + await readBody(req) + res.writeHead(200, { 'content-type': 'application/json' }) + res.end(RESPONSES_RESPONSE) + }) + + try { + const model = await openai('test-key', 'm', { + baseUrl: url, + headers: JSON.stringify({ 'X-Custom-Header': 'custom-value' }), + }) + await model.generateText(JSON.stringify('Hello')) + + t.is(headers['x-custom-header'], 'custom-value') + } finally { + await closeServer(server) + } +}) + +// ── backward compatibility ─────────────────────────────────────────────────── + +test('factory accepts bare string baseUrl (backward compatible)', async (t) => { + const { server, url } = await startMockServer(async (req, res) => { + t.is(req.url, '/responses') + await readBody(req) + res.writeHead(200, { 'content-type': 'application/json' }) + res.end(RESPONSES_RESPONSE) + }) + + try { + // 3rd param is a plain string (old API) + const model = await openai('test-key', 'm', url) + const result = JSON.parse(await model.generateText(JSON.stringify('Hello'))) + t.is(result.text, 'Done.') + } finally { + await closeServer(server) + } +}) + +// ── maxRetries ──────────────────────────────────────────────────────────────── + +test('maxRetries: 0 disables retries (single request on 500)', async (t) => { + let requestCount = 0 + const { server, url } = await startMockServer(async (req, res) => { + requestCount++ + await readBody(req) + res.writeHead(500, { 'content-type': 'application/json' }) + res.end(JSON.stringify({ error: { message: 'Internal server error', type: 'server_error' } })) + }) + + try { + const model = await openai('test-key', 'm', url) + const opts = JSON.stringify({ max_retries: 0 }) + + await t.throwsAsync( + async () => model.generateText(JSON.stringify('Hello'), opts), + { message: /500|Internal|server_error/i }, + ) + + // With retries disabled, exactly 1 request should have been made. + t.is(requestCount, 1, 'should make exactly 1 request when maxRetries=0') + } finally { + await closeServer(server) + } +}) diff --git a/bindings/node/__test__/reasoning_replay.test.ts b/bindings/node/__test__/reasoning_replay.test.ts index 1983c6336..5c0a18454 100644 --- a/bindings/node/__test__/reasoning_replay.test.ts +++ b/bindings/node/__test__/reasoning_replay.test.ts @@ -17,7 +17,7 @@ import test from 'ava' import { createServer, type Server } from 'node:http' -import { openai } from '../src/native.ts' +import { deepseek } from '../src/native.ts' function startMockServer(handler: (req: any, res: any) => void): Promise<{ server: Server; url: string }> { return new Promise((resolve) => { @@ -63,7 +63,7 @@ test('issue 1: tool_result with legacy `output` field is accepted and carries to }) try { - const model = await openai('test-key', 'deepseek-v4-flash', url) + const model = await deepseek('test-key', 'deepseek-v4-flash', url) // The exact shape from the user's failing case: a tool message whose // ContentPart uses the legacy `output` field (Vercel AI SDK / 0.1.0 TS). @@ -114,7 +114,7 @@ test('issue 2: assistant reasoning + tool_call lifts reasoning to reasoning_cont }) try { - const model = await openai('test-key', 'deepseek-v4-flash', url) + const model = await deepseek('test-key', 'deepseek-v4-flash', url) // Reproduces the user's table row: assistant // [{reasoning},{tool_call}] + tool result. The first turn produced a @@ -166,7 +166,7 @@ test('issue 2: assistant reasoning + text (no tool calls) emits reasoning_conten }) try { - const model = await openai('test-key', 'deepseek-v4-flash', url) + const model = await deepseek('test-key', 'deepseek-v4-flash', url) const messages = [ { role: 'user', content: 'What is 2+2?' }, @@ -201,7 +201,7 @@ test('issue 2: assistant without reasoning omits reasoning_content', async (t) = }) try { - const model = await openai('test-key', 'deepseek-v4-flash', url) + const model = await deepseek('test-key', 'deepseek-v4-flash', url) // The tool call has its result: a prompt that leaves a tool call without // one is a MissingToolResultsError (AI SDK `convertToLanguageModelPrompt`). diff --git a/bindings/node/__test__/session.test.ts b/bindings/node/__test__/session.test.ts index 242773883..1cc0b8a45 100644 --- a/bindings/node/__test__/session.test.ts +++ b/bindings/node/__test__/session.test.ts @@ -9,7 +9,7 @@ import test from 'ava' import { createServer, type Server } from 'node:http' import { - openai, + deepseek, generateText, initSessionStore, initSessionInfer, @@ -54,7 +54,7 @@ test.serial('sessionId groups calls and query API returns typed results', async initSessionStore() initSessionInfer(false) // explicit ids only - const model = await openai('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) + const model = await deepseek('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) // Two calls in the same session. await generateText(model, 'first', { session_id: 'sess-1' }) @@ -91,7 +91,7 @@ test.serial('opt-in inferer groups prefix continuations into auto sessions', asy initSessionStore() initSessionInfer(true) - const model = await openai('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) + const model = await deepseek('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) await generateText(model, 'u1') await generateText(model, [{ role: 'user', content: 'u1' }, { role: 'assistant', content: 'a1' }, { role: 'user', content: 'u2' }]) diff --git a/bindings/node/__test__/trace.test.ts b/bindings/node/__test__/trace.test.ts index d52fc8ea5..69e58f915 100644 --- a/bindings/node/__test__/trace.test.ts +++ b/bindings/node/__test__/trace.test.ts @@ -8,7 +8,7 @@ import test from 'ava' import { createServer, type Server } from 'node:http' import { - openai, + deepseek, generateText, type GenerateTextResult, } from '../src/index.ts' @@ -26,7 +26,10 @@ const MOCK_RESPONSE = JSON.stringify({ prompt_tokens: 200, completion_tokens: 5, total_tokens: 205, - prompt_tokens_details: { cached_tokens: 128 }, + // DeepSeek reports cache hits in its own fields (@ai-sdk/deepseek reads + // prompt_cache_hit_tokens, not prompt_tokens_details.cached_tokens). + prompt_cache_hit_tokens: 128, + prompt_cache_miss_tokens: 72, }, }) @@ -47,7 +50,7 @@ test('trace() records calls and query API returns typed results', async (t) => { const { server, url } = await startMockServer() t.teardown(() => server.close()) - const raw = await openai('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) + const raw = await deepseek('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) const traced = raw.traceAudited(true) // strict mode // > 4 KiB user message so block-aligned prefixes actually match. @@ -69,7 +72,7 @@ test('trace() records calls and query API returns typed results', async (t) => { const statsJson = traced.traceAggregate() const stats = JSON.parse(statsJson) as any[] t.is(stats.length, 1) - t.is(stats[0].provider, 'openai') + t.is(stats[0].provider, 'deepseek.chat') t.is(stats[0].requests, 2) t.truthy(stats[0].reported_hit_rate) t.truthy(stats[0].client_upper_bound_hit_rate) @@ -99,7 +102,7 @@ test('default streams carry no Raw parts (RFC-0016 M2 contract)', async (t) => { const { server, url } = await startMockServer() t.teardown(() => server.close()) - const model = await openai('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) + const model = await deepseek('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) const gen = await model.streamText(JSON.stringify('hello'), undefined) const parts: string[] = [] for await (const json of gen) { @@ -119,7 +122,7 @@ test('non-traced model still generates normally', async (t) => { const { server, url } = await startMockServer() t.teardown(() => server.close()) - const model = await openai('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) + const model = await deepseek('sk-test-fake-key', 'gpt-4o-mini', { baseUrl: url }) const result: GenerateTextResult = await generateText(model, 'hello') t.truthy(result.text) t.is(result.usage.input_tokens, 200) diff --git a/bindings/node/__test__/wrapper.test.ts b/bindings/node/__test__/wrapper.test.ts index b6eb8f272..f101b9c4c 100644 --- a/bindings/node/__test__/wrapper.test.ts +++ b/bindings/node/__test__/wrapper.test.ts @@ -9,9 +9,9 @@ import test from 'ava' import { createServer, type Server } from 'node:http' -// The wrapper under test. It re-exports the raw `openai`/`anthropic` factories +// The wrapper under test. It re-exports the raw `deepseek`/`anthropic` factories // and the ts-rs types, so a consumer needs only this one import. -import { openai, anthropic, generateText, streamText } from '../src/index.ts' +import { deepseek, anthropic, generateText, streamText } from '../src/index.ts' import type { GenerateTextResult, TextStreamPart, @@ -144,7 +144,7 @@ test('wrapper: generateText returns a typed object (.text / .usage / .finish_rea }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) // Typed input (string prompt) + typed output — no manual JSON.stringify/parse. const result: GenerateTextResult = await generateText(model, 'What is Rust?') @@ -172,7 +172,7 @@ test('wrapper: generateText parses tool_calls + content ToolCall', async (t) => }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) // Typed options object: tools pass through to the provider. const result = await generateText(model, "What's the weather in Tokyo?", { tools: [weatherTool], @@ -202,7 +202,7 @@ test('wrapper: generateText passes tool_choice through to the provider', async ( const { server, url } = await startMockServer(cap.handler) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const toolChoice: ToolChoice = 'required' await generateText(model, 'Hello', { tools: [weatherTool], @@ -221,7 +221,7 @@ test('wrapper: generateText accepts typed multi-role messages', async (t) => { const { server, url } = await startMockServer(cap.handler) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) // Typed ModelMessage[] prompt — no manual JSON.stringify. const messages: ModelMessage[] = [ { role: 'system', content: 'You are a helpful assistant.' }, @@ -267,7 +267,7 @@ test('wrapper: streamText yields typed TextDelta parts', async (t) => { }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const parts: TextStreamPart[] = [] // streamText is an async generator of typed TextStreamParts — no JSON.parse. for await (const part of streamText(model, 'Say hello')) { @@ -295,7 +295,7 @@ test('wrapper: streamText emits Raw parts when include_raw_chunks set', async (t }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const parts: TextStreamPart[] = [] for await (const part of streamText(model, 'Say hello', { include_raw_chunks: true, @@ -328,7 +328,7 @@ test('wrapper: no Raw parts by default', async (t) => { }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const parts: TextStreamPart[] = [] for await (const part of streamText(model, 'Say hello')) { parts.push(part) @@ -346,7 +346,7 @@ test('wrapper: streamText yields typed ToolCall parts', async (t) => { }) try { - const model = await openai('test-key', 'gpt-4o', url) + const model = await deepseek('test-key', 'gpt-4o', url) const parts: TextStreamPart[] = [] for await (const part of streamText(model, "What's the weather?", { tools: [weatherTool], diff --git a/bindings/node/index.d.ts b/bindings/node/index.d.ts index 4e4b2055a..14c294876 100644 --- a/bindings/node/index.d.ts +++ b/bindings/node/index.d.ts @@ -391,10 +391,9 @@ export declare function googleVideo(apiKey: string, modelId: string, baseUrl?: s export declare function initLogging(level: string): void /** - * Set the global proxy configuration (M6, RFC-0016). Must be called before - * the first `generateText` / `streamText` call; a no-op if the shared HTTP - * client is already initialised (the shared client is lazily built on first - * use and locked for the process lifetime). + * Replace the default transport with one routed by this proxy configuration + * (M6, RFC-0016). Requests started afterwards use it; calling it again + * replaces the previous configuration. * * `configJson` shape: `{ "http_url": "...", "https_url": "...", "all_url": * "...", "no_proxy": "..." }` (all fields optional; omitting all is equivalent @@ -496,13 +495,23 @@ export interface ProviderConfig { organization?: string /** OpenAI project ID (sent via `OpenAI-Project` header). */ project?: string - /** Override the provider's retry count. `0` disables retries. */ + /** + * Values of a preset's template parameters, named as the registry + * declares them (e.g. `{ account_id: '…' }`), for providers created by + * name (`provider()` / `createProvider()`). Only parameters the preset + * declares are accepted. + */ + params?: Record + /** + * Rejected with `InvalidArgumentError`: retry is a per-call setting + * (`maxRetries` in the call options). The field exists only so passing it + * is reported instead of ignored. + */ maxRetries?: number /** - * Provider-level request body overrides as a JSON string (deep-merged - * into every request). Per-call `bodyOverrides` in GenerateTextOptions - * takes precedence. Pass a JSON object string, e.g. - * `'{"enable_thinking": false}'`. + * Rejected with `InvalidArgumentError`: request-body overrides were + * removed. The field exists only so passing it is reported instead of + * ignored. */ bodyOverrides?: string } diff --git a/bindings/node/src/error.rs b/bindings/node/src/error.rs index 8cfaa7f24..4356b3142 100644 --- a/bindings/node/src/error.rs +++ b/bindings/node/src/error.rs @@ -25,6 +25,8 @@ const ERROR_CLASS_NAMES: &[&str] = &[ "ToolCallRepairError", "InvalidArgumentError", "InvalidPromptError", + "LoadAPIKeyError", + "LoadSettingError", "TokenExpiredError", "UnsupportedFunctionalityError", "NoSuchModelError", @@ -191,6 +193,8 @@ fn aimux_error_class_name(error: &AiMuxError) -> &'static str { AiMuxError::InvalidPrompt(_) | AiMuxError::MissingToolResults { .. } => { "InvalidPromptError" } + AiMuxError::LoadApiKey { .. } => "LoadAPIKeyError", + AiMuxError::LoadSetting { .. } => "LoadSettingError", AiMuxError::TokenExpired(_) => "TokenExpiredError", AiMuxError::UnsupportedFunctionality(_) => "UnsupportedFunctionalityError", AiMuxError::NoSuchModel { .. } => "NoSuchModelError", @@ -297,8 +301,27 @@ fn aimux_error_object<'env>(env: &'env Env, error: &AiMuxError) -> NapiResult { + AiMuxError::NoSuchProvider { + provider_id, + model_id, + model_type, + available_providers, + } => { obj.set("providerId", provider_id.as_str())?; + obj.set("modelId", model_id.as_str())?; + obj.set("modelType", model_type.as_str())?; + obj.set("availableProviders", available_providers.clone())?; + } + AiMuxError::LoadApiKey { + env_var, + description, + } => { + obj.set("envVar", env_var.as_str())?; + obj.set("description", description.as_str())?; + } + AiMuxError::LoadSetting { env_var, name } => { + obj.set("envVar", env_var.as_str())?; + obj.set("settingName", name.as_str())?; } AiMuxError::NoSuchTool { tool_name, @@ -322,7 +345,7 @@ fn aimux_error_object<'env>(env: &'env Env, error: &AiMuxError) -> NapiResult { - obj.set("originalError", &aimux_error_object(env, original_error)?)?; + obj.set("originalError", aimux_error_object(env, original_error)?)?; } _ => {} } diff --git a/bindings/node/src/error.ts b/bindings/node/src/error.ts index 74ff79358..d5066936a 100644 --- a/bindings/node/src/error.ts +++ b/bindings/node/src/error.ts @@ -102,6 +102,26 @@ export class ToolCallRepairError extends AimuxError { } export class InvalidArgumentError extends AimuxError {} export class InvalidPromptError extends AimuxError {} +/** + * No API key was passed and the fallback environment variable is unset (AI SDK + * `LoadAPIKeyError`). Raised when the request is made, before any HTTP call. + */ +export class LoadAPIKeyError extends AimuxError { + /** The environment variable consulted as the fallback, e.g. `'OPENAI_API_KEY'`. */ + declare readonly envVar: string + /** What the key is for, e.g. `'OpenAI'`. */ + declare readonly description: string +} +/** + * A required provider setting was not passed and its fallback environment + * variable is unset (AI SDK `LoadSettingError`), e.g. the AWS region. + */ +export class LoadSettingError extends AimuxError { + /** The environment variable consulted as the fallback. */ + declare readonly envVar: string + /** The setting's parameter name, e.g. `'region'`. */ + declare readonly settingName: string +} export class TokenExpiredError extends AimuxError { /** Token expiry is produced only from an observed 401 response. */ declare readonly status: 401 diff --git a/bindings/node/src/index.ts b/bindings/node/src/index.ts index 2d410fa21..68fcee0bc 100644 --- a/bindings/node/src/index.ts +++ b/bindings/node/src/index.ts @@ -109,6 +109,8 @@ export { ToolCallRepairError, InvalidArgumentError, InvalidPromptError, + LoadAPIKeyError, + LoadSettingError, TokenExpiredError, UnsupportedFunctionalityError, NoSuchModelError, diff --git a/bindings/node/src/lib.rs b/bindings/node/src/lib.rs index a3e9b7cb3..d3560a429 100644 --- a/bindings/node/src/lib.rs +++ b/bindings/node/src/lib.rs @@ -35,6 +35,7 @@ use aimux_core::message::ModelPrompt; use aimux_core::openai_output::OpenAiStreamOptions; use aimux_core::parse_tool_call::{ToolCallRepairReply, tool_call_repair_inputs}; use aimux_core::tool::ToolCall; +use aimux_providers::provider as providers; use napi::bindgen_prelude::*; use napi_derive::napi; @@ -573,8 +574,6 @@ impl Model { } } -type StreamReceiver = tokio::sync::mpsc::Receiver>; - /// AsyncGenerator that yields StreamPart JSON strings. /// /// The stream is started eagerly in `stream_text()` — a tokio task drives @@ -582,9 +581,16 @@ type StreamReceiver = tokio::sync::mpsc::Receiver>>, + rx: SharedStreamReceiver, } +/// The channel end `next()` drains; `None` once the stream is exhausted. +type SharedStreamReceiver = std::sync::Arc< + tokio::sync::Mutex< + Option>>, + >, +>; + #[napi] impl AsyncGenerator for StreamTextGenerator { type Yield = StreamItem; @@ -635,42 +641,81 @@ pub struct ProviderConfig { pub organization: Option, /// OpenAI project ID (sent via `OpenAI-Project` header). pub project: Option, - /// Override the provider's retry count. `0` disables retries. + /// Values of a preset's template parameters, named as the registry + /// declares them (e.g. `{ account_id: '…' }`), for providers created by + /// name (`provider()` / `createProvider()`). Only parameters the preset + /// declares are accepted. + pub params: Option>, + /// Rejected with `InvalidArgumentError`: retry is a per-call setting + /// (`maxRetries` in the call options). The field exists only so passing it + /// is reported instead of ignored. pub max_retries: Option, - /// Provider-level request body overrides as a JSON string (deep-merged - /// into every request). Per-call `bodyOverrides` in GenerateTextOptions - /// takes precedence. Pass a JSON object string, e.g. - /// `'{"enable_thinking": false}'`. + /// Rejected with `InvalidArgumentError`: request-body overrides were + /// removed. The field exists only so passing it is reported instead of + /// ignored. pub body_overrides: Option, } -/// Convert a ProviderConfig (from JS) into OpenAIConfig builder steps. -fn apply_provider_config_openai( - mut config: aimux_providers::openai::OpenAIConfig, - cfg: &ProviderConfig, -) -> MResult { - if let Some(url) = &cfg.base_url { - config = config.with_base_url(url); - } - if let Some(ref json_str) = cfg.headers { - let h: std::collections::HashMap = - parse_wire_json("config.headers", json_str)?; - config = config.with_headers(h); - } - if let Some(ref org) = cfg.organization { - config = config.with_org_id(org); - } - if let Some(ref proj) = cfg.project { - config = config.with_project(proj); +impl ProviderConfig { + /// Fail on the two removed keys instead of dropping them silently. + fn reject_removed_keys(&self) -> MResult<()> { + providers::reject_removed_provider_options( + self.max_retries.is_some(), + self.body_overrides.is_some(), + ) + .map_err(|e| AiMuxBindingError::from(&e)) } - if let Some(max) = cfg.max_retries { - config.retry_config.max_retries = max; +} + +/// What the native constructors take from their 3rd argument (a base URL +/// string, or a [`ProviderConfig`]). Every native constructor goes through +/// here, so `maxRetries` / `bodyOverrides` are rejected for all of them and +/// none of them drops a key silently. +struct NativeConfig { + base_url: Option, + headers: Option, + /// OpenAI only (`OpenAI-Organization` / `OpenAI-Project`). + organization: Option, + project: Option, +} + +fn native_config(config: Option>) -> MResult { + let mut native = NativeConfig { + base_url: None, + headers: None, + organization: None, + project: None, + }; + match config { + None => {} + Some(Either::A(url)) => native.base_url = Some(url), + Some(Either::B(cfg)) => { + cfg.reject_removed_keys()?; + native.base_url = cfg.base_url; + if let Some(json_str) = &cfg.headers { + let h: std::collections::HashMap = + parse_wire_json("config.headers", json_str)?; + native.headers = Some(h.into_iter().map(|(k, v)| (k, Some(v))).collect()); + } + native.organization = cfg.organization; + native.project = cfg.project; + } } - if let Some(ref json_str) = cfg.body_overrides { - let overrides: serde_json::Value = parse_wire_json("config.bodyOverrides", json_str)?; - config = config.with_body_overrides(overrides); + Ok(native) +} + +/// OpenAI provider settings for an explicit key handed over by the host, with +/// an optional base URL. The key is an explicit value (`""` included), so it +/// never falls back to `OPENAI_API_KEY`. +fn openai_settings( + api_key: String, + base_url: Option, +) -> aimux_providers::openai::OpenAIProviderSettings { + aimux_providers::openai::OpenAIProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() } - Ok(config) } /// 初始化全局日志（RFC-0014）。幂等：多次调用无副作用；宿主已自建 @@ -722,7 +767,7 @@ pub fn register_providers(config_json: String) -> error::AimuxResult<()> { // schema mismatch as `JsonParse`, which is core's InvalidArgument // here (the text already parsed) — JsonParse is for provider responses. let _: serde_json::Value = parse_wire_json("config_json", &config_json)?; - aimux_providers::load_providers_from_json(&config_json).map_err(|e| match e { + providers::load_providers_from_json(&config_json).map_err(|e| match e { AiMuxError::JsonParse(m) => crate::error::AiMuxBindingError::from( &AiMuxError::InvalidArgument(format!("config_json: {m}")), ), @@ -731,10 +776,9 @@ pub fn register_providers(config_json: String) -> error::AimuxResult<()> { })()) } -/// Set the global proxy configuration (M6, RFC-0016). Must be called before -/// the first `generateText` / `streamText` call; a no-op if the shared HTTP -/// client is already initialised (the shared client is lazily built on first -/// use and locked for the process lifetime). +/// Replace the default transport with one routed by this proxy configuration +/// (M6, RFC-0016). Requests started afterwards use it; calling it again +/// replaces the previous configuration. /// /// `configJson` shape: `{ "http_url": "...", "https_url": "...", "all_url": /// "...", "no_proxy": "..." }` (all fields optional; omitting all is equivalent @@ -747,9 +791,11 @@ pub fn init_proxy(config_json: String) -> error::AimuxResult<()> { AimuxResult((|| -> error::MResult<()> { let config: aimux_provider_utils::ProxyConfig = parse_wire_json("config_json", &config_json)?; - // `init_proxy` returns false when the shared client is already up; treat - // that as success (idempotent) so callers don't reason about ordering. - let _ = aimux_provider_utils::init_proxy(config); + // Replaces the default transport for later requests (see + // `aimux_init_proxy`). + aimux_provider_utils::set_default_transport(aimux_provider_utils::ReqwestFetch::new( + config, + )); Ok(()) })()) } @@ -951,19 +997,15 @@ struct RouterFfiConfig { /// or no store is registered. #[napi] pub fn session_calls(session_id: String) -> AimuxResult { - AimuxResult({ - let calls = aimux_core::session::session_calls(&session_id); - serialize_result(&calls) - }) + AimuxResult(serialize_result(&aimux_core::session::session_calls( + &session_id, + ))) } /// Query: all known sessions (RFC-0024), as a JSON-serialized `SessionView[]`. #[napi] pub fn list_sessions() -> AimuxResult { - AimuxResult({ - let views = aimux_core::session::list_sessions(); - serialize_result(&views) - }) + AimuxResult(serialize_result(&aimux_core::session::list_sessions())) } /// Create an OpenAI model instance. @@ -976,24 +1018,19 @@ pub async fn openai( AimuxResult({ let __r: crate::error::MResult = async { use aimux_core::provider::Provider; - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - - let mut cfg = OpenAIConfig::new(api_key); - match config { - Some(Either::A(url)) => { - cfg = cfg.with_base_url(url); - } - Some(Either::B(opts)) => { - cfg = apply_provider_config_openai(cfg, &opts)?; - } - None => {} - } - let provider = OpenAIProvider::new(cfg); - let model = provider - .language_model(&model_id) - .map_err(|e| AiMuxBindingError::from(&e))?; + use aimux_providers::openai::create_openai; + + let mut settings = openai_settings(api_key, None); + let native = native_config(config)?; + settings.base_url = native.base_url; + settings.headers = native.headers; + settings.organization = native.organization; + settings.project = native.project; + let provider = create_openai(settings).map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: provider + .language_model(&model_id) + .map_err(|e| AiMuxBindingError::from(&e))?, trace_store: None, }) } @@ -1012,39 +1049,21 @@ pub async fn anthropic( AimuxResult({ let __r: crate::error::MResult = async { use aimux_core::provider::Provider; - use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider}; + use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; - let mut cfg = AnthropicConfig::new(api_key); - match config { - Some(Either::A(url)) => { - cfg = cfg.with_base_url(url); - } - Some(Either::B(opts)) => { - if let Some(url) = &opts.base_url { - cfg = cfg.with_base_url(url); - } - if let Some(ref json_str) = opts.headers { - let h: std::collections::HashMap = - parse_wire_json("config.headers", json_str)?; - cfg = cfg.with_headers(h); - } - if let Some(max) = opts.max_retries { - cfg.retry_config.max_retries = max; - } - if let Some(ref json_str) = opts.body_overrides { - let overrides: serde_json::Value = - parse_wire_json("config.bodyOverrides", json_str)?; - cfg = cfg.with_body_overrides(overrides); - } - } - None => {} - } - let provider = AnthropicProvider::new(cfg); + let mut settings = AnthropicProviderSettings { + api_key: Some(api_key), + ..Default::default() + }; + let native = native_config(config)?; + settings.base_url = native.base_url; + settings.headers = native.headers; + let provider = create_anthropic(settings).map_err(|e| AiMuxBindingError::from(&e))?; let model = provider .language_model(&model_id) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1063,10 +1082,10 @@ pub async fn deepseek( AimuxResult({ let __r: crate::error::MResult = async { let options = provider_options_from_config(config)?; - let model = aimux_providers::provider("deepseek", Some(api_key), &model_id, options) + let model = providers::provider("deepseek", Some(api_key), &model_id, options) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1089,26 +1108,21 @@ pub async fn google( AimuxResult({ let __r: crate::error::MResult = async { use aimux_core::provider::Provider; - use aimux_providers::google::{GoogleConfig, GoogleProvider}; + use aimux_providers::google::{GoogleProviderSettings, create_google}; - let mut cfg = GoogleConfig::new(api_key); - match config { - Some(Either::A(url)) => { - cfg = cfg.with_base_url(url); - } - Some(Either::B(opts)) => { - if let Some(url) = &opts.base_url { - cfg = cfg.with_base_url(url); - } - } - None => {} - } - let provider = GoogleProvider::new(cfg); + let mut settings = GoogleProviderSettings { + api_key: Some(api_key), + ..Default::default() + }; + let native = native_config(config)?; + settings.base_url = native.base_url; + settings.headers = native.headers; + let provider = create_google(settings).map_err(|e| AiMuxBindingError::from(&e))?; let model = provider .language_model(&model_id) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1127,26 +1141,21 @@ pub async fn cohere( AimuxResult({ let __r: crate::error::MResult = async { use aimux_core::provider::Provider; - use aimux_providers::cohere::{CohereConfig, CohereProvider}; + use aimux_providers::cohere::{CohereProviderSettings, create_cohere}; - let mut cfg = CohereConfig::new(api_key); - match config { - Some(Either::A(url)) => { - cfg = cfg.with_base_url(url); - } - Some(Either::B(opts)) => { - if let Some(url) = &opts.base_url { - cfg = cfg.with_base_url(url); - } - } - None => {} - } - let provider = CohereProvider::new(cfg); + let mut cfg = CohereProviderSettings { + api_key: Some(api_key), + ..Default::default() + }; + let native = native_config(config)?; + cfg.base_url = native.base_url; + cfg.headers = native.headers; + let provider = create_cohere(cfg).map_err(|e| AiMuxBindingError::from(&e))?; let model = provider .language_model(&model_id) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1165,26 +1174,21 @@ pub async fn mistral( AimuxResult({ let __r: crate::error::MResult = async { use aimux_core::provider::Provider; - use aimux_providers::mistral::{MistralConfig, MistralProvider}; + use aimux_providers::mistral::{MistralProviderSettings, create_mistral}; - let mut cfg = MistralConfig::new(api_key); - match config { - Some(Either::A(url)) => { - cfg = cfg.with_base_url(url); - } - Some(Either::B(opts)) => { - if let Some(url) = &opts.base_url { - cfg = cfg.with_base_url(url); - } - } - None => {} - } - let provider = MistralProvider::new(cfg); + let mut cfg = MistralProviderSettings { + api_key: Some(api_key), + ..Default::default() + }; + let native = native_config(config)?; + cfg.base_url = native.base_url; + cfg.headers = native.headers; + let provider = create_mistral(cfg).map_err(|e| AiMuxBindingError::from(&e))?; let model = provider .language_model(&model_id) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1203,26 +1207,21 @@ pub async fn xai( AimuxResult({ let __r: crate::error::MResult = async { use aimux_core::provider::Provider; - use aimux_providers::xai::{XAIConfig, XAIProvider}; + use aimux_providers::xai::{XAIProviderSettings, create_xai}; - let mut cfg = XAIConfig::new(api_key); - match config { - Some(Either::A(url)) => { - cfg = cfg.with_base_url(url); - } - Some(Either::B(opts)) => { - if let Some(url) = &opts.base_url { - cfg = cfg.with_base_url(url); - } - } - None => {} - } - let provider = XAIProvider::new(cfg); + let mut cfg = XAIProviderSettings { + api_key: Some(api_key), + ..Default::default() + }; + let native = native_config(config)?; + cfg.base_url = native.base_url; + cfg.headers = native.headers; + let provider = create_xai(cfg).map_err(|e| AiMuxBindingError::from(&e))?; let model = provider .language_model(&model_id) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1243,27 +1242,30 @@ pub async fn bedrock( AimuxResult({ let __r: crate::error::MResult = async { use aimux_core::provider::Provider; - use aimux_providers::bedrock::{BedrockProvider, BedrockProviderConfig}; - - let mut cfg = BedrockProviderConfig::new(access_key_id, secret_access_key, region); - if let Some(cfg_config) = config { - match cfg_config { - Either::A(url) => { - cfg = cfg.with_base_url(url); - } - Either::B(opts) => { - if let Some(url) = &opts.base_url { - cfg = cfg.with_base_url(url); - } - } - } - } - let provider = BedrockProvider::new(cfg); + use aimux_providers::bedrock::{AmazonBedrockProviderSettings, create_amazon_bedrock}; + + let mut settings = AmazonBedrockProviderSettings { + credential_provider: Some(aimux_provider_utils::Resolvable::Value( + aimux_provider_utils::AwsCredentials { + access_key_id, + secret_access_key, + session_token: None, + region: region.clone(), + }, + )), + region: Some(region), + ..Default::default() + }; + let native = native_config(config)?; + settings.base_url = native.base_url; + settings.headers = native.headers; + let provider = + create_amazon_bedrock(settings).map_err(|e| AiMuxBindingError::from(&e))?; let model = provider .language_model(&model_id) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1284,27 +1286,24 @@ pub async fn vertex( AimuxResult({ let __r: crate::error::MResult = async { use aimux_core::provider::Provider; - use aimux_providers::vertex::{VertexProvider, VertexProviderConfig}; + use aimux_providers::vertex::{VertexProviderSettings, create_google_vertex}; - let mut cfg = VertexProviderConfig::new(access_token, project, location); - if let Some(cfg_config) = config { - match cfg_config { - Either::A(url) => { - cfg = cfg.with_base_url(url); - } - Either::B(opts) => { - if let Some(url) = &opts.base_url { - cfg = cfg.with_base_url(url); - } - } - } - } - let provider = VertexProvider::new(cfg); + let mut settings = VertexProviderSettings { + access_token: Some(aimux_provider_utils::Resolvable::Value(access_token)), + project: Some(project), + location: Some(location), + ..Default::default() + }; + let native = native_config(config)?; + settings.base_url = native.base_url; + settings.headers = native.headers.map(aimux_provider_utils::Resolvable::Value); + let provider = + create_google_vertex(settings).map_err(|e| AiMuxBindingError::from(&e))?; let model = provider .language_model(&model_id) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1325,28 +1324,26 @@ pub async fn anthropic_aws( let __r: crate::error::MResult = async { use aimux_core::provider::Provider; use aimux_providers::anthropic_aws::{ - AnthropicAwsProvider, AnthropicAwsProviderConfig, + AnthropicAwsAuth, AnthropicAwsProviderSettings, create_anthropic_aws, }; - let mut cfg = AnthropicAwsProviderConfig::with_api_key(api_key, region); - if let Some(cfg_config) = config { - match cfg_config { - Either::A(url) => { - cfg = cfg.with_base_url(url); - } - Either::B(opts) => { - if let Some(url) = &opts.base_url { - cfg = cfg.with_base_url(url); - } - } - } - } - let provider = AnthropicAwsProvider::new(cfg); + let mut settings = AnthropicAwsProviderSettings { + region: Some(region), + auth: Some(AnthropicAwsAuth::ApiKey( + aimux_provider_utils::Resolvable::Value(api_key), + )), + ..Default::default() + }; + let native = native_config(config)?; + settings.base_url = native.base_url; + settings.headers = native.headers; + let provider = + create_anthropic_aws(settings).map_err(|e| AiMuxBindingError::from(&e))?; let model = provider .language_model(&model_id) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1369,34 +1366,31 @@ pub async fn azure( AimuxResult({ let __r: crate::error::MResult = async { use aimux_core::provider::Provider; - use aimux_providers::azure::{AzureConfig, AzureProvider}; + use aimux_providers::azure::{AzureOpenAIProviderSettings, create_azure}; - let mut cfg = AzureConfig::new().with_api_key(api_key); - match config { - Some(Either::A(url)) => { - cfg = cfg.with_base_url(url); - } - Some(Either::B(opts)) => { - if let Some(url) = &opts.base_url { - cfg = cfg.with_base_url(url); - } - } - None => {} - } + let mut cfg = AzureOpenAIProviderSettings { + api_key: Some(aimux_provider_utils::Resolvable::Value(api_key)), + ..Default::default() + }; + let native = native_config(config)?; + cfg.base_url = native.base_url; + cfg.headers = native.headers; + // A dated `api_version` belongs to the deployment URL form. if let Some(version) = api_version && !version.is_empty() { - cfg = cfg.with_api_version(version); + cfg.api_version = Some(version); + cfg.use_deployment_based_urls = true; } if !resource_name.is_empty() { - cfg = cfg.with_resource_name(resource_name); + cfg.resource_name = Some(resource_name); } - let provider = AzureProvider::new(cfg).map_err(|e| AiMuxBindingError::from(&e))?; + let provider = create_azure(cfg).map_err(|e| AiMuxBindingError::from(&e))?; let model = provider .language_model(&deployment) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1421,10 +1415,10 @@ pub async fn provider( Some(cfg) => provider_options_from_config(Some(Either::B(cfg)))?, None => None, }; - let model = aimux_providers::provider(&name, api_key, &model_id, options) + let model = providers::provider(&name, api_key, &model_id, options) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(model), + inner: model, trace_store: None, }) } @@ -1477,6 +1471,12 @@ impl ProviderHandle { let __r: crate::error::MResult = async { let models = self .inner + .discovery() + .ok_or_else(|| { + AiMuxBindingError::from(AiMuxError::UnsupportedFunctionality( + "provider does not support listing models".into(), + )) + })? .list_models() .await .map_err(|e| AiMuxBindingError::from(&e))?; @@ -1498,7 +1498,7 @@ impl ProviderHandle { .language_model(&model_id) .map_err(|e| AiMuxBindingError::from(&e))?; Ok(Model { - inner: Arc::from(m), + inner: m, // A model built from a provider handle starts untraced; enable // probing via `Model::trace()` (same as Python binding). trace_store: None, @@ -1529,11 +1529,9 @@ pub async fn create_provider( Some(cfg) => provider_options_from_config(Some(Either::B(cfg)))?, None => None, }; - let p = aimux_providers::provider_handle(&name, api_key, options) + let inner = providers::provider_handle(&name, api_key, options) .map_err(|e| AiMuxBindingError::from(&e))?; - Ok(ProviderHandle { - inner: Arc::from(p), - }) + Ok(ProviderHandle { inner }) } .await; __r @@ -1560,15 +1558,16 @@ pub async fn get_model_specs(source_url: Option) -> AimuxResult /// Build `ProviderOptions` from a Node `ProviderConfig` (3rd factory arg). fn provider_options_from_config( config: Option>, -) -> MResult> { +) -> MResult> { let opts = match config { None => None, - Some(Either::A(url)) => Some(aimux_providers::ProviderOptions { + Some(Either::A(url)) => Some(providers::ProviderOptions { base_url: Some(url), ..Default::default() }), Some(Either::B(cfg)) => { - let mut o = aimux_providers::ProviderOptions::default(); + cfg.reject_removed_keys()?; + let mut o = providers::ProviderOptions::default(); if let Some(url) = cfg.base_url { o.base_url = Some(url); } @@ -1581,12 +1580,7 @@ fn provider_options_from_config( if let Some(proj) = cfg.project { o.project = Some(proj); } - if let Some(max) = cfg.max_retries { - o.max_retries = Some(max); - } - if let Some(ref json_str) = cfg.body_overrides { - o.body_overrides = Some(parse_wire_json("config.bodyOverrides", json_str)?); - } + o.params = cfg.params; Some(o) } }; diff --git a/bindings/node/src/multimodal.rs b/bindings/node/src/multimodal.rs index f3818c936..a9a9f401e 100644 --- a/bindings/node/src/multimodal.rs +++ b/bindings/node/src/multimodal.rs @@ -451,13 +451,10 @@ pub async fn openai_embedding( ) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); - let model = provider.embedding_model(&model_id); + use aimux_providers::openai::create_openai; + let provider = create_openai(crate::openai_settings(api_key, base_url)) + .map_err(|e| AiMuxBindingError::from(&e))?; + let model = provider.embedding(&model_id); Ok(EmbeddingModel { inner: Arc::new(model), }) @@ -476,12 +473,9 @@ pub async fn openai_speech( ) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); + use aimux_providers::openai::create_openai; + let provider = create_openai(crate::openai_settings(api_key, base_url)) + .map_err(|e| AiMuxBindingError::from(&e))?; let model = provider.speech(&model_id); Ok(SpeechModel { inner: Arc::new(model), @@ -501,12 +495,9 @@ pub async fn openai_image( ) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); + use aimux_providers::openai::create_openai; + let provider = create_openai(crate::openai_settings(api_key, base_url)) + .map_err(|e| AiMuxBindingError::from(&e))?; let model = provider.image(&model_id); Ok(ImageModel { inner: Arc::new(model), @@ -526,12 +517,9 @@ pub async fn openai_transcription( ) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); + use aimux_providers::openai::create_openai; + let provider = create_openai(crate::openai_settings(api_key, base_url)) + .map_err(|e| AiMuxBindingError::from(&e))?; let model = provider.transcription(&model_id); Ok(TranscriptionModel { inner: Arc::new(model), @@ -547,12 +535,9 @@ pub async fn openai_transcription( pub async fn openai_files(api_key: String, base_url: Option) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); + use aimux_providers::openai::create_openai; + let provider = create_openai(crate::openai_settings(api_key, base_url)) + .map_err(|e| AiMuxBindingError::from(&e))?; let files = provider.files(); Ok(Files { inner: Arc::new(files), @@ -572,13 +557,14 @@ pub async fn cohere_embedding( ) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::cohere::{CohereConfig, CohereProvider}; - let mut config = CohereConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = CohereProvider::new(config); - let model = provider.embedding_model(&model_id); + use aimux_providers::cohere::{CohereProviderSettings, create_cohere}; + let provider = create_cohere(CohereProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) + .map_err(|e| AiMuxBindingError::from(&e))?; + let model = provider.embedding(&model_id); Ok(EmbeddingModel { inner: Arc::new(model), }) @@ -597,13 +583,14 @@ pub async fn cohere_reranking( ) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::cohere::{CohereConfig, CohereProvider}; - let mut config = CohereConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = CohereProvider::new(config); - let model = provider.reranking_model(&model_id); + use aimux_providers::cohere::{CohereProviderSettings, create_cohere}; + let provider = create_cohere(CohereProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) + .map_err(|e| AiMuxBindingError::from(&e))?; + let model = provider.reranking(&model_id); Ok(RerankingModel { inner: Arc::new(model), }) @@ -622,13 +609,14 @@ pub async fn google_embedding( ) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::google::{GoogleConfig, GoogleProvider}; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = GoogleProvider::new(config); - let model = provider.embedding_model(&model_id); + use aimux_providers::google::{GoogleProviderSettings, create_google}; + let provider = create_google(GoogleProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) + .map_err(|e| AiMuxBindingError::from(&e))?; + let model = provider.embedding(&model_id); Ok(EmbeddingModel { inner: Arc::new(model), }) @@ -647,12 +635,13 @@ pub async fn google_image( ) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::google::{GoogleConfig, GoogleProvider}; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = GoogleProvider::new(config); + use aimux_providers::google::{GoogleProviderSettings, create_google}; + let provider = create_google(GoogleProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) + .map_err(|e| AiMuxBindingError::from(&e))?; let model = provider.image(&model_id); Ok(ImageModel { inner: Arc::new(model), @@ -672,12 +661,13 @@ pub async fn google_video( ) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::google::{GoogleConfig, GoogleProvider}; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = GoogleProvider::new(config); + use aimux_providers::google::{GoogleProviderSettings, create_google}; + let provider = create_google(GoogleProviderSettings { + api_key: Some(api_key), + base_url, + ..Default::default() + }) + .map_err(|e| AiMuxBindingError::from(&e))?; let model = provider.video(&model_id); Ok(VideoModel { inner: Arc::new(model), @@ -693,12 +683,13 @@ pub async fn google_video( pub async fn tavily_search(api_key: String, base_url: Option) -> AimuxResult { AimuxResult({ let __r: crate::error::MResult = async { - use aimux_providers::tavily::{TavilyConfig, TavilyProvider}; - let mut config = TavilyConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = TavilyProvider::new(config); + use aimux_providers::tavily::{TavilyProviderSettings, create_tavily}; + let provider = create_tavily(TavilyProviderSettings { + api_key: Some(aimux_provider_utils::Resolvable::Value(api_key)), + base_url, + ..Default::default() + }) + .map_err(|e| AiMuxBindingError::from(&e))?; let model = provider.search_model(); Ok(SearchModel { inner: Arc::new(model), @@ -830,7 +821,7 @@ pub async fn start_transcription_session( } }; match tx.try_send(Ok(json.clone())) { - Ok(()) => {}, + Ok(()) => {} Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { // Full: block on send (waits for // capacity) unless aborted. @@ -845,7 +836,6 @@ pub async fn start_transcription_session( return; } } - } } Err(e) => { @@ -853,18 +843,17 @@ pub async fn start_transcription_session( // try_send + abort-select (a full channel must not // stall; abort covers the session-drop path). match tx.try_send(Err(AiMuxBindingError::from(&e))) { - Ok(()) => {}, + Ok(()) => {} Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { tokio::select! { - _ = effective.cancelled() => return, + _ = effective.cancelled() => {} res = tx.send(Err(AiMuxBindingError::from(&e))) => { - if res.is_err() { return; } + let _ = res; } } } - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => return, + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {} } - } } }); diff --git a/bindings/node/src/types.ts b/bindings/node/src/types.ts index 4c777e8ac..a4079fce2 100644 --- a/bindings/node/src/types.ts +++ b/bindings/node/src/types.ts @@ -84,7 +84,6 @@ export type * from './types/LanguageModelV4ToolChoice' export type * from './types/Modality' export type * from './types/ModelCapabilities' export type * from './types/ModelCost' -export type * from './types/ModelId' export type * from './types/ModelLimits' export type * from './types/ModelMessage' export type * from './types/ModelModalities' @@ -140,6 +139,7 @@ export type * from './types/SessionSource' export type * from './types/SessionStepStat' export type * from './types/SessionView' export type * from './types/Size' +export type * from './types/SkillFile' export type * from './types/Source' export type * from './types/SpeechCallOptions' export type * from './types/SpeechRequest' @@ -188,6 +188,8 @@ export type * from './types/TypedToolResult' export type * from './types/UploadFileCallOptions' export type * from './types/UploadFileData' export type * from './types/UploadFileResult' +export type * from './types/UploadSkillCallOptions' +export type * from './types/UploadSkillResult' export type * from './types/Usage' export type * from './types/UsageSnapshot' export type * from './types/UserContent' diff --git a/bindings/node/src/types/AiMuxError.ts b/bindings/node/src/types/AiMuxError.ts index 93544dd57..e14125f85 100644 --- a/bindings/node/src/types/AiMuxError.ts +++ b/bindings/node/src/types/AiMuxError.ts @@ -15,9 +15,25 @@ export type AiMuxError = { "ApiCall": ApiCallError } | { "Retry": RetryError } | * Original argument text, when supplied by the provider. Absent in * older serialized errors and errors constructed without a call. */ -tool_input?: string | null, } } | { "InvalidToolInput": { tool_name: string, tool_input: string, cause: string, } } | { "ToolCallRepair": { original_error: AiMuxError, cause: AiMuxError, } } | { "InvalidArgument": string } | { "InvalidPrompt": string } | { "MissingToolResults": { tool_call_ids: Array, } } | { "ToolChoiceViolation": ToolChoiceViolationError } | { "TokenExpired": string } | { "UnsupportedFunctionality": string } | { "NoSuchModel": { model_id: string, +tool_input?: string | null, } } | { "InvalidToolInput": { tool_name: string, tool_input: string, cause: string, } } | { "ToolCallRepair": { original_error: AiMuxError, cause: AiMuxError, } } | { "InvalidArgument": string } | { "InvalidPrompt": string } | { "LoadApiKey": { +/** + * The environment variable consulted as the fallback. + */ +env_var: string, +/** + * What the key is for, e.g. `"OpenAI"`. + */ +description: string, } } | { "LoadSetting": { +/** + * The environment variable consulted as the fallback. + */ +env_var: string, +/** + * The setting's parameter name, e.g. `"region"`. + */ +name: string, } } | { "MissingToolResults": { tool_call_ids: Array, } } | { "ToolChoiceViolation": ToolChoiceViolationError } | { "TokenExpired": string } | { "UnsupportedFunctionality": string } | { "NoSuchModel": { model_id: string, /** * What kind of model was requested (`"languageModel"`, * `"imageModel"`, …), the AI SDK's `modelType`. */ -model_type: string, } } | { "NoSuchProvider": { provider_id: string, } } | { "Timeout": string } | { "Aborted": string } | { "Other": string }; +model_type: string, } } | { "NoSuchProvider": { provider_id: string, model_id: string, model_type: string, available_providers: Array, } } | { "Timeout": string } | { "Aborted": string } | { "Other": string }; diff --git a/bindings/node/src/types/AudioInput.ts b/bindings/node/src/types/AudioInput.ts index e513ffda3..eb4873bf2 100644 --- a/bindings/node/src/types/AudioInput.ts +++ b/bindings/node/src/types/AudioInput.ts @@ -1,6 +1,6 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. /** - * Audio input: raw bytes or a base64-encoded string. + * Audio input: raw bytes, a base64-encoded string, or a URL. */ -export type AudioInput = { "Binary": Array } | { "Base64": string }; +export type AudioInput = { "Binary": Array } | { "Base64": string } | { "Url": string }; diff --git a/bindings/node/src/types/CallOptions.ts b/bindings/node/src/types/CallOptions.ts index 93b7a271c..ba931b2ac 100644 --- a/bindings/node/src/types/CallOptions.ts +++ b/bindings/node/src/types/CallOptions.ts @@ -81,19 +81,15 @@ provider_options: { [key in string]: { [key in string]: JsonValue } } | null, */ reasoning: ReasoningEffort | null, /** - * Per-call request body overrides. Deep-merged into the provider-built - * request body (after any built-in vendor override) before sending. - * `null` values delete the corresponding key. See RFC-0017. - */ -body_overrides: JsonValue | null, -/** - * Emit raw provider stream chunks as `StreamPart::Raw` (debugging aid). - * When `Some(true)`, streaming providers yield one `Raw` part per JSON - * SSE event, carrying the **parsed JSON payload** of the event, emitted - * before the parsed parts. Excludes the `[DONE]` sentinel; unparsable - * chunks emit only `Error`. `null`/`Some(false)` = off. - * Currently honored by the OpenAI-compatible family (openai / azure / - * openai-compatible registry providers); other provider families ignore - * it for now (RFC-0016 M2). + * Emit raw provider stream chunks as `StreamPart::Raw` (the AI SDK's + * `includeRawChunks`). When `Some(true)`, a streaming model yields one + * `Raw` part per SSE event, carrying the event's parsed JSON, before the + * parts made from it; the `[DONE]` sentinel is excluded. A chunk that is + * not JSON yields `Raw` with `null` in Groq and Anthropic, as the AI SDK + * does, and only `Error` elsewhere. `None` / `Some(false)` = off. + * Honored by OpenAI chat, OpenAI-compatible chat and completion (registry + * presets and Azure included), Groq, DeepSeek, Anthropic Messages, and + * the OpenAI and ElevenLabs transcription streams; the other model + * families do not emit `Raw` yet. */ include_raw_chunks: boolean | null, }; diff --git a/bindings/node/src/types/EmbeddingCallOptions.ts b/bindings/node/src/types/EmbeddingCallOptions.ts index ca1caf1d2..d86e94f90 100644 --- a/bindings/node/src/types/EmbeddingCallOptions.ts +++ b/bindings/node/src/types/EmbeddingCallOptions.ts @@ -14,9 +14,7 @@ export type EmbeddingCallOptions = { values: Array, /** * Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - * retrying. `None` uses 2, the AI SDK default, unless the model's - * `retry_config()` sets another count (model-level retry settings are - * removed in the target design, docs/aisdk-architecture-alignment.md D14). + * retrying. `None` uses 2, the AI SDK default. */ max_retries: number | null, /** diff --git a/bindings/node/src/types/GenerateTextOptions.ts b/bindings/node/src/types/GenerateTextOptions.ts index a8682a834..aef7398a2 100644 --- a/bindings/node/src/types/GenerateTextOptions.ts +++ b/bindings/node/src/types/GenerateTextOptions.ts @@ -29,15 +29,9 @@ reasoning: ReasoningEffort | null, * System instructions prepended to the prompt. */ instructions: string | null, -/** - * Per-call request body overrides (deep-merged). See RFC-0017. - */ -body_overrides: JsonValue | null, /** * Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - * retrying. `None` uses 2, the AI SDK default, unless the model's - * `retry_config()` sets another count (model-level retry settings are - * removed in the target design, docs/aisdk-architecture-alignment.md D14). + * retrying. `None` uses 2, the AI SDK default. */ max_retries: number | null, /** diff --git a/bindings/node/src/types/ImageCallOptions.ts b/bindings/node/src/types/ImageCallOptions.ts index 412cfe027..bf89f2889 100644 --- a/bindings/node/src/types/ImageCallOptions.ts +++ b/bindings/node/src/types/ImageCallOptions.ts @@ -48,9 +48,7 @@ mask: ImageFile | null, provider_options: { [key in string]: { [key in string]: JsonValue } }, /** * Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - * retrying. `None` uses 2, the AI SDK default, unless the model's - * `retry_config()` sets another count (model-level retry settings are - * removed in the target design, docs/aisdk-architecture-alignment.md D14). + * retrying. `None` uses 2, the AI SDK default. */ max_retries: number | null, /** diff --git a/bindings/node/src/types/InputRecord.ts b/bindings/node/src/types/InputRecord.ts index 7b40111d6..b174b1070 100644 --- a/bindings/node/src/types/InputRecord.ts +++ b/bindings/node/src/types/InputRecord.ts @@ -1,4 +1,5 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { DecisionCapabilities } from "./DecisionCapabilities"; import type { LanguageModelMessage } from "./LanguageModelMessage"; import type { RecordingOperation } from "./RecordingOperation"; import type { JsonValue } from "./serde_json/JsonValue"; @@ -6,14 +7,20 @@ import type { JsonValue } from "./serde_json/JsonValue"; /** * ① 输入侧:完整调用参数,足以重建 generate_text 调用。 */ -export type InputRecord = { operation?: RecordingOperation, +export type InputRecord = { /** - * 完整 prompt(消息数组,含 ContentPart::Image 等多模态)。 + * Decision contract used for validation and deterministic offline replay. + * Connection settings remain outside the recording. + */ +decision_capabilities?: DecisionCapabilities | null, operation?: RecordingOperation, +/** + * 完整 prompt(消息数组,含 ContentPart::Image 等多模态)。Loads as + * empty when this version cannot read it ([`Self::unreadable_prompt`]). */ prompt: Array, /** * 序列化的 CallOptions(abort_signal 已 serde skip), * 加上所在操作的 `max_retries` / `timeout` / `session_id` - * (`OperationSettings`);headers/provider_options/body_overrides 已递归脱敏。 + * (`OperationSettings`);headers/provider_options 已递归脱敏。 */ options: JsonValue, }; diff --git a/bindings/node/src/types/ModelId.ts b/bindings/node/src/types/ModelId.ts deleted file mode 100644 index 3456afef4..000000000 --- a/bindings/node/src/types/ModelId.ts +++ /dev/null @@ -1,6 +0,0 @@ -// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. - -/** - * A model identifier in `"provider/model-name"` format, e.g. `"openai/gpt-4o"`. - */ -export type ModelId = { provider: string, model: string, }; diff --git a/bindings/node/src/types/ModelSpec.ts b/bindings/node/src/types/ModelSpec.ts index 58eb9bd31..69ead413e 100644 --- a/bindings/node/src/types/ModelSpec.ts +++ b/bindings/node/src/types/ModelSpec.ts @@ -12,8 +12,7 @@ import type { JsonValue } from "./serde_json/JsonValue"; * The full static portrait for a model, supplemented from community knowledge. * * All fields are advisory — aimux never auto-applies them in the request path. - * The caller reads them to decide how to configure `GenerateTextOptions` / - * `bodyOverrides`. + * The caller reads them to decide how to configure `GenerateTextOptions`. */ export type ModelSpec = { display_name?: string | null, type: ModelType, limits: ModelLimits, modalities: ModelModalities, capabilities: ModelCapabilities, reasoning?: ReasoningSpec | null, cost?: ModelCost | null, source: CatalogueSource, /** diff --git a/bindings/node/src/types/ProviderRecord.ts b/bindings/node/src/types/ProviderRecord.ts index 4af13e674..bbb3c1082 100644 --- a/bindings/node/src/types/ProviderRecord.ts +++ b/bindings/node/src/types/ProviderRecord.ts @@ -1,31 +1,23 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -import type { JsonValue } from "./serde_json/JsonValue"; /** * ② 配置侧:provider 身份(请求回放重建用)。 + * + * 只记身份,不记配置:回放时按 `provider_id` + `model_id` 重新构建 model, + * 凭证与 base URL 只来自环境/registry(RFC-0036 §3.3)。 */ -export type ProviderRecord = { +export type ProviderRecord = { /** - * `model.provider()`,如 "openai"。 + * provider 名(registry 名),回放重建的查找键,如 "openai"、"deepseek"。 + * Empty in a recording written by aimux 0.5 or earlier, which has none + * ([`Self::registry_id`]). */ -provider: string, +provider_id: string, /** - * `model.model_id()`,如 "gpt-4o"。 - */ -model_id: string, -/** - * provider 的 base_url。 - */ -base_url: string | null, -/** - * api_key 来源(不存明文):"env:OPENAI_API_KEY" / "explicit" / "none" / "unknown"。 - */ -api_key_source: string, -/** - * OpenAICompatProfile(能力差异)。 + * `model.provider()`,如 "openai.chat"。 */ -profile: JsonValue | null, +provider: string, /** - * ProviderOptions(headers/org/project/...;已脱敏)。 + * `model.model_id()`,如 "gpt-4o"。 */ -provider_options: JsonValue | null, }; +model_id: string, }; diff --git a/bindings/node/src/types/Recording.ts b/bindings/node/src/types/Recording.ts index 3bd5bc35c..f1ffb9286 100644 --- a/bindings/node/src/types/Recording.ts +++ b/bindings/node/src/types/Recording.ts @@ -7,50 +7,55 @@ import type { ProviderRecord } from "./ProviderRecord"; /** * 一次完整调用的录制记录(三层 + call_id 关联)。 */ -export type Recording = { +export type Recording = { /** - * 格式版本。 + * 格式版本。Written as [`RECORDING_SCHEMA`]; schemas 1 and 2 are read. + * Recordings written by aimux 0.5 and earlier also carry a configuration + * snapshot under `provider` (base URL, key source, profile, provider + * options), which is ignored. */ -schema: number, +schema: number, /** * 暴露逻辑调用 ID,关联三层。 */ -call_id: string, +call_id: string, /** * ISO 8601 时间戳。 */ -recorded_at: string, +recorded_at: string, /** * ① 输入侧:调用参数(prompt + options)。 */ -input: InputRecord, +input: InputRecord, /** * ② 配置侧:provider 身份与配置。 */ -provider: ProviderRecord, +provider: ProviderRecord, /** * ③ HTTP 侧:实际 wire 交换(per-attempt 一条)。 */ -exchanges: Array, +exchanges: Array, /** * 最终结果摘要(状态 + finish_reason + usage)。 */ -outcome: OutcomeRecord, +outcome: OutcomeRecord, /** - * wire 是否完整(流式 body 未补全时为 false;由 writer 标记)。 + * wire 是否完整(流式 body 未补全时为 false;由 writer 标记)。A call + * that lost an event (full queue, writer gone) or was marked + * inconsistent (duplicate attempt, unmatched update) is never complete. */ -complete: boolean, +complete: boolean, /** * 传输层封闭信号:层 B 声明"不再有 exchange"(P1:层 A 收尾自动发; * P2:由层 B 发送)。false = 仍可能来 exchange(记录暂不可定稿)。 */ -transport_closed: boolean, +transport_closed: boolean, /** * 会话归组(RFC-0024 P3):所在会话 id。None = 未归组(无 session_id 且 * 推断关闭)。由 `Recorder::record_session` 填充,写入 InputRecord 之前 * 或之后均可(writer 端按 call_id 合并)。 */ -session_id?: string | null, +session_id?: string | null, /** * 会话内步号(0 起,由 SessionStore 分配)。与 `session_id` 同生命周期。 */ diff --git a/bindings/node/src/types/RerankingCallOptions.ts b/bindings/node/src/types/RerankingCallOptions.ts index 1c0b79d1d..1aa545d7e 100644 --- a/bindings/node/src/types/RerankingCallOptions.ts +++ b/bindings/node/src/types/RerankingCallOptions.ts @@ -23,9 +23,7 @@ query: string, top_n: number | null, /** * Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - * retrying. `None` uses 2, the AI SDK default, unless the model's - * `retry_config()` sets another count (model-level retry settings are - * removed in the target design, docs/aisdk-architecture-alignment.md D14). + * retrying. `None` uses 2, the AI SDK default. */ max_retries: number | null, /** diff --git a/bindings/node/src/types/SearchCallOptions.ts b/bindings/node/src/types/SearchCallOptions.ts index 9da40b33f..b4cc779c2 100644 --- a/bindings/node/src/types/SearchCallOptions.ts +++ b/bindings/node/src/types/SearchCallOptions.ts @@ -36,9 +36,7 @@ include_domains: Array | null, exclude_domains: Array | null, /** * Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - * retrying. `None` uses 2, the AI SDK default, unless the model's - * `retry_config()` sets another count (model-level retry settings are - * removed in the target design, docs/aisdk-architecture-alignment.md D14). + * retrying. `None` uses 2, the AI SDK default. */ max_retries: number | null, /** diff --git a/bindings/node/src/types/SkillFile.ts b/bindings/node/src/types/SkillFile.ts new file mode 100644 index 000000000..8807afee2 --- /dev/null +++ b/bindings/node/src/types/SkillFile.ts @@ -0,0 +1,4 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { UploadFileData } from "./UploadFileData"; + +export type SkillFile = { path: string, data: UploadFileData, }; diff --git a/bindings/node/src/types/SpeechCallOptions.ts b/bindings/node/src/types/SpeechCallOptions.ts index ed08278d7..56294ef8e 100644 --- a/bindings/node/src/types/SpeechCallOptions.ts +++ b/bindings/node/src/types/SpeechCallOptions.ts @@ -41,9 +41,7 @@ language: string | null, provider_options: { [key in string]: { [key in string]: JsonValue } } | null, /** * Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - * retrying. `None` uses 2, the AI SDK default, unless the model's - * `retry_config()` sets another count (model-level retry settings are - * removed in the target design, docs/aisdk-architecture-alignment.md D14). + * retrying. `None` uses 2, the AI SDK default. */ max_retries: number | null, /** diff --git a/bindings/node/src/types/TranscriptionCallOptions.ts b/bindings/node/src/types/TranscriptionCallOptions.ts index 3eaa29126..eb4dc86a0 100644 --- a/bindings/node/src/types/TranscriptionCallOptions.ts +++ b/bindings/node/src/types/TranscriptionCallOptions.ts @@ -10,11 +10,14 @@ import type { JsonValue } from "./serde_json/JsonValue"; */ export type TranscriptionCallOptions = { /** - * Audio data to transcribe (raw bytes or a base64-encoded string). + * Audio to transcribe: raw bytes, a base64-encoded string, or a URL + * that [`transcribe`] downloads. */ audio: AudioInput, /** - * The IANA media type of the audio data, e.g. `"audio/mp3"`. + * The IANA media type of the audio data, e.g. `"audio/mp3"`. For a URL, + * the media type the download reports replaces it, as the AI SDK uses + * the downloaded media type (`transcribe.ts:227-233`). */ media_type: string, /** @@ -23,9 +26,7 @@ media_type: string, provider_options: { [key in string]: { [key in string]: JsonValue } } | null, /** * Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - * retrying. `None` uses 2, the AI SDK default, unless the model's - * `retry_config()` sets another count (model-level retry settings are - * removed in the target design, docs/aisdk-architecture-alignment.md D14). + * retrying. `None` uses 2, the AI SDK default. */ max_retries: number | null, /** diff --git a/bindings/node/src/types/UploadSkillCallOptions.ts b/bindings/node/src/types/UploadSkillCallOptions.ts new file mode 100644 index 000000000..7d6e3aa38 --- /dev/null +++ b/bindings/node/src/types/UploadSkillCallOptions.ts @@ -0,0 +1,5 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { SkillFile } from "./SkillFile"; +import type { JsonValue } from "./serde_json/JsonValue"; + +export type UploadSkillCallOptions = { files: Array, display_title: string | null, provider_options: { [key in string]: { [key in string]: JsonValue } } | null, }; diff --git a/bindings/node/src/types/UploadSkillResult.ts b/bindings/node/src/types/UploadSkillResult.ts new file mode 100644 index 000000000..17a61f981 --- /dev/null +++ b/bindings/node/src/types/UploadSkillResult.ts @@ -0,0 +1,5 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { Warning } from "./Warning"; +import type { JsonValue } from "./serde_json/JsonValue"; + +export type UploadSkillResult = { provider_reference: { [key in string]: string }, display_title: string | null, name: string | null, description: string | null, latest_version: string | null, provider_metadata: { [key in string]: { [key in string]: JsonValue } } | null, warnings: Array, }; diff --git a/bindings/node/src/types/VideoCallOptions.ts b/bindings/node/src/types/VideoCallOptions.ts index 48304854e..5586d33d0 100644 --- a/bindings/node/src/types/VideoCallOptions.ts +++ b/bindings/node/src/types/VideoCallOptions.ts @@ -12,73 +12,71 @@ import type { JsonValue } from "./serde_json/JsonValue"; * * Aligned with V4 `VideoModelV4CallOptions`. */ -export type VideoCallOptions = { +export type VideoCallOptions = { /** * Text prompt for the video generation. `None` when not required. */ -prompt: string | null, +prompt: string | null, /** * Number of videos to generate. Default `1`; most models only support * `n = 1` due to computational cost. */ -n: number, +n: number, /** * Aspect ratio, in `{width}:{height}` format (e.g. `"16:9"`). */ -aspect_ratio: AspectRatio | null, +aspect_ratio: AspectRatio | null, /** * Resolution, in `{width}x{height}` format (e.g. `"1280x720"`). */ -resolution: Size | null, +resolution: Size | null, /** * Duration of the video in seconds. Typically 3–10 seconds. */ -duration: number | null, +duration: number | null, /** * Frames per second. Common values: 24, 30, 60. */ -fps: number | null, +fps: number | null, /** * Seed for deterministic generation. `None` uses a random seed. */ -seed: number | null, +seed: number | null, /** * Input image for image-to-video generation (the starting frame). */ -image: VideoFile | null, +image: VideoFile | null, /** * Role-tagged image inputs for first-last-frame generation. */ -frame_images: Array | null, +frame_images: Array | null, /** * Reference inputs for reference-to-video generation (images or videos). */ -input_references: Array | null, +input_references: Array | null, /** * Whether the model should generate audio alongside the video. */ -generate_audio: boolean | null, +generate_audio: boolean | null, /** * Additional provider-specific options, keyed by provider name. */ -provider_options: { [key in string]: { [key in string]: JsonValue } }, +provider_options: { [key in string]: { [key in string]: JsonValue } }, /** * Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables - * retrying. `None` uses 2, the AI SDK default, unless the model's - * `retry_config()` sets another count (model-level retry settings are - * removed in the target design, docs/aisdk-architecture-alignment.md D14). + * retrying. `None` uses 2, the AI SDK default. */ -max_retries: number | null, +max_retries: number | null, /** * Poll pacing for the start/status flow, the AI SDK's `generateVideo` * `poll` option. Unset fields use the AI SDK defaults (5000 ms between * status checks, 600000 ms in total). */ -poll: VideoPollOptions | null, +poll: VideoPollOptions | null, /** * Per-call operation timeout. */ -timeout: TimeoutConfiguration | null, +timeout: TimeoutConfiguration | null, /** * Additional HTTP headers to send with the request. */ diff --git a/bindings/python/python/aimux/__init__.py b/bindings/python/python/aimux/__init__.py index 9e581fa13..5dd1f0fa6 100644 --- a/bindings/python/python/aimux/__init__.py +++ b/bindings/python/python/aimux/__init__.py @@ -21,6 +21,8 @@ ToolCallRepairError, InvalidArgumentError, InvalidPromptError, + LoadAPIKeyError, + LoadSettingError, TokenExpiredError, UnsupportedFunctionalityError, NoSuchModelError, @@ -100,6 +102,8 @@ "ToolCallRepairError", "InvalidArgumentError", "InvalidPromptError", + "LoadAPIKeyError", + "LoadSettingError", "TokenExpiredError", "UnsupportedFunctionalityError", "NoSuchModelError", @@ -217,7 +221,9 @@ def provider( model_id: Model ID. base_url: Base-URL override (wins over config["base_url"]). config: Full ProviderOptions dict — base_url / headers / organization / - project / max_retries / body_overrides. + project / params (a preset's template parameters, e.g. + ``{"account_id": "..."}``). ``max_retries`` (a call-level option) and + ``body_overrides`` are rejected with ``InvalidArgumentError``. """ config_json = json.dumps(config) if config is not None else None return _native_provider(name, api_key, model_id, base_url, config_json) @@ -239,7 +245,9 @@ def create_provider( api_key: API key; None reads the provider's env var. base_url: Base-URL override (wins over config["base_url"]). config: Full ProviderOptions dict — base_url / headers / organization / - project / max_retries / body_overrides. + project / params (a preset's template parameters, e.g. + ``{"account_id": "..."}``). ``max_retries`` (a call-level option) and + ``body_overrides`` are rejected with ``InvalidArgumentError``. """ config_json = json.dumps(config) if config is not None else None return _native_create_provider(name, api_key, base_url, config_json) diff --git a/bindings/python/python/aimux/wrapper.py b/bindings/python/python/aimux/wrapper.py index 4999b4d6b..f903431a4 100644 --- a/bindings/python/python/aimux/wrapper.py +++ b/bindings/python/python/aimux/wrapper.py @@ -1183,7 +1183,6 @@ class GenerateTextOptions(BaseModel): provider_options: Optional[Dict[str, Any]] = None reasoning: Optional[ReasoningEffort] = None instructions: Optional[str] = None - body_overrides: Optional[Any] = None max_retries: Optional[int] = None session_id: Optional[str] = None """Session identifier (RFC-0024): groups consecutive calls into a session. diff --git a/bindings/python/src/error.rs b/bindings/python/src/error.rs index a44ca9840..44e660d29 100644 --- a/bindings/python/src/error.rs +++ b/bindings/python/src/error.rs @@ -64,6 +64,18 @@ create_exception!( ); create_exception!(aimux, InvalidArgumentError, AimuxError, "Invalid argument"); create_exception!(aimux, InvalidPromptError, AimuxError, "Invalid prompt"); +create_exception!( + aimux, + LoadAPIKeyError, + AimuxError, + "No API key was passed and the fallback environment variable is unset" +); +create_exception!( + aimux, + LoadSettingError, + AimuxError, + "A required provider setting was not passed and its environment variable is unset" +); create_exception!(aimux, TokenExpiredError, AimuxError, "Access token expired"); create_exception!( aimux, @@ -254,6 +266,8 @@ fn variant_instance<'py>(py: Python<'py>, e: &AiMuxError) -> PyResult { py.get_type_bound::() } + AiMuxError::LoadApiKey { .. } => py.get_type_bound::(), + AiMuxError::LoadSetting { .. } => py.get_type_bound::(), AiMuxError::TokenExpired(_) => py.get_type_bound::(), AiMuxError::UnsupportedFunctionality(_) => { py.get_type_bound::() @@ -332,8 +346,27 @@ fn variant_instance<'py>(py: Python<'py>, e: &AiMuxError) -> PyResult { + AiMuxError::NoSuchProvider { + provider_id, + model_id, + model_type, + available_providers, + } => { inst.setattr("provider_id", provider_id.as_str())?; + inst.setattr("model_id", model_id.as_str())?; + inst.setattr("model_type", model_type.as_str())?; + inst.setattr("available_providers", available_providers.clone())?; + } + AiMuxError::LoadApiKey { + env_var, + description, + } => { + inst.setattr("env_var", env_var.as_str())?; + inst.setattr("description", description.as_str())?; + } + AiMuxError::LoadSetting { env_var, name } => { + inst.setattr("env_var", env_var.as_str())?; + inst.setattr("setting_name", name.as_str())?; } AiMuxError::NoSuchTool { tool_name, @@ -399,6 +432,8 @@ pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { "InvalidPromptError", py.get_type_bound::(), )?; + m.add("LoadAPIKeyError", py.get_type_bound::())?; + m.add("LoadSettingError", py.get_type_bound::())?; m.add( "TokenExpiredError", py.get_type_bound::(), diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 98771f790..970aec831 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -6,7 +6,12 @@ // pyo3 0.22 macros generate unsafe-op-in-unsafe-fn calls that trigger // edition-2024 lint. Suppress until pyo3 0.23+ lands. #![allow(unsafe_op_in_unsafe_fn)] -// PyO3 0.22 wraps PyResult returns in an identity conversion. +// The same macros test the `gil-refs` feature of the *calling* crate, which +// this crate does not declare (rustc's check-cfg flags it). +#![allow(unexpected_cfgs)] +// pyo3 0.22's `#[pyfunction]` / `#[pymethods]` expansion converts a `PyErr` into +// itself on every `PyResult` return (fixed in 0.23); clippy blames the +// signature. #![allow(clippy::useless_conversion)] mod decision; @@ -439,7 +444,7 @@ impl StreamIterator { let item = py.allow_threads(|| runtime().block_on(self.rx.recv())); match item { - Some(Ok(json)) => Ok(Some(json.to_object(py))), + Some(Ok(json)) => Ok(Some(json.to_object(py).into())), Some(Err(f)) => Err(f.to_py_err()), None => Ok(None), // stream finished } @@ -450,23 +455,32 @@ impl StreamIterator { // Provider constructors // ───────────────────────────────────────────────────────────────────────────── +/// An OpenAI provider for an explicit key handed over by the host, with an +/// optional base URL. The key is an explicit value (`""` included), so it never +/// falls back to `OPENAI_API_KEY`. +pub(crate) fn openai_provider( + api_key: &str, + base_url: Option<&str>, +) -> PyResult { + aimux_providers::openai::create_openai(aimux_providers::openai::OpenAIProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e)) +} + /// Create an OpenAI model instance. #[pyfunction] #[pyo3(signature = (api_key, model_id, base_url=None))] fn openai(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); - let model = provider - .language_model(model_id) - .map_err(|e| to_py_err(&e))?; + let provider = openai_provider(api_key, base_url)?; Ok(Model { - inner: Arc::from(model), + inner: provider + .language_model(model_id) + .map_err(|e| to_py_err(&e))?, trace_store: None, }) } @@ -476,13 +490,14 @@ fn openai(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider}; + use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; - let mut config = AnthropicConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = AnthropicProvider::new(config); + let provider = create_anthropic(AnthropicProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider .language_model(model_id) .map_err(|e| to_py_err(&e))?; @@ -496,12 +511,17 @@ fn anthropic(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult< #[pyfunction] #[pyo3(signature = (api_key, model_id, base_url=None))] fn deepseek(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult { - let options = base_url.map(|url| aimux_providers::ProviderOptions { + let options = base_url.map(|url| aimux_providers::provider::ProviderOptions { base_url: Some(url.to_string()), ..Default::default() }); - let model = aimux_providers::provider("deepseek", Some(api_key.to_string()), model_id, options) - .map_err(|e| to_py_err(&e))?; + let model = aimux_providers::provider::provider( + "deepseek", + Some(api_key.to_string()), + model_id, + options, + ) + .map_err(|e| to_py_err(&e))?; Ok(Model { inner: Arc::from(model), trace_store: None, @@ -514,13 +534,14 @@ fn deepseek(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::google::{GoogleConfig, GoogleProvider}; + use aimux_providers::google::{GoogleProviderSettings, create_google}; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = GoogleProvider::new(config); + let provider = create_google(GoogleProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider .language_model(model_id) .map_err(|e| to_py_err(&e))?; @@ -535,13 +556,14 @@ fn google(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::cohere::{CohereConfig, CohereProvider}; + use aimux_providers::cohere::{CohereProviderSettings, create_cohere}; - let mut config = CohereConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = CohereProvider::new(config); + let provider = create_cohere(CohereProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider .language_model(model_id) .map_err(|e| to_py_err(&e))?; @@ -556,13 +578,14 @@ fn cohere(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::mistral::{MistralConfig, MistralProvider}; + use aimux_providers::mistral::{MistralProviderSettings, create_mistral}; - let mut config = MistralConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = MistralProvider::new(config); + let provider = create_mistral(MistralProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider .language_model(model_id) .map_err(|e| to_py_err(&e))?; @@ -577,13 +600,14 @@ fn mistral(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::xai::{XAIConfig, XAIProvider}; + use aimux_providers::xai::{XAIProviderSettings, create_xai}; - let mut config = XAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = XAIProvider::new(config); + let provider = create_xai(XAIProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider .language_model(model_id) .map_err(|e| to_py_err(&e))?; @@ -604,13 +628,22 @@ fn bedrock( base_url: Option<&str>, ) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::bedrock::{BedrockProvider, BedrockProviderConfig}; - - let mut config = BedrockProviderConfig::new(access_key_id, secret_access_key, region); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = BedrockProvider::new(config); + use aimux_providers::bedrock::{AmazonBedrockProviderSettings, create_amazon_bedrock}; + + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + credential_provider: Some(aimux_provider_utils::Resolvable::Value( + aimux_provider_utils::AwsCredentials { + access_key_id: access_key_id.to_string(), + secret_access_key: secret_access_key.to_string(), + session_token: None, + region: region.to_string(), + }, + )), + region: Some(region.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider .language_model(model_id) .map_err(|e| to_py_err(&e))?; @@ -631,13 +664,18 @@ fn vertex( base_url: Option<&str>, ) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::vertex::{VertexProvider, VertexProviderConfig}; - - let mut config = VertexProviderConfig::new(access_token, project, location); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = VertexProvider::new(config); + use aimux_providers::vertex::{VertexProviderSettings, create_google_vertex}; + + let provider = create_google_vertex(VertexProviderSettings { + access_token: Some(aimux_provider_utils::Resolvable::Value( + access_token.to_string(), + )), + project: Some(project.to_string()), + location: Some(location.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider .language_model(model_id) .map_err(|e| to_py_err(&e))?; @@ -657,13 +695,19 @@ fn anthropic_aws( base_url: Option<&str>, ) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::anthropic_aws::{AnthropicAwsProvider, AnthropicAwsProviderConfig}; + use aimux_providers::anthropic_aws::{ + AnthropicAwsAuth, AnthropicAwsProviderSettings, create_anthropic_aws, + }; - let mut config = AnthropicAwsProviderConfig::with_api_key(api_key, region); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = AnthropicAwsProvider::new(config); + let provider = create_anthropic_aws(AnthropicAwsProviderSettings { + region: Some(region.to_string()), + auth: Some(AnthropicAwsAuth::ApiKey( + aimux_provider_utils::Resolvable::Value(api_key.to_string()), + )), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider .language_model(model_id) .map_err(|e| to_py_err(&e))?; @@ -686,21 +730,24 @@ fn azure( base_url: Option<&str>, ) -> PyResult { use aimux_core::provider::Provider; - use aimux_providers::azure::{AzureConfig, AzureProvider}; + use aimux_providers::azure::{AzureOpenAIProviderSettings, create_azure}; - let mut config = AzureConfig::new().with_api_key(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } + let mut config = AzureOpenAIProviderSettings { + api_key: Some(aimux_provider_utils::Resolvable::Value(api_key.to_string())), + base_url: base_url.map(str::to_string), + ..Default::default() + }; + // A dated `api_version` belongs to the deployment URL form. if let Some(version) = api_version && !version.is_empty() { - config = config.with_api_version(version); + config.api_version = Some(version.to_string()); + config.use_deployment_based_urls = true; } if !resource_name.is_empty() { - config = config.with_resource_name(resource_name); + config.resource_name = Some(resource_name.to_string()); } - let provider = AzureProvider::new(config).map_err(|e| to_py_err(&e))?; + let provider = create_azure(config).map_err(|e| to_py_err(&e))?; let model = provider .language_model(deployment) .map_err(|e| to_py_err(&e))?; @@ -713,8 +760,9 @@ fn azure( /// Create a language model from the built-in registry by provider name /// (RFC-0017 phase 4). `api_key=None` reads the provider's env var. /// `config_json` is a serialized `ProviderOptions` object (`base_url` / -/// `headers` / `organization` / `project` / `max_retries` / -/// `body_overrides`); the `base_url` parameter wins over the JSON field. +/// `headers` / `organization` / `project` / `params`; `max_retries` and +/// `body_overrides` are rejected as invalid arguments); the `base_url` +/// parameter wins over the JSON field. #[pyfunction] #[pyo3(signature = (name, api_key, model_id, base_url=None, config_json=None))] fn provider( @@ -724,15 +772,15 @@ fn provider( base_url: Option<&str>, config_json: Option<&str>, ) -> PyResult { - let mut options: Option = match config_json { + let mut options: Option = match config_json { Some(s) if !s.trim().is_empty() && s.trim() != "null" => Some(wire_json("config_json", s)?), _ => None, }; if let Some(url) = base_url { options.get_or_insert_with(Default::default).base_url = Some(url.to_string()); } - let model = - aimux_providers::provider(name, api_key, model_id, options).map_err(|e| to_py_err(&e))?; + let model = aimux_providers::provider::provider(name, api_key, model_id, options) + .map_err(|e| to_py_err(&e))?; Ok(Model { inner: Arc::from(model), trace_store: None, @@ -766,9 +814,14 @@ impl ProviderHandle { /// List models available on this provider (runtime discovery + anya2a spec). /// Returns a JSON array of RuntimeModel. fn list_models(&self, py: Python<'_>) -> PyResult { + let discovery = self.inner.discovery().ok_or_else(|| { + to_py_err(&AiMuxError::UnsupportedFunctionality( + "provider does not support model listing".into(), + )) + })?; let rt = runtime(); let models = py - .allow_threads(|| rt.block_on(async { self.inner.list_models().await })) + .allow_threads(|| rt.block_on(discovery.list_models())) .map_err(|e| to_py_err(&e))?; serialize_result(&models) } @@ -798,17 +851,16 @@ fn create_provider( base_url: Option<&str>, config_json: Option<&str>, ) -> PyResult { - let mut options: Option = match config_json { + let mut options: Option = match config_json { Some(s) if !s.trim().is_empty() && s.trim() != "null" => Some(wire_json("config_json", s)?), _ => None, }; if let Some(url) = base_url { options.get_or_insert_with(Default::default).base_url = Some(url.to_string()); } - let p = aimux_providers::provider_handle(name, api_key, options).map_err(|e| to_py_err(&e))?; - Ok(ProviderHandle { - inner: Arc::from(p), - }) + let inner = aimux_providers::provider::provider_handle(name, api_key, options) + .map_err(|e| to_py_err(&e))?; + Ok(ProviderHandle { inner }) } /// Fetch the community model catalogue (RFC-0027) and return it as a JSON @@ -853,7 +905,7 @@ fn init_logging(level: &str) { #[pyfunction] fn register_providers(config_json: &str) -> PyResult<()> { let _: serde_json::Value = wire_json("config_json", config_json)?; - aimux_providers::load_providers_from_json(config_json).map_err(|e| match e { + aimux_providers::provider::load_providers_from_json(config_json).map_err(|e| match e { AiMuxError::JsonParse(m) => { to_py_err(&AiMuxError::InvalidArgument(format!("config_json: {m}"))) } @@ -861,9 +913,9 @@ fn register_providers(config_json: &str) -> PyResult<()> { }) } -/// Set the global proxy configuration (M6, RFC-0016). Must be called before the -/// first `generate_text` / `stream_text` call; a no-op if the shared HTTP -/// client is already initialised. +/// Replace the default transport with one routed by this proxy configuration +/// (M6, RFC-0016). Requests started afterwards use it; calling it again +/// replaces the previous configuration. /// /// `config_json` shape: `{ "http_url": "...", "https_url": "...", "all_url": /// "...", "no_proxy": "..." }` (all fields optional). Raises `ValueError` on @@ -871,9 +923,11 @@ fn register_providers(config_json: &str) -> PyResult<()> { #[pyfunction] fn init_proxy(config: &str) -> PyResult<()> { let proxy_config: aimux_provider_utils::ProxyConfig = wire_json("config", config)?; - // `init_proxy` returns false when the shared client is already up; treat - // that as success (idempotent). - let _ = aimux_provider_utils::init_proxy(proxy_config); + // Replaces the default transport for later requests (see + // `aimux_init_proxy`). + aimux_provider_utils::set_default_transport(aimux_provider_utils::ReqwestFetch::new( + proxy_config, + )); Ok(()) } diff --git a/bindings/python/src/multimodal.rs b/bindings/python/src/multimodal.rs index 9a2406a09..f8562a7e2 100644 --- a/bindings/python/src/multimodal.rs +++ b/bindings/python/src/multimodal.rs @@ -425,13 +425,8 @@ pub fn openai_embedding( model_id: &str, base_url: Option<&str>, ) -> PyResult { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); - let model = provider.embedding_model(model_id); + let provider = crate::openai_provider(api_key, base_url)?; + let model = provider.embedding(model_id); Ok(EmbeddingModel { inner: Arc::new(model), }) @@ -445,12 +440,7 @@ pub fn openai_speech( model_id: &str, base_url: Option<&str>, ) -> PyResult { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); + let provider = crate::openai_provider(api_key, base_url)?; let model = provider.speech(model_id); Ok(SpeechModel { inner: Arc::new(model), @@ -461,12 +451,7 @@ pub fn openai_speech( #[pyfunction] #[pyo3(signature = (api_key, model_id, base_url=None))] pub fn openai_image(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); + let provider = crate::openai_provider(api_key, base_url)?; let model = provider.image(model_id); Ok(ImageModel { inner: Arc::new(model), @@ -481,12 +466,7 @@ pub fn openai_transcription( model_id: &str, base_url: Option<&str>, ) -> PyResult { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); + let provider = crate::openai_provider(api_key, base_url)?; let model = provider.transcription(model_id); Ok(TranscriptionModel { inner: Arc::new(model), @@ -497,12 +477,7 @@ pub fn openai_transcription( #[pyfunction] #[pyo3(signature = (api_key, base_url=None))] pub fn openai_files(api_key: &str, base_url: Option<&str>) -> PyResult { - use aimux_providers::openai::{OpenAIConfig, OpenAIProvider}; - let mut config = OpenAIConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = OpenAIProvider::new(config); + let provider = crate::openai_provider(api_key, base_url)?; let files = provider.files(); Ok(Files { inner: Arc::new(files), @@ -521,13 +496,14 @@ pub fn cohere_embedding( model_id: &str, base_url: Option<&str>, ) -> PyResult { - use aimux_providers::cohere::{CohereConfig, CohereProvider}; - let mut config = CohereConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = CohereProvider::new(config); - let model = provider.embedding_model(model_id); + use aimux_providers::cohere::{CohereProviderSettings, create_cohere}; + let provider = create_cohere(CohereProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; + let model = provider.embedding(model_id); Ok(EmbeddingModel { inner: Arc::new(model), }) @@ -541,13 +517,14 @@ pub fn cohere_reranking( model_id: &str, base_url: Option<&str>, ) -> PyResult { - use aimux_providers::cohere::{CohereConfig, CohereProvider}; - let mut config = CohereConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = CohereProvider::new(config); - let model = provider.reranking_model(model_id); + use aimux_providers::cohere::{CohereProviderSettings, create_cohere}; + let provider = create_cohere(CohereProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; + let model = provider.reranking(model_id); Ok(RerankingModel { inner: Arc::new(model), }) @@ -565,13 +542,14 @@ pub fn google_embedding( model_id: &str, base_url: Option<&str>, ) -> PyResult { - use aimux_providers::google::{GoogleConfig, GoogleProvider}; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = GoogleProvider::new(config); - let model = provider.embedding_model(model_id); + use aimux_providers::google::{GoogleProviderSettings, create_google}; + let provider = create_google(GoogleProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; + let model = provider.embedding(model_id); Ok(EmbeddingModel { inner: Arc::new(model), }) @@ -581,12 +559,13 @@ pub fn google_embedding( #[pyfunction] #[pyo3(signature = (api_key, model_id, base_url=None))] pub fn google_image(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult { - use aimux_providers::google::{GoogleConfig, GoogleProvider}; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = GoogleProvider::new(config); + use aimux_providers::google::{GoogleProviderSettings, create_google}; + let provider = create_google(GoogleProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider.image(model_id); Ok(ImageModel { inner: Arc::new(model), @@ -597,12 +576,13 @@ pub fn google_image(api_key: &str, model_id: &str, base_url: Option<&str>) -> Py #[pyfunction] #[pyo3(signature = (api_key, model_id, base_url=None))] pub fn google_video(api_key: &str, model_id: &str, base_url: Option<&str>) -> PyResult { - use aimux_providers::google::{GoogleConfig, GoogleProvider}; - let mut config = GoogleConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = GoogleProvider::new(config); + use aimux_providers::google::{GoogleProviderSettings, create_google}; + let provider = create_google(GoogleProviderSettings { + api_key: Some(api_key.to_string()), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider.video(model_id); Ok(VideoModel { inner: Arc::new(model), @@ -613,12 +593,13 @@ pub fn google_video(api_key: &str, model_id: &str, base_url: Option<&str>) -> Py #[pyfunction] #[pyo3(signature = (api_key, base_url=None))] pub fn tavily_search(api_key: &str, base_url: Option<&str>) -> PyResult { - use aimux_providers::tavily::{TavilyConfig, TavilyProvider}; - let mut config = TavilyConfig::new(api_key); - if let Some(url) = base_url { - config = config.with_base_url(url); - } - let provider = TavilyProvider::new(config); + use aimux_providers::tavily::{TavilyProviderSettings, create_tavily}; + let provider = create_tavily(TavilyProviderSettings { + api_key: Some(aimux_provider_utils::Resolvable::Value(api_key.to_string())), + base_url: base_url.map(str::to_string), + ..Default::default() + }) + .map_err(|e| to_py_err(&e))?; let model = provider.search_model(); Ok(SearchModel { inner: Arc::new(model), @@ -759,13 +740,13 @@ pub fn start_transcription_session( Ok(()) => {} Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { tokio::select! { - _ = effective.cancelled() => return, + _ = effective.cancelled() => {} res = tx.send(Err(e.clone().into())) => { - if res.is_err() { return; } + let _ = res; } } } - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => return, + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {} } } } diff --git a/bindings/python/tests/test_aimux.py b/bindings/python/tests/test_aimux.py index 34b177d03..60b3397f9 100644 --- a/bindings/python/tests/test_aimux.py +++ b/bindings/python/tests/test_aimux.py @@ -51,14 +51,36 @@ def test_provider_accepts_full_config(): "headers": {"X-Custom": "1"}, "organization": "org-1", "project": "proj-1", - "max_retries": 0, - "body_overrides": {"temperature": 0.1}, }, ) assert model is not None assert hasattr(model, "generate_text") +@pytest.mark.parametrize( + "key,value", + [("max_retries", 0), ("body_overrides", {"temperature": 0.1})], +) +def test_provider_rejects_removed_config_keys(key, value): + """max_retries is call-level and body_overrides is gone: neither is silently ignored.""" + from aimux import provider + + with pytest.raises(Exception, match=key): + provider("groq", "sk-test-fake-key", "llama-3.3-70b", config={key: value}) + + +@pytest.mark.parametrize( + "key,value", + [("max_retries", 0), ("body_overrides", {"temperature": 0.1})], +) +def test_create_provider_rejects_removed_config_keys(key, value): + """create_provider() goes through the same ProviderOptions check.""" + from aimux import InvalidArgumentError, create_provider + + with pytest.raises(InvalidArgumentError, match=key): + create_provider("groq", "sk-test-fake-key", config={key: value}) + + def test_provider_base_url_param_wins_over_config(): """Explicit base_url parameter overrides config["base_url"].""" from aimux import provider @@ -82,16 +104,49 @@ def test_provider_invalid_config_raises(): "groq", "sk-test-fake-key", "llama-3.3-70b", - config={"max_retries": "not-a-number"}, + config={"headers": "not-a-map"}, ) -def test_provider_missing_env_key_raises(): - """provider() with api_key=None reads the env var and fails clearly when unset.""" - from aimux import provider +def test_provider_missing_env_key_raises(monkeypatch): + """provider() with api_key=None reads the env var when a request is made, and + fails clearly when it is unset (the AI SDK loads the key lazily).""" + from aimux import LoadAPIKeyError, generate_text, provider + + monkeypatch.delenv("ABACUS_API_KEY", raising=False) + model = provider("abacus", None, "m") + with pytest.raises(LoadAPIKeyError, match="(?i)api key"): + generate_text(model, "hi") + + +def test_provider_params_fill_a_preset_template(monkeypatch): + """config["params"] reaches the preset's template parameters.""" + from aimux import InvalidArgumentError, generate_text, provider - with pytest.raises(Exception, match="(?i)api key"): - provider("abacus", None, "m") + monkeypatch.delenv("CLOUDFLARE_ACCOUNT_ID", raising=False) + # Without the parameter the preset cannot build its base URL (reported at call time). + unfilled = provider("cloudflare_workers_ai", "sk-test-fake-key", "m") + with pytest.raises(InvalidArgumentError, match="account_id"): + generate_text(unfilled, "hi", {"max_retries": 0}) + model = provider( + "cloudflare_workers_ai", + "sk-test-fake-key", + "m", + config={"params": {"account_id": "acct123"}}, + ) + assert model is not None + + +def test_provider_params_the_preset_does_not_declare_are_rejected(): + from aimux import InvalidArgumentError, provider + + with pytest.raises(InvalidArgumentError, match="no template parameter `account_id`"): + provider( + "abacus", + "sk-test-fake-key", + "m", + config={"params": {"account_id": "acct123"}}, + ) def test_openai_creates_model(): diff --git a/bindings/python/tests/test_cassette.py b/bindings/python/tests/test_cassette.py index be85d0487..124b8be45 100644 --- a/bindings/python/tests/test_cassette.py +++ b/bindings/python/tests/test_cassette.py @@ -7,7 +7,7 @@ import json import pytest -from aimux import openai, anthropic, generate_text, stream_text +from aimux import openai, deepseek, anthropic, generate_text, stream_text from tests.cassette_replay import CassetteServer @@ -59,7 +59,7 @@ class TestDeepSeekCassette: def test_generate(self): with CassetteServer("deepseek") as srv: assert srv.count > 0 - model = openai("test-key", "deepseek-chat", srv.url) + model = deepseek("test-key", "deepseek-chat", srv.url) result = generate_text(model, "Hello") assert "error" not in result, f"unexpected error: {result.get('error')}" assert isinstance(result["text"], str) @@ -70,7 +70,7 @@ def test_generate_step_usage_raw_carries_vendor_fields(self): The result's ``usage`` is summed over the steps and, as the AI SDK's addLanguageModelUsage, carries no raw object.""" with CassetteServer("deepseek") as srv: - model = openai("test-key", "deepseek-chat", srv.url) + model = deepseek("test-key", "deepseek-chat", srv.url) result = generate_text(model, "Hello") raw = result["steps"][0]["usage"].get("raw") @@ -87,7 +87,7 @@ class TestGroqCassette: def test_generate(self): with CassetteServer("groq") as srv: assert srv.count > 0 - model = openai("test-key", "llama-3.3-70b-versatile", f"{srv.url}/openai/v1") + model = deepseek("test-key", "llama-3.3-70b-versatile", f"{srv.url}/openai/v1") result = generate_text(model, "Hello") assert "error" not in result, f"unexpected error: {result.get('error')}" assert isinstance(result["text"], str) @@ -98,7 +98,7 @@ class TestMistralCassette: def test_generate(self): with CassetteServer("mistral") as srv: assert srv.count > 0 - model = openai("test-key", "ministral-8b-latest", f"{srv.url}/v1") + model = deepseek("test-key", "ministral-8b-latest", f"{srv.url}/v1") result = generate_text(model, "Hello") assert "error" not in result, f"unexpected error: {result.get('error')}" assert isinstance(result["text"], str) @@ -109,7 +109,7 @@ class TestOllamaCassette: def test_generate(self): with CassetteServer("ollama") as srv: assert srv.count > 0 - model = openai("test-key", "qwen3:4b", f"{srv.url}/v1") + model = deepseek("test-key", "qwen3:4b", f"{srv.url}/v1") result = generate_text(model, "Hello") assert "error" not in result, f"unexpected error: {result.get('error')}" assert isinstance(result["text"], str) @@ -120,7 +120,7 @@ class TestPerplexityCassette: def test_generate(self): with CassetteServer("perplexity") as srv: assert srv.count > 0 - model = openai("test-key", "sonar", srv.url) + model = deepseek("test-key", "sonar", srv.url) result = generate_text(model, "Hello") assert "error" not in result, f"unexpected error: {result.get('error')}" assert isinstance(result["text"], str) diff --git a/bindings/python/tests/test_cassette_exhaustive.py b/bindings/python/tests/test_cassette_exhaustive.py index 046e912ee..138656d34 100644 --- a/bindings/python/tests/test_cassette_exhaustive.py +++ b/bindings/python/tests/test_cassette_exhaustive.py @@ -15,7 +15,7 @@ import pytest -from aimux import openai, generate_text, stream_text +from aimux import deepseek, generate_text, stream_text CASSETTE_BASE = Path(__file__).resolve().parents[3] / "aimux-providers" / "tests" / "cassettes" @@ -154,7 +154,7 @@ def test_all_chat_cassettes_replayed(self): prompt = _extract_prompt(cass) try: - model = openai("test-key", model_id, base_url) + model = deepseek("test-key", model_id, base_url) if cass["is_stream"]: parts = list(stream_text(model, prompt)) if len(parts) == 0: diff --git a/bindings/python/tests/test_e2e.py b/bindings/python/tests/test_e2e.py index ea6cb50f8..0deea9111 100644 --- a/bindings/python/tests/test_e2e.py +++ b/bindings/python/tests/test_e2e.py @@ -19,7 +19,7 @@ import pytest -from aimux import openai, anthropic, generate_text, stream_text +from aimux import deepseek, anthropic, generate_text, stream_text # ── Mock server in a separate process ─────────────────────────────────────── @@ -143,14 +143,14 @@ class TestOpenAIE2E: def test_generate_text(self): with MockServer(OPENAI_CHAT) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) result = generate_text(model, "What is Rust?") assert result["text"] == "Rust is a systems programming language." assert result["usage"] is not None def test_stream_text(self): with MockServer(OPENAI_STREAM, content_type='text/event-stream') as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) parts = list(stream_text(model, "Say hello")) assert len(parts) > 0 text = "".join( @@ -160,7 +160,7 @@ def test_stream_text(self): def test_generate_text_with_options(self): with MockServer(OPENAI_CHAT) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) generate_text(model, "Hello", {"max_output_tokens": 100, "temperature": 0.5}) # Options should reach the HTTP request (tested implicitly — if they # don't serialize, the Rust side would error) @@ -344,7 +344,7 @@ class TestStructuredContent: def test_generate_text_parses_tool_calls(self): with RecordingMockServer(OPENAI_TOOL_CALL) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) opts = { "tools": [{ "type": "function", @@ -376,7 +376,7 @@ def test_generate_text_parses_tool_calls(self): def test_multi_role_messages_reach_provider(self): with RecordingMockServer(OPENAI_CHAT) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) prompt = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is Rust?"}, @@ -394,7 +394,7 @@ def test_multi_role_messages_reach_provider(self): def test_tool_choice_reaches_provider(self): with RecordingMockServer(OPENAI_TOOL_CALL) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) opts = { "tools": [{ "type": "function", @@ -420,7 +420,7 @@ def test_stream_text_parses_tool_call_parts(self): 'data: [DONE]\n\n' ) with MockServer(sse_body, content_type='text/event-stream') as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) opts = { "tools": [{ "type": "function", @@ -463,7 +463,7 @@ def test_tool_call_round_trip(self): """ # The mock answers request #1 with a tool call and request #2 with text. with SequencedMockServer([OPENAI_TOOL_CALL, OPENAI_CHAT]) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) opts = { "tools": [{ "type": "function", diff --git a/bindings/python/tests/test_error.py b/bindings/python/tests/test_error.py index cc6adc58e..c76b09851 100644 --- a/bindings/python/tests/test_error.py +++ b/bindings/python/tests/test_error.py @@ -163,3 +163,49 @@ def test_malformed_wire_json_is_a_value_error_not_a_core_error(): model.generate_text("{not json") assert not isinstance(excinfo.value, AimuxError) + + +def test_missing_api_key_is_load_api_key_error(monkeypatch): + """No key passed and the env var unset: a typed LoadAPIKeyError at call time, not InvalidArgument.""" + from aimux import AimuxError, InvalidArgumentError, LoadAPIKeyError, generate_text, provider + + monkeypatch.delenv("GROQ_API_KEY", raising=False) + model = provider("groq", None, "some-model") + with pytest.raises(LoadAPIKeyError) as excinfo: + generate_text(model, "hi", {"max_retries": 0}) + err = excinfo.value + + assert type(err) is LoadAPIKeyError + assert isinstance(err, AimuxError) + assert not isinstance(err, InvalidArgumentError) + assert err.env_var == "GROQ_API_KEY" + assert err.description == "Groq" + assert "GROQ_API_KEY" in str(err) + # Pre-HTTP: the ApiCall-only attributes do not exist. + for attr in ("status", "retryable", *_API_CALL_ONLY): + assert not hasattr(err, attr), f"{attr} must not exist on LoadAPIKeyError" + + +def test_explicit_empty_api_key_is_not_a_load_error(): + """The AI SDK sends an explicit empty key verbatim; only None consults the env.""" + from aimux import provider + + assert provider("groq", "", "some-model") is not None + + +def test_missing_setting_is_load_setting_error_at_call_time(monkeypatch): + """A required setting unset in the environment fails the call, typed.""" + from aimux import AimuxError, LoadSettingError, azure, generate_text + + monkeypatch.delenv("AZURE_RESOURCE_NAME", raising=False) + model = azure("test-key", "", "some-deployment") + with pytest.raises(LoadSettingError) as excinfo: + generate_text(model, "hi", {"max_retries": 0}) + err = excinfo.value + + assert type(err) is LoadSettingError + assert isinstance(err, AimuxError) + assert err.env_var == "AZURE_RESOURCE_NAME" + assert err.setting_name == "resourceName" + for attr in ("status", "retryable", *_API_CALL_ONLY): + assert not hasattr(err, attr) diff --git a/bindings/python/tests/test_gil.py b/bindings/python/tests/test_gil.py index acb53e1e6..252b93c52 100644 --- a/bindings/python/tests/test_gil.py +++ b/bindings/python/tests/test_gil.py @@ -14,11 +14,11 @@ import pytest -from aimux import generate_text, openai, openai_embedding +from aimux import deepseek, generate_text, openai_embedding CHAT = json.dumps({ "id": "chatcmpl-gil", - "model": "gpt-4o", + "model": "deepseek-chat", "choices": [{ "message": {"role": "assistant", "content": "served by a Python thread"}, "finish_reason": "stop", @@ -74,7 +74,8 @@ def __exit__(self, *args): def test_generate_text_lets_other_python_threads_run(): with _ThreadServer(CHAT) as server: - model = openai("test-key", "gpt-4o", server.url) + # A chat-completions provider: openai() uses the Responses API. + model = deepseek("test-key", "deepseek-chat", server.url) started = time.monotonic() result = generate_text(model, "hi", LIMITS) assert result["text"] == "served by a Python thread" diff --git a/bindings/python/tests/test_session.py b/bindings/python/tests/test_session.py index f3f802986..e04ee87c9 100644 --- a/bindings/python/tests/test_session.py +++ b/bindings/python/tests/test_session.py @@ -8,7 +8,7 @@ import pytest from aimux import ( - openai, + deepseek, init_session_store, init_session_infer, session_calls, @@ -27,7 +27,7 @@ def test_explicit_session_groups_calls(self): init_session_infer(False) with MockServer(OPENAI_CHAT) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) generate_text(model, "first", GenerateTextOptions(session_id="sess-1")) generate_text(model, "second", GenerateTextOptions(session_id="sess-1")) @@ -58,7 +58,7 @@ def test_opt_in_inferer_groups_prefix_continuations(self): init_session_infer(True) with MockServer(OPENAI_CHAT) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) generate_text(model, "u1") generate_text( diff --git a/bindings/python/tests/test_trace.py b/bindings/python/tests/test_trace.py index c4f1b3389..513b97f2b 100644 --- a/bindings/python/tests/test_trace.py +++ b/bindings/python/tests/test_trace.py @@ -7,7 +7,7 @@ import json -from aimux import openai +from aimux import deepseek from aimux.wrapper import GenerateTextOptions, generate_text from test_e2e import MockServer, OPENAI_CHAT @@ -16,7 +16,7 @@ class TestTraceProbe: def test_trace_records_calls_and_query_api(self): with MockServer(OPENAI_CHAT) as mock: - raw = openai("test-key", "gpt-4o", mock.url) + raw = deepseek("test-key", "gpt-4o", mock.url) traced = raw.trace_audited(True) # > 4 KiB user message so block-aligned prefixes actually match. @@ -39,7 +39,7 @@ def test_trace_records_calls_and_query_api(self): stats = json.loads(traced.trace_aggregate()) assert len(stats) == 1 assert stats[0]["requests"] == 2 - assert stats[0]["provider"] == "openai" + assert stats[0]["provider"] == "deepseek.chat" assert "verdict_counts" in stats[0] chain = json.loads(traced.trace_session_chain("sess-1")) @@ -58,7 +58,7 @@ def test_trace_records_calls_and_query_api(self): def test_untraced_model_rejects_query_api(self): with MockServer(OPENAI_CHAT) as mock: - raw = openai("test-key", "gpt-4o", mock.url) + raw = deepseek("test-key", "gpt-4o", mock.url) try: raw.trace_aggregate() assert False, "untraced model must reject trace_aggregate" @@ -67,6 +67,6 @@ def test_untraced_model_rejects_query_api(self): def test_non_traced_model_generates_normally(self): with MockServer(OPENAI_CHAT) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) result = generate_text(model, "hello") assert result.text diff --git a/bindings/python/tests/test_wrapper.py b/bindings/python/tests/test_wrapper.py index dbd9b72fc..013aa6798 100644 --- a/bindings/python/tests/test_wrapper.py +++ b/bindings/python/tests/test_wrapper.py @@ -11,7 +11,7 @@ import pytest -from aimux import openai, anthropic +from aimux import deepseek, anthropic from aimux.wrapper import ( GenerateTextOptions, GenerateTextResult, @@ -45,7 +45,7 @@ class TestGenerateText: def test_returns_typed_result_with_text(self): with MockServer(OPENAI_CHAT) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) result = generate_text(model, "What is Rust?") assert isinstance(result, GenerateTextResult) @@ -58,7 +58,7 @@ def test_returns_typed_result_with_text(self): def test_parses_tool_calls(self): with RecordingMockServer(OPENAI_TOOL_CALL) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) opts = GenerateTextOptions( tools=[ FunctionTool( @@ -92,7 +92,7 @@ def test_parses_tool_calls(self): def test_content_is_text_variant(self): with MockServer(OPENAI_CHAT) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) result = generate_text(model, "What is Rust?") # content[0] is a typed Text content part. @@ -107,7 +107,7 @@ class TestStreamText: def test_yields_text_delta_dicts(self): with MockServer(OPENAI_STREAM, content_type="text/event-stream") as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) parts = list(stream_text(model, "Say hello")) assert len(parts) > 0 @@ -117,7 +117,7 @@ def test_yields_text_delta_dicts(self): def test_stream_part_can_be_typed(self): """A yielded dict round-trips through parse_stream_part into a model.""" with MockServer(OPENAI_STREAM, content_type="text/event-stream") as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) first_delta = next( p for p in stream_text(model, "Say hello") if "TextDelta" in p ) @@ -135,7 +135,7 @@ def test_include_raw_chunks_emits_raw_parts(self): """include_raw_chunks=True surfaces one Raw part per JSON SSE event, before the parsed parts ([DONE] excluded).""" with MockServer(OPENAI_STREAM, content_type="text/event-stream") as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) parts = list( stream_text(model, "Say hello", GenerateTextOptions(include_raw_chunks=True)) @@ -160,7 +160,7 @@ def test_include_raw_chunks_emits_raw_parts(self): def test_raw_chunks_off_by_default(self): """Default options emit no Raw parts.""" with MockServer(OPENAI_STREAM, content_type="text/event-stream") as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) parts = list(stream_text(model, "Say hello")) assert all("Raw" not in p for p in parts) @@ -194,7 +194,7 @@ def test_provider_tool_deferred_results_roundtrip(self, supported): def test_tools_and_tool_choice_reach_provider(self): with RecordingMockServer(OPENAI_TOOL_CALL) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) opts = GenerateTextOptions( tools=[ FunctionTool( @@ -229,7 +229,7 @@ def test_tool_choice_tool_variant_forces_tool_at_provider(self): """The ``tool`` tool-choice variant forces the specific tool at the provider (OpenAI renders it as ``{"type":"function","function":{...}}``).""" with RecordingMockServer(OPENAI_TOOL_CALL) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) opts = GenerateTextOptions( tools=[FunctionTool(name="get_weather", input_schema={"type": "object"})], tool_choice=ToolChoiceTool(tool_name="get_weather"), @@ -267,7 +267,7 @@ class TestPromptMessages: def test_multi_role_messages_reach_provider(self): with RecordingMockServer(OPENAI_CHAT) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) prompt = [ ModelMessage(role="system", content="You are a helpful assistant."), ModelMessage(role="user", content="What is Rust?"), @@ -332,7 +332,7 @@ def test_round_trip_with_typed_content_parts(self): ) with SequencedMockServer([OPENAI_TOOL_CALL, OPENAI_CHAT]) as mock: - model = openai("test-key", "gpt-4o", mock.url) + model = deepseek("test-key", "gpt-4o", mock.url) # Step 1: model requests a tool call. result = generate_text(model, "What's the weather in Tokyo?", opts) diff --git a/bindings/swift/Sources/Aimux/Aimux.swift b/bindings/swift/Sources/Aimux/Aimux.swift index be79f0507..f59d55fd0 100644 --- a/bindings/swift/Sources/Aimux/Aimux.swift +++ b/bindings/swift/Sources/Aimux/Aimux.swift @@ -11,7 +11,7 @@ import Foundation // // Every fallible C function returns `aimux_error_t *` (`OpaquePointer?`): // NULL = success (result in the trailing out-param), non-NULL = failure. The -// unified code is AiMuxError (1...18), RecordingError (100...105), or a C ABI +// unified code is AiMuxError (1...20), RecordingError (100...105), or a C ABI // failure (200...206). The three `expect*` decoders copy the relevant fields, release // it with `aimux_error_free` (exactly once) and return the Swift error // to throw. Errors are not handles: never `aimux_drop_handle` one. @@ -44,7 +44,7 @@ func expectFfiError(_ e: OpaquePointer, context: String) -> any Error { return invariant("aimux ffi: \(context): \(message)") } -/// Decode a returned error from an `[AiMuxError]` call: 1...18 becomes +/// Decode a returned error from an `[AiMuxError]` call: 1...20 becomes /// `AimuxError`; 200...206 is decoded by `expectFfiError`. Frees `e` once. func expectAimuxError(_ e: OpaquePointer, context: String) -> any Error { let code = aimux_error_code(e) @@ -179,6 +179,13 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl /// was repairing, as externally-tagged wire JSON (the same encoding as /// `ToolCall.error`). case toolCallRepair(message: String, status: Int, retryMs: Int64, retryable: Bool, originalError: String) + /// No API key was passed and the fallback environment variable is unset + /// (`AIMUX_E_LOAD_API_KEY`, the AI SDK's `LoadAPIKeyError`). No request + /// was made; `envVar` names the variable that was consulted. + case loadApiKey(message: String, status: Int, retryMs: Int64, retryable: Bool, envVar: String) + /// A required provider setting was not passed and its fallback environment + /// variable is unset (`AIMUX_E_LOAD_SETTING`, `LoadSettingError`). + case loadSetting(message: String, status: Int, retryMs: Int64, retryable: Bool, envVar: String) case other(message: String, status: Int, retryMs: Int64, retryable: Bool) // MARK: Accessors @@ -201,6 +208,8 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl .noSuchTool(let m, let s, let r, let t, _, _), .invalidToolInput(let m, let s, let r, let t, _, _), .toolCallRepair(let m, let s, let r, let t, _), + .loadApiKey(let m, let s, let r, let t, _), + .loadSetting(let m, let s, let r, let t, _), .other(let m, let s, let r, let t): return (m, s, r, t) case .retry(let m, _, _): @@ -230,6 +239,8 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl case .noSuchTool: c = AIMUX_E_NO_SUCH_TOOL case .invalidToolInput: c = AIMUX_E_INVALID_TOOL_INPUT case .toolCallRepair: c = AIMUX_E_TOOL_CALL_REPAIR + case .loadApiKey: c = AIMUX_E_LOAD_API_KEY + case .loadSetting: c = AIMUX_E_LOAD_SETTING case .other: c = AIMUX_E_OTHER } return Int32(bitPattern: c.rawValue) @@ -362,6 +373,17 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl return nil } + /// `.loadApiKey` / `.loadSetting` only: the environment variable that was + /// consulted (e.g. `"OPENAI_API_KEY"`). + public var envVar: String? { + switch self { + case .loadApiKey(_, _, _, _, let v), .loadSetting(_, _, _, _, let v): + return v + default: + return nil + } + } + public var description: String { message } public var errorDescription: String? { @@ -451,6 +473,13 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl case AIMUX_E_TOOL_CALL_REPAIR: return .toolCallRepair(message: message, status: status, retryMs: retryMs, retryable: retryable, originalError: takeCString(aimux_error_original_error(h)) ?? "") + case AIMUX_E_LOAD_API_KEY: + // The consulted environment variable rides the provider-code getter. + return .loadApiKey(message: message, status: status, retryMs: retryMs, retryable: retryable, + envVar: takeCString(aimux_error_provider_code(h)) ?? "") + case AIMUX_E_LOAD_SETTING: + return .loadSetting(message: message, status: status, retryMs: retryMs, retryable: retryable, + envVar: takeCString(aimux_error_provider_code(h)) ?? "") case AIMUX_E_OTHER: return .other(message: message, status: status, retryMs: retryMs, retryable: retryable) default: @@ -723,8 +752,10 @@ public final class Model: @unchecked Sendable { /// registry entry. /// - modelId: Model id. /// - configJson: Optional JSON object of `ProviderOptions` - /// (`{"base_url": "...", "headers": {...}, "max_retries": 0, - /// "body_overrides": {...}}`); `nil` for defaults. + /// (`{"base_url": "...", "headers": {...}, "organization": "...", + /// "project": "...", "params": {...}}`); `nil` for defaults. + /// `max_retries` (a per-call option) and `body_overrides` (removed) are + /// rejected as `.invalidArgument`. public static func provider( name: String, apiKey: String? = nil, modelId: String, configJson: String? = nil ) throws -> Model { @@ -884,9 +915,9 @@ public final class Model: @unchecked Sendable { if let e = aimux_register_providers(configJSON) { throw expectAimuxError(e, context: "registerProviders") } } - /// Set the global proxy configuration (M6, RFC-0016). Must be called before - /// the first `generateText` / `streamText` call; a no-op (returns without - /// error) if the shared HTTP client is already initialised. + /// Replace the default transport with one routed by this proxy configuration + /// (M6, RFC-0016). Requests started afterwards use it; calling it again + /// replaces the previous configuration. /// /// `configJSON` shape: `{ "http_url", "https_url", "all_url", "no_proxy" }` /// (all fields optional; omitting all is equivalent to relying on the diff --git a/bindings/swift/Sources/Aimux/Types.swift b/bindings/swift/Sources/Aimux/Types.swift index 88337156e..a57b04317 100644 --- a/bindings/swift/Sources/Aimux/Types.swift +++ b/bindings/swift/Sources/Aimux/Types.swift @@ -1067,12 +1067,7 @@ public enum GenerateContent: Codable, Equatable { } } -public struct GenerateRequest: Codable, Equatable { - public var body: JSONValue? - public init(body: JSONValue? = nil) { self.body = body } -} - -public struct GenerateResponseMetadata: Codable, Equatable { +public struct ResponseInfo: Codable, Equatable { public var id: String? public var timestamp: String? public var modelId: String? @@ -1081,13 +1076,17 @@ public struct GenerateResponseMetadata: Codable, Equatable { enum CodingKeys: String, CodingKey { case id, timestamp, headers, body, modelId = "model_id" } - public init(id: String? = nil, timestamp: String? = nil, modelId: String? = nil, - headers: [String: String]? = nil, body: JSONValue? = nil) { - self.id = id; self.timestamp = timestamp; self.modelId = modelId - self.headers = headers; self.body = body + public init(id: String? = nil, timestamp: String? = nil, modelId: String? = nil, headers: [String: String]? = nil, body: JSONValue? = nil) { + self.id = id; self.timestamp = timestamp; self.modelId = modelId; self.headers = headers; self.body = body } } +public struct RequestInfo: Codable, Equatable { + public var body: JSONValue? + + public init(body: JSONValue? = nil) { self.body = body } +} + /// Raw provider result. public struct GenerateResult: Codable, Equatable { public var content: [GenerateContent] @@ -1095,21 +1094,20 @@ public struct GenerateResult: Codable, Equatable { public var usage: Usage? public var warnings: [Warning]? public var providerMetadata: JSONValue? - public var response: GenerateResponseMetadata? - public var request: GenerateRequest? + public var response: ResponseInfo? + public var request: RequestInfo? enum CodingKeys: String, CodingKey { case content case finishReason = "finish_reason" case usage, warnings case providerMetadata = "provider_metadata" - case response - case request + case response, request } public init(content: [GenerateContent], finishReason: FinishReason? = nil, usage: Usage? = nil, warnings: [Warning]? = nil, providerMetadata: JSONValue? = nil, - response: GenerateResponseMetadata? = nil, request: GenerateRequest? = nil) { + response: ResponseInfo? = nil, request: RequestInfo? = nil) { self.content = content; self.finishReason = finishReason; self.usage = usage self.warnings = warnings; self.providerMetadata = providerMetadata self.response = response; self.request = request @@ -1292,7 +1290,6 @@ public struct GenerateTextOptions: Codable, Equatable { public var providerOptions: JSONValue? public var reasoning: ReasoningEffort? public var instructions: String? - public var bodyOverrides: JSONValue? public var maxRetries: UInt32? public var timeout: TimeoutConfiguration? public var includeRawChunks: Bool? @@ -1346,7 +1343,6 @@ public struct GenerateTextOptions: Codable, Equatable { case headers case providerOptions = "provider_options" case reasoning, instructions - case bodyOverrides = "body_overrides" case maxRetries = "max_retries" case timeout case includeRawChunks = "include_raw_chunks" @@ -1360,7 +1356,7 @@ public struct GenerateTextOptions: Codable, Equatable { tools: [Tool]? = nil, toolChoice: ToolChoice? = nil, headers: [String: String]? = nil, providerOptions: JSONValue? = nil, reasoning: ReasoningEffort? = nil, instructions: String? = nil, - bodyOverrides: JSONValue? = nil, maxRetries: UInt32? = nil, + maxRetries: UInt32? = nil, timeout: TimeoutConfiguration? = nil, includeRawChunks: Bool? = nil, sessionId: String? = nil, @@ -1371,7 +1367,7 @@ public struct GenerateTextOptions: Codable, Equatable { self.output = output; self.seed = seed; self.tools = tools self.toolChoice = toolChoice; self.headers = headers; self.providerOptions = providerOptions self.reasoning = reasoning; self.instructions = instructions - self.bodyOverrides = bodyOverrides; self.maxRetries = maxRetries; self.timeout = timeout + self.maxRetries = maxRetries; self.timeout = timeout self.includeRawChunks = includeRawChunks self.sessionId = sessionId self.repairToolCallBox = repairToolCall.map(RepairToolCallBox.init(run:)) diff --git a/bindings/swift/Tests/AimuxTests/AimuxTests.swift b/bindings/swift/Tests/AimuxTests/AimuxTests.swift index 89058965c..6c6f20fb4 100644 --- a/bindings/swift/Tests/AimuxTests/AimuxTests.swift +++ b/bindings/swift/Tests/AimuxTests/AimuxTests.swift @@ -186,6 +186,25 @@ final class AimuxTests: XCTestCase { Int32(AIMUX_E_TIMEOUT.rawValue)) } + /// 19 / 20: a missing API key / provider setting have their own codes and + /// carry the consulted environment variable. + func testLoadErrorCodes() { + XCTAssertEqual(AIMUX_E_NO_OUTPUT_GENERATED.rawValue, 18) + XCTAssertEqual(AimuxError.noOutputGenerated(message: "no output", status: -1, + retryMs: -1, retryable: false).code, 18) + XCTAssertEqual(AIMUX_E_LOAD_API_KEY.rawValue, 19) + XCTAssertEqual(AIMUX_E_LOAD_SETTING.rawValue, 20) + let key = AimuxError.loadApiKey(message: "no key", status: -1, retryMs: -1, retryable: false, + envVar: "OPENAI_API_KEY") + XCTAssertEqual(key.code, Int32(AIMUX_E_LOAD_API_KEY.rawValue)) + XCTAssertEqual(key.envVar, "OPENAI_API_KEY") + XCTAssertNil(key.status) + let setting = AimuxError.loadSetting(message: "no region", status: -1, retryMs: -1, retryable: false, + envVar: "AWS_REGION") + XCTAssertEqual(setting.code, Int32(AIMUX_E_LOAD_SETTING.rawValue)) + XCTAssertEqual(setting.envVar, "AWS_REGION") + } + /// HTTP-shaped failures all arrive as `.apiCall`; the classification is the /// `status` field (there are no per-status cases any more). func testApiCallClassifiesByStatus() { @@ -363,9 +382,7 @@ final class AimuxTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let resultJson = try model.generateText(prompt: jsonEncodeString("What is Rust?")) let r = parseJSON(resultJson) XCTAssertEqual(r["text"] as? String, "Rust is a systems programming language.") @@ -380,9 +397,7 @@ final class AimuxTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let opts = jsonEncode([ "tools": [[ "type": "function", @@ -422,9 +437,7 @@ final class AimuxTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let prompt = jsonEncode([ ["role": "system", "content": "You are a helpful assistant."], ["role": "user", "content": "What is Rust?"], @@ -448,9 +461,7 @@ final class AimuxTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let opts = jsonEncode([ "tools": [[ "type": "function", @@ -476,9 +487,7 @@ final class AimuxTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) var parts: [String] = [] var streamErr: (any Error)? model.streamText( @@ -504,9 +513,7 @@ final class AimuxTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let opts = jsonEncode([ "tools": [[ "type": "function", @@ -558,7 +565,7 @@ final class AimuxTests: XCTestCase { let model = try Model.anthropic( apiKey: "test-key", modelId: "claude-3-5-sonnet-20241022", - baseUrl: server.baseURL + baseUrl: server.baseURL + "/v1" ) let r = parseJSON(try model.generateText(prompt: jsonEncodeString("Hello"))) XCTAssertEqual(r["text"] as? String, "Hello from Claude!") @@ -583,9 +590,7 @@ final class AimuxTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let opts = jsonEncode([ "tools": [[ "type": "function", diff --git a/bindings/swift/Tests/AimuxTests/MockHTTPServer.swift b/bindings/swift/Tests/AimuxTests/MockHTTPServer.swift index f77fdb474..5adc2b90e 100644 --- a/bindings/swift/Tests/AimuxTests/MockHTTPServer.swift +++ b/bindings/swift/Tests/AimuxTests/MockHTTPServer.swift @@ -11,12 +11,23 @@ // close` so each FFI call maps to exactly one connection. import Foundation +@testable import Aimux #if canImport(Glibc) import Glibc #elseif canImport(Darwin) import Darwin #endif +/// A model that speaks the OpenAI chat-completions wire format at `baseURL`: +/// the registry's OpenAI-compatible preset (`Model.openai` talks to the +/// Responses API). +func chatModel(_ baseURL: String) throws -> Model { + try Model.provider( + name: "deepseek", apiKey: "test-key", modelId: "gpt-4o", + configJson: #"{"base_url":"\#(baseURL)"}"# + ) +} + /// A preset HTTP response returned by `MockHTTPServer`. struct MockResponse { let status: Int diff --git a/bindings/swift/Tests/AimuxTests/WrapperTests.swift b/bindings/swift/Tests/AimuxTests/WrapperTests.swift index 97f9d0077..77edee6e4 100644 --- a/bindings/swift/Tests/AimuxTests/WrapperTests.swift +++ b/bindings/swift/Tests/AimuxTests/WrapperTests.swift @@ -21,9 +21,7 @@ final class WrapperTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let result = try model.generateText(prompt: .text("What is Rust?")) XCTAssertEqual(result.text, "Rust is a systems programming language.") @@ -85,9 +83,7 @@ final class WrapperTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let options = GenerateTextOptions(tools: [ .function(FunctionTool( name: "get_weather", @@ -123,9 +119,7 @@ final class WrapperTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let result = try model.generateText(prompt: .messages([ .system("You are a helpful assistant."), .user("What is Rust?"), @@ -151,9 +145,7 @@ final class WrapperTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let options = GenerateTextOptions( tools: [.function(FunctionTool( name: "get_weather", @@ -202,9 +194,7 @@ final class WrapperTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) var parts: [StreamPart] = [] var streamErr: (any Error)? model.streamText( @@ -243,9 +233,7 @@ final class WrapperTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let options = GenerateTextOptions(tools: [ .function(FunctionTool( name: "get_weather", @@ -289,9 +277,7 @@ final class WrapperTests: XCTestCase { try server.start() defer { server.stop() } - let model = try Model.openai( - apiKey: "test-key", modelId: "gpt-4o", baseUrl: server.baseURL - ) + let model = try chatModel(server.baseURL) let parts = try await collect(model.streamTextAsync(prompt: .text("Say hello"))) XCTAssertFalse(parts.isEmpty) let text = parts.reduce(into: "") { acc, part in diff --git a/docs/API.md b/docs/API.md index 12ca1bd86..b2c2bc2b1 100644 --- a/docs/API.md +++ b/docs/API.md @@ -60,19 +60,28 @@ OpenAI, Anthropic, Google, Bedrock, Vertex, Azure, Cohere, Mistral, xAI, Anthropic-AWS — one constructor per provider: `openai(apiKey, model, baseUrl?)` / `anthropic(apiKey, model, baseUrl?)` (Node, Python), `NewOpenAI(apiKey, model)` (Go), `Model.openai(apiKey, modelId)` (Java, Kotlin, Flutter), -`Aimux.openai(apiKey:modelId:)` (Swift), `OpenAIProvider::new(..)` (Rust). +`Aimux.openai(apiKey:modelId:)` (Swift), `create_openai(..)` (Rust). Multimodal, local-inference and search providers have their own constructors too — full list: [reference.md](api/reference.md). ### OpenAI-compatible (registry-backed) -One function in every binding: +Rust uses `aimux_providers::create_provider(name, PresetSettings)` for vendor +packages and all 282 presets, then `Provider::language_model(model_id)`. +`default_providers()` supplies a map for +`aimux_core::provider_registry::create_provider_registry(providers, options)`, +whose model ids default to `provider:model`. `provider_names()` lists accepted +names; `Provider::discovery()` exposes runtime model listing. + +The language bindings retain: ```text -provider(name, api_key?, model_id, config?) // all languages - name — registry provider name as a string +provider(name, api_key?, model_id, config?) // language bindings + name — vendor package or registry preset name as a string api_key — optional; omitted/None reads the provider's env var - config — optional overrides (base_url / headers / maxRetries / body_overrides) + config — optional overrides (base_url / headers / organization / project / params) + `maxRetries` (a per-call option) and `bodyOverrides` (removed) are + rejected with InvalidArgument, not ignored ``` - Full list (name / env var / base URL): [providers.md](api/providers.md) @@ -100,9 +109,9 @@ Examples: [Node.js](api/node.md#text-generation) · [Python](api/python.md#text- | `tool_choice` | `ToolChoice?` | Tool selection strategy | | `instructions` | `string?` | System instructions | | `reasoning` | `ReasoningEffort?` | Reasoning effort | -| `max_retries` | `number?` | Per-call retry override; `0` disables retries (`None` = provider default, 2) | +| `max_retries` | `number?` | Per-call retry override; `0` disables retries (`None` = 2) | | `timeout` | `TimeoutConfiguration?` | Per-call timeouts (total / step / first-chunk / chunk idle) — see [Timeouts](#timeouts) | -| `body_overrides` | `object?` | Per-call request-body overrides, deep-merged; `null` values delete keys | +| `provider_options` | `object?` | Per-provider options keyed by namespace; for OpenAI-compatible providers, unknown fields of the provider's own namespace go to the request body as given | | `headers` | `object?` | Extra HTTP headers | Node.js additionally accepts an `AbortSignal` as the 4th argument of @@ -202,6 +211,10 @@ when the call is aborted. | `Error` | Stream error | | `Raw` | Provider raw chunk (for debugging, when `include_raw_chunks` is set) | +`stream_text` emits the call-layer `TextStreamPart` type. Provider-layer +`StreamPart::ResponseMetadata` events are consumed internally and do not +appear in this stream. + ### Request Cancellation (abort) Calls can be cancelled mid-flight. Cancellation covers the whole request @@ -269,7 +282,7 @@ Per-call timeout limits, JSON-serializable in every binding: `None`/absent disables the corresponding limit. On expiry the call fails with `AiMuxError::Timeout` (not retryable); streaming timeouts surface as a -`StreamPart::Error` item (`"first chunk timeout"` / `"chunk idle timeout"` / +`TextStreamPart::Error` item (`"first chunk timeout"` / `"chunk idle timeout"` / `"total timeout"`). Unrepresentable values (e.g. `u64::MAX` ms on narrower platforms) are rejected with `AiMuxError::InvalidArgument` instead of panicking. When both abort and a deadline are in play, abort wins. @@ -353,7 +366,7 @@ Examples: [Node.js](api/node.md#image-generation) · [Python](api/python.md#imag ### Video Generation -Video generation typically returns a URL (not binary). +Most video models return a URL; the video operation downloads it, as the AI SDK's `generateVideo` does, and returns the bytes. Examples: [Node.js](api/node.md#video-generation) · [Python](api/python.md#video-generation) · [Rust](api/rust.md#video-generation) · [Go](api/go.md#video-generation) · [C ABI](api/c.md#video-generation) diff --git a/docs/PROJECT-OVERVIEW.md b/docs/PROJECT-OVERVIEW.md index b29c96998..24b73979e 100644 --- a/docs/PROJECT-OVERVIEW.md +++ b/docs/PROJECT-OVERVIEW.md @@ -19,7 +19,7 @@ aimux does not do agent loops, RAG, or orchestration — it focuses solely on un | Metric | Value | |------|------| | Rust code | 144,500+ lines | -| AI providers | 327 (251 registry-backed OpenAI-compatible + 76 typed; counts by category in [api/providers.md](api/providers.md), generated) | +| AI providers | 327 accepted names (282 registry-backed presets + 45 vendor packages; listed by `provider_names()`) | | Modality traits | 8 (text/embedding/image/video/speech/transcription/reranking/search) | | Test cassettes | 2,650 recorded replays | | Test files | 118 | @@ -28,10 +28,10 @@ aimux does not do agent loops, RAG, or orchestration — it focuses solely on un | Type definitions | 79 (auto-generated by ts-rs) | > **Provider access (RFC-0017 phase 4)**: all OpenAI-compatible providers are -> registry-backed — `provider-registry.json` is the single source of truth, and -> every binding exposes the string-keyed `provider(name, ...)`. The -> per-provider shell types (`XxxConfig`/`XxxProvider`) -> were retired in phase 4. +> represented by presets or vendor packages. Rust uses +> `create_provider(name, PresetSettings)`; bindings retain `provider(name, ...)`. +> Presets live in `provider_registry.json`; Groq and DeepSeek have their own +> vendor packages and chat models. Per-preset shell types were retired. --- @@ -118,14 +118,15 @@ Counts by category live in [api/providers.md](api/providers.md#totals) | Type | Representatives | |------|------| -| Native protocol | OpenAI, Anthropic, Google, Bedrock, Vertex, Azure, Cohere, Mistral, xAI, Codex, OpenRouter | -| OpenAI-compatible (registry) | Groq, Fireworks, Together, Perplexity, Ollama, Alibaba Tongyi, Zhipu, Baidu, Tencent, Moonshot AI, SiliconFlow… | +| Native protocol | OpenAI, Anthropic, Google, Bedrock, Vertex, Azure, Cohere, Mistral, Groq, DeepSeek, xAI | +| OpenAI-compatible (registry) | Fireworks, Together, Perplexity, Ollama, Alibaba Tongyi, Zhipu, Baidu, Tencent, Moonshot AI, SiliconFlow… | | Voice/transcription | ElevenLabs, Deepgram, AssemblyAI, Cartesia… | | Image/video | Black Forest Labs, Replicate, Fal, KlingAI… | | Search | Tavily, Exa, Serper, Firecrawl… | See [api/providers.md](api/providers.md) for the full list with counts. -All registry providers are accessed via the unified `provider(name, ...)` entry point +Rust providers are accessed via `create_provider(name, PresetSettings)`; +bindings retain their `provider(name, ...)` entry point — see [API.md](API.md#providers). ### Data model (aligned with AI SDK V4) @@ -163,7 +164,7 @@ aimux/ │ └── RerankingModel / SearchModel ├── aimux-providers # provider implementations (canonical counts: docs/api/providers.md) │ ├── native protocols # standalone model + convert, handles provider-specific differences -│ ├── OpenAI compatible # registry-backed: provider-registry.json + provider(name, ...) entry (RFC-0017 phase 4) +│ ├── OpenAI compatible # registry-backed: provider_registry.json + create_provider(name, PresetSettings) (RFC-0017 phase 4) │ └── modalities/search # voice / image / video / search implementations ├── aimux-stream # SSE parsing ├── aimux-provider-utils # One-exchange HTTP helpers, response handlers, API-key loading, streamed tool-call tracking @@ -194,21 +195,22 @@ pub trait LanguageModel: Send + Sync { This means `Box` can be swapped between OpenAI / Anthropic / Google — **switching providers only requires changing the provider construction, and the model usage stays exactly the same**. -#### 2. OpenAICompatProfile configuration descriptor struct +#### 2. Provider factories and registry presets -A thin wrapper does not lose differences. Each OpenAI-compatible service has subtle differences (some support top_k, some do not support tools, some have a different streaming usage format): +A thin wrapper does not lose differences. Each OpenAI-compatible service has subtle differences (some do not accept `top_k`, some use `max_completion_tokens`, some have a different streaming usage format). They are expressed once, as data and dialect hooks, instead of one hand-written shell type per vendor: ```rust -pub struct OpenAICompatProfile { - pub supports_top_k: bool, // Groq supports it, OpenAI does not - pub supports_tools: bool, // some services do not support it - pub supports_response_format: bool, - pub streaming_usage_format: UsageFormat, // streaming usage in data line vs chunk line - pub post_process_request: Option, -} +// Every package follows the AI SDK factory shape. +let provider = create_openai_compatible(OpenAICompatibleProviderSettings { + name: "acme".into(), // model.provider() == "acme.chat" + base_url: "https://api.acme.example/v1".into(), + api_key: Some("sk-...".to_string().into()), // resolved per request; None sends no Authorization + ..Default::default() +})?; +let model = provider.chat("acme-large"); ``` -Use a profile descriptor struct to express the differences, rather than writing an independent model for each provider — this is why the 145 compatible providers only need a thin wrapper. +The 282 presets are rows of `provider_registry.json` (base URL, key variable, `auth: api_key | none`, template parameters, `profile.max_tokens_key`); the embedded JSON is parsed once into a runtime descriptor table. `create_provider(name, PresetSettings)` creates presets or vendor packages by name. There is no `family` field or generated presets source. Groq and DeepSeek have their own chat models; xAI and Hugging Face use Responses for text. #### 3. JSON string FFI boundary @@ -360,12 +362,15 @@ for part in result: ```rust use aimux_core::prelude::*; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::{OpenAIProviderSettings, create_openai}; #[tokio::main] async fn main() -> Result<(), AiMuxError> { - let provider = OpenAIProvider::new(OpenAIConfig::new("sk-...")); - let model = provider.model("gpt-4o"); + let provider = create_openai(OpenAIProviderSettings { + api_key: Some("sk-...".to_string().into()), + ..Default::default() + })?; + let model = provider.chat("gpt-4o"); let result = generate_text( &model, @@ -410,20 +415,15 @@ Not a single trick, but systematic design choices: ### 3. The way to unify all providers -Don't write an independent model for each provider — that would explode. Use `OpenAICompatProfile` to describe the differences: +Don't write an independent model for each provider — that would explode. A compatible chat model plus per-vendor preset rows handles compatible services; Groq and DeepSeek have their own vendor chat models: ```rust -// one profile describes a compatible provider's differences -let profile = OpenAICompatProfile { - supports_top_k: true, // Groq supports top_k - supports_tools: true, // supports function calling - streaming_usage_format: UsageFormat::InDataLine, - post_process_request: Some(strip_unsupported_fields), - ..Default::default() -}; +// Create a vendor provider by name, then select its model. +let provider = aimux_providers::create_provider("groq", Default::default())?; +let model = provider.language_model("llama-3.3-70b-versatile")?; ``` -The native protocol providers have independent models + convert (handling differences such as Anthropic message format / Google generateContent / Bedrock SigV4), while the OpenAI-compatible providers are registry-backed (provider-registry.json + unified provider(name, ...) entry, RFC-0017 phase 4). +The native protocol providers have independent models + convert (handling differences such as Anthropic message format / Google generateContent / Bedrock SigV4), while compatible presets use provider_registry.json and a runtime descriptor table. `create_provider(name, PresetSettings)` selects either a preset or a vendor package. ### 4. Recorded testing with 2650 cassettes @@ -439,7 +439,7 @@ aimux's type design directly targets Vercel AI SDK V4 provider types — the `Ge ### Completed -- [x] 327 providers integrated (registry-backed OpenAI-compatible + native/modality typed; counts by category in [providers.md](api/providers.md), generated) +- [x] 327 provider names accepted by `create_provider` (282 presets + 45 vendor packages; `provider_names()` lists them) - [x] 8 modality traits (text/embedding/image/video/speech/transcription/reranking/search) - [x] 7 language bindings (Node/Python/Swift/Kotlin/Flutter/C/Rust) - [x] 2650 cassette recorded tests diff --git a/docs/ai-sdk-request-pipeline.md b/docs/ai-sdk-request-pipeline.md index c6f9c1570..2ecaf9315 100644 --- a/docs/ai-sdk-request-pipeline.md +++ b/docs/ai-sdk-request-pipeline.md @@ -88,7 +88,7 @@ Core user operation | error context 安全 | AI SDK 传原始 request context；Aimux 写入 public error 前使用统一白名单/脱敏/大小限制，binary 与大型 data URL 只保留摘要 | | 错误 body 上限 | AI SDK 全量读取错误 body（仅受 2 GiB 防 OOM 上限，超限直接抛 `DownloadError` 替换原错误）；Aimux 的 `ApiCallError` 会跨 FFI 序列化并写入 recording，因此错误 body 采用 best-effort 截断读：public `response_body` 上限 64 KiB（lossy 解码后按字符边界执行，带 `…(truncated)` 标记），解析上限 1 MiB 保证超大但合法的错误 JSON 仍能进 provider mapper，读取中途连接死亡保留已收到的部分 | | 成功 JSON body 上限 | AI SDK 对成功 body 复用同一个 2 GiB 下载上限；Aimux 的成功 JSON body 会同时以 bytes + `T` 的形式驻留（`raw_value` 是 best-effort 的第二次解析，不再额外持有 `Value` 克隆），2 GiB 上限会放大峰值内存，因此 `create_json_response_handler` 改用独立更小的默认值 `DEFAULT_MAX_JSON_RESPONSE_SIZE`（64 MiB），可通过 `HttpRequest::max_json_response_bytes` 按请求覆盖；`create_binary_response_handler` 不受影响，仍使用 `DEFAULT_MAX_DOWNLOAD_SIZE`（2 GiB） | -| provider 默认 retry | 保留既有 `RetryConfig` 源码与行为兼容；Core 读取 model config，per-call 只覆盖 `max_retries`（目标设计已撤销：[aisdk-architecture-alignment.md](aisdk-architecture-alignment.md) D14、§9.4 删除 model/provider 级默认 retry，组合模型改为给每个 child 配预算；删除它的那层合入前，这里仍是主线现状） | +| provider 默认 retry | 调用级常量默认（`max_retries` 2 / `initial_delay` 2000 ms / `backoff_factor` 2）；provider 没有 retry 设置，`prepare_retries(max_retries, abort)` 是唯一读取点，per-call 只传 `max_retries`（RFC-0036 D-e） | | jitter | AI SDK 默认不抖动但留 `getDelayInMs` 注入点；Aimux 用同一注入点默认注入 RFC-0009 Full Jitter，只作用于 exponential delay，server hint 精确遵守（§6.5） | | 第一条 SSE error | AI SDK OpenAI provider 会扫描到首个 semantic output；Aimux 保留 RFC-0016 的更窄 first-event peek，使立即到达的 error 成为 retry 边界内的 attempt 失败（§8.3） | | 默认超时 | 与 AI SDK 一致无默认 `total_ms`；Aimux 有意保留非流式 exchange 的 30s whole-response 上限，但从 shared-client 全局配置下移到单次 exchange，流式 exchange 豁免（§5.4） | @@ -296,7 +296,7 @@ redaction helper。不得依赖每个 Provider 自己记得脱敏。 它们遵循 AI SDK 的单次 fetch attempt 契约： -- 不接受 `RetryConfig`； +- 不接受 retry 配置； - 不接受 `TimeoutConfiguration`； - 不执行 backoff； - 不调用自身或另一个 API helper 进行 retry； @@ -357,17 +357,18 @@ fetch 无默认 exchange timeout 的有意差异，CHANGELOG 必须明确。 | `mergeAbortSignals` | 不逐字移植；Core 的 `tokio::select!` 同时观察 caller signal、deadline 和 operation future（§8.0） | | `setAbortTimeout` | `timeout::OperationTimeout` 保存 deadline，驱动 future 直接 `sleep_until`（§8.0） | -Core operation 把 per-call override、既有 model `RetryConfig` 和 caller abort 交给同名函数： +Core operation 把 per-call override 和 caller abort 交给同名函数： ```text -prepare_retries(max_retries, retry_config, abort_signal) +prepare_retries(max_retries, abort_signal) ``` -`max_retries=None` 使用 `retry_config.max_retries`（默认 2）；per-call `Some(n)` 只覆盖 -该计数，`initial_delay` 和 `backoff_factor` 保持既有配置。`RetryConfig` 的 canonical -定义移到 `aimux-core::retry`，`aimux-provider-utils::RetryConfig` 和 -`aimux_provider_utils::retry::RetryConfig` 都 re-export 同一类型。旧 -`retry_config` / `with_retry_config` 保留，不增加第二套 provider config 命名。 +默认值是调用级常量（`DEFAULT_MAX_RETRIES` = 2、`DEFAULT_INITIAL_DELAY_MS` = 2000、 +`DEFAULT_BACKOFF_FACTOR` = 2，与 AI SDK 一致）：`max_retries=None` 用 2，per-call +`Some(n)` 只覆盖该计数，`initial_delay` 和 `backoff_factor` 是常量。**provider 没有 retry +设置**：不再有 provider 级的 retry 配置类型、配置字段、builder 方法或 model trait 上的只读 +getter，`prepare_retries` 是 retry 配置的唯一读取点（RFC-0036 D-e）。jitter、`Retry-After` +和错误历史不变。 两层函数的边界与 AI SDK 相同，**不得合并成一个 APICall-aware 的 primitive**： @@ -395,7 +396,7 @@ async fn retry_with_exponential_backoff_respecting_retry_headers<..>( } ``` -`prepare_retries(max_retries, retry_config, abort_signal)` 返回 `PreparedRetries { max_retries, retry }` +`prepare_retries(max_retries, abort_signal)` 返回 `PreparedRetries { max_retries, retry }` 的 Rust 形式是 `PreparedRetries::retry(&self, op)`，与 AI SDK 的 `{ maxRetries, retry }` 一致； 不另起额外的 retry 类型。 @@ -788,6 +789,10 @@ start/status 流），Core 拥有 poll 循环并分别 retry 两个阶段： 设置；`poll.timeout` 是整个轮询的预算：每次 `do_status`（含其重试）与剩余时间 竞争，超时即丢弃挂起的请求，以 `Video generation timed out after {timeoutMs}ms.` 失败（与上游一致，ai 7.0.127 `generate-video/generate-video.ts:543-549`、`:586-625`）； +- 所有批次完成后，Core 按顺序下载以 URL 返回的视频，返回 `VideoData::Binary` + （上游 `generateVideo` 返回 `GeneratedFile`，ai 7.0.127 + `generate-video/generate-video.ts:383-410`）；传输为 `VideoCallOptions::download`， + 缺省为受控下载 `create_download(None)`。下载不重试，计入 operation timeout； - Provider-specific polling delay 不是通用 exponential retry； - 耗尽产生的 `RetryError` 原样透传，外层不得再次 submit。 - **修订**（RFC-0031 review）：`options.n` 按 `VideoModel::max_videos_per_call` @@ -822,15 +827,14 @@ Composite 外层不重放；重试放在已有语义边界内： ### 9.4 Files -`Files::upload_file` 在本 RFC 中直接作为公开 operation，没有对应的 -`do_upload_file`；本轮不为它增加第二层。（上游另有用户操作 `uploadFile`，直接调用 -provider 的 `files().uploadFile` 且不重试，见 ai 7.0.127 `upload-file/upload-file.ts:37`、 -`:94`。V4 分层中删除 provider 级 `RetryConfig` 的那层会去掉这里的重试，与上游一致； -该层合入前，下面描述的重试仍是主线现状。）**修订**（RFC-0031 review）：`upload_file` -自己用 Provider 已配置的 `RetryConfig` 重试（与 `execute_list_models` 同一 -primitive）。这是“创建资源类 POST 不自动重试”一般规则的显式例外，前提是 -upload 不计费、失败的 create-file 请求不会返回可复用的 file id。多阶段的 -Google Files 按 exchange 分别重试，而不是包住整个 `upload_file`。 +`Files::upload_file` 按 AI SDK Provider SPI 本身就是公开 operation，没有对应的 +`do_upload_file`；本轮不为它虚构第二层。**修订**（RFC-0031 review）：`upload_file` +自己重试（RFC-0031 review 的修订）。**RFC-0036 撤销该修订**：provider 没有 retry +设置，`upload_file`（OpenAI、Anthropic、Google Files 的多阶段上传同样）与 +`list_models` 一样是单次 exchange，不重试——“创建资源类 POST 不自动重试”的一般规则 +不再有例外；需要重试的调用方自己包一层。上游的用户操作 `uploadFile` 同样直接调用 +provider 的 `files().uploadFile` 且不重试（ai 7.0.127 `upload-file/upload-file.ts:37`、 +`:94`）。 --- @@ -856,9 +860,9 @@ search 所有 binding 调用这些 user operations，不得直接调用 provider-facing `do_*`。`do_*` 继续作为 Provider SPI。 -（目标设计已撤销：[aisdk-architecture-alignment.md](aisdk-architecture-alignment.md) D14、§9.4 删除 model/provider 级默认 retry，组合模型改为给每个 child 配预算；删除它的那层合入前，这里仍是主线现状）为保留现有 model-level 配置，各 model trait 提供只读 `retry_config()`；只有 -Core user operation 读取它。Provider 的 `do_*` 不得重试整个 operation；§9.1 的 -已创建 job 安全 exchange 例外除外。 +retry 只有调用级一个入口（`prepare_retries(max_retries, abort)`），model trait 上没有 +任何 retry getter。Provider 的 `do_*` 不得重试整个 operation；§9.1 的已创建 job 安全 +exchange 例外除外（轮询阶段按 provider 包内常量有界执行，`shared::poll`）。 ### 10.2 Recording @@ -919,9 +923,9 @@ Node/Python/Go/Java/Kotlin/Swift/Flutter。 | `parse_provider_error` | 删除；由 failed response handler 替代 | | `TimeoutBodyStream` | 删除；由 Core semantic stream timeout 替代 | | `get_retry_delay_ms_with_jitter` | 删除；Full Jitter 改为 `aimux-core::retry` 里 `get_delay_ms` hook 的默认实现，且不再抖动 server hint（§6.5） | -| Provider `resolve_retry_config` | 删除；effective max retries 由 Core 解析 | -| Provider-owned retry execution | 删除；旧 `RetryConfig` 类型/配置 API 保留，由 Core 执行 | -| `aimux-provider-utils/src/retry.rs` | 仅保留 `RetryConfig` 兼容 re-export；retry 实现只在 `aimux-core::retry` | +| Provider 内的 retry 解析函数 | 删除；effective max retries 由 Core 解析（`prepare_retries`） | +| Provider-owned retry execution | 删除；provider 没有 retry 设置，Core 以调用级常量默认执行 | +| `aimux-provider-utils/src/retry.rs` | 删除；retry 实现只在 `aimux-core::retry` | | `logging.rs::is_sensitive_key` | 删除；统一用 `aimux_core::recording::is_sensitive_key` | | `shared.rs::AbortSignal` | 实现迁到 `abort_signal.rs`，旧路径 re-export；timeout 不塞进 signal | diff --git a/docs/aisdk-architecture-alignment.md b/docs/aisdk-architecture-alignment.md index 2b387b89d..8660a277f 100644 --- a/docs/aisdk-architecture-alignment.md +++ b/docs/aisdk-architecture-alignment.md @@ -2,9 +2,9 @@ > **性质**: [RFC-0036](../rfc/0036-positioning-and-layered-architecture.md)(定位与分层架构)之下的实施设计文档,不是独立 RFC,不占 RFC 编号。随本文件合入生效;§0.4 的四项决定、§0.7 对 ROADMAP 既有承诺的保留 / 调整 / 后置清单,以本文件为准,后续变更直接修改本文件并在 CHANGELOG 记录 > -> **Date**: 2026-09-29 · 2026-10-04 修订（撤销切换后的 ABI 共存承诺、登记 xAI/HF chat 扩展、修复附录 §6.4-14 裁定缺失等评审问题）· 2026-10-12 修订（§9.1 / Q3 / D31 改为逐层合入主线、全部合完前不发布；新 C ABI、生成链、完整录制 schema 3 与绑定迁移排在最后）· 2026-10-13 修订（删除未实现且不再采用的 aimux 自有设计：LanguageModelRef / ResolvedModel / RegistryModelReference、注册 sidecar 元数据、SessionInferer trait、session affinity middleware；体积与性能门禁、能力矩阵移到 ROADMAP；§9.1 登记已知差距）· 2026-10-13 修订（[4/8] 结果与流：对象生成改为 `output` 的 deprecated 包装（§3.5）；异步 job 跟随上游 start/status 与轮询预算（§4.5、§9.3）；登记视频多批 idempotency-key、流中途重试、首块超时起点（§5）；结果便捷字段 deprecated 与回放规则补充（第二部分）） +> **Date**: 2026-09-29 · 2026-10-04 修订（撤销切换后的 ABI 共存承诺、xAI/HF chat 扩展登记后撤回、preset 改为运行时 descriptor 表、重定向规则收归 helper、修复附录 §6.4-14 裁定缺失等评审问题） · 2026-10-12 修订（§9.1 / Q3 / D31 改为逐层合入主线、全部合完前不发布；新 C ABI、生成链、完整录制 schema 3 与绑定迁移排在最后）· 2026-10-13 修订（删除未实现且不再采用的 aimux 自有设计：LanguageModelRef / ResolvedModel / RegistryModelReference、注册 sidecar 元数据、SessionInferer trait、session affinity middleware；体积与性能门禁、能力矩阵移到 ROADMAP；§9.1 登记已知差距）· 2026-10-13 修订（[4/8] 结果与流：对象生成改为 `output` 的 deprecated 包装（§3.5）；异步 job 跟随上游 start/status 与轮询预算（§4.5、§9.3）；登记视频多批 idempotency-key、流中途重试、首块超时起点（§5）；结果便捷字段 deprecated 与回放规则补充（第二部分）） > -> **Reference baseline**: `reference/aisdk-review/node_modules`(本地安装、带 `src/`):`ai` 7.0.122、`@ai-sdk/provider` 4.0.19、`@ai-sdk/provider-utils` 5.0.51、openai 4.0.80、anthropic 4.0.68、google 4.0.85、google-vertex 5.0.98、amazon-bedrock 5.0.100、azure 4.0.84、xai 5.0.12、mistral 4.0.54、cohere 4.0.52、deepseek 3.0.56、groq 4.0.52、openai-compatible 3.0.59、gateway 4.0.100 +> **Reference baseline**: `reference/aisdk-review/node_modules`(本地安装、带 `src/`):`ai` 7.0.122、`@ai-sdk/provider` 4.0.19、`@ai-sdk/provider-utils` 5.0.51、openai 4.0.80、anthropic 4.0.68、google 4.0.85、google-vertex 5.0.98、amazon-bedrock 5.0.100、azure 4.0.84、xai 5.0.12、mistral 4.0.54、cohere 4.0.52、deepseek 3.0.56、groq 4.0.52、openai-compatible 3.0.59、gateway 4.0.100。这是调研与正文 file:line 证据对应的版本;实现用于行为对照的协议 fixture 由随实现提交入库的 `fixtures/aisdk/VERSIONS.json` 单独锁定,可晚于此处,升级时以该文件为准 > > **Scope**: provider 创建与设置求值、model 获取与解析(registry / 默认 provider)、传输(fetch/WS)、调用运行时(retry / timeout / telemetry / scope)、消息协议(用户输入 / provider 输入 / provider 输出 / core 输出四层)、recording / replay / trace / session / catalogue、组合模型、FFI 与 8 个绑定、CLI / Web 工具 > @@ -12,7 +12,7 @@ > > **Supersedes**: 早期仅对齐 provider 工厂外观的草案(本文件的前身);[RFC-0008](../rfc/0008-multimodal-bindings.md)「Alternative: Unified Provider Object」的 Decision;`docs/internal/handoff/HANDOFF.md` §11 > -> **Amends**: [ROADMAP](../ROADMAP.md) 与 [RFC-0036](../rfc/0036-positioning-and-layered-architecture.md) 中与切换方式、L2 选型相关的承诺(逐项见 §0.7,旧文档中对应段落已加注指向本文);[RFC-0031](ai-sdk-request-pipeline.md)(provider 默认 retry、composite 零重试、helper 录制等有意差异被撤销或改写,见第一部分 §9.4)、[RFC-0017](../rfc/0017-provider-config-dx.md)(通用 body override 路线撤销)、[RFC-0020](../rfc/0020-external-provider-config.md)(全局 overlay 撤销)、[RFC-0023](../rfc/0023-runtime-request-recording.md)(录制 schema 3)、[RFC-0027](../rfc/0027-model-catalogue-and-list-api.md)(one-shot 工厂承诺撤销) +> **Amends**: [ROADMAP](../ROADMAP.md) 与 [RFC-0036](../rfc/0036-positioning-and-layered-architecture.md) 中与切换方式、L2 选型相关的承诺(逐项见 §0.7,旧文档中对应段落已加注指向本文);[RFC-0031](ai-sdk-request-pipeline.md)(provider 默认 retry、composite 零重试、helper 录制等有意差异被撤销或改写,见第一部分 §9.4)、[RFC-0017](../rfc/0017-provider-config-dx.md)(通用 body override 路线撤销)、[RFC-0020](../rfc/0020-external-provider-config.md)(全局 overlay 撤销)、[RFC-0023](../rfc/0023-runtime-request-recording.md)(录制 schema 3)、[RFC-0027](../rfc/0027-model-catalogue-and-list-api.md)(one-shot 工厂承诺撤销)、[RFC-0016](../rfc/0016-align-with-aisdk.md)(L2 tool-call input 维持 `Value`、L7 不做 metadataExtractor 等结论由第二部分与 D15 取代) > > **Appendix**: [aisdk-architecture-impact-map.md](aisdk-architecture-impact-map.md) —— 全仓库 12 模块、201 条经对抗核验的发现与跨模块耦合链路 @@ -48,8 +48,8 @@ AI SDK 的做法:provider 把 settings 编译成闭包(`url()`、`headers()`(可 aimux-provider V4 规范:Provider / 各模态 Model trait、V4 prompt / content / stream part、错误、middleware 规范 ↑ (对应 @ai-sdk/provider;无运行时依赖,即第二部分所说的"协议叶子 crate") aimux-provider-utils Fetch / WS / Resolvable / load_* / headers / 单次 HTTP helper / 下载守卫 / 传输诊断 / ModelMessage 用户层类型 - ↑ ← aimux-stream (对应 @ai-sdk/provider-utils;aimux-stream 保持独立的 SSE / NDJSON 解析库) -aimux-providers create_xxx 工厂、私有 model config、协议转换、鉴权组合、preset(生成)、catalogue 摄取 + ↑ ← aimux-stream (对应 @ai-sdk/provider-utils;aimux-stream 保持独立的 SSE 解析库) +aimux-providers create_xxx 工厂、私有 model config、协议转换、鉴权组合、preset(运行时注册表)、catalogue 摄取 aimux 用户操作、resolve / registry / custom / wrap、prompt 标准化与下载、retry / timeout、telemetry、operation scope、组合模型 aimux-devtools recording(订阅事件 + 传输诊断)、replay_fetch / replay_web_socket、trace middleware、缓存审计、session store FFI / Node / Python / 5 个 C-ABI 绑定 / CLI / Web 由 descriptor + manifest 单一生成链驱动 @@ -112,7 +112,7 @@ FFI / Node / Python / 5 个 C-ABI 绑定 / CLI / Web 由 descriptor + manifest | 0.6 版本语义"binding wire 与 C ABI 不变";0.8 C1 "旧导出保留共存(#166 要求共存至少一个 minor 版本)";1.0 C4 删除旧导出 | 调整 | 本文的切换(§9.1 P1,逐层合入主线,全部合完后统一发布)是一次 **breaking 发布**:对象模型、wire、C ABI 同时替换,不提供旧导出共存、deprecated 别名或转发层(D31、§5"不采用")。发布为一个 0.x minor,CHANGELOG 给迁移说明。此后 ops 协议(见下一行)引入时的旧符号同样按切换门一次替换,不再承诺"共存一个 minor"(与 [RFC-0036 §12.2](../rfc/0036-positioning-and-layered-architecture.md#122-明确替代清单) 的 C2 / C4 行一致;切换门由 [RFC-0039](../rfc/0039-ops-bindings-contract.md) 定义) | | ops 协议(op 表 / 错误信封 / 二进制帧 / 版本协商)、stdio CLI 入口(RFC-0036 §4;ROADMAP 0.6 schema、0.8 C1) | 保留,后置到本文切换之后 | 本文 §6 替换的是现有 123 个 `extern "C"` 的句柄对象模型与错误传输;`dispatch(op, handle, json)` 应建立在 §6.1 的句柄种类与 §6.5 的 manifest 之上,因此排在 P1 之后。本文不设计 ops 协议 | | L0 native passthrough(RFC-0036 §3 原则 4;ROADMAP 0.6) | 保留,后置到本文切换之后 | 落点是 §3.2 的 `Fetch` / `WsConnector` 注入与 leaf 诊断:passthrough 复用同一传输边界即可获得 auth / 重试 / 录制。本文不实现它 | -| B 轨 protocol registry = L1 协议层(RFC-0032,B1 "protocol 列 + `from_resolved`",B2–B8 "17 个 LanguageModel → ~7 个协议、33 个 wrapper 退役") | 保留方向,形式调整 | "协议"对应本文 §3.3 / §4.1 的模型实现家族(OpenAI chat / Responses、Anthropic Messages、Google、compat 等);251 个 registry preset 由 descriptor 生成独立厂商工厂并复用 compat 实现(D12、D20),wrapper 退役由此完成。RFC-0032 的 registry 列保留为 descriptor 的数据来源;`from_resolved` 式构造由 `create_xxx` 工厂替代;C2 "FFI 构造器转发 shim"不再需要(§6 一次切换) | +| B 轨 protocol registry = L1 协议层(RFC-0032,B1 "protocol 列 + `from_resolved`",B2–B8 "17 个 LanguageModel → ~7 个协议、33 个 wrapper 退役") | 保留方向,形式调整 | "协议"对应本文 §3.3 / §4.1 的模型实现家族(OpenAI chat / Responses、Anthropic Messages、Google、compat 等);registry preset 由内嵌的 registry 在运行时解析成 descriptor 表,按名字创建并复用 compat 实现(D12),不为每一行生成 Rust 源码,wrapper 退役由此完成。RFC-0032 的 registry 列保留为 descriptor 的数据来源;`from_resolved` 式构造由 `create_xxx` 工厂替代;C2 "FFI 构造器转发 shim"不再需要(§6 一次切换) | | #174 Auth L1(registry `auth` schema + `apply_auth()`)、#175 Auth L2(Credential 解析序 + CredentialStore + TokenRefresher) | 调整 | 鉴权按 D11 在各厂商包内以 headers 闭包 / fetch 装饰器 / Rust 内建 credential resolver 实现(§3.3、§4.1),registry `auth` 列作为 preset descriptor 数据保留,`apply_auth()` 即生成 headers 闭包。CredentialStore / TokenRefresher 涉及宿主回调与有状态刷新,与 Q4 一并后置;RFC-0018 的无状态刷新边界保留(§4.1 Codex 行) | | #167 transport-level replay("mock 挂 HTTP 层跑真协议代码";`ProviderRecord` = registry 行 + protocol)、#179 replay 子命令 | 保留,机制按本文 | 即 D4 / D5 的 `replay_fetch` / `replay_web_socket`(§4.3):录制与回放共用 provider-utils 的 leaf 边界,真实 provider 编码器与解析器参与回放。"`ProviderRecord` = registry 行 + protocol"撤销:录制 schema 3 不保存可重建的 ProviderRecord,live replay 由宿主 registry 与 operation 目标引用驱动(D24、§4.4)。#179 子命令保留在 tools | | #185 "删 `StreamingToolCallTracker`(纯删)→ `ToolInput{Raw,Parsed}`" | 调整 | tracker 不删,改为对齐 `@ai-sdk/provider-utils` 并移入 aimux-provider-utils,由 OpenAI chat 流驱动(#204);对齐基线是 provider-utils 5.0.55,含该版新增的 `hasEmptyArguments` 修复(同一流里新 id、同名、首片段参数为空的并行调用不再被并入前一个调用;`streaming-tool-call-tracker.ts:140-141, 275-285`);`ToolInput` 的 Raw / Parsed 区分由第二部分的四层消息协议承担(V4 tool-call `input` 为原文 String,core 层 `parse_tool_call` 产出解析值) | @@ -156,15 +156,15 @@ FFI / Node / Python / 5 个 C-ABI 绑定 / CLI / Web 由 descriptor + manifest | D9 | providerOptions 规则逐模型实现声明 | OpenAI、Anthropic、Google、compat 的规则不同；允许包内按 SDK 规则选择命名空间，禁止全局统一推导 | | D10 | 设置求值时间逐包对齐 | 区分工厂、取模型、请求三个阶段，不统一成“所有设置都懒加载” | | D11 | 特殊鉴权维持包内组合 | Azure token、Vertex express、Bedrock 签名使用 fetch 装饰器；Vertex ADC 使用异步凭证能力 | -| D12 | preset 生成独立厂商工厂，复用 compat 内部模型实现 | preset 的 key 懒加载是明确的产品扩展；不以原生 OpenAI model 加 profile 模拟所有厂商 | -| D13 | V4 Provider 与厂商扩展分开 | `chat`、`responses`、`getAvailableModels`、tools 等通过注册条目的扩展能力访问；registry、wrap 不丢失这些能力 | +| D12 | preset 按名字从运行时 descriptor 表创建，复用 compat 内部模型实现 | registry JSON 是唯一数据来源，首次使用时解析并校验一次；不为每一行生成 `create_` 源码（同一份数据不存两遍）。preset 的 key 懒加载是明确的产品扩展；不以原生 OpenAI model 加 profile 模拟所有厂商 | +| D13 | V4 Provider 与厂商扩展分开 | `chat`、`responses`、`getAvailableModels`、tools 等通过注册条目的扩展能力访问；registry、wrap 不丢失这些能力。上游只有 gateway 提供 `getAvailableModels`（`@ai-sdk/gateway` 4.0.100 `gateway-provider.ts:93`），其他厂商的 discovery 是 aimux 扩展（§5） | | D14 | 删除所有 model/provider 默认 retry 配置 | retry 属于操作；router/MoA 为每个 child 配置预算，最终组合失败不得触发整组重跑 | | D15 | 删除通用 `body_overrides` | compat 保留 SDK 的 providerOptions 透传与 `transformRequestBody`；原生包不增加统一请求体改写契约 | | D16 | `aimux` 不依赖厂商包、不内建 Gateway | `aimux_providers::default_providers()` 返回注册条目集合；上层调用 `create_provider_registry()` 组装 | | D17 | 用 crate 依赖强制分层 | 规范、工具、厂商、调用运行时、devtools 分开；`aimux-stream` 保持独立 | | D18 | C ABI 保持同步；本文不加入宿主控制流回调 | 所有绑定只注入 Rust 内建能力句柄；Node/Python 的 TSFN / GIL 宿主回调桥接与 C 宿主回调一并后置（§0.4 Q4） | | D19 | 保留单调递增 `u64` ID 与有类型枚举的句柄表 | 当前实现不是代际 slab；扩展种类、错误校验与耗尽处理，不重新发明句柄系统 | -| D20 | descriptor 与 codegen 是唯一生成链 | preset Rust 源码、manifest、绑定与文档提交入库；不使用 `build.rs` 隐式生成另一套产物 | +| D20 | descriptor 与 codegen 是唯一生成链 | manifest、绑定与文档提交入库；preset 不生成 Rust 源码（见 D12）；不使用 `build.rs` 隐式生成另一套产物 | | D21 | 区分模型身份、registry 引用、包元数据 | `package_id` 不是模型的第四个运行身份，也不靠 endpoint 字符串反解 | | D22 | 实时转写采用 `TranscriptionModelV4.doStream?` | Rust 用可选流能力表达“没有方法”；不以一次 `UnsupportedFunctionality` 调用结果代替能力发现 | | D23 | provider-defined tools 是一等能力 | 工厂暴露 tools；独立搜索服务继续使用 aimux 的 Search 扩展 | @@ -311,8 +311,8 @@ FFI / Node / Python / tools ```rust #[async_trait] pub trait LanguageModel: Send + Sync { - fn specification_version(&self) -> SpecificationVersion { - SpecificationVersion::V4 + fn specification_version(&self) -> &'static str { + "v4" } fn provider(&self) -> &str; @@ -347,7 +347,9 @@ pub trait LanguageModel: Send + Sync { ```rust pub trait Provider: Send + Sync { - fn specification_version(&self) -> SpecificationVersion; + fn specification_version(&self) -> &'static str { + "v4" + } fn language_model(&self, id: &str) -> Result, AiMuxError>; @@ -402,7 +404,6 @@ pub struct FetchRequest { pub url: Url, pub headers: Headers, pub body: Bytes, - pub redirect: RedirectPolicy, pub signal: Option, } @@ -427,14 +428,15 @@ pub type FetchFunction = Arc; 关键规则： - multipart 在 helper 内编码成字节，签名处理最终发送字节。 +- 重定向由 helper 逐跳处理，transport（含注入的 `Fetch`）不跟随重定向、原样返回 `3xx`。普通 API 调用跟随所有重定向（同上游 `postToApi` 在 fetch 默认 `redirect: "follow"` 下的行为，provider-utils@5.0.53 `post-to-api.ts:98-107`）；一旦离开请求的 origin，只保留 `User-Agent`，其余请求头全部丢弃，并以 `RequestCredentials::Omit` 通知 transport（同 `fetch-with-validated-redirects.ts:195-206`），请求头与 fetch 装饰器注入的凭证都不会到达另一个 origin；每一跳重新经过请求的 `Fetch`，签名装饰器对实际发送的 URL 重新签名。校验下载（D26）沿用逐跳校验、pin 与离开 credentialed origin 后剥离调用方头的规则。 - 默认 HTTP client 按 Tokio runtime 分片，保留现有连接池生命周期约束。 - `settings.fetch == None` 时，每次请求读取当前默认 leaf；不在工厂创建时冻结全局 fetch。 -- WebSocket 使用独立 connector，和默认 HTTP leaf 共用显式 `TransportSettings`。 +- WebSocket 使用独立 connector（`TungsteniteConnector::new(ProxyConfig)`）；未注入 connector 时，连接时读取默认 HTTP leaf 的 `ProxyConfig`。 - 代理与 TLS 设置属于 transport 实例，不保留初始化时机敏感的 `init_proxy`。 - 30 秒非流式 whole-response 上限留在 helper，覆盖注入 fetch 后的完整响应处理；streaming 不套用该上限。 - helper 不执行完整 model operation 的 retry。 -设置解析： +设置解析按包定义，不把动态 key、动态 headers 或 body-transform hook 作为所有 provider 的通用设置。OpenAI、Anthropic、Google 的 key 是固定字符串，headers 是固定 map；自定义传输使用 `fetch`。OpenAI-compatible 包保留上游的 `transformRequestBody` 对应 hook。工具层的 `Resolvable` 仅用于实际暴露该类型的设置： ```rust pub enum Resolvable { @@ -450,7 +452,7 @@ pub enum Resolvable { `Future` 表示一次 Promise 值；`AsyncFn` 表示每次重新求值。两者不能混为一谈。`Shared` 所需的 Clone/Send 约束落实在类型定义和构造器中。 -Headers 采用大小写不敏感的覆盖语义： +Headers 采用大小写不敏感的覆盖语义（有意差异，见 §5：上游 `combineHeaders` 合并时大小写敏感，`normalizeHeaders` 之后才转小写并丢弃 undefined，见 `@ai-sdk/provider-utils` 5.0.53 `combine-headers.ts:1-11`、`normalize-headers.ts:8-38`；因此上游 `{a: undefined}` 不会删除先前的 `A`）： 1. 名称统一小写； 2. 后者覆盖前者； @@ -495,7 +497,7 @@ pub fn create_openai( pub fn openai() -> &'static OpenAIProvider; ``` -具体 provider 暴露 `.call()`、`.chat()`、`.responses()` 和实际支持的模态方法。OpenAI 默认 LM 为 Responses；xAI、Hugging Face 的 V4 入口按本地源码只提供 Responses LM，不新增伪装成 SDK 能力的 `.chat()`；既有 Chat Completions 实现保留为 `Provider` trait 之外的显式扩展方法 `chat_completions(id)`（上游没有该入口，登记为产品差异，见 §5）。 +具体 provider 暴露 `.call()`、`.chat()`、`.responses()` 和实际支持的模态方法。OpenAI 默认 LM 为 Responses；xAI、Hugging Face 按本地源码只提供 Responses LM，不提供 `.chat()`；两者既有的 Chat Completions 实现随切换删除，不保留为 trait 之外的扩展入口（相对 master 的能力变化，记入 CHANGELOG；仍需该端点的调用方用 `create_openai_compatible` 指向对应 base URL）。 config 只对该包及明确的 `internal` 复用接口开放。不得提供公开 config getter。 @@ -735,7 +737,7 @@ session 推断沿用现有的 opt-in 前缀推断实现（`SessionInferer`）， | Bedrock | 明确 Converse、Anthropic InvokeModel、Mantle 三条能力；不把 Anthropic-AWS 当作 Bedrock-Anthropic | | Vertex 子路径 | 包含 anthropic、maas、xai；本地源码存在的子路径不能遗漏 | | Anthropic-AWS | 保留 Claude Platform on AWS 的独立端点与签名服务名 | -| preset 与本地服务 | 生成工厂和 descriptor；本地服务允许 `auth=none`；不继续使用 `XxxConfig(OpenAIConfig)` 薄包装 | +| preset 与本地服务 | 运行时 descriptor 表，按名字创建；本地服务允许 `auth=none`；不继续使用 `XxxConfig(OpenAIConfig)` 薄包装 | | Codex | 保留无状态刷新；401 映射 TokenExpired；宿主持久化和刷新后重建 provider | | 单模态厂商 | 与其他包一样使用工厂、模态方法、helper 和 transport 注入 | @@ -749,7 +751,7 @@ providerOptions 必须描述**查找、合并、回退、回写**，而非仅列 | Google/Vertex | 包内保留 Vertex 分支；使用 `googleVertex`，保留功能性的 `google` fallback | | Bedrock | `amazonBedrock`；另外独立读取 `anthropic` 选项，二者不是同一条 fallback 链 | | compat | `openaiCompatible` 与 name 的规范 camelCase key；优先级及未知字段透传按包定义 | -| aimux preset | 由生成 descriptor 明确命名空间，不通过宿主 registry alias 改变 | +| aimux preset | 由 descriptor 明确命名空间，不通过宿主 registry alias 改变 | 按本文的无兼容立场，SDK 中明确标记为历史兼容的 `vertex`、`bedrock`、`openai-compatible` 等旧键及对应双写/告警转发不移植。该偏差进入 §5，不能声称逐字复制 SDK 所有历史行为。 @@ -1084,23 +1086,26 @@ RFC-0014 的 `tracing` span 保留为统一事件的内建日志适配器。HTTP |---|---|---| | 语言适配 | 可调用 provider | Rust 使用 `.call()`；支持的绑定恢复调用语法 | | 语言适配 | 可选方法 | `Option>` 或可选 future，区分能力缺失与执行失败 | -| 语言适配 | 同步/异步设置联合类型 | `Resolvable` | +| 语言适配 | 同步/异步设置联合类型 | 仅在包的 settings 实际暴露时使用 `Resolvable`；OpenAI、Anthropic、Google 的 key 是 `Option`，headers 是固定 map，不是动态回调 | | 语言适配 | WHATWG Request/Response | Rust FetchRequest/FetchResponse；请求体以 bytes 表达 | | 语言适配 | AsyncLocalStorage | task-local 加显式 spawn/stream 生命周期传播 | | 语言适配 | 泛型 telemetry wrapper | `Box` 擦除结果类型，保留所有权 | | 语言适配 | 工厂抛异常 | `Result<_, AiMuxError>` | | 产品选择 | 不可失败的默认 provider 访问 | 延后失败配置的解析；显式工厂仍返回 Result | +| 产品选择 | headers 合并大小写不敏感 | 名称先转小写再按后者覆盖，`None` 删除同名头；上游 `combineHeaders` 按原样合并、`normalizeHeaders` 之后才转小写（`@ai-sdk/provider-utils` 5.0.53 `combine-headers.ts:1-11`、`normalize-headers.ts:8-38`），`{a: undefined}` 不删除先前的 `A`（§3.3） | +| 语言适配 | 全局 fetch | Rust 没有 `globalThis.fetch`：代理属于 transport 实例（`ReqwestFetch::new(ProxyConfig)` 作为 provider 的 `fetch`）；未给 `fetch` 的 provider 每次请求读取进程默认 transport，`set_default_transport` 替换它，内建 WebSocket connector 与下载跳共用其代理设置（§3.3） | | 产品选择 | 默认字符串模型 | 绑定/工具安装本地 registry；不隐式请求 Vercel Gateway | | 产品选择 | 观测范围 | 为 object 和非 LM 模态增加 operation/attempt 事件 | | 产品选择 | telemetry 关闭 | 关闭观测，保留功能性 scope 与 session 推断 | | 产品选择 | recording/replay/audit/session | devtools 能力，不进入 model trait | -| 产品选择 | preset 数量与构建 | 单 crate 模块、feature 和提交入库的代码生成 | +| 产品选择 | preset 数量与构建 | 单 crate 内嵌 registry，运行时解析成 descriptor 表；没有逐厂商的 Rust 函数，按名字创建 | | 产品选择 | preset 能力配置 | `max_tokens_key`、`supports_*` 等仅为显式 preset 适配，不冒充 SDK compat settings | | 产品选择 | preset key | 请求时懒加载；不同于通用 compat 的固定 key | | 产品选择 | 配置 secret 引用 | 配置解析层支持结构化 env 引用 | | 产品选择 | discovery/search/composite | aimux 扩展，有独立 descriptor 和测试 | | 产品选择 | retry 细节 | 保留 jitter、Retry-After 处理和既定错误历史 | | 产品选择 | helper 超时 | 非流式 30 秒上限，与注入 fetch 无关 | +| 产品选择 | API 请求的 redirect | 默认跟随，同上游 `postToApi`（provider-utils@5.0.55 `post-to-api.ts:98-107`）；离开请求的 origin 后只保留 `User-Agent`，transport 收到 `RequestCredentials::Omit`，跳回原 origin 也不恢复；最多 10 跳。比上游严：fetch 规范跨 origin 只去掉 `Authorization`、最多 20 跳（`fetch-with-validated-redirects.ts:89-92`）。与 RFC-0040 §3.2 原文“默认不跟随 redirect”不同，该节已加修订说明（#253） | | 产品选择 | SSE 预读 | 保留首事件错误预读，明确计时影响 | | 产品选择 | 视频多批的 idempotency-key | 调用方给 key 时，多批按批次派生 `:batch:`，单批保持原键；上游每个批次的 doStart 带同一个调用方 key（ai 7.0.127 `A/generate-video/generate-video.ts:338-372`、`:520-536`），provider 可能据此把几个批次去重成一个任务 | | 尚未对齐 | 流中途重试 | 上游 `streamText` 有 `streamRetries`（不设时不重试；设为 n 时流中途的 provider error 最多重新调用模型 n 次，`onError` 也可请求重试，见 ai 7.0.127 `A/generate-text/stream-text.ts:374`、`:1948-1953`、`:2495-2665`）。aimux 待实现，现状与不设时相同：流返回后不重试（RFC-0031 §8.3；§9.1 已知差距） | @@ -1110,6 +1115,16 @@ RFC-0014 的 `tracing` span 保留为统一事件的内建日志适配器。HTTP | 不采用 | 历史版本和 deprecated 兼容 | 不实现 V2/V3 适配、旧 namespace 双写、旧方法转发、旧录制读取 | | 本期范围外 | Evaluation/Realtime/SpeechTranslation/Batch 等完整链路 | 不注册虚假能力；后续需完整接通操作与观测 | | ABI 范围外 | C 宿主异步 fetch/model/headers/telemetry | 由独立异步 ABI 设计处理；本期提供 Rust 内建句柄 | +| 上游对齐 | `specification_version` | `Provider` 与各模型 trait 保留 `specification_version()`，默认返回 `"v4"`，与上游 `ProviderV4` 及各模型 V4 接口的 `specificationVersion: 'v4'` 一致 | +| 语言适配 | `Provider.name()` | `Provider` 不带 `name()`（D-c）：注册名归持有 registry 的一方，身份由模型自己的 `provider()`（`"{name}.{method}"`）承担 | +| 语言适配 | Rust `.call()` 的具体形式 | 每个包的 provider 提供 `call(model_id) -> Arc`，与 `Provider::language_model` 返回同一个模型；Rust 无可调用对象，绑定侧在语言允许处恢复调用语法 | +| 上游对齐 | OpenAI 默认 `language_model` | `language_model` 和 `call` 返回 Responses（`openai.responses`），与上游 `openai(id)` 一致；`chat(id)` 显式选择 Chat Completions。xAI、Hugging Face、Azure 的默认也为 Responses | +| 产品选择 | Bedrock 上的 Anthropic InvokeModel | aimux 没有 `bedrock.anthropic.messages` 对应的 provider（`anthropic_aws` 是 Claude Platform on AWS，不是 Bedrock InvokeModel）；属新增能力，不在本期 | +| 产品选择 | Azure 未移植的模型 | deepseek、completion、MAI speech 端点和 Foundry item type 未移植 | +| 上游对齐 | 视频轮询节奏 | provider 和 `VideoModel` 都没有 poll 设置（`VideoModel::poll_config()` 已删除，见 §8.1）；调用级 `VideoCallOptions.poll` 逐字段覆盖 AI SDK 默认值 5000 ms / 600000 ms（ai 7.0.127 `generate-video/generate-video.ts:543-544`） | +| 上游对齐 | 操作级下载传输 | `generate_text` / `stream_text` 的 `download`（`DownloadFunction`，上游 `experimental_download`）以及 `transcribe`、`generate_video` 的 `download`（`Download`）可注入，缺省为受控下载 `create_download(None)`（D26）。prompt 只下载模型 `supported_urls()` 不覆盖的 URL；`transcribe` 下载 `AudioInput::Url`；`generate_video` 下载 URL 结果并返回 `VideoData::Binary`（ai 7.0.127 `prompt/convert-to-language-model-prompt.ts:36-55`、`transcribe/transcribe.ts:56`、`:192-199`、`generate-video/generate-video.ts:136`、`:383-410`） | +| 产品选择 | 下载后的媒体类型 | Core 不做字节嗅探，省略上游的 `detectMediaType`：URL 音频取下载报告的媒体类型，否则沿用调用方的 `media_type`；视频依次取自身、下载报告的类型、`video/mp4`；prompt 文件只在原类型不完整时取下载报告的类型。`data:` URL 不进入下载计划，原样发送（上游先把它转成内联数据） | +| 产品选择 | 录制重建原生协议 provider | `rebuild_provider` 只走 registry 与 overlay；原生协议包返回 `NoSuchProvider`，回放时改用 `replay_with_model` 传入模型 | 本地没有安装源码的厂商包，不能仅凭命名推测宣称“已对齐 AI SDK”。descriptor 记录 `verified_sdk_source` 或 `aimux_extension`；发布前按实际支持承诺完成协议 fixture 验证。 @@ -1289,7 +1304,6 @@ aimux_error_t* aimux_stream_text( ```text 规范类型/Wire DTO + 包 descriptor + preset 输入 - → aimux-codegen 生成 preset Rust 源码 → aimux-manifest 汇总 → manifest.json → 八种语言类型/工厂、FFI 分发、文档、contract vectors @@ -1466,7 +1480,7 @@ providerOptions 的 namespace 不跟随 registry alias。将 `openai` 注册为 | P1：分层切换 | 按依赖顺序拆成若干层（core 类型 → provider 工厂与请求层 → 厂商包与其余模型种类 → 公共 JSON → 下文所列的最后几层），每层一个 PR，直接合入主线 | 每层在同一 PR 内改完它删除或改动的符号在整个依赖闭包（FFI、八种接口、三个工具）中的全部消费方；合入时 §9.2 中已具备的检查全部通过 | | P2：协议与生命周期验收 | 完成逐包差分、离线回放、并发、取消、认证、流、生成物检查和文档核对 | 全部发布门槛通过后才发布 | -每层直接写成目标设计，不引入之后要删除的过渡代码、deprecated 别名或旧 API 转发（D31）。某一层暂时达不到本文的某项要求时，在该层 PR 描述中写明缺什么、由哪一层补齐；与该层代码直接相关的本文修订随该层 PR 一起提交。例外：新 C ABI 与句柄对象模型（§6.1–§6.3）、descriptor/manifest/codegen（§6.5）、录制 schema 3 的完整形态（§4.3、§4.4），以及八种接口与工具的迁移（§6.4、§4.6），排在其余各层之后、发布之前。在此之前，现有按厂商命名的 C 构造器与按名字创建 provider 的入口保持对外签名，内部改调新工厂；它们随绑定迁移那一层一次替换，不另设 deprecated 别名，任何发布版本中都不会出现。每层 PR 描述还列出产物改前改后的体积（体积与性能门禁按 RFC-0043，见 §0.7）。 +每层直接写成目标设计，不引入之后要删除的过渡代码、deprecated 别名或旧 API 转发（D31）。某一层暂时达不到本文的某项要求时，在该层 PR 描述中写明缺什么、由哪一层补齐；与该层代码直接相关的本文修订随该层 PR 一起提交。例外：新 C ABI 与句柄对象模型（§6.1–§6.3）、descriptor/manifest/codegen（§6.5）、录制 schema 3（§4.3、§4.4），以及八种接口与工具的迁移（§6.4、§4.6），排在其余各层之后、发布之前。在此之前，现有按厂商命名的 C 构造器与按名字创建 provider 的入口保持对外签名，内部改调新工厂；它们随绑定迁移那一层一次替换，不另设 deprecated 别名，任何发布版本中都不会出现。每层 PR 描述还列出产物改前改后的体积（体积与性能门禁按 RFC-0043，见 §0.7）。 全部层合入并通过 P2 之前，主线处于迁移中途，**不从主线发布**。此期间已发布版本的修复，从对应 release 标签拉出维护分支发布（如 0.5.x 从 `v0.5.0` 拉出）。 @@ -1487,6 +1501,7 @@ providerOptions 的 namespace 不跟随 registry alias。将 `openai` 注册为 - feature 组合矩阵 CI（§9.2）。 - 协议 fixture 扩到四个家族以外，逐字段比对结果与流（§9.3）。 - CLI 探测加 `--method`；现在 OpenAI 只能探测 Responses（§4.6、RFC-0025）。 +- 录制 schema 3（§4.3、§4.4）；随最后几层做。在此之前各层仍写主线的 schema 2，`ProviderRecord` 只记身份；主线写的录制可以读取（只有主线写的字段忽略，只有新层写的字段读时可缺省），回放不了的部分在回放时报错说明。 - OpenAI 兼容输出的流带回上游的 id 与 model（RFC-0026）；这是回归，不是新功能。 ### 9.2 必须通过的构建检查 @@ -1578,7 +1593,7 @@ providerOptions 的 namespace 不跟随 registry alias。将 `openai` 注册为 | 未覆盖项：补回 onStepFinish 转发 | **不接受历史兼容转发。**只保留 onStepEnd；Evaluation 等能力按完整范围处理，不为消除表面缺项增加空事件 | | 影响面 S0-2/S0-3：统一 `.` 和首段 namespace | **源码反例足够，不能作为全局规则。**Bedrock 身份不符合该模板，Google/Vertex 与 Responses 有包内分支，Anthropic 有 canonical/custom 合并。保留逐包规则 | | 影响面 S1-4：必须使用 RecordingFetch 装饰器 | **机制不必照搬。**需要的是可替换传输与同一录制/回放边界；provider-utils 的统一 leaf 诊断满足要求，且不要求用户手工排列录制包装器 | -| 影响面 S4-7：xAI 同时保留 chat/responses | **作为 SDK 入口不接受；作为显式扩展保留。**本地 xAI 工厂只暴露 Responses LM，没有 chat LM 工厂，`languageModel` / `.call()` 不提供 chat；既有 Chat Completions 实现只作 trait 外的 `chat_completions(id)` 扩展并登记为产品差异（§3.3、§5）。OpenAI 的 chat 选择继续提供 | +| 影响面 S4-7：xAI 同时保留 chat/responses | **不符合本地基线。**本地 xAI 工厂只暴露 Responses LM，没有 chat LM 工厂；xAI 与 Hugging Face 既有的 Chat Completions 实现删除，不作为 trait 外扩展保留（§3.3）。OpenAI 的 chat 选择继续提供 | | 影响面 S5 的旧数据映射、重建和归一化 | **全部排除。**schema 3 仅处理新录制；live replay 依赖宿主目标引用，不保存或重建 ProviderRecord | | 影响面 S7-3/S8-2：本期必须加入所有语言 credential callback ABI | **需求接受，指定机制不接受。**动态凭证可以由 Rust 内建 resolver 表达；任意 C 宿主回调涉及异步 ABI 和重入契约，不能靠函数指针加 ctx 宣称已经解决 | @@ -1855,7 +1870,7 @@ pub enum Optional { 生成的 JSON TypeScript DTO 对缺失属性使用 `?`。上游运行时的 `number | undefined` 可在 SDK facade 中恢复，但不能要求 JSON 对象携带不存在的 undefined。 -provider options/metadata 的内部键完全透传，不执行 camelCase 改写，不从 `model.provider()` 拆字符串推导 namespace。 +provider options/metadata 的内部键完全透传，不执行 camelCase 改写。从 `model.provider()` 拆字符串推导 namespace 不作为通用规则；个别厂商包按上游各自的做法从 provider 名取 options 名（例如 `@ai-sdk/openai-compatible` 3.0.62 `openai-compatible-chat-language-model.ts:157-159`、`@ai-sdk/anthropic` 4.0.71 `anthropic-language-model.ts:373`），见下文衔接第 5 条。 **用户输入 `ModelMessage` 应按 role 建模。** diff --git a/docs/api/c.md b/docs/api/c.md index 77e8f37a7..c9557ce22 100644 --- a/docs/api/c.md +++ b/docs/api/c.md @@ -71,7 +71,7 @@ The code range identifies the source: | Range | Meaning | |---|---| -| `1..17` | `AiMuxError` (4 retired, 14 = `Retry`) | +| `1..20` | `AiMuxError` (4 retired, 14 = `Retry`) | | `100..105` | `RecordingError` | | `200..206` | failure detected by the C ABI | @@ -139,10 +139,11 @@ it with these same getters — it can itself be any AiMuxError code, including `aimux_error_free` independently of the parent (free order is unconstrained). -The unified `aimux_error_code_t` maps AiMuxError variants to values 1–17 +The unified `aimux_error_code_t` maps AiMuxError variants to values 1–20 (`AIMUX_E_RETRY` 14, `AIMUX_E_NO_SUCH_TOOL` 15, `AIMUX_E_INVALID_TOOL_INPUT` 16, -`AIMUX_E_TOOL_CALL_REPAIR` 17; 4 is retired — the legacy catch-all `Tool` -variant, never produced), adds RecordingError values +`AIMUX_E_TOOL_CALL_REPAIR` 17, `AIMUX_E_NO_OUTPUT_GENERATED` 18, +`AIMUX_E_LOAD_API_KEY` 19, `AIMUX_E_LOAD_SETTING` 20; +4 is retired — the legacy catch-all `Tool` variant, never produced), adds RecordingError values 100–105, and assigns C ABI failures 200–206: `NULL_POINTER`, `INVALID_UTF8`, `INVALID_WIRE_JSON`, `INVALID_HANDLE`, `REENTRANT_CALL`, `RESULT_SERIALIZATION`, and `CALLBACK_FAILURE`. Values are never renumbered @@ -154,6 +155,18 @@ ones: `aimux_error_tool_name` (codes 15/16), `aimux_error_available_tools` original failure as the same externally-tagged JSON used by `ToolCall.error`). All returned strings are caller-owned. +`AIMUX_E_LOAD_API_KEY` (19) and `AIMUX_E_LOAD_SETTING` (20) are the AI SDK's +`LoadAPIKeyError` / `LoadSettingError`: no key was passed and the fallback +environment variable is unset, or a required setting (AWS region, Azure +resource name, Vertex project, …) is. Provider keys resolve on every request, +so these arrive from the call (or from `aimux_provider_new` for registry +providers), never from a request that was sent. `aimux_error_provider_code` +returns the environment variable that was consulted; `aimux_error_provider_message` +the key's description (`"OpenAI"`) or the setting name (`"region"`). A provider +`config_json` is `{"base_url", "headers", "organization", "project", "params"}`; +`max_retries` (a per-call option) and `body_overrides` (removed) are rejected +with `AIMUX_E_INVALID_ARGUMENT`. + Internal panics abort the process in this workspace's **release** profile (`panic = "abort"`); in a `panic=unwind` build a Rust callback's panic is caught inside the C ABI and reported as a C ABI failure instead. Either @@ -313,8 +326,8 @@ high-level range. Any fallible call can additionally return 200..206. | `aimux_vertex_new(access_token, project, location, model_id, out_handle)` / `_with_base(…, base_url, out_handle)` | Vertex AI | | `aimux_anthropic_aws_new(api_key, region, model_id, out_handle)` / `_with_base(…, base_url, out_handle)` | Anthropic on AWS | | `aimux_azure_new(api_key, resource_name, deployment, api_version, out_handle)` / `_with_base(api_key, base_url, deployment, api_version, out_handle)` | Azure OpenAI (`_with_base` requires `base_url`) | -| `aimux_provider_new(name, api_key, model_id, config_json, uint64_t *out_handle)` | [AiMuxError] Registry provider (`api_key` NULL → env) | -| `aimux_provider_from_env(name, model_id, uint64_t *out_handle)` | [AiMuxError] Registry + env API key | +| `aimux_provider_new(name, api_key, model_id, config_json, uint64_t *out_handle)` | [AiMuxError] Vendor package or preset (`api_key` NULL → env at request time) | +| `aimux_provider_from_env(name, model_id, uint64_t *out_handle)` | [AiMuxError] By-name provider + env API key at request time | | `aimux_provider_handle_new(name, api_key, config_json, uint64_t *out_handle)` | [AiMuxError] RFC-0027 provider handle | | `aimux_provider_list_models(handle, char **out_models_json)` | [AiMuxError] JSON `RuntimeModel[]` | | `aimux_provider_model(handle, model_id, uint64_t *out_handle)` | [AiMuxError] Model from a provider handle | @@ -473,7 +486,7 @@ Pass `"session_id": "..."` inside `opts_json` of generate/stream to group a call |------|------| | `aimux_init_logging(level)` | [C ABI] `"off"…"trace"`, NULL = `"warn"` | | `aimux_register_providers(config_json)` | [AiMuxError] External OpenAI-compatible providers (RFC-0020) | -| `aimux_init_proxy(config_json)` | [AiMuxError] Global proxy (before first call) | +| `aimux_init_proxy(config_json)` | [AiMuxError] Replace the default transport's proxy (later requests) | | `aimux_router_new(const uint64_t *handles, size_t len, config_json, uint64_t *out_handle)` | [AiMuxError] RouterModel (RFC-0021) | | `aimux_moa_new(const uint64_t *reference_handles, size_t ref_len, aggregator, config_json, uint64_t *out_handle)` | [AiMuxError] MoaModel (RFC-0022) | diff --git a/docs/api/decision.md b/docs/api/decision.md index c463ab8d6..7a41b15ae 100644 --- a/docs/api/decision.md +++ b/docs/api/decision.md @@ -24,9 +24,10 @@ console.log(result.answers.refund) // boolean 概率或 { type: 'refusal' } ``` Python 对应 `create_provider('openai', key).decision_model('gpt-6-luna')`，随后 -调用下文的 `decide`。Rust 使用 `OpenAIProvider::new(OpenAIConfig::new(key))` -和 `Provider::decision_model`。服务地址为 `{base_url}/decisions`，默认使用 -`https://api.openai.com/v1`；继承 provider 的认证、headers、organization/project 和重试配置。 +调用下文的 `decide`。Rust 使用 `create_openai(settings)` 返回的 provider 的 +`decision(model_id)` 或 `Provider::decision_model`。服务地址为 `{base_url}/decisions`，默认使用 +`https://api.openai.com/v1`；继承 provider 的认证、headers、organization/project 和传输， +重试次数由调用的 `max_retries` 控制。 文本 state 原样发送；JSON 对象／数组序列化为文本。Boolean 映射 predicate。 Choice 支持 2–255 项，`label` 为规范化结果的 map key，可选 `value` 是原生 string @@ -53,7 +54,7 @@ const result = await decide(model, { 最多 128 张内嵌图，detail 支持 low/high/auto/original。也可使用 `data: { Binary: [137, 80, ...] }`，与已有 FileBytes contract 一致。 -对象/数组 instructions、Boolean criteria 和生成接口 body overrides 返回 Unsupported。 +对象/数组 instructions 和 Boolean criteria 返回 Unsupported。 `openai` options 仅接受可选 safety_identifier（最多 128 字符）。 原生概率/confidence 保留原值，不套用 Jev 两位小数规则；usage 缺失或合计不一致 不影响有效答案。本地 fixture 基于[官方 reference](https://developers.openai.com/api/reference/resources/decisions/methods/create)，并非新增服务的线上录制。 @@ -321,10 +322,11 @@ finally: ``` 标准 `recordings.jsonl` 每行记录一次逻辑调用：`input.operation = "decision"`， -`input.options` 保留 state/questions/images 以及协议、编码 options；provider 快照含 endpoint 和 capabilities； +`input.options` 保留 state/questions/images 以及协议、编码 options，`input.decision_capabilities` 保留能力声明； +provider 只记录 provider_id/provider/model_id，endpoint 与凭证来自调用方配置； exchanges 记录各重试 attempt 的 HTTP 请求、响应和时延； `outcome.decision_result` 保存规范化结果。凭据沿用统一脱敏规则。 -旧语言录制缺少 operation 时仍按 language model 解析，schema 2 保持兼容。 +录制使用 schema 2，与 aimux 0.5 相同；语言录制缺少 operation 时仍按 language model 解析。aimux 0.5 写的 decision 录制可以读取，但其能力声明在 `provider.profile` 中，本版本不读，离线回放会报错说明。 CLI 离线回放不需要 key： @@ -337,7 +339,7 @@ Rust 可用 `aimux_core::replay::MockDecisionReplayModel::from_jsonl` 加载录 再通过 `decide` 或 `replay_decision_with_model` 离线回放。匹配 provider/model、 state/questions/images、headers 和 provider_options；timeout/retry 控制不参与匹配。 输入未命中或记录不完整时返回错误。请求回放可通过 -`aimux_providers::rebuild_decision_provider` 重建所有上述官方 provider，再调用 +`aimux_providers::rebuild_decision_provider` 从调用方 registry 重建上述任一官方 provider 的 decision model，再调用 `replay_decision_with_model` 发出真实请求；CLI 不带 `--mock` 时也走这条路径。 实测录制在 `aimux-providers/tests/fixtures/jev_systemone_live.jsonl`。 diff --git a/docs/api/flutter.md b/docs/api/flutter.md index 1046ed087..8a75aea82 100644 --- a/docs/api/flutter.md +++ b/docs/api/flutter.md @@ -1,6 +1,6 @@ # aimux · Flutter/Dart API -> Unified LLM service access layer — one API to access 325 AI providers +> Unified LLM service access layer — one API to access AI providers Flutter/Dart wraps the Rust core through the `aimux-ffi` C ABI (via `dart:ffi`). @@ -47,11 +47,10 @@ model.close(); ## Providers -All 251 registry-backed OpenAI-compatible providers are reachable by string name: +The vendor packages and 282 registry-backed OpenAI-compatible providers are reachable by string name: -> **Scope:** `provider(name)` covers only the 251 registry OpenAI-compatible -> providers; Anthropic/Google/multimodal/local → typed factories -> (`Model.anthropic(apiKey, modelId)`); custom endpoints → base-URL variant. +> **Scope:** `provider(name)` reaches vendor packages and the 282 preset rows; +> typed factories remain available. Custom endpoints use the base-URL variant. > Full list: [providers.md](providers.md). ```dart @@ -73,7 +72,7 @@ shares a base with the other (both just `implements Exception`): | Source | Rust | Dart | C code | |---|---|---|---| -| AiMux | `AiMuxError` | `AimuxException` hierarchy | 1..17 (4 retired) | +| AiMux | `AiMuxError` | `AimuxException` hierarchy | 1..20 (4 retired) | | recorder | `RecordingError` | `RecordingException` | 100..105 | Every fallible C call returns an opaque `aimux_error_t *` (`NULL` = @@ -81,7 +80,7 @@ success, result in the trailing out-parameter). The binding has one decoder (`errors.dart`): `expectAimuxError(e, context)` for model calls, `expectRecordingError(e, context)` for `initRecording` / `recordingTryFlush`, `expectFfiError(e, context)` for utilities that can only fail in the C ABI. -One unified code selects 1..17, 100..105, or 200..206; each decoder copies the +One unified code selects 1..20, 100..105, or 200..206; each decoder copies the relevant fields, releases the error with `aimux_error_free` exactly once, and throws the matching `AimuxException` subclass / `RecordingException`. Codes 200..206 throw the native @@ -105,6 +104,9 @@ Exception (implements) ├── NoSuchToolError // code 15; tool not in the supplied tool set ├── InvalidToolInputError // code 16; tool arguments failed to parse/validate ├── ToolCallRepairError // code 17; a repairToolCall hook itself failed + ├── NoOutputGeneratedError // code 18; no output was generated + ├── LoadAPIKeyError // code 19; no API key, fallback env var unset (envVar) + ├── LoadSettingError // code 20; a required setting is missing (envVar) └── OtherError ``` diff --git a/docs/api/go.md b/docs/api/go.md index 01e93827b..a83c4a8d3 100644 --- a/docs/api/go.md +++ b/docs/api/go.md @@ -1,6 +1,6 @@ # aimux · Go API -> Unified LLM service access layer — one API to access 325 AI providers +> Unified LLM service access layer — one API to access AI providers The Go binding goes through the `aimux-ffi` C ABI: cgo statically links `libaimux_ffi.a`, so the Rust core is compiled into the executable and the @@ -63,11 +63,13 @@ if err != nil { | `AvailableTools` | `CodeNoSuchTool` payload: `[]string` of the tools the call offered; `nil` when not reported | | `ToolInput` | `CodeInvalidToolInput` payload: the raw argument text the model produced | | `OriginalError` | `CodeToolCallRepair` payload: the pre-repair failure as `json.RawMessage` wire JSON (same encoding as `ToolCall.Error`) | +| `EnvVar` | `CodeLoadAPIKey` / `CodeLoadSetting` payload: the environment variable that was consulted | -`Code` values 1..17 mirror aimux-core's `AiMuxError` variants; 4 is retired +`Code` values 1..20 mirror aimux-core's `AiMuxError` variants; 4 is retired (the legacy `Tool` variant) and 14 is `CodeRetry`. The tool-call variants arrive as `CodeNoSuchTool` (15), `CodeInvalidToolInput` (16), and -`CodeToolCallRepair` (17). A code outside the enum is a header/library +`CodeToolCallRepair` (17); `CodeNoOutputGenerated` is 18; a missing API key / provider setting are +`CodeLoadAPIKey` (19) / `CodeLoadSetting` (20). A code outside the enum is a header/library mismatch and fails with a `panic`, not an error type. Recording failures are a separate type, as in Rust (`recording::RecordingError` @@ -92,7 +94,7 @@ no Go type of their own; the binding maps them to native Go errors: |---------|----| | bad raw JSON in `promptJson` / `optsJson` / `configJSON` | plain `error` naming the parameter, e.g. `aimux: prompt_json: invalid JSON` — checked in Go before the C call with `json.Valid` **plus** a surrogate-pairing scan, because `json.Valid` accepts unpaired `\uD800`–`\uDFFF` escapes that serde_json rejects (required parameters reject `""`; optional `""` = default; JSONL by line) | | a string parameter that is not valid UTF-8, or contains a NUL | plain `error`: `aimux: : must be valid UTF-8` / `aimux: : must not contain NUL` — a Go string is an arbitrary byte sequence and `C.CString` passes it through verbatim (a NUL would truncate the argument silently), so every user-supplied string is checked before the C call. Two exceptions, both deliberate: the five `mustNew` constructors panic instead (see below), and `InitLogging` has no error channel and falls back to `"warn"` | -| a typed option struct (`ProviderConfig`, `EmbeddingCallOptions`, `SpeechCallOptions`, … , `TranscriptionSessionOpts`) whose raw JSON field carries bad bytes | plain `error` naming the marshalled parameter, e.g. `aimux: opts: invalid JSON: unpaired high surrogate \uD800`, `aimux: config_json: must be valid UTF-8` — `json.Marshal` is **not** self-validating: it coerces invalid UTF-8 in a Go *string* field to U+FFFD and escapes NUL, but a `json.RawMessage` field (`ProviderOptions`, `BodyOverrides`, `Audio`, `Documents`, `Files`, `Mask`, `Data`, `InputSchema`) is emitted through `compact()`, which checks JSON *syntax only*. Raw non-UTF-8 bytes and lone `\uD800`–`\uDFFF` escapes therefore survive marshalling, so every marshalled C argument gets the same `checkJSON` as a raw-string parameter | +| a typed option struct (`ProviderConfig`, `EmbeddingCallOptions`, `SpeechCallOptions`, … , `TranscriptionSessionOpts`) whose raw JSON field carries bad bytes | plain `error` naming the marshalled parameter, e.g. `aimux: opts: invalid JSON: unpaired high surrogate \uD800`, `aimux: config_json: must be valid UTF-8` — `json.Marshal` is **not** self-validating: it coerces invalid UTF-8 in a Go *string* field to U+FFFD and escapes NUL, but a `json.RawMessage` field (`ProviderOptions`, `Audio`, `Documents`, `Files`, `Mask`, `Data`, `InputSchema`) is emitted through `compact()`, which checks JSON *syntax only*. Raw non-UTF-8 bytes and lone `\uD800`–`\uDFFF` escapes therefore survive marshalling, so every marshalled C argument gets the same `checkJSON` as a raw-string parameter | | use-after-close (`Model`, `ProviderHandle`, multimodal models, `TranscriptionSession`) | `errors.Is(err, aimux.ErrClosed)` — guarded in Go before the C call | | a `nil` `*Model` handed to a composite constructor — a `NewRouter` child, a `NewMoa` reference, or the `NewMoa` aggregator | plain `error` naming the position: `aimux: router: models[2] is nil`, `aimux: moa: references[1] is nil`, `aimux: moa: aggregator is nil` — checked in Go before the C call, so a nil element is a returned error rather than the nil-pointer dereference that would otherwise take the process down | | trace query (`TraceAggregate`, `TraceSessionChain`, `TraceExportJsonl`, `TraceClear`) on a model that never went through `Trace` / `TraceAudited` | `errors.Is(err, aimux.ErrNotTraced)` — guarded in Go before the C call, ahead of argument validation; the trace store is keyed on the wrapper handle, so C can only report it as a missing handle | @@ -100,7 +102,7 @@ no Go type of their own; the binding maps them to native Go errors: Decoder: every fallible C call returns an opaque `aimux_error_t *` (NULL = success, result in the out-parameter). One `aimux_error_code()` distinguishes -`AiMuxError` (1–17), `RecordingError` (100–105), and C ABI failures (200–206). +`AiMuxError` (1–20), `RecordingError` (100–105), and C ABI failures (200–206). `expectAimuxError`, `expectRecordingError`, and `expectFfiError` enforce the range expected by each call; the first two restore `*Error` and `*RecordingError`, while 200–206 becomes a plain `error`. Every path frees the @@ -117,7 +119,7 @@ design** — that one is opt-in, and every one of its five entry points has a |------------|--------------------------| | `aimux.go` `mustNew` — behind `OpenAI` / `OpenAIWithBase` / `Anthropic` / `AnthropicWithBase` / `DeepSeek` | **Yes, by design.** `regexp.MustCompile` convention: an `apiKey` / `modelID` / `baseURL` that is not valid UTF-8 or contains a NUL panics, as does any AiMuxError failure. Use `NewOpenAI` / `NewOpenAIWithBase` / `NewAnthropic` / `NewAnthropicWithBase` / `NewDeepSeek` for anything caller-supplied | | `aimux.go` `InitLogging` — `expectFfiError` returned an error | No. `level` is coerced first: empty, non-UTF-8, or NUL-bearing falls back to `"warn"`, which is what aimux-core does with an unparseable level anyway (`AIMUX_LOG` / `AIMUX_LOG_LEVEL` outrank it regardless). That leaves no documented failure for `aimux_init_logging`, so a non-nil error here is a header/library mismatch | -| `aimux.go` `expectAimuxError` — `aimux_error_code_t` outside 1..17 | No. Header/library version mismatch | +| `aimux.go` `expectAimuxError` — `aimux_error_code_t` outside 1..20 | No. Header/library version mismatch | | `aimux.go` `expectRecordingError` — `aimux_error_code_t` outside the enum | No. Header/library version mismatch | | `multimodal.go` `TranscriptionSession.NextPart` — unknown `aimux_transcription_next_part` state | No. Header/library version mismatch | @@ -155,11 +157,10 @@ The producer closes `Parts()` after the blocking native call returns. ## Providers -All 251 registry-backed OpenAI-compatible providers are reachable by string name: +The vendor packages and 282 registry-backed OpenAI-compatible providers are reachable by string name: -> **Scope:** `provider(name)` covers only the 251 registry OpenAI-compatible -> providers; Anthropic/Google/multimodal/local → typed constructors -> (`NewAnthropic(apiKey, model)`); custom endpoints → `WithBase` variant. +> **Scope:** `provider(name)` reaches vendor packages and the 282 preset rows; +> typed factories remain available. Custom endpoints use the `WithBase` variant. > Full list: [providers.md](providers.md). ```go @@ -171,15 +172,25 @@ defer model.Close() model2, err := aimux.ProviderWithBase("groq", "sk-...", "llama-3.3-70b", "https://relay.example/v1") defer model2.Close() -// 完整 ProviderOptions(base_url / headers / organization / project / -// max_retries / body_overrides): +// Full ProviderOptions (base_url / headers / organization / project / params). +// Retry is a per-call option (CallOptions.MaxRetries) and request-body overrides +// no longer exist: a config_json carrying max_retries or body_overrides is +// rejected as CodeInvalidArgument, so ProviderConfig has neither. model3, err := aimux.ProviderWithConfig("groq", "sk-...", "llama-3.3-70b", &aimux.ProviderConfig{ Headers: map[string]string{"X-Custom": "1"}, }) defer model3.Close() + +// A preset with template parameters: +model4, err := aimux.ProviderWithConfig("cloudflare_workers_ai", "sk-...", "@cf/meta/llama-3.1-8b-instruct", + &aimux.ProviderConfig{Params: map[string]string{"account_id": "acct123"}}) +defer model4.Close() + +// A missing key or setting is CodeLoadAPIKey / CodeLoadSetting; Error.EnvVar +// names the environment variable that was consulted. ``` -`NewDeepSeek` / `DeepSeek` remain as shortcuts (registry-backed). Unknown +`NewDeepSeek` / `DeepSeek` remain as by-name shortcuts to the vendor package. Unknown names return an error. ## Text Generation @@ -374,7 +385,8 @@ if err != nil { ## Video Generation -Video generation typically returns a URL (not binary). +Most video models return a URL; `Generate` downloads it, as the AI SDK's +`generateVideo` does, and returns the video's bytes. ```go prompt := "A cat playing piano" @@ -403,7 +415,7 @@ result, err := aimux.ParseVideoResult(resultJSON) if err != nil { log.Fatal(err) } -// result.Videos[0].Url.URL — video URL +// result.Videos[0].Binary.Data — the downloaded video ``` ## Reranking diff --git a/docs/api/java.md b/docs/api/java.md index 37619ec6d..8dc876472 100644 --- a/docs/api/java.md +++ b/docs/api/java.md @@ -1,6 +1,6 @@ # aimux · Java API -> Unified LLM service access layer — one API to access 325 AI providers +> Unified LLM service access layer — one API to access AI providers The Java binding goes through the `aimux-ffi` C ABI via JNA — no native toolchain is needed at build time. Artifact: `ai.arcships:aimux-java:0.3.0`, @@ -33,11 +33,10 @@ try (Model model = Model.openaiWithBase("sk-...", "gpt-4o", "http://localhost:30 ## Providers -All 251 registry-backed OpenAI-compatible providers are reachable by string name: +The vendor packages and 282 registry-backed OpenAI-compatible providers are reachable by string name: -> **Scope:** `provider(name)` covers only the 251 registry OpenAI-compatible -> providers; Anthropic/Google/multimodal/local → typed factories -> (`Model.anthropic(apiKey, modelId)`); custom endpoints → base-URL variant. +> **Scope:** `provider(name)` reaches vendor packages and the 282 preset rows; +> typed factories remain available. Custom endpoints use the base-URL variant. > Full list: [providers.md](providers.md). ```java @@ -53,7 +52,7 @@ try (Model model = Model.provider("groq", "sk-...", "llama-3.3-70b", null)) { } ``` -`deepseek(apiKey, modelId)` remains as a shortcut (registry-backed). +`deepseek(apiKey, modelId)` remains as a by-name shortcut to the vendor package. Unknown names throw `NoSuchProviderError` naming the requested provider. ## Errors @@ -75,6 +74,9 @@ RuntimeException ├── NoSuchToolError // code 15: the tool the model called is not in the tool set ├── InvalidToolInputError // code 16: the tool arguments failed to parse or validate ├── ToolCallRepairError // code 17: tool-call repair itself failed + ├── NoOutputGeneratedError // code 18: no output was generated + ├── LoadAPIKeyError // code 19: no API key, fallback env var unset (getEnvVar()) + ├── LoadSettingError // code 20: a required setting is missing (getEnvVar()) └── OtherError ``` @@ -83,7 +85,7 @@ Every instance has: | Field | Meaning | |-------|---------| | `getMessage()` | human-readable text from C | -| `getCode()` | `aimux_error_code_t` value 1–17 (4 retired, 14 = `Retry`; matches `aimux-error.h`) | +| `getCode()` | `aimux_error_code_t` value 1–20 (4 retired, 14 = `Retry`; matches `aimux-error.h`) | | `getStatusCode()` | HTTP status, or `-1` | | `getRetryMs()` | rate-limit hint, or `-1` (`0` = retry now) | | `isRetryable()` | the `AiMuxError` retry verdict (not derivable from status) | @@ -149,7 +151,7 @@ types and share no base beyond `RuntimeException`. Transport: every fallible C call returns an opaque `aimux_error_t *` (JNA `Pointer`) — `null` on success with the result in a trailing out-parameter (`LongByReference` handle / `PointerByReference` JSON), non-null on failure. -`AimuxResult` reads one unified code: 1–17 restores the matching +`AimuxResult` reads one unified code: 1–20 restores the matching `AimuxException` subclass, 100–105 restores `RecordingException`, and 200–206 becomes `IllegalStateException("aimux ffi: …")`. Payload getters are read only under their owning AiMuxError code; a `RetryError`'s attempt errors are new diff --git a/docs/api/kotlin.md b/docs/api/kotlin.md index 4601c413d..310a70f47 100644 --- a/docs/api/kotlin.md +++ b/docs/api/kotlin.md @@ -1,6 +1,6 @@ # aimux · Kotlin API -> Unified LLM service access layer — one API to access 325 AI providers +> Unified LLM service access layer — one API to access AI providers Kotlin wraps the Rust core through the `aimux-ffi` C ABI (via JNA). @@ -27,11 +27,10 @@ Model.openai("sk-...", "gpt-4o", "http://localhost:3000").use { model -> ## Providers -All 251 registry-backed OpenAI-compatible providers are reachable by string name: +The vendor packages and 282 registry-backed OpenAI-compatible providers are reachable by string name: -> **Scope:** `provider(name)` covers only the 251 registry OpenAI-compatible -> providers; Anthropic/Google/multimodal/local → typed factories -> (`Model.anthropic(apiKey, modelId)`); custom endpoints → base-URL variant. +> **Scope:** `provider(name)` reaches vendor packages and the 282 preset rows; +> typed factories remain available. Custom endpoints use the base-URL variant. > Full list: [providers.md](providers.md). ```kotlin @@ -53,7 +52,7 @@ Two aimux exception types, each mirroring its own Rust type — **AiMux** (`AimuxException`) and **recorder** (`RecordingException`). They share no base beyond `RuntimeException`; catch each on its own. Every fallible C call returns an `aimux_error_t *` (null = success, result in the out-parameter). The -binding reads one unified code: 1..17 restores an `AimuxException` subclass, +binding reads one unified code: 1..20 restores an `AimuxException` subclass, 100..105 restores `RecordingException`, and 200..206 becomes `IllegalStateException("aimux ffi: …")`. Payload getters are read only under their owning AiMuxError code. @@ -82,6 +81,9 @@ RuntimeException ├── NoSuchToolError // toolName + availableTools (null = no tool set supplied) ├── InvalidToolInputError // toolName + toolInput (the raw argument text) ├── ToolCallRepairError // originalError (wire JSON, same encoding as ToolCall.error) + ├── NoOutputGeneratedError // code 18: no output was generated + ├── LoadAPIKeyError // code 19: no API key, fallback env var unset (envVar) + ├── LoadSettingError // code 20: a required setting is missing (envVar) └── OtherError ``` @@ -104,7 +106,7 @@ try { | Field | Meaning | |-------|---------| -| `code` | `AIMUX_E_*` matching C `aimux_error_code_t` (1..17; 1 is the catch-all `Other`; 4 is retired — the legacy `Tool` variant; 14 = `Retry`) | +| `code` | `AIMUX_E_*` matching C `aimux_error_code_t` (1..20; 1 is the catch-all `Other`; 4 is retired — the legacy `Tool` variant; 14 = `Retry`) | | `status` | HTTP status when known; otherwise `-1` | | `retryMs` | Rate-limit hint in ms; `-1` if none; `0` = retry immediately | diff --git a/docs/api/node.md b/docs/api/node.md index 3f43ccf0a..bc087f930 100644 --- a/docs/api/node.md +++ b/docs/api/node.md @@ -1,6 +1,6 @@ # aimux · Node.js API -> Unified LLM service access layer — one API to access 325 AI providers +> Unified LLM service access layer — one API to access AI providers Shared reference — parameter tables, result shapes, factory functions, and the feature coverage matrix — lives in the [API overview](../API.md). @@ -21,9 +21,8 @@ console.log(result.text) ## Providers -All 251 built-in OpenAI-compatible providers are registry-backed. Look them up -by string name; the current list is generated from `provider_registry.json` in -[providers.md](providers.md): +The vendor packages and 282 registry-backed OpenAI-compatible providers are +reachable by string name; see [providers.md](providers.md) for the provider list: ```typescript import { provider, generateText } from '@arcships/aimux' @@ -31,18 +30,24 @@ import { provider, generateText } from '@arcships/aimux' const model = await provider('groq', undefined, 'llama-3.3-70b') const relay = await provider('groq', 'sk-...', 'llama-3.3-70b', { baseUrl: 'https://relay.example/v1', - maxRetries: 0, }) const result = await generateText(model, 'Hello') ``` -`openai` / `anthropic` / `deepseek` factories remain (deepseek is now -registry-backed). For custom providers not in the registry, build from the -base classes with `createProvider`-style config via the base-URL override. +`openai` / `anthropic` / `deepseek` factories remain; DeepSeek uses its own vendor +chat model. Custom endpoints use a base-URL override with a compatible provider. +`openai` and by-name `provider('openai', …)` select the Responses API. -> **Scope:** `provider(name)` covers only the 251 registry OpenAI-compatible -> providers; Anthropic/Google/multimodal/local → typed factories -> (`anthropic(apiKey, model)`); custom endpoints → `baseUrl` override. +The 3rd argument of every constructor is a base URL string or a +`ProviderConfig` (`baseUrl`, `headers` as a JSON string, `organization`, +`project`, and `params` — a `Record` filling a preset's template +parameters, e.g. `{ account_id: '…' }` for `cloudflare_workers_ai`). +`maxRetries` and `bodyOverrides` are kept in the type only so that passing them +throws `InvalidArgumentError`; retry is a per-call option +(`max_retries` in the call options). + +> **Scope:** `provider(name)` reaches vendor packages and the 282 preset rows; +> typed factories remain available. Custom endpoints use the `baseUrl` override. > Full list: [providers.md](providers.md). ## Desktop and Electron compatibility @@ -91,6 +96,7 @@ Error ├── JSONParseError / InvalidResponseDataError ├── NoSuchToolError / InvalidToolInputError / ToolCallRepairError // tool-contract errors ├── InvalidArgumentError / InvalidPromptError + ├── LoadAPIKeyError / LoadSettingError // no API key / required setting: envVar (+ description / settingName) ├── TokenExpiredError ├── UnsupportedFunctionalityError ├── NoSuchModelError / NoSuchProviderError @@ -132,8 +138,9 @@ that carries them: `APICallError` adds `retryable` and optional `status` / error — or `'errorNotRetryable'` — a later attempt failed non-retryably), `errors` — the per-attempt history, oldest first, each itself an error from this hierarchy — and `lastError`; `TokenExpiredError` carries `status: 401`; -`NoSuchModelError` adds `modelId` / `modelType`; and `NoSuchProviderError` -adds `providerId`. Missing HTTP status and retry hints are absent rather than +`NoSuchModelError` adds `modelId` / `modelType`; `NoSuchProviderError` +adds `providerId`, `modelId`, `modelType` and `availableProviders`; and `LoadAPIKeyError` / `LoadSettingError` add `envVar` +(the environment variable consulted) plus `description` / `settingName`. Missing HTTP status and retry hints are absent rather than represented by `-1`. ```typescript @@ -160,7 +167,7 @@ try { ``` The ts-rs wire type `AiMuxError` is only for payload unions inside -`StreamPart`, not for throws. +`TextStreamPart`, not for throws. ## Text Generation @@ -411,7 +418,8 @@ controller.abort() // cancels the image call ## Video Generation -Video generation typically returns a URL (not binary). +Most video models return a URL; `generate` downloads it, as the AI SDK's +`generateVideo` does, and returns the video's bytes (`Binary`). ```typescript const { googleVideo } = require('@arcships/aimux/raw') @@ -425,9 +433,10 @@ const resultJson = await videor.generate(JSON.stringify({ })) const result = JSON.parse(resultJson) -// result.videos is usually [{ Url: { url, media_type } }] -if (result.videos[0].Url) { - console.log('Video URL:', result.videos[0].Url.url) +// result.videos is usually [{ Binary: { data, media_type } }] +const video = result.videos[0].Binary +if (video) { + require('fs').writeFileSync('video.mp4', Buffer.from(video.data)) } ``` @@ -508,7 +517,7 @@ The `@arcships/aimux` package has two layers: | `RerankingModel` | `cohereReranking` | `rerank(query, docsJson, optsJson?)` | | `SearchModel` | `tavilySearch` | `search(query, optsJson?)` | | `Files` | `openaiFiles(apiKey, baseUrl?)` | `uploadFile(dataBase64, mediaType, optsJson?)` | -| `StreamTextGenerator` | returned by `Model.streamText` | async iterable of `StreamPart` JSON strings | +| `StreamTextGenerator` | returned by `Model.streamText` | async iterable of `TextStreamPart` JSON strings | All factories return a `Promise` and accept an optional `baseUrl` as the last parameter. All native methods take and return JSON strings — the typed wrapper @@ -567,6 +576,9 @@ export type TextStreamPart = | { Abort: ... } | { Error: ... } | { Raw: ... } ``` +Provider-layer `StreamPart` also has `ResponseMetadata`; the call layer +consumes those events internally. + The full declarations live in `bindings/node/src/types/` — `GenerateTextOptions.ts`, `ModelMessage.ts`, `Tool.ts`, `ToolChoice.ts`, `ContentPart.ts`, `GenerateContent.ts`, `GenerateResult.ts`, and the `types/` directory of the diff --git a/docs/api/providers.md b/docs/api/providers.md index c970a359e..2c8d116b5 100644 --- a/docs/api/providers.md +++ b/docs/api/providers.md @@ -9,282 +9,309 @@ | category | count | |----------|-------| -| Registry-backed OpenAI-compatible (`provider(name, ...)`) | 251 | +| Registry-backed OpenAI-compatible (`provider(name, ...)`) | 282 | | Non-registry: Native protocol providers | 14 | -| Non-registry: OpenAI-compatible thin wrappers (second batch) | 5 | | Non-registry: Speech-only providers (TTS) | 4 | | Non-registry: Transcription-only providers (STT) | 5 | | Non-registry: Image-only providers | 4 | | Non-registry: Video-only providers | 1 | | Non-registry: Generic Responses API wrapper | 1 | -| Non-registry: Bulk-generated thin-wrapper providers | 16 | | Non-registry: Modality-specific providers (non-language, e.g. rerank-only) | 1 | | Non-registry: AWS Polly speech (TTS) provider — SigV4 authenticated, speech modality only | 1 | | Non-registry: Recraft image provider (OpenAI Images-compatible + Recraft extension fields) | 1 | | Non-registry: Stability image provider (image modality only) | 1 | | Non-registry: Video-only provider (runwayml) | 1 | -| Non-registry: P1 thin-wrapper providers (provider-research batch) | 1 | -| Non-registry: Vertex AI MaaS partner-model providers (OpenAI-compatible thin wrappers) | 10 | | Non-registry: Search-only providers (web search modality) | 11 | -| **Total providers** | **328** | - -**251 registry-backed OpenAI-compatible providers** (construct via `provider(name, ...)`) + **77 non-registry providers** (construct via the typed factories listed below). - -## Registry-backed (OpenAI-compatible) — 251 - -| name | display | env var | base_url | -|------|---------|---------|----------| -| `abacus` | Abacus | `ABACUS_API_KEY` | `https://routellm.abacus.ai/v1` | -| `abliteration_ai` | Abliteration AI | `ABLIT_KEY` | `https://api.abliteration.ai/v1` | -| `ai21` | AI21 Labs | `AI21_API_KEY` | `https://api.ai21.ai/v1` | -| `ai302` | 302.AI | `AI302_API_KEY` | `https://api.302.ai/v1` | -| `ai_router` | AI-ROUTER | `AI_ROUTER_API_KEY` | `https://api.ai-router.dev/v1` | -| `aiand` | AIand | `AIAND_API_KEY` | `https://api.aiand.com/v1` | -| `aibadgr` | AI Badgr | `AIBADGR_API_KEY` | `https://api.aibadgr.com/v1` | -| `aigc2d` | AIGC2D | `AIGC2D_API_KEY` | `https://api.aigc2d.com/v1` | -| `aihubmix` | AIHubMix | `AIHUBMIX_API_KEY` | `https://aihubmix.com/v1` | -| `ails` | AILS | `AILS_API_KEY` | `https://api.caipacity.com/v1` | -| `aiml` | AI/ML API | `AIML_API_KEY` | `https://api.aimlapi.com/v1` | -| `aki_io` | AKI.IO | `AKI_IO_API_KEY` | `https://aki.io/openai/v1` | -| `albert` | Albert | `ALBERT_API_KEY` | `https://api.albert.ai/v1` | -| `alibaba` | Alibaba Cloud (DashScope) | `ALIBABA_API_KEY` | `https://dashscope-intl.aliyuncs.com/compatible-mode/v1` | -| `alibaba_coding_plan` | Alibaba Coding Plan | `ALIBABA_CODING_PLAN_API_KEY` | `https://coding-intl.dashscope.aliyuncs.com/v1` | -| `alibaba_coding_plan_cn` | Alibaba Coding Plan (China) | `ALIBABA_CODING_PLAN_API_KEY` | `https://coding.dashscope.aliyuncs.com/v1` | -| `alibaba_token_plan` | Alibaba Token Plan | `ALIBABA_TOKEN_PLAN_API_KEY` | `https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1` | -| `alibaba_token_plan_cn` | Alibaba Token Plan (China) | `ALIBABA_TOKEN_PLAN_API_KEY` | `https://token-plan.cn-beijing.maas.aliyuncs.com/compatible-mode/v1` | -| `ambient` | Ambient | `AMBIENT_API_KEY` | `https://api.ambient.xyz/v1` | -| `anyapi` | AnyAPI | `ANYAPI_KEY` | `https://api.anyapi.ai/v1` | -| `anyscale` | Anyscale | `ANYSCALE_API_KEY` | `https://api.endpoints.anyscale.com/v1` | -| `apertis` | Apertis | `STIMA_API_KEY` | `https://api.stima.tech/v1` | -| `api2d` | API2D | `API2D_API_KEY` | `https://oa.api2d.net/v1` | -| `api2gpt` | API2GPT | `API2GPT_API_KEY` | `https://api.api2gpt.com/v1` | -| `apiserpent` | API Serpent | `APISERPENT_API_KEY` | `https://api.apiserpent.com/v1` | -| `atlascloud` | AtlasCloud | `ATLASCLOUD_API_KEY` | `https://api.atlascloud.com/v1` | -| `atomic_chat` | Atomic Chat | `ATOMIC_CHAT_API_KEY` | `http://127.0.0.1:1337/v1` | -| `auriko` | Auriko | `AURIKO_API_KEY` | `https://api.auriko.ai/v1` | -| `azure_ai` | Azure AI | `AZURE_AI_API_KEY` | `https://models.inference.ai.azure.com` | -| `baichuan` | Baichuan AI | `BAICHUAN_API_KEY` | `https://api.baichuan-ai.com/v1` | -| `baidu` | Baidu (文心/ERNIE) | `BAIDU_API_KEY` | `https://qianfan.baidubce.com/v2` | -| `baidu_v2` | BaiduV2 | `QIANFAN_API_KEY` | `https://qianfan.baidubce.com/v2` | -| `bailing` | Bailing | `BAILING_API_TOKEN` | `https://api.ant-ling.com/v1` | -| `baseten` | Baseten | `BASETEN_API_KEY` | `https://inference.baseten.co/v1` | -| `berget` | Berget.AI | `BERGET_API_KEY` | `https://api.berget.ai/v1` | -| `bigmodel` | BigModel (智谱) | `BIGMODEL_API_KEY` | `https://open.bigmodel.cn/api/paas/v4` | -| `blueclaw` | Blue Claw | `BLUECLAW_API_KEY` | `https://openai.blueclaw.network/v1` | -| `bytedance` | ByteDance | `ARK_API_KEY` | `https://ark.cn-beijing.volces.com/api/v3` | -| `byteplus` | BytePlus (Volcano) | `BYTEPLUS_API_KEY` | `https://ark.bytepluses.com/api/v3` | -| `bytez` | Bytez | `BYTEZ_API_KEY` | `https://api.bytez.com/v2` | -| `canopywave` | Canopywave | `CANOPYWAVE_API_KEY` | `https://api.canopywave.com/v1` | -| `cerebras` | Cerebras | `CEREBRAS_API_KEY` | `https://api.cerebras.ai/v1` | -| `chatgpt` | ChatGPT (订阅) | `CHATGPT_API_KEY` | `https://chatgpt.com/backend-api/codex` | -| `cherryin` | cherryin | `CHERRYIN_API_KEY` | `https://open.cherryin.net` | -| `chutes` | Chutes | `CHUTES_API_KEY` | `https://llm.chutes.ai/v1` | -| `clarifai` | Clarifai | `CLARIFAI_API_KEY` | `https://api.clarifai.com/v2/ext/openai/v1` | -| `claudinio` | Claudinio | `CLAUDINIO_API_KEY` | `https://api.claudin.io` | -| `cline_pass` | Cline | `CLINE_API_KEY` | `https://api.cline.bot/v1` | -| `closeai` | CloseAI | `CLOSEAI_API_KEY` | `https://api.closeai-proxy.xyz/v1` | -| `cloudferro_sherlock` | CloudFerro Sherlock | `CLOUDFERRO_SHERLOCK_API_KEY` | `https://api-sherlock.cloudferro.com/openai/v1` | -| `cloudflare` | Cloudflare | `CLOUDFLARE_API_KEY` | `https://api.cloudflare.com/client/v4/accounts/{CLOUDFLARE_ACCOUNT_ID}/ai/v1` | -| `cloudflare_workers_ai` | Cloudflare Workers AI | `CLOUDFLARE_API_KEY` | `https://api.cloudflare.com/client/v4/accounts/{CLOUDFLARE_ACCOUNT_ID}/ai/v1` | -| `codestral` | Codestral (Mistral) | `CODESTRAL_API_KEY` | `https://api.mistral.ai/v1` | -| `cometapi` | CometAPI | `COMETAPI_API_KEY` | `https://api.cometapi.com/v1` | -| `commandcode` | CommandCode | `COMMANDCODE_API_KEY` | `https://api.commandcode.com/v1` | -| `compactifai` | CompactifAI | `COMPACTIFAI_API_KEY` | `https://api.compactif.ai/v1` | -| `copilot` | GitHub Copilot | `COPILOT_API_KEY` | `https://api.githubcopilot.com` | -| `cortecs` | Cortecs | `CORTECS_API_KEY` | `https://api.cortecs.ai/v1/` | -| `coze` | Coze (扣子) | `COZE_API_KEY` | `https://api.coze.cn/v1` | -| `crof` | CrofAI | `CROF_API_KEY` | `https://crof.ai/v1` | -| `crossmodel` | CrossModel | `CROSSMODEL_API_KEY` | `https://api.crossmodel.ai/v1` | -| `crusoe` | Crusoe | `CRUSOE_API_KEY` | `https://api.inference.crusoecloud.com/v1` | -| `daoxe` | DaoXE | `DAOXE_API_KEY` | `https://daoxe.com/v1` | -| `darkbloom` | Darkbloom | `DARKBLOOM_API_KEY` | `https://api.darkbloom.dev/v1` | -| `databricks` | Databricks | `DATABRICKS_API_KEY` | `https://databricks.com/serving-endpoints` | -| `datarobot` | DataRobot | `DATAROBOT_API_TOKEN` | `https://app.datarobot.com/api/v2` | -| `deepbricks` | DeepBricks | `DEEPBRICKS_API_KEY` | `https://api.deepbricks.ai/v1` | -| `deepinfra` | DeepInfra | `DEEPINFRA_API_KEY` | `https://api.deepinfra.com/v1/openai` | -| `deepseek` | DeepSeek | `DEEPSEEK_API_KEY` | `https://api.deepseek.com/v1` | -| `digitalocean` | DigitalOcean | `DIGITALOCEAN_ACCESS_TOKEN` | `https://inference.do-ai.run` | -| `dinference` | DInference | `DINFERENCE_API_KEY` | `https://api.dinference.com/v1` | -| `doubao` | Doubao | `ARK_API_KEY` | `https://ark.cn-beijing.volces.com/api/v3` | -| `doubleword` | Doubleword | `DOUBLEWORD_API_KEY` | `https://api.doubleword.ai/v1` | -| `drun` | D.Run (China) | `DRUN_API_KEY` | `https://chat.d.run/v1` | -| `ebcloud` | EBCloud | `EBCLOUD_API_KEY` | `https://maas-api.ebcloud.com/v1` | -| `embercloud` | Embercloud | `EMBERCLOUD_API_KEY` | `https://api.embercloud.com/v1` | -| `empiriolabs` | EmpirioLabs AI | `EMPIRIOLABS_API_KEY` | `https://api.empiriolabs.ai/v1` | -| `evroc` | evroc | `EVROC_API_KEY` | `https://models.think.evroc.com/v1` | -| `fastcrw` | FastCRW | `FASTCRW_API_KEY` | `https://fastcrw.com/api/v1` | -| `fastgpt` | FastGPT | `FASTGPT_API_KEY` | `https://api.fastgpt.in/v1` | -| `fastrouter` | FastRouter | `FASTROUTER_API_KEY` | `https://api.fastrouter.ai/v1` | -| `featherless_ai` | Featherless AI | `FEATHERLESS_API_KEY` | `https://api.featherless.ai/v1` | -| `firepass` | Fireworks (Firepass) | `FIREWORKS_API_KEY` | `https://api.fireworks.ai/inference/v1` | -| `fireworks` | Fireworks | `FIREWORKS_API_KEY` | `https://api.fireworks.ai/inference/v1` | -| `freemodel` | FreeModel | `FREEMODEL_API_KEY` | `https://api.freemodel.dev/v1` | -| `friendliai` | FriendliAI | `FRIENDLIAI_API_KEY` | `https://inference.friendli.ai/v1` | -| `frogbot` | FrogBot | `FROGBOT_API_KEY` | `https://app.frogbot.ai/api/v1` | -| `galadriel` | Galadriel | `GALADRIEL_API_KEY` | `https://api.galadriel.com/v1` | -| `gdc` | GDC | `GDC_API_KEY` | `https://api.gdc.ai/v1` | -| `gigachat` | GigaChat (Sberbank) | `GIGACHAT_API_KEY` | `https://gigachat.devices.sberbank.ru/api/v1` | -| `github` | GitHub Models | `GITHUB_TOKEN` | `https://models.inference.ai.azure.com` | -| `gmi` | GMI | `GMI_API_KEY` | `https://api.gmi-serving.com/v1` | -| `gmicloud` | GMI Cloud | `GMI_API_KEY` | `https://api.gmi-serving.com/v1` | -| `gonka24` | Gonka24 | `GONKA24_API_KEY` | `https://api.gonka24.com/v1` | -| `gradient_ai` | Gradient AI | `GRADIENT_API_KEY` | `https://inference.do-ai.run/v1` | -| `groq` | Groq | `GROQ_API_KEY` | `https://api.groq.com/openai/v1` | -| `helicone` | Helicone | `HELICONE_API_KEY` | `https://api.helicone.ai/v1` | -| `heroku` | Heroku AI | `HEROKU_API_KEY` | `https://api.heroku.com/inference/v1` | -| `hetzner` | Hetzner | `HETZNER_VLLM_API_KEY` | `https://inference.hetzner.com/api/v1` | -| `hosted_vllm` | Hosted vLLM | `HOSTED_VLLM_API_KEY` | `https://hosted-vllm-api.com/v1` | -| `hpc_ai` | HPC-AI | `INFERENCE_API_KEY` | `https://api.hpc-ai.com/inference/v1` | -| `hyperbolic` | Hyperbolic | `HYPERBOLIC_API_KEY` | `https://api.hyperbolic.xyz/v1` | -| `iflowcn` | iFlow | `IFLOW_API_KEY` | `https://apis.iflow.cn/v1` | -| `inception` | Inception Labs | `INCEPTION_API_KEY` | `https://api.inceptionlabs.ai/v1` | -| `inceptron` | Inceptron | `INCEPTRON_API_KEY` | `https://api.inceptron.io/v1` | -| `inference_net` | Inference.net | `INFERENCE_NET_API_KEY` | `https://api.inference.net/v1` | -| `inferencehub` | InferenceHub | `INFERENCEHUB_API_KEY` | `https://app.inferencehub.tech/v1` | -| `inferx` | InferX | `INFERX_API_KEY` | `https://model.inferx.net/v1` | -| `infinity` | Infinity AI | `INFINITY_API_KEY` | `https://infinity.ai/api/v1` | -| `io_net` | IO.NET | `IOINTELLIGENCE_API_KEY` | `https://api.intelligence.io.solutions/api/v1` | -| `jiekou` | Jiekou.AI | `JIEKOU_API_KEY` | `https://api.highwayapi.ai/openai` | -| `kenari` | Kenari | `KENARI_API_KEY` | `https://kenari.id/v1` | -| `kilo` | Kilo | `KILO_API_KEY` | `https://api.kilo.ai/v1` | -| `kimi` | Kimi | `MOONSHOT_API_KEY` | `https://api.moonshot.ai/v1` | -| `kimi_for_coding` | Kimi For Coding | `KIMI_API_KEY` | `https://api.kimi.com/coding/v1` | -| `kiro` | Kiro | `KIRO_API_KEY` | `https://api.kiro.dev/v1` | -| `kluster_ai` | Kluster AI | `KLUSTER_API_KEY` | `https://api.kluster.ai/v1` | -| `krutrim` | Krutrim | `KRUTRIM_API_KEY` | `https://api.krutrim.ai/v1` | -| `kuae_cloud_coding_plan` | KUAE Cloud Coding Plan | `KUAE_API_KEY` | `https://coding-plan-endpoint.kuaecloud.net/v1` | -| `lambda_ai` | Lambda AI | `LAMBDA_API_KEY` | `https://api.lambda.ai/v1` | -| `lemonade` | Lemonade | `LEMONADE_API_KEY` | `http://localhost:13305/v1` | -| `lemonfox_ai` | Lemonfox AI | `LEMONFOX_API_KEY` | `https://api.lemonfox.ai/v1` | -| `libertai` | Libertai | `LIBERTAI_API_KEY` | `https://api.libertai.io/v1` | -| `lilac` | Lilac | `LILAC_API_KEY` | `https://api.getlilac.com/v1` | -| `lingyiwanwu` | Lingyiwanwu (零一万物) | `LINGYIWANWU_API_KEY` | `https://api.lingyiwanwu.com/v1` | -| `llama` | Llama | `LLAMA_API_KEY` | `https://api.llama.com/compat/v1/` | -| `llamagate` | Llamagate | `LLAMAGATE_API_KEY` | `https://api.llamagate.dev/v1` | -| `llmgateway` | LLM Gateway | `LLM_GATEWAY_API_KEY` | `https://api.llmgateway.io/v1` | -| `llmtr` | LLMTR | `LLMTR_API_KEY` | `https://llmtr.com/v1` | -| `longcat` | LongCat | `LONGCAT_API_KEY` | `https://api.longcat.chat/v1` | -| `lucidquery` | LucidQuery | `LUCIDQUERY_API_KEY` | `https://api.lucidquery.com/v1` | -| `lynkr` | Lynkr | `LYNKR_API_KEY` | `http://localhost:8081/v1` | -| `matterai` | Matter AI | `MATTERAI_API_KEY` | `https://api.matterai.com/v1` | -| `meganova` | Meganova | `MEGANOVA_API_KEY` | `https://api.meganova.ai/v1` | -| `merge_gateway` | Merge Gateway | `MERGE_GATEWAY_API_KEY` | `https://api-gateway.merge.dev/v1/openai` | -| `meta` | Meta | `MODEL_API_KEY` | `https://api.meta.ai/v1` | -| `meta_llama` | Meta Llama API | `LLAMA_API_KEY` | `https://api.llama.com/compat/v1` | -| `mimo` | Mimo | `MIMO_API_KEY` | `https://api.xiaomimimo.com/v1` | -| `minimax` | MiniMax | `MINIMAX_API_KEY` | `https://api.minimax.io/v1` | -| `minimax_cn` | MiniMax (minimaxi.com) | `MINIMAX_API_KEY` | `https://api.minimaxi.com/v1` | -| `minimax_cn_coding_plan` | MiniMax Token Plan (minimaxi.com) | `MINIMAX_API_KEY` | `https://api.minimaxi.com/v1` | -| `minimax_coding_plan` | MiniMax Token Plan (minimax.io) | `MINIMAX_API_KEY` | `https://api.minimax.io/anthropic/v1` | -| `mira` | Mira | `MIRA_API_KEY` | `https://api.mira.so/v1` | -| `mixlayer` | Mixlayer | `MIXLAYER_API_KEY` | `https://models.mixlayer.ai/v1` | -| `moark` | Moark | `MOARK_API_KEY` | `https://api.moark.com/v1` | -| `modal` | Modal | `MODAL_API_KEY` | `https://modal.com/v1` | -| `model_oracle_ai` | Model Oracle AI | `MODEL_ORACLE_API_KEY` | `https://api.modeloracle.com/api/v1` | -| `modelscope` | ModelScope | `MODELSCOPE_API_KEY` | `https://api-inference.modelscope.cn/v1` | -| `moonshotai` | Moonshot AI | `MOONSHOT_API_KEY` | `https://api.moonshot.cn/v1` | -| `moonshotai_cn` | Moonshot AI (China) | `MOONSHOT_API_KEY` | `https://api.moonshot.cn/anthropic/v1` | -| `morph` | Morph LLM | `MORPH_API_KEY` | `https://api.morphllm.com/v1` | -| `nanogpt` | NanoGPT | `NANOGPT_API_KEY` | `https://api.nanogpt.com/v1` | -| `ncompass` | Ncompass | `NCOMPASS_API_KEY` | `https://api.ncompass.tech/v1` | -| `nearai` | NEAR AI Cloud | `NEARAI_API_KEY` | `https://cloud-api.near.ai/v1` | -| `nebius` | Nebius AI | `NEBIUS_API_KEY` | `https://api.studio.nebius.ai/v1` | -| `neon` | Neon | `NEON_AI_GATEWAY_TOKEN` | `https:///v1` | -| `neuralwatt` | Neuralwatt | `NEURALWATT_API_KEY` | `https://api.neuralwatt.com/v1` | -| `nextbit` | NextBit | `NEXTBIT_API_KEY` | `https://api.nextbit.ai/v1` | -| `nlp_cloud` | NLP Cloud | `NLPCLOUD_API_KEY` | `https://api.nlpcloud.io/v1` | -| `nous_research` | Nous Research | `NOUS_API_KEY` | `https://api.nousresearch.com/v1` | -| `novita` | Novita AI | `NOVITA_API_KEY` | `https://api.novita.ai/v1` | -| `nscale` | Nscale | `NSCALE_API_KEY` | `https://inference.api.nscale.com/v1` | -| `nvidia_nim` | NVIDIA NIM | `NVIDIA_API_KEY` | `https://integrate.api.nvidia.com/v1` | -| `oci` | OCI | `OCI_API_KEY` | `https://inference.generativeai.${region}.oci.oraclecloud.com/openai/v1` | -| `ofox` | OfoxAI | `OFOX_API_KEY` | `https://api.ofox.ai/v1` | -| `ohmygpt` | OhMyGPT | `OHMYGPT_API_KEY` | `https://api.ohmygpt.com/v1` | -| `ollama_cloud` | Ollama Cloud | `OLLAMA_CLOUD_API_KEY` | `https://api.ollama.com/v1` | -| `openaimax` | OpenAIMax | `OPENAIMAX_API_KEY` | `https://api.openaimax.com/v1` | -| `openaisb` | OpenAI-SB | `OPENAISB_API_KEY` | `https://api.openaisb.com/v1` | -| `opencode` | OpenCode Zen | `OPENCODE_API_KEY` | `https://api.opencode.zen/v1` | -| `opencode_go` | OpenCode Go | `OPENCODE_GO_API_KEY` | `https://api.opencode.dev/v1` | -| `opencode_zen` | OpenCode Zen | `OPENCODE_ZEN_API_KEY` | `https://api.opencode.zen/v1` | -| `orcarouter` | OrcaRouter | `ORCAROUTER_API_KEY` | `https://api.orcarouter.com/v1` | -| `ovhcloud` | OVHcloud AI | `OVHCLOUD_API_KEY` | `https://oai.endpoints.kepler.ai.cloud.ovh.net/v1` | -| `parasail` | Parasail | `PARASAIL_API_KEY` | `https://api.parasail.io/v1` | -| `perfxcloud` | PerfXCloud | `PERFXCLOUD_API_KEY` | `https://api.perfxcloud.com/v1` | -| `perplexity` | Perplexity | `PERPLEXITY_API_KEY` | `https://api.perplexity.ai` | -| `perplexity_agent` | Perplexity Agent | `PERPLEXITY_API_KEY` | `https://api.perplexity.ai/v1` | -| `petals` | Petals | `PETALS_API_KEY` | `https://api.petals.dev/v1` | -| `pinstripes` | Pinstripes | `PINSTRIPES_API_KEY` | `https://api.pinstripes.io/v1` | -| `pioneer` | Pioneer | `PIONEER_API_KEY` | `https://api.pioneer.ai/v1` | -| `poe` | Poe | `POE_API_KEY` | `https://api.poe.com/v1` | -| `poolside` | Poolside | `POOLSIDE_API_KEY` | `https://inference.poolside.ai/v1` | -| `portkey` | Portkey Gateway | `PORTKEY_API_KEY` | `https://api.portkey.ai/v1` | -| `ppinfra` | PPInfra（PPIO 派欧云） | `PPIO_API_KEY` | `https://api.ppio.com/openai` | -| `predibase` | Predibase | `PREDIBASE_API_KEY` | `https://serving.app.predibase.com/v1` | -| `privatemode_ai` | Privatemode AI | `PRIVATEMODE_API_KEY` | `http://localhost:8080/v1` | -| `publicai` | Publicai | `PUBLICAI_API_KEY` | `https://platform.publicai.co/v1` | -| `qihang_ai` | QiHang（启航 AI） | `QIHANG_API_KEY` | `https://api.qhaigc.net/v1` | -| `qihoo360` | 360 AI | `AI360_API_KEY` | `https://api.360.cn/v1` | -| `qiniu_ai` | Qiniu AI | `QINIU_API_KEY` | `https://api.qiniu.com/v1` | -| `regolo_ai` | Regolo AI | `REGOLO_API_KEY` | `https://api.regolo.ai/v1` | -| `reka_ai` | Reka AI | `REKA_API_KEY` | `https://api.reka.ai/v1` | -| `requesty` | Requesty | `REQUESTY_API_KEY` | `https://api.requesty.ai/v1` | -| `reve` | Reve | `REVE_API_KEY` | `https://api.reve.ai/v1` | -| `routing_run` | routing.run | `ROUTING_RUN_API_KEY` | `https://api.routing.run/v1` | -| `sakana` | Sakana AI | `SAKANA_API_KEY` | `https://api.sakana.ai/v1` | -| `sambanova` | SambaNova | `SAMBANOVA_API_KEY` | `https://api.sambanova.ai/v1` | -| `sarvam` | Sarvam AI | `SARVAM_API_KEY` | `https://api.sarvam.ai/v1` | -| `scaleway` | Scaleway AI | `SCALEWAY_API_KEY` | `https://api.scaleway.ai/v1` | -| `scx_ai` | SCX AI | `SCX_AI_API_KEY` | `https://api.scx.ai/v1` | -| `siliconflow` | SiliconFlow | `SILICONFLOW_API_KEY` | `https://api.siliconflow.cn/v1` | -| `snowflake` | Snowflake | `SNOWFLAKE_PAT` | `https://.snowflakecomputing.com/api/v2/cortex/v1` | -| `snowflake_cortex` | Snowflake Cortex | `SNOWFLAKE_CORTEX_PAT` | `https://.snowflakecomputing.com/api/v2/cortex/v1` | -| `stackit` | STACKIT | `STACKIT_API_KEY` | `https://api.openai-compat.model-serving.eu01.onstackit.cloud/v1` | -| `stepfun` | StepFun (阶跃星辰) | `STEPFUN_API_KEY` | `https://api.stepfun.com/v1` | -| `stepfun_ai_step_plan` | StepFun Step Plan (Global) | `STEPFUN_API_KEY` | `https://api.stepfun.ai/step_plan/v1` | -| `stepfun_step_plan` | StepFun Step Plan (China) | `STEPFUN_API_KEY` | `https://api.stepfun.com/step_plan/v1` | -| `subconscious` | Subconscious | `SUBCONSCIOUS_API_KEY` | `https://api.subconscious.dev/v1` | -| `submodel` | SubModel | `SUBMODEL_API_KEY` | `https://api.submodel.com/v1` | -| `synthetic` | Synthetic | `SYNTHETIC_API_KEY` | `https://api.synthetic.new/openai/v1` | -| `tencent` | Tencent (混元/Hunyuan) | `TENCENT_API_KEY` | `https://api.hunyuan.cloud.tencent.com/v1` | -| `tencent_coding_plan` | Tencent Coding Plan (China) | `TENCENT_CODING_PLAN_API_KEY` | `https://api.lkeap.cloud.tencent.com/coding/v3` | -| `tencent_token_plan` | Tencent Token Plan | `TENCENT_TOKEN_PLAN_API_KEY` | `https://api.lkeap.cloud.tencent.com/plan/v3` | -| `tencent_token_plan_enterprise_auto` | 腾讯云 Token Plan / Token Plan 企业版轻享套餐 | `TENCENT_TOKEN_PLAN_ENTERPRISE_API_KEY` | `https://tokenhub.tencentmaas.com/plan/v3` | -| `tencent_token_plan_enterprise_pro` | 腾讯云 Token Plan / Token Plan 企业版专业套餐 | `TENCENT_TOKEN_PLAN_ENTERPRISE_API_KEY` | `https://tokenhub.tencentmaas.com/plan/v3` | -| `tencent_token_plan_general_personal` | 腾讯云 Token Plan / 通用 Token Plan（个人版） | `TENCENT_TOKEN_PLAN_API_KEY` | `https://api.lkeap.cloud.tencent.com/plan/v3` | -| `tencent_token_plan_hy_personal` | 腾讯云 Token Plan / Hy Token Plan（个人版） | `TENCENT_TOKEN_PLAN_API_KEY` | `https://api.lkeap.cloud.tencent.com/plan/v3` | -| `tencent_tokenhub` | Tencent TokenHub | `TENCENT_TOKENHUB_API_KEY` | `https://tokenhub.tencentmaas.com/v1` | -| `tensormesh` | Tensormesh | `YOUR_API_KEY` | `https://serverless.tensormesh.ai` | -| `the_grid_ai` | The Grid AI | `THEGRIDAI_API_KEY` | `https://api.thegrid.ai/v1` | -| `thinkingmachines` | Thinking Machines | `TINKER_API_KEY` | `https://tinker.thinkingmachines.dev/services/tinker-prod/oai/api/v1` | -| `tinfoil` | Tinfoil | `TINFOIL_API_KEY` | `https://inference.tinfoil.sh/v1` | -| `togetherai` | Together AI | `TOGETHER_API_KEY` | `https://api.together.xyz/v1` | -| `tokenflux` | Tokenflux | `TOKENFLUX_API_KEY` | `https://tokenflux.ai/v1` | -| `tokenpony` | TokenPony | `TOKENPONY_API_KEY` | `https://api.tokenpony.com/v1` | -| `trustedrouter` | TrustedRouter | `TRUSTEDROUTER_API_KEY` | `https://api.trustedrouter.com/v1` | -| `tundra` | Tundra | `TUNDRA_API_KEY` | `https://api.tundra.ai/v1` | -| `umans_ai` | Umans AI | `UMANS_AI_API_KEY` | `https://api.code.umans.ai/v1` | -| `unorouter` | UnoRouter | `UNOROUTER_API_KEY` | `https://unorouter.com/en` | -| `upstage` | Upstage | `UPSTAGE_API_KEY` | `https://api.upstage.ai/v1` | -| `v0` | v0 (Vercel) | `V0_API_KEY` | `https://api.v0.dev/v1` | -| `venice` | Venice | `VENICE_API_KEY` | `https://api.venice.ai/api/v1` | -| `vercel` | Vercel | `VERCEL_API_KEY` | `https://api.v0.dev/v1` | -| `vivgrid` | Vivgrid | `VIVGRID_API_KEY` | `https://api.vivgrid.com/v1` | -| `volc_engine` | VolcEngine | `ARK_API_KEY` | `https://ark.cn-beijing.volces.com/api/v3` | -| `vultr` | Vultr | `VULTR_API_KEY` | `https://api.vultrinference.com/v1` | -| `wafer` | Wafer | `WAFER_API_KEY` | `https://api.wafer.ai/v1` | -| `wandb` | Weights & Biases | `WANDB_API_KEY` | `https://api.inference.wandb.ai/v1` | -| `xiaomi_token_plan_ams` | Xiaomi Token Plan (Europe) | `MIMO_API_KEY` | `https://token-plan-ams.xiaomimimo.com/v1` | -| `xiaomi_token_plan_cn` | Xiaomi Token Plan (China) | `MIMO_API_KEY` | `https://token-plan-cn.xiaomimimo.com/v1` | -| `xiaomi_token_plan_sgp` | Xiaomi Token Plan (Singapore) | `MIMO_API_KEY` | `https://token-plan-sgp.xiaomimimo.com/v1` | -| `xiaomimimo` | Xiaomi MiMo | `XIAOMI_API_KEY` | `https://mimo.xiaomi.com/v1` | -| `xpersona` | Xpersona | `XPERSONA_API_KEY` | `https://www.xpersona.co/v1` | -| `xunfei` | Xunfei | `XUNFEI_API_PASSWORD` | `https://spark-api-open.xf-yun.com/v1` | -| `zai` | Zai | `ZAI_API_KEY` | `https://api.z.ai/api/paas/v4` | -| `zai_coding_plan` | Z.AI Coding Plan | `ZHIPU_API_KEY` | `https://api.z.ai/api/anthropic` | -| `zeldoc` | Zeldoc | `ZELDOC_API_KEY` | `https://api.zeldoc.ai/v1` | -| `zenmux` | ZenMux | `ZENMUX_API_KEY` | `https://zenmux.ai/api/v1` | -| `zhipu_v4` | ZhipuV4 | `ZHIPU_API_KEY` | `https://open.bigmodel.cn/api/paas/v4` | -| `zhipuai_coding_plan` | Zhipu AI Coding Plan | `ZHIPU_API_KEY` | `https://open.bigmodel.cn/api/coding/paas/v4` | +| **Total providers** | **327** | + +**282 registry-backed OpenAI-compatible providers** (create presets by name via `provider(name, ...)`) + **45 non-registry providers** (construct via the typed factories listed below). + +## Registry-backed (OpenAI-compatible) — 282 + +| name | display | auth | env var | base_url | +|------|---------|------|---------|----------| +| `abacus` | Abacus | api_key | `ABACUS_API_KEY` | `https://routellm.abacus.ai/v1` | +| `abliteration_ai` | Abliteration AI | api_key | `ABLIT_KEY` | `https://api.abliteration.ai/v1` | +| `ai21` | AI21 Labs | api_key | `AI21_API_KEY` | `https://api.ai21.ai/v1` | +| `ai302` | 302.AI | api_key | `AI302_API_KEY` | `https://api.302.ai/v1` | +| `ai_router` | AI-ROUTER | api_key | `AI_ROUTER_API_KEY` | `https://api.ai-router.dev/v1` | +| `aiand` | AIand | api_key | `AIAND_API_KEY` | `https://api.aiand.com/v1` | +| `aibadgr` | AI Badgr | api_key | `AIBADGR_API_KEY` | `https://api.aibadgr.com/v1` | +| `aigc2d` | AIGC2D | api_key | `AIGC2D_API_KEY` | `https://api.aigc2d.com/v1` | +| `aihubmix` | AIHubMix | api_key | `AIHUBMIX_API_KEY` | `https://aihubmix.com/v1` | +| `ails` | AILS | api_key | `AILS_API_KEY` | `https://api.caipacity.com/v1` | +| `aiml` | AI/ML API | api_key | `AIML_API_KEY` | `https://api.aimlapi.com/v1` | +| `aki_io` | AKI.IO | api_key | `AKI_IO_API_KEY` | `https://aki.io/openai/v1` | +| `albert` | Albert | api_key | `ALBERT_API_KEY` | `https://api.albert.ai/v1` | +| `alibaba` | Alibaba Cloud (DashScope) | api_key | `ALIBABA_API_KEY` | `https://dashscope-intl.aliyuncs.com/compatible-mode/v1` | +| `alibaba_coding_plan` | Alibaba Coding Plan | api_key | `ALIBABA_CODING_PLAN_API_KEY` | `https://coding-intl.dashscope.aliyuncs.com/v1` | +| `alibaba_coding_plan_cn` | Alibaba Coding Plan (China) | api_key | `ALIBABA_CODING_PLAN_API_KEY` | `https://coding.dashscope.aliyuncs.com/v1` | +| `alibaba_token_plan` | Alibaba Token Plan | api_key | `ALIBABA_TOKEN_PLAN_API_KEY` | `https://token-plan.ap-southeast-1.maas.aliyuncs.com/compatible-mode/v1` | +| `alibaba_token_plan_cn` | Alibaba Token Plan (China) | api_key | `ALIBABA_TOKEN_PLAN_API_KEY` | `https://token-plan.cn-beijing.maas.aliyuncs.com/compatible-mode/v1` | +| `ambient` | Ambient | api_key | `AMBIENT_API_KEY` | `https://api.ambient.xyz/v1` | +| `anyapi` | AnyAPI | api_key | `ANYAPI_KEY` | `https://api.anyapi.ai/v1` | +| `anyscale` | Anyscale | api_key | `ANYSCALE_API_KEY` | `https://api.endpoints.anyscale.com/v1` | +| `apertis` | Apertis | api_key | `STIMA_API_KEY` | `https://api.stima.tech/v1` | +| `api2d` | API2D | api_key | `API2D_API_KEY` | `https://oa.api2d.net/v1` | +| `api2gpt` | API2GPT | api_key | `API2GPT_API_KEY` | `https://api.api2gpt.com/v1` | +| `apiserpent` | API Serpent | api_key | `APISERPENT_API_KEY` | `https://api.apiserpent.com/v1` | +| `atlascloud` | AtlasCloud | api_key | `ATLASCLOUD_API_KEY` | `https://api.atlascloud.com/v1` | +| `atomic_chat` | Atomic Chat | api_key | `ATOMIC_CHAT_API_KEY` | `http://127.0.0.1:1337/v1` | +| `auriko` | Auriko | api_key | `AURIKO_API_KEY` | `https://api.auriko.ai/v1` | +| `azure_ai` | Azure AI | api_key | `AZURE_AI_API_KEY` | `https://models.inference.ai.azure.com` | +| `baichuan` | Baichuan AI | api_key | `BAICHUAN_API_KEY` | `https://api.baichuan-ai.com/v1` | +| `baidu` | Baidu (文心/ERNIE) | api_key | `BAIDU_API_KEY` | `https://qianfan.baidubce.com/v2` | +| `baidu_v2` | BaiduV2 | api_key | `QIANFAN_API_KEY` | `https://qianfan.baidubce.com/v2` | +| `bailing` | Bailing | api_key | `BAILING_API_TOKEN` | `https://api.ant-ling.com/v1` | +| `baseten` | Baseten | api_key | `BASETEN_API_KEY` | `https://inference.baseten.co/v1` | +| `bedrock_mantle` | Bedrock Mantle | api_key | `BEDROCK_MANTLE_API_KEY` | `https://bedrock-mantle.{region}.api.aws/v1` (params: `region`) | +| `berget` | Berget.AI | api_key | `BERGET_API_KEY` | `https://api.berget.ai/v1` | +| `bigmodel` | BigModel (智谱) | api_key | `BIGMODEL_API_KEY` | `https://open.bigmodel.cn/api/paas/v4` | +| `blueclaw` | Blue Claw | api_key | `BLUECLAW_API_KEY` | `https://openai.blueclaw.network/v1` | +| `bytedance` | ByteDance | api_key | `ARK_API_KEY` | `https://ark.cn-beijing.volces.com/api/v3` | +| `byteplus` | BytePlus (Volcano) | api_key | `BYTEPLUS_API_KEY` | `https://ark.bytepluses.com/api/v3` | +| `bytez` | Bytez | api_key | `BYTEZ_API_KEY` | `https://api.bytez.com/v2` | +| `canopywave` | Canopywave | api_key | `CANOPYWAVE_API_KEY` | `https://api.canopywave.com/v1` | +| `cerebras` | Cerebras | api_key | `CEREBRAS_API_KEY` | `https://api.cerebras.ai/v1` | +| `chatgpt` | ChatGPT (订阅) | api_key | `CHATGPT_API_KEY` | `https://chatgpt.com/backend-api/codex` | +| `cherryin` | cherryin | api_key | `CHERRYIN_API_KEY` | `https://open.cherryin.net` | +| `chutes` | Chutes | api_key | `CHUTES_API_KEY` | `https://llm.chutes.ai/v1` | +| `clarifai` | Clarifai | api_key | `CLARIFAI_API_KEY` | `https://api.clarifai.com/v2/ext/openai/v1` | +| `claudinio` | Claudinio | api_key | `CLAUDINIO_API_KEY` | `https://api.claudin.io` | +| `cline_pass` | Cline | api_key | `CLINE_API_KEY` | `https://api.cline.bot/v1` | +| `closeai` | CloseAI | api_key | `CLOSEAI_API_KEY` | `https://api.closeai-proxy.xyz/v1` | +| `cloudferro_sherlock` | CloudFerro Sherlock | api_key | `CLOUDFERRO_SHERLOCK_API_KEY` | `https://api-sherlock.cloudferro.com/openai/v1` | +| `cloudflare` | Cloudflare | api_key | `CLOUDFLARE_API_KEY` | `https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1` (params: `account_id`) | +| `cloudflare_workers_ai` | Cloudflare Workers AI | api_key | `CLOUDFLARE_API_KEY` | `https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1` (params: `account_id`) | +| `codestral` | Codestral (Mistral) | api_key | `CODESTRAL_API_KEY` | `https://api.mistral.ai/v1` | +| `cometapi` | CometAPI | api_key | `COMETAPI_API_KEY` | `https://api.cometapi.com/v1` | +| `commandcode` | CommandCode | api_key | `COMMANDCODE_API_KEY` | `https://api.commandcode.com/v1` | +| `compactifai` | CompactifAI | api_key | `COMPACTIFAI_API_KEY` | `https://api.compactif.ai/v1` | +| `copilot` | GitHub Copilot | api_key | `COPILOT_API_KEY` | `https://api.githubcopilot.com` | +| `cortecs` | Cortecs | api_key | `CORTECS_API_KEY` | `https://api.cortecs.ai/v1/` | +| `coze` | Coze (扣子) | api_key | `COZE_API_KEY` | `https://api.coze.cn/v1` | +| `crof` | CrofAI | api_key | `CROF_API_KEY` | `https://crof.ai/v1` | +| `crossmodel` | CrossModel | api_key | `CROSSMODEL_API_KEY` | `https://api.crossmodel.ai/v1` | +| `crusoe` | Crusoe | api_key | `CRUSOE_API_KEY` | `https://api.inference.crusoecloud.com/v1` | +| `cybertron` | Cybertron | none | — | `http://127.0.0.1:8080/v1` (URL from `CYBERTRON_BASE_URL`) | +| `daoxe` | DaoXE | api_key | `DAOXE_API_KEY` | `https://daoxe.com/v1` | +| `darkbloom` | Darkbloom | api_key | `DARKBLOOM_API_KEY` | `https://api.darkbloom.dev/v1` | +| `databricks` | Databricks | api_key | `DATABRICKS_API_KEY` | `https://databricks.com/serving-endpoints` | +| `datarobot` | DataRobot | api_key | `DATAROBOT_API_TOKEN` | `https://app.datarobot.com/api/v2` | +| `deepbricks` | DeepBricks | api_key | `DEEPBRICKS_API_KEY` | `https://api.deepbricks.ai/v1` | +| `deepinfra` | DeepInfra | api_key | `DEEPINFRA_API_KEY` | `https://api.deepinfra.com/v1/openai` | +| `digitalocean` | DigitalOcean | api_key | `DIGITALOCEAN_ACCESS_TOKEN` | `https://inference.do-ai.run` | +| `dinference` | DInference | api_key | `DINFERENCE_API_KEY` | `https://api.dinference.com/v1` | +| `docker_model_runner` | Docker Model Runner | none | — | `http://model-runner.docker.internal/engines/llama.cpp/v1` (URL from `DOCKER_MODEL_RUNNER_BASE_URL`) | +| `doubao` | Doubao | api_key | `ARK_API_KEY` | `https://ark.cn-beijing.volces.com/api/v3` | +| `doubleword` | Doubleword | api_key | `DOUBLEWORD_API_KEY` | `https://api.doubleword.ai/v1` | +| `drun` | D.Run (China) | api_key | `DRUN_API_KEY` | `https://chat.d.run/v1` | +| `ebcloud` | EBCloud | api_key | `EBCLOUD_API_KEY` | `https://maas-api.ebcloud.com/v1` | +| `embercloud` | Embercloud | api_key | `EMBERCLOUD_API_KEY` | `https://api.embercloud.com/v1` | +| `empiriolabs` | EmpirioLabs AI | api_key | `EMPIRIOLABS_API_KEY` | `https://api.empiriolabs.ai/v1` | +| `evroc` | evroc | api_key | `EVROC_API_KEY` | `https://models.think.evroc.com/v1` | +| `fastcrw` | FastCRW | api_key | `FASTCRW_API_KEY` | `https://fastcrw.com/api/v1` | +| `fastgpt` | FastGPT | api_key | `FASTGPT_API_KEY` | `https://api.fastgpt.in/v1` | +| `fastrouter` | FastRouter | api_key | `FASTROUTER_API_KEY` | `https://api.fastrouter.ai/v1` | +| `featherless_ai` | Featherless AI | api_key | `FEATHERLESS_API_KEY` | `https://api.featherless.ai/v1` | +| `firepass` | Fireworks (Firepass) | api_key | `FIREWORKS_API_KEY` | `https://api.fireworks.ai/inference/v1` | +| `fireworks` | Fireworks | api_key | `FIREWORKS_API_KEY` | `https://api.fireworks.ai/inference/v1` | +| `freemodel` | FreeModel | api_key | `FREEMODEL_API_KEY` | `https://api.freemodel.dev/v1` | +| `friendliai` | FriendliAI | api_key | `FRIENDLIAI_API_KEY` | `https://inference.friendli.ai/v1` | +| `frogbot` | FrogBot | api_key | `FROGBOT_API_KEY` | `https://app.frogbot.ai/api/v1` | +| `galadriel` | Galadriel | api_key | `GALADRIEL_API_KEY` | `https://api.galadriel.com/v1` | +| `gaudi` | Intel Gaudi | none | — | `http://127.0.0.1:8080/v1` (URL from `GAUDI_BASE_URL`) | +| `gdc` | GDC | api_key | `GDC_API_KEY` | `https://api.gdc.ai/v1` | +| `gigachat` | GigaChat (Sberbank) | api_key | `GIGACHAT_API_KEY` | `https://gigachat.devices.sberbank.ru/api/v1` | +| `github` | GitHub Models | api_key | `GITHUB_TOKEN` | `https://models.inference.ai.azure.com` | +| `gmi` | GMI | api_key | `GMI_API_KEY` | `https://api.gmi-serving.com/v1` | +| `gmicloud` | GMI Cloud | api_key | `GMI_API_KEY` | `https://api.gmi-serving.com/v1` | +| `gonka24` | Gonka24 | api_key | `GONKA24_API_KEY` | `https://api.gonka24.com/v1` | +| `gradient_ai` | Gradient AI | api_key | `GRADIENT_API_KEY` | `https://inference.do-ai.run/v1` | +| `helicone` | Helicone | api_key | `HELICONE_API_KEY` | `https://api.helicone.ai/v1` | +| `heroku` | Heroku AI | api_key | `HEROKU_API_KEY` | `https://api.heroku.com/inference/v1` | +| `hetzner` | Hetzner | api_key | `HETZNER_VLLM_API_KEY` | `https://inference.hetzner.com/api/v1` | +| `hosted_vllm` | Hosted vLLM | api_key | `HOSTED_VLLM_API_KEY` | `https://hosted-vllm-api.com/v1` | +| `hpc_ai` | HPC-AI | api_key | `INFERENCE_API_KEY` | `https://api.hpc-ai.com/inference/v1` | +| `hyperbolic` | Hyperbolic | api_key | `HYPERBOLIC_API_KEY` | `https://api.hyperbolic.xyz/v1` | +| `iflowcn` | iFlow | api_key | `IFLOW_API_KEY` | `https://apis.iflow.cn/v1` | +| `inception` | Inception Labs | api_key | `INCEPTION_API_KEY` | `https://api.inceptionlabs.ai/v1` | +| `inceptron` | Inceptron | api_key | `INCEPTRON_API_KEY` | `https://api.inceptron.io/v1` | +| `inference_net` | Inference.net | api_key | `INFERENCE_NET_API_KEY` | `https://api.inference.net/v1` | +| `inferencehub` | InferenceHub | api_key | `INFERENCEHUB_API_KEY` | `https://app.inferencehub.tech/v1` | +| `inferx` | InferX | api_key | `INFERX_API_KEY` | `https://model.inferx.net/v1` | +| `infinity` | Infinity AI | api_key | `INFINITY_API_KEY` | `https://infinity.ai/api/v1` | +| `io_net` | IO.NET | api_key | `IOINTELLIGENCE_API_KEY` | `https://api.intelligence.io.solutions/api/v1` | +| `jiekou` | Jiekou.AI | api_key | `JIEKOU_API_KEY` | `https://api.highwayapi.ai/openai` | +| `jlama` | Jlama | none | — | `http://127.0.0.1:8080/v1` (URL from `JLAMA_BASE_URL`) | +| `kenari` | Kenari | api_key | `KENARI_API_KEY` | `https://kenari.id/v1` | +| `kilo` | Kilo | api_key | `KILO_API_KEY` | `https://api.kilo.ai/v1` | +| `kimi` | Kimi | api_key | `MOONSHOT_API_KEY` | `https://api.moonshot.ai/v1` | +| `kimi_for_coding` | Kimi For Coding | api_key | `KIMI_API_KEY` | `https://api.kimi.com/coding/v1` | +| `kiro` | Kiro | api_key | `KIRO_API_KEY` | `https://api.kiro.dev/v1` | +| `kluster_ai` | Kluster AI | api_key | `KLUSTER_API_KEY` | `https://api.kluster.ai/v1` | +| `krutrim` | Krutrim | api_key | `KRUTRIM_API_KEY` | `https://api.krutrim.ai/v1` | +| `kuae_cloud_coding_plan` | KUAE Cloud Coding Plan | api_key | `KUAE_API_KEY` | `https://coding-plan-endpoint.kuaecloud.net/v1` | +| `lambda_ai` | Lambda AI | api_key | `LAMBDA_API_KEY` | `https://api.lambda.ai/v1` | +| `laya` | Laya | none | `LAYA_API_KEY` | `http://127.0.0.1:8000/v1` (URL from `LAYA_BASE_URL`) | +| `lemonade` | Lemonade | api_key | `LEMONADE_API_KEY` | `http://localhost:13305/v1` | +| `lemonfox_ai` | Lemonfox AI | api_key | `LEMONFOX_API_KEY` | `https://api.lemonfox.ai/v1` | +| `libertai` | Libertai | api_key | `LIBERTAI_API_KEY` | `https://api.libertai.io/v1` | +| `lilac` | Lilac | api_key | `LILAC_API_KEY` | `https://api.getlilac.com/v1` | +| `lingyiwanwu` | Lingyiwanwu (零一万物) | api_key | `LINGYIWANWU_API_KEY` | `https://api.lingyiwanwu.com/v1` | +| `litellm_proxy` | LiteLLM Proxy | none | — | `http://127.0.0.1:4000/v1` (URL from `LITELLM_PROXY_BASE_URL`) | +| `llama` | Llama | api_key | `LLAMA_API_KEY` | `https://api.llama.com/compat/v1/` | +| `llamacpp` | llama.cpp | none | `LLAMACPP_API_KEY` | `http://127.0.0.1:8080/v1` (URL from `LLAMACPP_BASE_URL`) | +| `llamafile` | Llamafile | none | — | `http://127.0.0.1:8080/v1` (URL from `LLAMAFILE_BASE_URL`) | +| `llamagate` | Llamagate | api_key | `LLAMAGATE_API_KEY` | `https://api.llamagate.dev/v1` | +| `llmgateway` | LLM Gateway | api_key | `LLM_GATEWAY_API_KEY` | `https://api.llmgateway.io/v1` | +| `llmtr` | LLMTR | api_key | `LLMTR_API_KEY` | `https://llmtr.com/v1` | +| `lmstudio` | LM Studio | none | — | `http://127.0.0.1:1234/v1` (URL from `LMSTUDIO_BASE_URL`) | +| `local` | Local LLM | none | — | `http://127.0.0.1:8080/v1` (URL from `LOCAL_LLM_BASE_URL`) | +| `localai` | LocalAI | none | `LOCALAI_API_KEY` | `http://127.0.0.1:8080/v1` (URL from `LOCALAI_BASE_URL`) | +| `longcat` | LongCat | api_key | `LONGCAT_API_KEY` | `https://api.longcat.chat/v1` | +| `lucidquery` | LucidQuery | api_key | `LUCIDQUERY_API_KEY` | `https://api.lucidquery.com/v1` | +| `lynkr` | Lynkr | api_key | `LYNKR_API_KEY` | `http://localhost:8081/v1` | +| `matterai` | Matter AI | api_key | `MATTERAI_API_KEY` | `https://api.matterai.com/v1` | +| `meganova` | Meganova | api_key | `MEGANOVA_API_KEY` | `https://api.meganova.ai/v1` | +| `merge_gateway` | Merge Gateway | api_key | `MERGE_GATEWAY_API_KEY` | `https://api-gateway.merge.dev/v1/openai` | +| `meta` | Meta | api_key | `MODEL_API_KEY` | `https://api.meta.ai/v1` | +| `meta_llama` | Meta Llama API | api_key | `LLAMA_API_KEY` | `https://api.llama.com/compat/v1` | +| `mimo` | Mimo | api_key | `MIMO_API_KEY` | `https://api.xiaomimimo.com/v1` | +| `minimax` | MiniMax | api_key | `MINIMAX_API_KEY` | `https://api.minimax.io/v1` | +| `minimax_cn` | MiniMax (minimaxi.com) | api_key | `MINIMAX_API_KEY` | `https://api.minimaxi.com/v1` | +| `minimax_cn_coding_plan` | MiniMax Token Plan (minimaxi.com) | api_key | `MINIMAX_API_KEY` | `https://api.minimaxi.com/v1` | +| `minimax_coding_plan` | MiniMax Token Plan (minimax.io) | api_key | `MINIMAX_API_KEY` | `https://api.minimax.io/anthropic/v1` | +| `mira` | Mira | api_key | `MIRA_API_KEY` | `https://api.mira.so/v1` | +| `mistralrs` | Mistral.rs | none | — | `http://127.0.0.1:8080/v1` (URL from `MISTRALRS_BASE_URL`) | +| `mixlayer` | Mixlayer | api_key | `MIXLAYER_API_KEY` | `https://models.mixlayer.ai/v1` | +| `mlx` | MLX (Apple Silicon) | none | — | `http://127.0.0.1:8080/v1` (URL from `MLX_BASE_URL`) | +| `moark` | Moark | api_key | `MOARK_API_KEY` | `https://api.moark.com/v1` | +| `modal` | Modal | api_key | `MODAL_API_KEY` | `https://modal.com/v1` | +| `model_oracle_ai` | Model Oracle AI | api_key | `MODEL_ORACLE_API_KEY` | `https://api.modeloracle.com/api/v1` | +| `modelscope` | ModelScope | api_key | `MODELSCOPE_API_KEY` | `https://api-inference.modelscope.cn/v1` | +| `moonshotai` | Moonshot AI | api_key | `MOONSHOT_API_KEY` | `https://api.moonshot.cn/v1` | +| `moonshotai_cn` | Moonshot AI (China) | api_key | `MOONSHOT_API_KEY` | `https://api.moonshot.cn/anthropic/v1` | +| `morph` | Morph LLM | api_key | `MORPH_API_KEY` | `https://api.morphllm.com/v1` | +| `nanogpt` | NanoGPT | api_key | `NANOGPT_API_KEY` | `https://api.nanogpt.com/v1` | +| `ncompass` | Ncompass | api_key | `NCOMPASS_API_KEY` | `https://api.ncompass.tech/v1` | +| `nearai` | NEAR AI Cloud | api_key | `NEARAI_API_KEY` | `https://cloud-api.near.ai/v1` | +| `nebius` | Nebius AI | api_key | `NEBIUS_API_KEY` | `https://api.studio.nebius.ai/v1` | +| `neon` | Neon | api_key | `NEON_AI_GATEWAY_TOKEN` | `https://{branch_host}/v1` (params: `branch_host`) | +| `neuralwatt` | Neuralwatt | api_key | `NEURALWATT_API_KEY` | `https://api.neuralwatt.com/v1` | +| `nextbit` | NextBit | api_key | `NEXTBIT_API_KEY` | `https://api.nextbit.ai/v1` | +| `nlp_cloud` | NLP Cloud | api_key | `NLPCLOUD_API_KEY` | `https://api.nlpcloud.io/v1` | +| `nous_research` | Nous Research | api_key | `NOUS_API_KEY` | `https://api.nousresearch.com/v1` | +| `novita` | Novita AI | api_key | `NOVITA_API_KEY` | `https://api.novita.ai/v1` | +| `nscale` | Nscale | api_key | `NSCALE_API_KEY` | `https://inference.api.nscale.com/v1` | +| `nvidia_nim` | NVIDIA NIM | api_key | `NVIDIA_API_KEY` | `https://integrate.api.nvidia.com/v1` | +| `oci` | OCI | api_key | `OCI_API_KEY` | `https://inference.generativeai.{region}.oci.oraclecloud.com/openai/v1` (params: `region`) | +| `ofox` | OfoxAI | api_key | `OFOX_API_KEY` | `https://api.ofox.ai/v1` | +| `ohmygpt` | OhMyGPT | api_key | `OHMYGPT_API_KEY` | `https://api.ohmygpt.com/v1` | +| `ollama` | Ollama | none | `OLLAMA_API_KEY` | `http://127.0.0.1:11434/v1` (URL from `OLLAMA_BASE_URL`) | +| `ollama_cloud` | Ollama Cloud | api_key | `OLLAMA_CLOUD_API_KEY` | `https://api.ollama.com/v1` | +| `omlx` | OMLX / MLX LM | none | — | `http://127.0.0.1:8080/v1` (URL from `OMLX_BASE_URL`) | +| `onnx` | ONNX Runtime | none | — | `http://127.0.0.1:8080/v1` (URL from `ONNX_BASE_URL`) | +| `oobabooba` | Oobabooga Text Generation WebUI | none | — | `http://127.0.0.1:5000/v1` (URL from `OOBABOOBA_BASE_URL`) | +| `openaimax` | OpenAIMax | api_key | `OPENAIMAX_API_KEY` | `https://api.openaimax.com/v1` | +| `openaisb` | OpenAI-SB | api_key | `OPENAISB_API_KEY` | `https://api.openaisb.com/v1` | +| `opencode` | OpenCode Zen | api_key | `OPENCODE_API_KEY` | `https://api.opencode.zen/v1` | +| `opencode_go` | OpenCode Go | api_key | `OPENCODE_GO_API_KEY` | `https://api.opencode.dev/v1` | +| `opencode_zen` | OpenCode Zen | api_key | `OPENCODE_ZEN_API_KEY` | `https://api.opencode.zen/v1` | +| `openrouter` | OpenRouter | api_key | `OPENROUTER_API_KEY` | `https://openrouter.ai/api/v1` | +| `openvino` | OpenVINO | none | — | `http://127.0.0.1:8080/v1` (URL from `OPENVINO_BASE_URL`) | +| `orcarouter` | OrcaRouter | api_key | `ORCAROUTER_API_KEY` | `https://api.orcarouter.com/v1` | +| `ovhcloud` | OVHcloud AI | api_key | `OVHCLOUD_API_KEY` | `https://oai.endpoints.kepler.ai.cloud.ovh.net/v1` | +| `parasail` | Parasail | api_key | `PARASAIL_API_KEY` | `https://api.parasail.io/v1` | +| `perfxcloud` | PerfXCloud | api_key | `PERFXCLOUD_API_KEY` | `https://api.perfxcloud.com/v1` | +| `perplexity` | Perplexity | api_key | `PERPLEXITY_API_KEY` | `https://api.perplexity.ai` | +| `perplexity_agent` | Perplexity Agent | api_key | `PERPLEXITY_API_KEY` | `https://api.perplexity.ai/v1` | +| `petals` | Petals | api_key | `PETALS_API_KEY` | `https://api.petals.dev/v1` | +| `pinstripes` | Pinstripes | api_key | `PINSTRIPES_API_KEY` | `https://api.pinstripes.io/v1` | +| `pioneer` | Pioneer | api_key | `PIONEER_API_KEY` | `https://api.pioneer.ai/v1` | +| `poe` | Poe | api_key | `POE_API_KEY` | `https://api.poe.com/v1` | +| `poolside` | Poolside | api_key | `POOLSIDE_API_KEY` | `https://inference.poolside.ai/v1` | +| `portkey` | Portkey Gateway | api_key | `PORTKEY_API_KEY` | `https://api.portkey.ai/v1` | +| `ppinfra` | PPInfra（PPIO 派欧云） | api_key | `PPIO_API_KEY` | `https://api.ppio.com/openai` | +| `predibase` | Predibase | api_key | `PREDIBASE_API_KEY` | `https://serving.app.predibase.com/v1` | +| `privatemode_ai` | Privatemode AI | api_key | `PRIVATEMODE_API_KEY` | `http://localhost:8080/v1` | +| `publicai` | Publicai | api_key | `PUBLICAI_API_KEY` | `https://platform.publicai.co/v1` | +| `qihang_ai` | QiHang（启航 AI） | api_key | `QIHANG_API_KEY` | `https://api.qhaigc.net/v1` | +| `qihoo360` | 360 AI | api_key | `AI360_API_KEY` | `https://api.360.cn/v1` | +| `qiniu_ai` | Qiniu AI | api_key | `QINIU_API_KEY` | `https://api.qiniu.com/v1` | +| `regolo_ai` | Regolo AI | api_key | `REGOLO_API_KEY` | `https://api.regolo.ai/v1` | +| `reka_ai` | Reka AI | api_key | `REKA_API_KEY` | `https://api.reka.ai/v1` | +| `requesty` | Requesty | api_key | `REQUESTY_API_KEY` | `https://api.requesty.ai/v1` | +| `reve` | Reve | api_key | `REVE_API_KEY` | `https://api.reve.ai/v1` | +| `routing_run` | routing.run | api_key | `ROUTING_RUN_API_KEY` | `https://api.routing.run/v1` | +| `sakana` | Sakana AI | api_key | `SAKANA_API_KEY` | `https://api.sakana.ai/v1` | +| `sambanova` | SambaNova | api_key | `SAMBANOVA_API_KEY` | `https://api.sambanova.ai/v1` | +| `sarvam` | Sarvam AI | api_key | `SARVAM_API_KEY` | `https://api.sarvam.ai/v1` | +| `scaleway` | Scaleway AI | api_key | `SCALEWAY_API_KEY` | `https://api.scaleway.ai/v1` | +| `scx_ai` | SCX AI | api_key | `SCX_AI_API_KEY` | `https://api.scx.ai/v1` | +| `sglang` | SGLang | none | `SGLANG_API_KEY` | `http://127.0.0.1:30000/v1` (URL from `SGLANG_BASE_URL`) | +| `siliconflow` | SiliconFlow | api_key | `SILICONFLOW_API_KEY` | `https://api.siliconflow.cn/v1` | +| `snowflake` | Snowflake | api_key | `SNOWFLAKE_PAT` | `https://{account_identifier}.snowflakecomputing.com/api/v2/cortex/v1` (params: `account_identifier`) | +| `snowflake_cortex` | Snowflake Cortex | api_key | `SNOWFLAKE_CORTEX_PAT` | `https://{account_identifier}.snowflakecomputing.com/api/v2/cortex/v1` (params: `account_identifier`) | +| `stackit` | STACKIT | api_key | `STACKIT_API_KEY` | `https://api.openai-compat.model-serving.eu01.onstackit.cloud/v1` | +| `stepfun` | StepFun (阶跃星辰) | api_key | `STEPFUN_API_KEY` | `https://api.stepfun.com/v1` | +| `stepfun_ai_step_plan` | StepFun Step Plan (Global) | api_key | `STEPFUN_API_KEY` | `https://api.stepfun.ai/step_plan/v1` | +| `stepfun_step_plan` | StepFun Step Plan (China) | api_key | `STEPFUN_API_KEY` | `https://api.stepfun.com/step_plan/v1` | +| `subconscious` | Subconscious | api_key | `SUBCONSCIOUS_API_KEY` | `https://api.subconscious.dev/v1` | +| `submodel` | SubModel | api_key | `SUBMODEL_API_KEY` | `https://api.submodel.com/v1` | +| `synthetic` | Synthetic | api_key | `SYNTHETIC_API_KEY` | `https://api.synthetic.new/openai/v1` | +| `tencent` | Tencent (混元/Hunyuan) | api_key | `TENCENT_API_KEY` | `https://api.hunyuan.cloud.tencent.com/v1` | +| `tencent_coding_plan` | Tencent Coding Plan (China) | api_key | `TENCENT_CODING_PLAN_API_KEY` | `https://api.lkeap.cloud.tencent.com/coding/v3` | +| `tencent_token_plan` | Tencent Token Plan | api_key | `TENCENT_TOKEN_PLAN_API_KEY` | `https://api.lkeap.cloud.tencent.com/plan/v3` | +| `tencent_token_plan_enterprise_auto` | 腾讯云 Token Plan / Token Plan 企业版轻享套餐 | api_key | `TENCENT_TOKEN_PLAN_ENTERPRISE_API_KEY` | `https://tokenhub.tencentmaas.com/plan/v3` | +| `tencent_token_plan_enterprise_pro` | 腾讯云 Token Plan / Token Plan 企业版专业套餐 | api_key | `TENCENT_TOKEN_PLAN_ENTERPRISE_API_KEY` | `https://tokenhub.tencentmaas.com/plan/v3` | +| `tencent_token_plan_general_personal` | 腾讯云 Token Plan / 通用 Token Plan（个人版） | api_key | `TENCENT_TOKEN_PLAN_API_KEY` | `https://api.lkeap.cloud.tencent.com/plan/v3` | +| `tencent_token_plan_hy_personal` | 腾讯云 Token Plan / Hy Token Plan（个人版） | api_key | `TENCENT_TOKEN_PLAN_API_KEY` | `https://api.lkeap.cloud.tencent.com/plan/v3` | +| `tencent_tokenhub` | Tencent TokenHub | api_key | `TENCENT_TOKENHUB_API_KEY` | `https://tokenhub.tencentmaas.com/v1` | +| `tensormesh` | Tensormesh | api_key | `YOUR_API_KEY` | `https://serverless.tensormesh.ai` | +| `the_grid_ai` | The Grid AI | api_key | `THEGRIDAI_API_KEY` | `https://api.thegrid.ai/v1` | +| `thinkingmachines` | Thinking Machines | api_key | `TINKER_API_KEY` | `https://tinker.thinkingmachines.dev/services/tinker-prod/oai/api/v1` | +| `tinfoil` | Tinfoil | api_key | `TINFOIL_API_KEY` | `https://inference.tinfoil.sh/v1` | +| `togetherai` | Together AI | api_key | `TOGETHER_API_KEY` | `https://api.together.xyz/v1` | +| `tokenflux` | Tokenflux | api_key | `TOKENFLUX_API_KEY` | `https://tokenflux.ai/v1` | +| `tokenpony` | TokenPony | api_key | `TOKENPONY_API_KEY` | `https://api.tokenpony.com/v1` | +| `trustedrouter` | TrustedRouter | api_key | `TRUSTEDROUTER_API_KEY` | `https://api.trustedrouter.com/v1` | +| `tundra` | Tundra | api_key | `TUNDRA_API_KEY` | `https://api.tundra.ai/v1` | +| `umans_ai` | Umans AI | api_key | `UMANS_AI_API_KEY` | `https://api.code.umans.ai/v1` | +| `unorouter` | UnoRouter | api_key | `UNOROUTER_API_KEY` | `https://unorouter.com/en` | +| `upstage` | Upstage | api_key | `UPSTAGE_API_KEY` | `https://api.upstage.ai/v1` | +| `v0` | v0 (Vercel) | api_key | `V0_API_KEY` | `https://api.v0.dev/v1` | +| `venice` | Venice | api_key | `VENICE_API_KEY` | `https://api.venice.ai/api/v1` | +| `vercel` | Vercel | api_key | `VERCEL_API_KEY` | `https://api.v0.dev/v1` | +| `vertex_ai_ai21_models` | AI21 models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vertex_ai_anthropic_models` | Anthropic Claude models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vertex_ai_deepseek_models` | DeepSeek models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vertex_ai_llama_models` | Meta Llama models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vertex_ai_minimax_models` | MiniMax models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vertex_ai_mistral_models` | Mistral models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vertex_ai_moonshot_models` | Moonshot AI models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vertex_ai_openai_models` | OpenAI models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vertex_ai_qwen_models` | Qwen models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vertex_ai_zai_models` | Z.AI models on Vertex AI | api_key | `GOOGLE_VERTEX_ACCESS_TOKEN` | `https://{host}/v1/projects/{project}/locations/{location}/endpoints/openapi` (params: `project`, `location`) | +| `vivgrid` | Vivgrid | api_key | `VIVGRID_API_KEY` | `https://api.vivgrid.com/v1` | +| `vllm` | vLLM | none | `VLLM_API_KEY` | `http://127.0.0.1:8000/v1` (URL from `VLLM_BASE_URL`) | +| `volc_engine` | VolcEngine | api_key | `ARK_API_KEY` | `https://ark.cn-beijing.volces.com/api/v3` | +| `vultr` | Vultr | api_key | `VULTR_API_KEY` | `https://api.vultrinference.com/v1` | +| `wafer` | Wafer | api_key | `WAFER_API_KEY` | `https://api.wafer.ai/v1` | +| `wandb` | Weights & Biases | api_key | `WANDB_API_KEY` | `https://api.inference.wandb.ai/v1` | +| `xiaomi_token_plan_ams` | Xiaomi Token Plan (Europe) | api_key | `MIMO_API_KEY` | `https://token-plan-ams.xiaomimimo.com/v1` | +| `xiaomi_token_plan_cn` | Xiaomi Token Plan (China) | api_key | `MIMO_API_KEY` | `https://token-plan-cn.xiaomimimo.com/v1` | +| `xiaomi_token_plan_sgp` | Xiaomi Token Plan (Singapore) | api_key | `MIMO_API_KEY` | `https://token-plan-sgp.xiaomimimo.com/v1` | +| `xiaomimimo` | Xiaomi MiMo | api_key | `XIAOMI_API_KEY` | `https://mimo.xiaomi.com/v1` | +| `xinference` | Xinference | none | — | `http://127.0.0.1:9997/v1` (URL from `XINFERENCE_BASE_URL`) | +| `xpersona` | Xpersona | api_key | `XPERSONA_API_KEY` | `https://www.xpersona.co/v1` | +| `xunfei` | Xunfei | api_key | `XUNFEI_API_PASSWORD` | `https://spark-api-open.xf-yun.com/v1` | +| `zai` | Zai | api_key | `ZAI_API_KEY` | `https://api.z.ai/api/paas/v4` | +| `zai_coding_plan` | Z.AI Coding Plan | api_key | `ZHIPU_API_KEY` | `https://api.z.ai/api/anthropic` | +| `zeldoc` | Zeldoc | api_key | `ZELDOC_API_KEY` | `https://api.zeldoc.ai/v1` | +| `zenmux` | ZenMux | api_key | `ZENMUX_API_KEY` | `https://zenmux.ai/api/v1` | +| `zhipu_v4` | ZhipuV4 | api_key | `ZHIPU_API_KEY` | `https://open.bigmodel.cn/api/paas/v4` | +| `zhipuai_coding_plan` | Zhipu AI Coding Plan | api_key | `ZHIPU_API_KEY` | `https://open.bigmodel.cn/api/coding/paas/v4` | ## Typed factories (non-registry) @@ -294,156 +321,104 @@ These providers are **not** name-addressable: `provider("anthropic", ...)` fails | module | typed entry points | |--------|--------------------| -| `anthropic` | `AnthropicConfig` / `AnthropicProvider` | -| `anthropic_aws` | `AnthropicAwsConfig` / `AnthropicAwsProvider` | -| `azure` | `AzureConfig` / `AzureProvider` | -| `bedrock` | `BedrockConfig` / `BedrockProvider` | -| `cohere` | `CohereConfig` / `CohereProvider` | -| `google` | `GoogleConfig` / `GoogleProvider` | +| `anthropic` | `AnthropicProvider` / `AnthropicProviderSettings` / `create_anthropic` | +| `anthropic_aws` | `AnthropicAwsProvider` / `AnthropicAwsProviderSettings` / `create_anthropic_aws` | +| `azure` | `AzureOpenAIProvider` / `AzureOpenAIProviderSettings` / `create_azure` | +| `bedrock` | `AmazonBedrockProvider` / `AmazonBedrockProviderSettings` / `create_amazon_bedrock` | +| `cohere` | `CohereProvider` / `CohereProviderSettings` / `create_cohere` | +| `google` | `GoogleProvider` / `GoogleProviderSettings` / `create_google` | | `jev` | `JevConfig` / `JevProvider` | -| `mistral` | `MistralConfig` / `MistralProvider` | -| `openai` | `OpenAIConfig` / `OpenAIProvider` | -| `vertex` | `VertexConfig` / `VertexProvider` | -| `voyage` | `VoyageConfig` / `VoyageProvider` | -| `codex` | `CodexConfig` / `CodexProvider` | -| `openrouter` | `OpenRouterConfig` / `OpenRouterProvider` | -| `xai` | `XAIConfig` / `XAIProvider` | - -### OpenAI-compatible thin wrappers (second batch) — 5 - -| module | typed entry points | -|--------|--------------------| -| `huggingface` | `HuggingFaceConfig` / `HuggingFaceProvider` | -| `llamafile` | `LlamafileConfig` / `LlamafileProvider` | -| `lmstudio` | `LmStudioConfig` / `LmStudioProvider` | -| `mistralrs` | `MistralrsConfig` / `MistralrsProvider` | -| `ollama` | `OllamaConfig` / `OllamaProvider` | +| `mistral` | `MistralProvider` / `MistralProviderSettings` / `create_mistral` | +| `openai` | `OpenAIProvider` / `OpenAIProviderSettings` / `create_openai` | +| `vertex` | `VertexProvider` / `VertexProviderSettings` / `create_google_vertex` | +| `voyage` | `VoyageProvider` / `VoyageProviderSettings` / `create_voyage` | +| `codex` | `CodexProvider` / `CodexProviderSettings` / `create_codex` | +| `xai` | `XAIProvider` / `XAIProviderSettings` / `create_xai` | +| `huggingface` | `HuggingFaceProvider` / `HuggingFaceProviderSettings` / `create_huggingface` | ### Speech-only providers (TTS) — 4 | module | typed entry points | |--------|--------------------| -| `cartesia` | `CartesiaConfig` / `CartesiaProvider` | -| `elevenlabs` | `ElevenLabsConfig` / `ElevenLabsProvider` | -| `hume` | `HumeConfig` / `HumeProvider` | -| `lmnt` | `LMNTConfig` / `LMNTProvider` | +| `cartesia` | `CartesiaProvider` / `CartesiaProviderSettings` / `create_cartesia` | +| `elevenlabs` | `ElevenLabsProvider` / `ElevenLabsProviderSettings` / `create_elevenlabs` | +| `hume` | `HumeProvider` / `HumeProviderSettings` / `create_hume` | +| `lmnt` | `LMNTProvider` / `LMNTProviderSettings` / `create_lmnt` | ### Transcription-only providers (STT) — 5 | module | typed entry points | |--------|--------------------| -| `assemblyai` | `AssemblyAIConfig` / `AssemblyAIProvider` | -| `deepgram` | `DeepgramConfig` / `DeepgramProvider` | -| `fal` | `FalConfig` / `FalProvider` | -| `gladia` | `GladiaConfig` / `GladiaProvider` | -| `revai` | `RevaiConfig` / `RevaiProvider` | +| `assemblyai` | `AssemblyAIProvider` / `AssemblyAIProviderSettings` / `create_assemblyai` | +| `deepgram` | `DeepgramProvider` / `DeepgramProviderSettings` / `create_deepgram` | +| `fal` | `FalProvider` / `FalProviderSettings` / `create_fal` | +| `gladia` | `GladiaProvider` / `GladiaProviderSettings` / `create_gladia` | +| `revai` | `RevaiProvider` / `RevaiProviderSettings` / `create_revai` | ### Image-only providers — 4 | module | typed entry points | |--------|--------------------| -| `black_forest_labs` | `BlackForestLabsConfig` / `BlackForestLabsProvider` | -| `luma` | `LumaConfig` / `LumaProvider` | -| `prodia` | `ProdiaConfig` / `ProdiaProvider` | -| `replicate` | `ReplicateConfig` / `ReplicateProvider` | +| `black_forest_labs` | `BlackForestLabsProvider` / `BlackForestLabsProviderSettings` / `create_black_forest_labs` | +| `luma` | `LumaProvider` / `LumaProviderSettings` / `create_luma` | +| `prodia` | `ProdiaProvider` / `ProdiaProviderSettings` / `create_prodia` | +| `replicate` | `ReplicateProvider` / `ReplicateProviderSettings` / `create_replicate` | ### Video-only providers — 1 | module | typed entry points | |--------|--------------------| -| `klingai` | `KlingAIConfig` / `KlingAIProvider` | +| `klingai` | `KlingAIProvider` / `KlingAIProviderSettings` / `create_klingai` | ### Generic Responses API wrapper — 1 | module | typed entry points | |--------|--------------------| -| `open_responses` | `OpenResponsesConfig` / `OpenResponsesProvider` | - -### Bulk-generated thin-wrapper providers — 16 - -| module | typed entry points | -|--------|--------------------| -| `cybertron` | `CybertronConfig` / `CybertronProvider` | -| `docker_model_runner` | `DockerModelRunnerConfig` / `DockerModelRunnerProvider` | -| `gaudi` | `GaudiConfig` / `GaudiProvider` | -| `jlama` | `JlamaConfig` / `JlamaProvider` | -| `litellm_proxy` | `LitellmProxyConfig` / `LitellmProxyProvider` | -| `llamacpp` | `LlamacppConfig` / `LlamacppProvider` | -| `local` | `LocalConfig` / `LocalProvider` | -| `localai` | `LocalaiConfig` / `LocalaiProvider` | -| `mlx` | `MlxConfig` / `MlxProvider` | -| `omlx` | `OmlxConfig` / `OmlxProvider` | -| `onnx` | `OnnxConfig` / `OnnxProvider` | -| `oobabooba` | `OobaboobaConfig` / `OobaboobaProvider` | -| `openvino` | `OpenvinoConfig` / `OpenvinoProvider` | -| `sglang` | `SglangConfig` / `SglangProvider` | -| `vllm` | `VllmConfig` / `VllmProvider` | -| `xinference` | `XinferenceConfig` / `XinferenceProvider` | +| `open_responses` | `OpenResponsesProvider` / `OpenResponsesProviderSettings` / `create_open_responses` | ### Modality-specific providers (non-language, e.g. rerank-only) — 1 | module | typed entry points | |--------|--------------------| -| `jina_ai` | `JinaAiConfig` / `JinaAiProvider` | +| `jina_ai` | `JinaAiProvider` / `JinaAiProviderSettings` / `create_jina_ai` | ### AWS Polly speech (TTS) provider — SigV4 authenticated, speech modality only — 1 | module | typed entry points | |--------|--------------------| -| `aws_polly` | `AwsPollyConfig` / `AwsPollyProvider` | +| `aws_polly` | `AwsPollyProvider` / `AwsPollyProviderSettings` / `create_aws_polly` | ### Recraft image provider (OpenAI Images-compatible + Recraft extension fields) — 1 | module | typed entry points | |--------|--------------------| -| `recraft` | `RecraftConfig` / `RecraftProvider` | +| `recraft` | `RecraftProvider` / `RecraftProviderSettings` / `create_recraft` | ### Stability image provider (image modality only) — 1 | module | typed entry points | |--------|--------------------| -| `stability` | `StabilityConfig` / `StabilityProvider` | +| `stability` | `StabilityProvider` / `StabilityProviderSettings` / `create_stability` | ### Video-only provider (runwayml) — 1 | module | typed entry points | |--------|--------------------| -| `runwayml` | `RunwaymlConfig` / `RunwaymlProvider` | - -### P1 thin-wrapper providers (provider-research batch) — 1 - -| module | typed entry points | -|--------|--------------------| -| `bedrock_mantle` | `BedrockMantleConfig` / `BedrockMantleProvider` | - -### Vertex AI MaaS partner-model providers (OpenAI-compatible thin wrappers). Each wraps the shared OpenAIProvider against the Vertex AI MaaS OpenAPI endpoint, authenticating with a Google Cloud Bearer token — 10 - -| module | typed entry points | -|--------|--------------------| -| `vertex_ai_ai21_models` | `VertexAiAi21ModelsConfig` / `VertexAiAi21ModelsProvider` | -| `vertex_ai_anthropic_models` | `VertexAiAnthropicModelsConfig` / `VertexAiAnthropicModelsProvider` | -| `vertex_ai_deepseek_models` | `VertexAiDeepseekModelsConfig` / `VertexAiDeepseekModelsProvider` | -| `vertex_ai_llama_models` | `VertexAiLlamaModelsConfig` / `VertexAiLlamaModelsProvider` | -| `vertex_ai_minimax_models` | `VertexAiMinimaxModelsConfig` / `VertexAiMinimaxModelsProvider` | -| `vertex_ai_mistral_models` | `VertexAiMistralModelsConfig` / `VertexAiMistralModelsProvider` | -| `vertex_ai_moonshot_models` | `VertexAiMoonshotModelsConfig` / `VertexAiMoonshotModelsProvider` | -| `vertex_ai_openai_models` | `VertexAiOpenaiModelsConfig` / `VertexAiOpenaiModelsProvider` | -| `vertex_ai_qwen_models` | `VertexAiQwenModelsConfig` / `VertexAiQwenModelsProvider` | -| `vertex_ai_zai_models` | `VertexAiZaiModelsConfig` / `VertexAiZaiModelsProvider` | +| `runwayml` | `RunwaymlProvider` / `RunwaymlProviderSettings` / `create_runwayml` | ### Search-only providers (web search modality) — 11 | module | typed entry points | |--------|--------------------| -| `dataforseo` | `DataforseoConfig` / `DataforseoProvider` | -| `exa_ai` | `ExaAiConfig` / `ExaAiProvider` | -| `firecrawl` | `FirecrawlConfig` / `FirecrawlProvider` | -| `google_pse` | `GooglePseConfig` / `GooglePseProvider` | -| `linkup` | `LinkupConfig` / `LinkupProvider` | -| `parallel_ai` | `ParallelAiConfig` / `ParallelAiProvider` | -| `searxng` | `SearxngConfig` / `SearxngProvider` | -| `serper` | `SerperConfig` / `SerperProvider` | -| `tavily` | `TavilyConfig` / `TavilyProvider` | -| `tinyfish` | `TinyfishConfig` / `TinyfishProvider` | -| `you_com` | `YouComConfig` / `YouComProvider` | +| `dataforseo` | `DataforseoProvider` / `DataforseoProviderSettings` / `create_dataforseo` | +| `exa_ai` | `ExaAiProvider` / `ExaAiProviderSettings` / `create_exa_ai` | +| `firecrawl` | `FirecrawlProvider` / `FirecrawlProviderSettings` / `create_firecrawl` | +| `google_pse` | `GooglePseProvider` / `GooglePseProviderSettings` / `create_google_pse` | +| `linkup` | `LinkupProvider` / `LinkupProviderSettings` / `create_linkup` | +| `parallel_ai` | `ParallelAiProvider` / `ParallelAiProviderSettings` / `create_parallel_ai` | +| `searxng` | `SearxngProvider` / `SearxngProviderSettings` / `create_searxng` | +| `serper` | `SerperProvider` / `SerperProviderSettings` / `create_serper` | +| `tavily` | `TavilyProvider` / `TavilyProviderSettings` / `create_tavily` | +| `tinyfish` | `TinyfishProvider` / `TinyfishProviderSettings` / `create_tinyfish` | +| `you_com` | `YouComProvider` / `YouComProviderSettings` / `create_you_com` | diff --git a/docs/api/python.md b/docs/api/python.md index 568a9dbc5..0cb14b383 100644 --- a/docs/api/python.md +++ b/docs/api/python.md @@ -1,6 +1,6 @@ # aimux · Python API -> Unified LLM service access layer — one API to access 325 AI providers +> Unified LLM service access layer — one API to access AI providers Shared reference — parameter tables, result shapes, factory functions, and the feature coverage matrix — lives in the [API overview](../API.md). @@ -30,20 +30,25 @@ model = provider("groq", None, "llama-3.3-70b") model = provider("groq", "sk-...", "llama-3.3-70b", "https://relay.example/v1") # Full ProviderOptions via config dict: model = provider("groq", "sk-...", "llama-3.3-70b", - config={"headers": {"X-Custom": "1"}, "max_retries": 0}) + config={"headers": {"X-Custom": "1"}}) +# A preset with template parameters: +model = provider("cloudflare_workers_ai", "sk-...", "@cf/meta/llama-3.1-8b-instruct", + config={"params": {"account_id": "acct123"}}) result = generate_text(model, "Hello") ``` -`provider(name, api_key, model_id, base_url=None, config=None)` covers all -251 built-in OpenAI-compatible providers. `config` takes the full +`provider(name, api_key, model_id, base_url=None, config=None)` covers the vendor packages and +282 registry-backed OpenAI-compatible providers. `config` takes the full `ProviderOptions` shape (`base_url` / `headers` / `organization` / `project` / -`max_retries` / `body_overrides`); the `base_url` parameter wins over -`config["base_url"]`. `openai` / `anthropic` / `deepseek` factories -remain (deepseek is now registry-backed). - -> **Scope:** `provider(name)` covers only the 251 registry OpenAI-compatible -> providers; Anthropic/Google/multimodal/local → typed factories -> (`anthropic(api_key, model)`); custom endpoints → `base_url` param. +`params`); the `base_url` parameter wins over `config["base_url"]`. +`max_retries` (a per-call option) and `body_overrides` (removed) in `config` +raise `InvalidArgumentError`. A missing key or setting raises `LoadAPIKeyError` +(`env_var`, `description`) / `LoadSettingError` (`env_var`, `setting_name`). `openai` / `anthropic` / `deepseek` factories +remain; DeepSeek now uses its own vendor chat model. `openai` and by-name +`provider("openai", …)` select the Responses API. + +> **Scope:** `provider(name)` reaches vendor packages and the 282 preset rows; +> typed factories remain available. Custom endpoints use the `base_url` param. > Full list: [providers.md](providers.md). ## Text Generation @@ -247,7 +252,8 @@ if "Base64" in result["images"]: ## Video Generation -Video generation typically returns a URL (not binary). +Most video models return a URL; `generate` downloads it, as the AI SDK's +`generateVideo` does, and returns the video's bytes (`Binary`). ```python from aimux import google_video @@ -260,9 +266,10 @@ result = json.loads(videor.generate(json.dumps({ "provider_options": {}, }))) -# result["videos"] is usually [{"Url": {"url": "...", "media_type": "..."}}] -if "Url" in result["videos"][0]: - print(result["videos"][0]["Url"]["url"]) +# result["videos"] is usually [{"Binary": {"data": [...], "media_type": "..."}}] +if "Binary" in result["videos"][0]: + with open("video.mp4", "wb") as f: + f.write(bytes(result["videos"][0]["Binary"]["data"])) ``` ## Reranking @@ -392,6 +399,7 @@ Exception ├── JSONParseError / InvalidResponseDataError ├── NoSuchToolError / InvalidToolInputError / ToolCallRepairError # tool-contract errors ├── InvalidArgumentError / InvalidPromptError + ├── LoadAPIKeyError / LoadSettingError # no API key / required setting: env_var (+ description / setting_name) ├── TokenExpiredError ├── UnsupportedFunctionalityError ├── NoSuchModelError / NoSuchProviderError @@ -430,7 +438,10 @@ permitted attempt failed with a retryable error — or `"errorNotRetryable"` — a later attempt failed non-retryably), `errors` — the per-attempt history, oldest first, each itself an exception from this hierarchy — and `last_error`. `TokenExpiredError` has `status == 401`, `NoSuchModelError` has -`model_id` / `model_type`, and `NoSuchProviderError` has `provider_id`. +`model_id` / `model_type`, `NoSuchProviderError` has `provider_id`, +`model_id`, `model_type` and `available_providers`, and +`LoadAPIKeyError` / `LoadSettingError` have `env_var` (the environment variable +consulted) plus `description` / `setting_name`. ```python from aimux import ( @@ -478,3 +489,6 @@ SDK's full stream: `Start`, `StartStep` … `FinishStep`, `Finish`), e.g. `{"TextDelta": {"id": ..., "text": ...}}`. Iterate dicts with `if "TextDelta" in part:` (as in [Streaming Generation](#streaming-generation)) or validate them with `parse_stream_part(part)` for attribute access. + +Provider-layer `StreamPart::ResponseMetadata` events are consumed internally +and are not emitted by `stream_text`. diff --git a/docs/api/reference.md b/docs/api/reference.md index 2631fb408..b01e9e13c 100644 --- a/docs/api/reference.md +++ b/docs/api/reference.md @@ -15,7 +15,7 @@ | `ResponseFormat` | [options.rs](../../aimux-core/src/options.rs) | [ResponseFormat.ts](../../bindings/node/src/types/ResponseFormat.ts) | text / json_object / json_schema | | `ReasoningEffort` | [types.rs](../../aimux-core/src/types.rs) | [ReasoningEffort.ts](../../bindings/node/src/types/ReasoningEffort.ts) | 7 levels, passed through verbatim | | `AbortSignal` | [abort_signal.rs](../../aimux-core/src/abort_signal.rs) | — (runtime handle) | Rust: options field; Node: `AbortBridge` + JS `AbortSignal` | -| `ProviderOptions` | [provider.rs](../../aimux-providers/src/provider.rs) | `ProviderConfig` (per binding) | `base_url` / `headers` / `organization` / `project` / `max_retries` / `body_overrides` | +| `aimux_providers::provider::ProviderOptions` (binding compatibility) | [provider.rs](../../aimux-providers/src/provider.rs) | `ProviderConfig` (per binding) | `base_url` / `headers` / `organization` / `project` / `params` (a `max_retries` or `body_overrides` key is rejected as `InvalidArgument`) | | `Output` | [output.rs](../../aimux-core/src/output.rs) | [Output.ts](../../bindings/node/src/types/Output.ts) | the `output` option: Text / Object / Array / Choice / Json | | `GenerateTextResult` | [generate.rs](../../aimux-core/src/generate.rs) | [GenerateTextResult.ts](../../bindings/node/src/types/GenerateTextResult.ts) | `steps`, `text`, `content`, `tool_calls`, `usage`, `finish_reason`, `output`, … | | `StepResult` | [step_result.rs](../../aimux-core/src/step_result.rs) | [StepResult.ts](../../bindings/node/src/types/StepResult.ts) | one model call: content, usage, performance, request, response | @@ -48,24 +48,23 @@ | Provider | Rust | Node / Python | Go / Java | Kotlin / Swift / Flutter | C ABI | |----------|------|--------------|-----------|--------------------------|-------| -| OpenAI | `OpenAIProvider::new(OpenAIConfig::new(key))` | `openai(apiKey, model, baseUrl?)` | `NewOpenAI(...)` / `Model.openai(...)` | `Model.openai(...)` / `Aimux.openai(apiKey:modelId:)` | `aimux_openai_new` | -| Anthropic | `AnthropicProvider::new(AnthropicConfig::new(key))` | `anthropic(apiKey, model, baseUrl?)` | `NewAnthropic(...)` / `Model.anthropic(...)` | `Model.anthropic(...)` / `Aimux.anthropic(apiKey:modelId:)` | `aimux_anthropic_new` | +| OpenAI | `create_openai(OpenAIProviderSettings { api_key: Some(key.into()), ..Default::default() })` | `openai(apiKey, model, baseUrl?)` | `NewOpenAI(...)` / `Model.openai(...)` | `Model.openai(...)` / `Aimux.openai(apiKey:modelId:)` | `aimux_openai_new` | +| Anthropic | `create_anthropic(AnthropicProviderSettings { api_key: Some(key.into()), ..Default::default() })` | `anthropic(apiKey, model, baseUrl?)` | `NewAnthropic(...)` / `Model.anthropic(...)` | `Model.anthropic(...)` / `Aimux.anthropic(apiKey:modelId:)` | `aimux_anthropic_new` | | Google | `GoogleProvider` | `google(apiKey, model, baseUrl?)` | — | — | — | -| DeepSeek (registry) | `provider("deepseek", ...)` | `deepseek(...)` | `NewDeepSeek(...)` / `Model.deepseek(...)` | — | `aimux_provider_new("deepseek", ...)` | +| DeepSeek | `aimux_providers::deepseek::create_deepseek(DeepSeekProviderSettings::default())` | `deepseek(...)` | `NewDeepSeek(...)` / `Model.deepseek(...)` | — | `aimux_provider_new("deepseek", ...)` | | Cohere | `CohereProvider` | `cohere(apiKey, model, baseUrl?)` | `NewCohere(...)` / `Model.cohere(...)` | `Model.cohere(...)` / `Aimux.cohere(...)` | `aimux_cohere_new` | | Mistral | `MistralProvider` | `mistral(apiKey, model, baseUrl?)` | `NewMistral(...)` / `Model.mistral(...)` | `Model.mistral(...)` / `Aimux.mistral(...)` | `aimux_mistral_new` | | xAI | `XAIProvider` | `xai(apiKey, model, baseUrl?)` | `NewXai(...)` / `Model.xai(...)` | `Model.xai(...)` / `Aimux.xai(...)` | `aimux_xai_new` | -| Bedrock | `BedrockProvider` | `bedrock(accessKeyId, secretKey, region, model, baseUrl?)` | `NewBedrock(...)` / `Model.bedrock(...)` | `Model.bedrock(...)` / `Aimux.bedrock(...)` | `aimux_bedrock_new` | +| Bedrock | `AmazonBedrockProvider` | `bedrock(accessKeyId, secretKey, region, model, baseUrl?)` | `NewBedrock(...)` / `Model.bedrock(...)` | `Model.bedrock(...)` / `Aimux.bedrock(...)` | `aimux_bedrock_new` | | Vertex | `VertexProvider` | `vertex(accessToken, project, location, model, baseUrl?)` | `NewVertex(...)` / `Model.vertex(...)` | `Model.vertex(...)` / `Aimux.vertex(...)` | `aimux_vertex_new` | | Anthropic-AWS | `AnthropicAwsProvider` | `anthropicAws(apiKey, region, model, baseUrl?)` | `NewAnthropicAws(...)` / `Model.anthropicAws(...)` | `Model.anthropicAws(...)` / `Aimux.anthropicAws(...)` | `aimux_anthropic_aws_new` | -| Azure | `AzureProvider` | `azure(apiKey, resource, deployment, apiVersion?, baseUrl?)` | `NewAzure(...)` / `Model.azure(...)` | `Model.azure(...)` / `Aimux.azure(...)` | `aimux_azure_new` | +| Azure | `AzureOpenAIProvider` | `azure(apiKey, resource, deployment, apiVersion?, baseUrl?)` | `NewAzure(...)` / `Model.azure(...)` | `Model.azure(...)` / `Aimux.azure(...)` | `aimux_azure_new` | `—` = not exposed in that binding yet (Google is Rust / Node / Python only). Every constructor has a base-URL variant (`WithBase` / `baseUrl` overload). -The remaining 65 providers (speech, transcription, image, video, local -inference, search, standalone OpenAI-compatible, Vertex-hosted) each have a -typed Rust constructor — full list, one entry per provider: +Other vendor packages (speech, transcription, image, video and search) have +typed Rust factories; presets use `create_provider(name, PresetSettings)`: [providers.md](providers.md). Their per-binding modality factories are in the function tables below. @@ -73,9 +72,11 @@ function tables below. | Function | Signature (abridged) | Notes | |----------|----------------------|-------| -| `provider` | `provider(name, api_key, model_id, options) -> Box` | registry lookup; unknown name → `NoSuchProvider` | -| `provider_from_env` | `provider_from_env(name, model_id, options)` | key read from the registry entry's env var | -| `provider_registry_entry` | `-> Option` | inspect a registry profile | +| `create_provider` | `(name, PresetSettings) -> Result, AiMuxError>` | vendor package or preset; unknown name → `NoSuchProvider`; keys read per request | +| `default_providers` | `() -> BTreeMap>` | all built-in providers with default settings | +| `provider_names` | `() -> Vec<&'static str>` | sorted names accepted by `create_provider` | +| `provider_registry::create_provider_registry` | `(providers, options) -> ProviderRegistry` | resolves `"provider:model"` ids; options set the separator | +| `Provider::discovery` | `() -> Option<&dyn ProviderDiscovery>` | optional model listing via `list_models().await` | | `generate_text` | `(model, prompt, options) -> GenerateTextResult` | non-streaming | | `stream_text` | `(model, prompt, options) -> StreamTextResult` | streaming | @@ -83,13 +84,17 @@ Per-modality entry points are trait methods on the model types (`EmbeddingModel`, `ImageModel`, `SpeechModel`, …) — see [API.md §Features](API.md#features). +The bindings keep their own compatibility API, implemented in +`aimux_providers::provider`; these helpers and option types are not Rust +crate-root exports. + ### Node (`@arcships/aimux`) | Function | Notes | |----------|-------| | `generateText(model, prompt, options?, signal?)` / `streamText(...)` | typed wrappers; `signal` = `AbortSignal` | -| `openai()` / `anthropic()` / `deepseek()` / `google()` | native typed path; `deepseek()` is registry-backed | -| `provider(name, apiKey?, modelId, config?)` | registry-backed; names listed in [providers.md](providers.md) | +| `openai()` / `anthropic()` / `deepseek()` / `google()` | native typed path; `deepseek()` uses the vendor chat model | +| `provider(name, apiKey?, modelId, config?)` | by-name vendor packages and presets; names listed in [providers.md](providers.md) | | `openaiEmbedding` / `cohereEmbedding` / `googleEmbedding` | embeddings | | `openaiSpeech` / `openaiTranscription` | speech / STT | | `openaiImage` / `googleImage` / `googleVideo` | image / video | @@ -101,7 +106,7 @@ Per-modality entry points are trait methods on the model types | Function | Notes | |----------|-------| | `generate_text(model, prompt, options?)` / `stream_text(...)` | typed wrappers | -| `openai` / `anthropic` / `deepseek` / `provider` | same semantics as Node; `deepseek` registry-backed | +| `openai` / `anthropic` / `deepseek` / `provider` | same semantics as Node; `deepseek` uses the vendor chat model | | `openai_embedding` / `cohere_embedding` / `google_embedding` | embeddings | | `openai_speech` / `openai_transcription` | speech / STT | | `openai_image` / `google_image` / `google_video` | image / video | @@ -111,7 +116,7 @@ Per-modality entry points are trait methods on the model types | Function | Notes | |----------|-------| -| `Provider(name, apiKey?, modelId, config?)` | registry-backed; names listed in [providers.md](providers.md) | +| `Provider(name, apiKey?, modelId, config?)` | by-name vendor packages and presets; names listed in [providers.md](providers.md) | | `NewOpenAI` / `NewAnthropic` / `NewDeepSeek` | native typed path (+ `WithBase` variants) | | `NewOpenAIEmbedding` / `NewCohereEmbedding` / `NewGoogleEmbedding` | embeddings | | `NewOpenAISpeech` / `NewOpenAITranscription` | speech / STT | @@ -122,14 +127,14 @@ Per-modality entry points are trait methods on the model types | Function | Notes | |----------|-------| -| `provider(name, apiKey?, modelId, config?)` | registry-backed | +| `provider(name, apiKey?, modelId, config?)` | by-name vendor packages and presets | | `openai(apiKey, modelId)` / `anthropic(apiKey, modelId)` | native typed path (+ `WithBase`/`baseUrl` variants; Swift: `Aimux.openai(apiKey:modelId:)`) | | modality classes | `EmbeddingModel` / `SpeechModel` / `ImageModel` / `TranscriptionModel` / `VideoModel` / `RerankingModel` / `SearchModel` / `Files`, each with `openai` / `cohere` / `google` … factories | ### C ABI (`aimux-ffi`) All functions communicate via JSON strings. Entry points are declared in -[aimux-ffi.h](../../aimux-ffi/aimux-ffi.h): `provider_new` (registry lookup), +[aimux-ffi.h](../../aimux-ffi/aimux-ffi.h): `provider_new` (by-name vendor package or preset), model constructors (`openai_new`, `anthropic_new`, …), `*_generate`, `*_stream` (push-callback), `*_embed` / `*_speech` / … per modality, and the `*_free` family. diff --git a/docs/api/rust.md b/docs/api/rust.md index 3f303d6bd..be4daa471 100644 --- a/docs/api/rust.md +++ b/docs/api/rust.md @@ -1,6 +1,6 @@ # aimux · Rust API -> Unified LLM service access layer — one API to access 325 AI providers +> Unified LLM service access layer — one API to access AI providers This is the core implementation language. Shared reference — parameter tables, result shapes, factory functions, and the feature coverage matrix — lives in @@ -10,12 +10,15 @@ the [API overview](../API.md). ```rust use aimux_core::prelude::*; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; +use aimux_providers::{OpenAIProviderSettings, create_openai}; #[tokio::main] async fn main() -> Result<(), AiMuxError> { - let provider = OpenAIProvider::new(OpenAIConfig::new("sk-...")); - let model = provider.model("gpt-4o"); + let provider = create_openai(OpenAIProviderSettings { + api_key: Some("sk-...".to_string()), + ..Default::default() + })?; + let model = provider.chat("gpt-4o"); let result = generate_text(&model, "What is Rust?", GenerateTextOptions::default()).await?; println!("{}", result.text); Ok(()) @@ -24,35 +27,93 @@ async fn main() -> Result<(), AiMuxError> { ## Providers -All 251 built-in OpenAI-compatible providers are registry-backed: no per-provider -`XxxConfig`/`XxxProvider` types. Look them up by string name: +Provider packages follow the AI SDK shape: `XxxProviderSettings`, +`create_xxx(settings)` and a default instance `xxx()` where the package +provides one. Required fields follow the package: OpenAI-compatible settings +require `name` and `base_url`. Models are taken from the provider (`provider.chat(id)`, +`provider.language_model(id)?`, `provider.embedding_model(id)?`, …). Model +identity follows the package (`openai.chat`, `anthropic.messages`); custom +name handling is package-specific. + +```rust +use aimux_providers::anthropic::{AnthropicProviderSettings, create_anthropic}; +use aimux_providers::openai::openai; + +// Default instance: OPENAI_API_KEY is read when a request is made; a missing +// key fails that call with `AiMuxError::LoadApiKey`, not the constructor. +let model = openai().chat("gpt-4o"); + +// Explicit settings. `api_key: Some(..)` is used as given (`""` included). +let anthropic = create_anthropic(AnthropicProviderSettings { + api_key: Some("sk-ant-...".to_string()), + base_url: Some("https://relay.example/v1".into()), + ..Default::default() +})?; +let model = anthropic.messages("claude-sonnet-4-5"); +``` -> **Scope:** `provider(name)` covers only the 251 registry OpenAI-compatible -> providers; Anthropic/Google/multimodal/local → typed factories -> (`AnthropicProvider::new(..)`); custom endpoints → -> `ProviderOptions.base_url` / `OpenAIConfig::with_base_url`. -> Full list: [providers.md](providers.md). +Settings are package-specific. OpenAI, Anthropic and Google take +`api_key: Option` and `headers: Option`; these are fixed +values, not callbacks. `fetch: Option` supplies a custom +transport. These packages have no `transform_request_body` setting; +OpenAI-compatible settings expose that hook, plus `supported_urls` and +`convert_usage` callbacks for chat models. Cohere and Mistral expose +`generate_id` for generated ids. Settings only expose `name` where the upstream +package does. Every request includes the package user-agent suffix +`ai-sdk-/`. +OpenAI defaults to Responses through `call` and `language_model`; `chat` +selects Chat Completions explicitly. +Retry is not a provider setting: `max_retries` on the call options (default +2, 2000 ms initial delay, factor 2). + +The 282 OpenAI-compatible presets use a runtime table parsed once from +`provider_registry.json`; there is no `family` field or generated presets source. +Groq and DeepSeek are vendor packages with their own chat models. xAI and +Hugging Face use the Responses API for language models, with no chat-completions +accessor. + +By-name creation covers both vendor packages and preset rows: ```rust -use aimux_providers::{provider, provider_from_env, ProviderOptions}; +use aimux_providers::{create_provider, default_providers, provider_names, PresetSettings}; +use aimux_core::provider_registry::create_provider_registry; -// Key from the provider's env var (GROQ_API_KEY etc.), base URL & profile from -// the registry — replaces the retired `XxxConfig::from_env()`. -let model = provider_from_env("groq", "llama-3.3-70b", None)?; +// The vendor's environment key is read when a request is made. +let provider = create_provider("groq", PresetSettings::default())?; +let model = provider.language_model("model-id")?; -// Explicit key + overrides — replaces the retired `XxxConfig::new().with_base_url(..)`. -let model = provider( - "groq", - Some("sk-...".to_string()), - "llama-3.3-70b", - Some(ProviderOptions { base_url: Some("https://relay.example/v1".into()), ..Default::default() }), -)?; +// Explicit key and base URL override. +let provider = create_provider("groq", PresetSettings { + api_key: Some("sk-...".to_string()), + base_url: Some("https://relay.example/v1".into()), + ..Default::default() +})?; + +// Combined ids use "provider:model" (":" is the default separator). +let registry = create_provider_registry(default_providers(), Default::default()); +let model = registry.language_model("groq:model-id")?; +let names = provider_names(); +// Model listing is optional: providers without discovery return None. +if let Some(discovery) = provider.discovery() { + let models = discovery.list_models().await?; +} ``` -Names come from `provider_registry.json`; the current list is in -[providers.md](providers.md). Unknown names fail with -`AiMuxError::NoSuchProvider { provider_id }`. +Unknown names fail with `AiMuxError::NoSuchProvider { provider_id, model_id, +model_type, available_providers }`. +Custom endpoints use `PresetSettings::base_url` or the vendor's settings. +The bindings retain their compatibility API in `aimux_providers::provider`; +its helpers and option types are not re-exported from the Rust crate root. + +## Middleware + +`wrap_language_model(model, middleware, model_id, provider_id)` applies the +first middleware as the outermost wrapper. Hooks can override provider id, +model id and supported URLs, transform call parameters, and wrap generation +or streaming with callbacks for both operations. Explicit id arguments take +precedence over middleware overrides. `wrap_image_model` provides the image +model hooks and the same explicit id overrides. ## Text Generation @@ -90,12 +151,22 @@ signal.abort(); // cancels connect, body reads, and streaming let result = task.await.unwrap()?; ``` +File URLs in the prompt — user `FileUrl` parts and `file` items of tool +outputs — are downloaded before the model call unless the model's +`supported_urls()` covers them, as the AI SDK's `generateText` does. Set +`download` (a `aimux_core::download::DownloadFunction`, the AI SDK's +`experimental_download`) to choose per URL; it receives every URL with +`is_url_supported_by_model` and returns the downloaded file or `None` to +send the URL. The default downloads with the guarded download. + > Parameters, return value, and the `content` variants are documented in > the [API overview](../API.md#text-generation). ## Streaming Generation -Returns generated content as a stream, output chunk by chunk. +Returns generated content as `TextStreamPart` items, output chunk by chunk. +Provider-layer `StreamPart::ResponseMetadata` events are consumed by the call +layer; they are not emitted by `stream_text`. ```rust use futures::StreamExt; @@ -115,7 +186,7 @@ Streaming honors the same `timeout` / `abort_signal` options as [text generation](#text-generation); streamed timeouts surface as an `Err(AiMuxError::Timeout(..))` stream item, and aborting the signal ends the stream with `Err(AiMuxError::Aborted(..))`. Provider-reported error events -remain `Ok(StreamPart::Error { .. })` data. +remain `Ok(TextStreamPart::Error { .. })` data. > Stream part variants are documented in the [API overview](../API.md#streaming-generation). @@ -172,7 +243,7 @@ Converts text into a vector representation. ```rust use aimux_core::embedding_model::{EmbeddingCallOptions, embed}; -let model = provider.embedding_model("text-embedding-3-small"); +let model = provider.embedding("text-embedding-3-small"); let opts = EmbeddingCallOptions::new("hello"); let result = embed(&model, opts).await?; // result.embeddings[0] is Vec @@ -207,6 +278,12 @@ let result = transcribe(&model, opts).await?; // result.text, result.segments, result.language ``` +`AudioInput::Url(url)` is downloaded by `transcribe` before the model call, +as the AI SDK's `transcribe` does. `opts.download` replaces the transport +(an `aimux_core::download::Download`); the default is +`create_download(None)`, the guarded download (public HTTP(S) URLs only, +2 GiB limit). + ## Image Generation ```rust @@ -220,7 +297,9 @@ let result = generate_image(&model, opts).await?; ## Video Generation -Video generation typically returns a URL (not binary). +Most video models return a URL. `generate_video` downloads it, as the AI +SDK's `generateVideo` does, with `opts.download` (default: the guarded +download), so the result holds the video's bytes. ```rust use aimux_core::video_model::{VideoCallOptions, generate_video}; @@ -228,7 +307,7 @@ use aimux_core::video_model::{VideoCallOptions, generate_video}; let model = provider.video("veo-3.0"); let opts = VideoCallOptions { prompt: Some("A cat".into()), n: 1, .. }; let result = generate_video(&model, opts).await?; -// result.videos[0] is VideoData::Url { url, media_type } +// result.videos[0] is VideoData::Binary { data, media_type } for a URL video ``` ## Reranking @@ -238,7 +317,7 @@ Reorders a document list by relevance. ```rust use aimux_core::reranking_model::{RerankingCallOptions, RerankingDocuments, rerank}; -let model = provider.reranking_model("rerank-v3.0"); +let model = provider.reranking("rerank-v3.0"); let opts = RerankingCallOptions::new("What is Rust?", docs); let result = rerank(&model, opts).await?; // result.ranking sorted by score @@ -251,7 +330,7 @@ Calls a search provider to obtain results. ```rust use aimux_core::search_model::{SearchCallOptions, search}; -let model = provider.search_model("tavily-search"); +let model = aimux_providers::tavily().search_model(); let opts = SearchCallOptions::new("What is Rust?"); let result = search(&model, opts).await?; // result.results is Vec @@ -276,11 +355,12 @@ let result = files.upload_file(opts).await?; ## Core Traits -The Rust core provides 10 traits/interfaces, implemented by each provider as needed: +The Rust core provides these traits/interfaces, implemented by each provider as needed: | Trait | Method | Semantics | |-------|------|------| -| `Provider` | `name`, `language_model` | Provider factory — holds API config, creates `LanguageModel` instances by model name | +| `Provider` | `language_model`, `embedding_model`, `image_model`, `discovery` | Provider factory — holds API config, creates `LanguageModel` instances by model name | +| `ProviderDiscovery` | `list_models` | Optional model listing through `Provider::discovery()` | | `LanguageModel` | `do_generate`, `do_stream` | Text generation | | `EmbeddingModel` | `do_embed` | Vector embedding | | `SpeechModel` | `do_generate` | Speech synthesis | @@ -293,14 +373,11 @@ The Rust core provides 10 traits/interfaces, implemented by each provider as nee The user-facing API consists of the `generate_text()` / `stream_text()` free functions, which internally call the trait methods. -> The multimodal accessors in the examples (`provider.embedding_model(...)`, -> `provider.speech(...)`, `provider.image(...)`, `provider.transcription(...)`, -> `provider.files()`, `provider.reranking_model(...)`, -> `provider.search_model(...)`, `provider.video(...)`) are **inherent methods -> on each provider struct**, not trait methods — they exist only on providers -> that support the feature (e.g. `OpenAIProvider` has `embedding_model` / -> `speech` / `image` / `transcription` / `files`; `CohereProvider` has -> `embedding_model` / `reranking_model`; `BedrockProvider` has `image`). +Concrete provider accessors such as `embedding`, `speech`, `image`, +`transcription`, `reranking` and `video` are inherent methods on the vendors +that offer them. Through `Provider`, required model methods return `Result`; +optional model methods return `Option>`, and `files()` returns +`Option>`. ## Types @@ -323,6 +400,7 @@ Rust types are the canonical definitions — one module per feature in | `types` | `Usage`, `TokenUsage`, `FinishReason`, `ResponseMetadata`, `Warning` | | `shared` | `FileBytes`, `FileData`, `Size`, `AspectRatio`, `ResponseInfo` | | `abort_signal` | `AbortSignal` | +| `download` | `Download`, `DownloadFunction`, `create_download`, `create_default_download_function` | | `embedding_model` | `EmbeddingModel`, `EmbeddingCallOptions`, `EmbeddingResult` | | `speech_model` | `SpeechModel`, `SpeechCallOptions`, `SpeechResult` | | `transcription_model` | `TranscriptionModel`, `TranscriptionCallOptions`, `TranscriptionResult`, `AudioInput` | diff --git a/docs/api/swift.md b/docs/api/swift.md index 04759832c..02b695518 100644 --- a/docs/api/swift.md +++ b/docs/api/swift.md @@ -1,6 +1,6 @@ # aimux · Swift API -> Unified LLM service access layer — one API to access 325 AI providers +> Unified LLM service access layer — one API to access AI providers Swift wraps the Rust core through the `aimux-ffi` C ABI (module `CAimuxFFI`), with ARC-managed model handles. @@ -31,11 +31,10 @@ print(result) ## Providers -All 251 registry-backed OpenAI-compatible providers are reachable by string name: +The vendor packages and 282 registry-backed OpenAI-compatible providers are reachable by string name: -> **Scope:** `provider(name)` covers only the 251 registry OpenAI-compatible -> providers; Anthropic/Google/multimodal/local → typed factories -> (`Model.anthropic(apiKey:modelId:)`); custom endpoints → base-URL variant. +> **Scope:** `provider(name)` reaches vendor packages and the 282 preset rows; +> typed factories remain available. Custom endpoints use the base-URL variant. > Full list: [providers.md](providers.md). ```swift @@ -150,7 +149,7 @@ have no Swift type — see "C ABI failures" below. Every fallible C function returns an opaque `aimux_error_t *` (`OpaquePointer?`): `NULL` = success (the result is in the trailing out-parameter), non-`NULL` = failure. One unified code selects `AimuxError` -(1...17; 4 retired), `RecordingError` (100...105), +(1...20; 4 retired), `RecordingError` (100...105), or a C ABI failure (200...206). The three decoders enforce the range expected by each call and restore the Swift error type; 200...206 collapses to `DecodingError.dataCorrupted`. @@ -180,6 +179,9 @@ and yields the invariant `DecodingError.dataCorrupted("aimux ffi: : | `.noSuchTool` | `AIMUX_E_NO_SUCH_TOOL` (15) | The model called a tool outside the supplied tool set; carries `toolName` and `availableTools` (`[String]?`, `nil` when no tool set was supplied) | | `.invalidToolInput` | `AIMUX_E_INVALID_TOOL_INPUT` (16) | Tool arguments failed to parse/validate; carries `toolName` and `toolInput` (the raw argument text) | | `.toolCallRepair` | `AIMUX_E_TOOL_CALL_REPAIR` (17) | A `repairToolCall` hook itself failed; carries `originalError` (the repaired-over error as wire JSON, the `ToolCall.error` encoding) | +| `.noOutputGenerated` | `AIMUX_E_NO_OUTPUT_GENERATED` (18) | No output was generated | +| `.loadApiKey` | `AIMUX_E_LOAD_API_KEY` (19) | No API key passed and the fallback environment variable is unset; carries `envVar` | +| `.loadSetting` | `AIMUX_E_LOAD_SETTING` (20) | A required provider setting is missing; carries `envVar` | | `.other` | `AIMUX_E_OTHER` (1) | Unclassified core error | There are no binding-local cases: only aimux-core produces an `AimuxError`. diff --git a/docs/codex-subscription-guide.md b/docs/codex-subscription-guide.md index 6cedaf48c..3389bdc0c 100644 --- a/docs/codex-subscription-guide.md +++ b/docs/codex-subscription-guide.md @@ -22,11 +22,17 @@ ## 2. API key 模式（Rust） ```rust -use aimux_providers::{CodexConfig, CodexProvider}; - -let config = CodexConfig::new("sk-..."); // 或 CodexConfig::from_env()（读 CODEX_API_KEY） -let model = CodexProvider::new(config).model("gpt-5.2-codex"); -// 之后就是标准 LanguageModel 用法：generate_text / stream_text 均可 +use aimux_providers::{CodexMode, CodexProviderSettings, codex, create_codex}; + +// 显式 key（`Some("")` 原样发送，不回退环境变量）： +let provider = create_codex(CodexProviderSettings { + mode: CodexMode::ApiKey(Some("sk-...".to_string().into())), + ..Default::default() +})?; +let model = provider.responses("gpt-5.2-codex"); +// 或默认实例：每次请求时读 CODEX_API_KEY，缺失则该次调用返回 `AiMuxError::LoadApiKey` +let model = codex().responses("gpt-5.2-codex"); +// 之后就是标准 LanguageModel 用法：generate_text / stream_text 均可；model.provider() == "codex.responses" ``` 可用模型（2026-08 核验）：`gpt-5.2-codex`、`gpt-5.1-codex`、`gpt-5.1-codex-mini`、`gpt-5-codex`。 @@ -85,12 +91,17 @@ async fn device_login(client_id: &str) -> Tokens { ```rust use aimux_core::error::AiMuxError; use aimux_core::generate::{generate_text, GenerateTextOptions}; -use aimux_providers::{CodexConfig, CodexProvider}; - -let config = CodexConfig::subscription(access_token) // OAuth 产物 - .with_chatgpt_account_id("acct_...") // 可选 - .with_originator("my-client"); // 可选，默认 "aimux" -let model = CodexProvider::new(config).model("gpt-5.2-codex"); +use aimux_providers::{CodexMode, CodexProviderSettings, create_codex}; + +let provider = create_codex(CodexProviderSettings { + mode: CodexMode::ChatGptAccount { + token: access_token.into(), // OAuth 产物；也可传 Resolvable::from_async_fn 每次请求取最新 token + account_id: Some("acct_...".to_string()), // 可选 + }, + originator: Some("my-client".to_string()), // 可选，默认 "aimux" + ..Default::default() // base_url 默认 https://chatgpt.com/backend-api/codex +})?; +let model = provider.responses("gpt-5.2-codex"); let result = generate_text(model, "Hello", GenerateTextOptions::default()).await?; ``` @@ -98,12 +109,20 @@ let result = generate_text(model, "Hello", GenerateTextOptions::default()).await ### 4.3 刷新编排（集成方） ```rust -use aimux_providers::{codex_refresh, CodexConfig, CodexProvider}; +use aimux_providers::{CodexMode, CodexProviderSettings, codex_refresh, create_codex}; + +/// 订阅通道的 provider：token 由调用方给（刷新后重新创建，或用 Resolvable::from_async_fn 动态取）。 +fn subscription(token: impl Into) -> Result { + create_codex(CodexProviderSettings { + mode: CodexMode::ChatGptAccount { token: token.into().into(), account_id: None }, + ..Default::default() + }) +} const CLIENT_ID: &str = "your-oauth-client-id"; async fn call_with_refresh(access: &str, refresh: &str, refresh_store: &mut String) -> Result<(), AiMuxError> { - let model = CodexProvider::new(CodexConfig::subscription(access)).model("gpt-5.2-codex"); + let model = subscription(access)?.responses("gpt-5.2-codex"); match generate_text(model, "Hello", Default::default()).await { Ok(_) => Ok(()), Err(AiMuxError::TokenExpired(_)) => { @@ -114,8 +133,7 @@ async fn call_with_refresh(access: &str, refresh: &str, refresh_store: &mut Stri *refresh_store = new_refresh.clone(); } // 3. 用新 access token 重试一次 - let model = CodexProvider::new(CodexConfig::subscription(tokens.access_token)) - .model("gpt-5.2-codex"); + let model = subscription(tokens.access_token)?.responses("gpt-5.2-codex"); generate_text(model, "Hello", Default::default()).await?; Ok(()) } diff --git a/docs/contributing/adding-a-provider.md b/docs/contributing/adding-a-provider.md index 080eb4793..fdb657d7a 100644 --- a/docs/contributing/adding-a-provider.md +++ b/docs/contributing/adding-a-provider.md @@ -14,59 +14,86 @@ simplify some steps below — noted per step). | Case | The vendor speaks… | You add | Example | |---|---|---|---| | 1. OpenAI-compatible | OpenAI Chat Completions wire format, own base URL | one registry row + derived cassettes | `deepinfra` | -| 2. A new protocol | its own HTTP API for text generation | a `src//` module implementing `do_generate` / `do_stream` | `cohere` | -| 3. Single modality | any API, but only one non-text modality | a typed shell + one modality model | `serper` (search), `lmnt` (speech) | +| 2. A new protocol | its own HTTP API for text generation | a `src//` package: `XxxProviderSettings` + `create_xxx` + a model implementing `do_generate` / `do_stream` | `cohere` | +| 3. Single modality | any API, but only one non-text modality | a `XxxProviderSettings` + `create_xxx` factory + one modality model | `serper` (search), `lmnt` (speech) | -If the vendor is OpenAI-compatible **and** needs quirks beyond a `profile` -flag (different auth header, non-standard error body, …), it is case 2, not -case 1. +If the vendor is OpenAI-compatible **and** needs quirks beyond what a row can +say (a different auth header, a non-standard error body, message conversion of +its own, ...), it is case 2: add its own package next to `groq/` and +`deepseek/`, which have their own chat models. The registry has no `family` field. ## Case 1 — OpenAI-compatible vendor (registry row) 1. Add one row to `aimux-providers/src/provider_registry.json` (file is - sorted by `name`; follow the flat 5-key shape): + sorted by `name`): ```json { "name": "example", "display": "Example", "env_var": "EXAMPLE_API_KEY", "base_url": "https://api.example.com/v1", "profile": {} } ``` - `profile` carries the known quirks — `supports_top_k`, `supports_tools`, - `supports_response_format`, `stream_usage_key`, `max_tokens_key` — see - the rows that already use one (e.g. `groq`). Anything beyond body-shape - flags rides in per-call `body_overrides` (RFC-0017); if the quirk changes - auth or error shapes, it is case 2. + Optional keys: `auth: "none"` (a local server: no key variable, no + `Authorization` header, no placeholder key) with `base_url_env` (the + variable holding the base URL); `params` for a templated `base_url` + (`{param}` placeholders, each declared with an optional `env` list and + `default`; a derived parameter takes `derive: {from, map, otherwise}`); + `profile.max_tokens_key` (`"max_tokens"` or `"max_completion_tokens"`: the + only max-token key the vendor accepts). Per-call body fields can use the + provider's options namespace. The OpenAI-compatible factory also exposes + its upstream `transform_request_body` hook; if the quirk changes auth or + error shapes, it is case 2. -2. Regenerate the provider documentation and commit its output: +2. Regenerate what is generated from the registry and commit its output: ```sh python scripts/gen_providers_doc.py # docs/api/providers.md (totals + list) ``` - CI runs it with `--check` in the `contract-tests` job and fails if the - committed output is stale. + The registry is embedded and parsed once into a runtime descriptor table. + Create presets with `create_provider(name, PresetSettings)`; there are no + per-name Rust functions or generated presets source. The table test validates + all rows and creates each with default settings. CI checks the generated documentation with `--check` + in the `contract-tests` job. 3. Derive replay cassettes: add a tuple to `PROVIDERS` in `scripts/generate_thin_wrapper_cassettes.py`, then run it. It derives `thin_wrapper_nonstream.json` / `thin_wrapper_stream.json` under `aimux-providers/tests/cassettes//` from **real OpenAI recordings** - (rewriting request path and model id — not fake data), because a + (rewriting request path and model id - not fake data), because a registry-backed provider's requests are byte-for-byte OpenAI shape. -4. Nothing else: `provider("example", ...)` now works in every binding, and - env-var key loading follows `env_var` automatically. +4. `create_provider("example", PresetSettings::default())` now works in Rust; + bindings retain `provider("example", ...)`. Keys are loaded from `env_var` + when a request is made. `provider_names()` lists accepted names. -> Roadmap note: #166 B2 turns the 33 standalone thin wrappers into registry -> rows and makes `conformance_test.rs` iterate the registry; the +> Roadmap note: #166 B2 makes `conformance_test.rs` iterate the registry; the > cassette-derivation step above is then replaced by that suite. ## Case 2 — a new protocol 1. New directory `aimux-providers/src//`: - - `model.rs` — `Provider` implementing `LanguageModel`: + - `mod.rs` — the factory, modelled on the vendor's `@ai-sdk/` + package: `ProviderSettings` with exactly that package's fields + (OpenAI and Anthropic use `api_key: Option`, + `headers: Option` and `fetch: Option`, + with no body-transform hook), `create_(settings)` (validates + package-specific settings, including base URLs and mutually exclusive + credentials), a default instance `()` where upstream provides one, and + `Provider` with one method per model (`chat(id)`, `embedding(id)`, + …) plus the `Provider` trait impl. For packages using the API-key loader, + `None` reads the environment variable **on every request**, + `Some("")` is sent verbatim and a missing key fails the call with + `AiMuxError::LoadApiKey`; a required setting that is missing fails the + call with `LoadSetting`. Model identity and custom-name handling follow + the package; do not impose a universal `"{name}.{method}"` rule. There is no `Config` type, no `from_env`, no `with_*` + builder, and no retry setting: the model reads a private config built by + the factory, and retry belongs to the caller (`max_retries` on the call); + - `model.rs` — the model implementing `LanguageModel`: `do_generate` / `do_stream` (plus other modality traits if the vendor offers them: embeddings, image, transcription…); + - `options.rs` — the providerOptions / providerMetadata namespace keys, + spelled once; - `convert.rs` — unified message format ⇄ vendor wire format, both directions, including the error-body shape; - `types.rs` — vendor request/response structs (`#[serde(default)]` on @@ -81,7 +108,7 @@ case 1. ```rust pub mod ; - pub use ::{Config, Provider}; + pub use ::{Provider, ProviderSettings, , create_}; ``` 3. Record real cassettes under `tests/cassettes//` (RFC-0003; no @@ -90,31 +117,37 @@ case 1. `do_stream_returns_parts` against those cassettes, mirroring the existing 20 modules. -4. Note: protocol providers are **not** name-addressable today — callers use - the typed factories (`CohereProvider::new(...)`), and - `provider("cohere", ...)` fails with `NoSuchProvider`. A registry row with - a `protocol` field arrives with #166 B1 (the `Protocol` enum + - `from_resolved`); until then, do not add protocol providers to the - registry JSON. +4. Add the package name and factory dispatch to `src/default_providers.rs`. + Rust callers can then use `create_provider("cohere", PresetSettings::default())` + or the typed factory `create_cohere(CohereProviderSettings {..})` / default + instance `cohere()`. `default_providers()` supplies the map for + `aimux_core::provider_registry::create_provider_registry(providers, options)`; + model ids use `provider:model` by default. Vendor packages do not need rows + in the compatible preset JSON. Model listing is exposed through + `Provider::discovery()`. ## Case 3 — single-modality vendor 1. One file `aimux-providers/src/.rs` (see `serper.rs` for search, `lmnt.rs` for speech; `cartesia/` for a vendor with two modalities): - - `Config { api_key, base_url }` with `new` / `with_base_url` / - `from_env`; - - `Provider` — a shell whose only job is returning the modality - model (`search_model()`, `speech_model()`, `image_model()`, …) and - reporting its `name()`; + - `ProviderSettings { api_key, base_url, headers, name, fetch, .. }` + (all optional), `create_(settings)` and an infallible default + instance `()` — the key is resolved per request, as in case 2; + - `Provider` — returns the modality model (`search_model()`, + `speech_model()`, `image_model()`, …); the model reports `provider()` as + `"{name}.{method}"` (`luma.image`, `tavily.search`); - the modality model implementing the trait's operation (`SearchModel::do_search`, `SpeechModel::do_speak`, …) as: build the request body → send through the provider-utils HTTP helpers → map the response into aimux types. Implement **only** the transform functions; transport, retry and timeouts are not your problem (Core owns them). -2. `pub mod` + `pub use` in `lib.rs` under the modality's section comment - (search-only, speech-only, image-only, video-only, …). +2. `pub mod` + `pub use` (`ProviderSettings`, `create_`, ``) in + `lib.rs` under the modality's section comment (search-only, speech-only, + image-only, video-only, …). Async jobs poll on package constants through + `shared/poll.rs`; `scripts/check_provider_boundaries.sh` keeps retry, + `options.max_retries` reads and builder-era names out of the package. 3. Cassettes + unit tests asserting the mapped result shapes. @@ -122,7 +155,8 @@ case 1. - **A generator may stay in `scripts/` only if its output carries a "GENERATED — do not edit" header and CI runs it with `--check`.** Today - that is `gen_providers_doc.py` (the `contract-tests` job). + that is `gen_providers_doc.py` (the `contract-tests` + job). `gen_ts_types.py` regenerates the ts-rs TypeScript types through `cargo test -p aimux-core --lib export` (the export tests are the gate). diff --git a/docs/error-model.md b/docs/error-model.md index 151e84e42..6a49bdbc7 100644 --- a/docs/error-model.md +++ b/docs/error-model.md @@ -20,7 +20,7 @@ - C ABI 成功时返回 `NULL` 并写入出参；失败时返回一个由调用方释放一次的 不透明 `aimux_error_t *`。 - Go、Java、Kotlin、Swift 和 Flutter 将 C 错误还原为本语言错误。错误码 - `1..17` 属于核心，`100..105` 属于录制，`200..206` 属于绑定边界。 + `1..20` 属于核心，`100..105` 属于录制，`200..206` 属于绑定边界。 结构化字段只放在真正拥有它的错误上。例如，HTTP 状态和响应头属于 `APICallError`，不放进通用基类（retry hint 与 request id 从 `response_headers` 读取）。 diff --git a/docs/external-provider-config-research.md b/docs/external-provider-config-research.md index 6676c5855..c3eb1c746 100644 --- a/docs/external-provider-config-research.md +++ b/docs/external-provider-config-research.md @@ -1,6 +1,7 @@ # 调研:允许外部提供 provider 配置的可行性与设计 > **Status**: 调研报告(非 RFC,供决策) +> **Historical (RFC-0036)**: 本报告写于 provider 工厂重构之前,下文的 `OpenAIConfig`、provider 配置上的 retry/body 覆盖字段等旧 API 已不存在,仅作决策记录保留;现行形态见 [docs/api/rust.md](api/rust.md#providers)(`create_provider(name, PresetSettings)`、`default_providers()` + `create_provider_registry`、各包 `XxxProviderSettings` + `create_xxx`)。旧 `provider()` / `ProviderOptions` / 注册与 JSON 加载入口只保留在 `aimux_providers::provider` 绑定兼容模块,不再从 Rust crate root 导出。 > **Date**: 2026-08-05 > **Scope**: 评估"运行时/宿主应用/配置文件提供 provider 配置"对 aimux 的可行性、形态、成本收益,给出明确建议 > **Related**: [RFC-0017](../rfc/0017-provider-config-dx.md) 配置 DX、[RFC-0019](../rfc/0019-session-affinity.md) 会话亲和 diff --git a/docs/provider-config-manual.md b/docs/provider-config-manual.md index 6de9b0915..5e6388449 100644 --- a/docs/provider-config-manual.md +++ b/docs/provider-config-manual.md @@ -1,28 +1,30 @@ # Provider Config 用户手册（思考开关 + 字段差异） -> **原则**：aimux **不内置任何厂商映射**（RFC-0017 v3）。用户只有两个入口： -> - `reasoning`（7 档，直传 `reasoning_effort`，厂商自决——支持则生效，不支持则忽略/报错） -> - `bodyOverrides`（provider 级 + per-call，JSON deep-merge，用户定义一切厂商差异） +> **原则**：配置按 provider 包解释；OpenAI-compatible 的透传规则不适用于所有原生厂商包。用户有两个调用级入口： +> - `reasoning`（由 provider 包映射为其 wire 参数；不支持的设置产生 warning） +> - `provider_options.`（per-call；OpenAI-compatible 家族里，厂商自己命名空间下 schema 不认识的字段**原样写入请求体**，用户定义一切厂商差异） +> +> 声明式 `bodyOverrides`（provider 级 + per-call 的 JSON deep-merge）已删除，`ProviderConfig` / `config_json` 里传 `body_overrides` 会报 `InvalidArgument`。OpenAI-compatible settings 保留上游的 `transform_request_body` 闭包；OpenAI、Anthropic、Google settings 没有这个字段。 > > 本手册是"知识"的归宿：各厂商的 wire 参数、配置示例、核实日期。知识会过期——以厂商官方文档为准，本手册仅作参考。 > 数据来源：[model-config-research/](internal/model-config-research/)（2026-08-01 全网调研，250 家）。 --- -## 1. reasoning：7 档直传 +## 1. reasoning：按包映射 ```ts await generateText(model, prompt, { reasoning: 'none' | 'minimal' | 'low' | 'medium' | 'high' | 'xhigh' }) // 注：枚举另有 provider-default（不传档位），实际可设档位 6 个 ``` -- aimux 把档位原样写成 `reasoning_effort` 字段（**无归一化**） +- OpenAI-compatible 默认把档位写成 `reasoning_effort`；原生 provider 按包规则映射。例如 Groq 将 minimal 映射为 low、xhigh 映射为 high；DeepSeek 将 minimal/low 映射为 low、medium/high 映射为 high、xhigh 映射为 max。 - 厂商认识哪个档位是厂商的事： - OpenAI 官方：minimal/low/medium/high 等（官方 API 枚举） - DeepSeek V4：官方接受 low/high/max/**xhigh**（自行映射，见 §2） - Perplexity：minimal/low/medium/high（四档） - Kimi k3：low/high/max（三档） - - 不认识的厂商：忽略或报错，由用户自行处理（或用 bodyOverrides 直接控制） + - 不认识的厂商：忽略或报错，由用户自行处理（或用 `provider_options.` 直接控制） ## 2. DeepSeek V4（官方核实 2026-08-02） @@ -37,43 +39,46 @@ await generateText(model, prompt, { reasoning: 'none' | 'minimal' | 'low' | 'med | 无效参数 | temperature/top_p/penalty 在思考模式下无效 | 不报错但无效果 | ```ts -// 关思考（v3 迁移：旧版 reasoning:'none' 自动注入已退役） -await generateText(model, p, { bodyOverrides: { thinking: { type: 'disabled' } } }) +// 显式关思考（优先于 reasoning） +await generateText(model, p, { provider_options: { deepseek: { thinking: { type: 'disabled' } } } }) // 开思考 + xhigh 档（官方样例：两参数独立） await generateText(model, p, { reasoning: 'xhigh', - bodyOverrides: { thinking: { type: 'enabled' } }, + provider_options: { deepseek: { thinking: { type: 'enabled' } } }, }) ``` -> ⚠️ 迁移说明：0.1.x 旧版对 DeepSeek 有内置特化（`reasoning:'none'` → 自动注入 `thinking:{type:"disabled"}`）。该内置已退役（2026-08），现在需要显式 `bodyOverrides`。 +> 当前 DeepSeek 包中，未显式指定 `thinking.type` 时，`reasoning: 'none'` 会关闭思考，其他显式档位会开启思考；`thinking.type` 优先。`reasoningEffort` 是该包的 option key，不能用 wire 字段名 `reasoning_effort` 替代。 ## 3. 思考开关配置示例（按调研，来源见各 batch 文件） +配置示例写在 `provider_options.` 里。OpenAI-compatible 的 `` 接受 provider 名及其 camelCase 形式（`zhipu_v4` → `zhipuV4`），原名含下划线时会收到 deprecation warning；原生包的 namespace 和 option key 按各包 schema 定义。 + | 厂商 | 关思考 wire | 配置示例 | 备注/来源 | |---|---|---|---| -| **GLM / Zhipu**（bigmodel/zai/zhipu_v4） | `thinking:{type:"disabled"}` | `bodyOverrides: { thinking: { type: 'disabled' } }` | 跨代稳定（batch-01/06） | -| **Qwen 系**（alibaba/baidu） | `enable_thinking: false`；预算 `thinking_budget` [100,16384] | `bodyOverrides: { enable_thinking: false }`；开思考 `{ enable_thinking: true, thinking_budget: 16384 }` | qwen3 混合（batch-01）；**纯思考版不可关**；新版混合用消息级 `/no_think` | -| **Kimi / Moonshot** | k3: `reasoning_effort`（low/high/max）；k2.5/k2.6: `thinking:{type:"disabled"}`；k2.7-code: **不可关** | `bodyOverrides: { thinking: { type: 'disabled' } }`（k2.5/k2.6 系） | by-model 三套（batch-03） | -| **MiniMax** | M3: `thinking:{type:"disabled"}`；M2.x: **不可关**（传 disabled 仍思考） | `bodyOverrides: { thinking: { type: 'disabled' } }` | 开启值 M3 用 `"adaptive"`（batch-04） | -| **方舟 / 豆包**（bytedance/byteplus） | `thinking:{type:"disabled"}`；带预算 `{type:"enabled",budget_tokens:N}` | `bodyOverrides: { thinking: { type: 'disabled' } }` | batch-01/02 | -| **DeepInfra** | `reasoning:{enabled:false}` | `bodyOverrides: { reasoning: { enabled: false } }` | 非 thinking 对象（batch-02） | -| **SiliconFlow** | `thinking_budget`（思维链 token 上限，Qwen3 系强制截断；**无 0=关 语义**，关思考走 qwen 系 `enable_thinking`） | `bodyOverrides: { thinking_budget: 1024 }`（调低预算） | batch-05 | -| **Perplexity** | `reasoning_effort` 四档 + `stream_mode`；推理 token **不可强制关闭** | `bodyOverrides: { stream_mode: 'concise' }` | batch-05 | -| **Groq** | `reasoning_format`（如 raw）+ effort **直传**（无归一化） | `bodyOverrides: { reasoning_format: 'raw' }` | batch-03 | -| **Heroku** | `extended_thinking:{enabled,budget_tokens,include_reasoning}`；未知参数需 `allow_ignored_params` | `bodyOverrides: { extended_thinking: { enabled: true, budget_tokens: 2000 } }`；非标准参数一并 `bodyOverrides: { allow_ignored_params: true, ... }` | batch-03 | -| **Hetzner** | `chat_template_kwargs:{enable_thinking:false}` | `bodyOverrides: { chat_template_kwargs: { enable_thinking: false } }` | 社区实测（batch-03） | -| **Venice** | `venice_parameters:{disable_thinking:true}` | `bodyOverrides: { venice_parameters: { disable_thinking: true } }` | 封闭字段（batch-06） | -| **腾讯 hy3**（TokenHub） | `thinking:{type:"enabled"}` + `reasoning_effort`（默认 low） | `bodyOverrides: { thinking: { type: 'enabled' }, reasoning_effort: 'low' }` | batch-06（C 级官方文档） | -| **StepFun** | `reasoning_format: "general"/"deepseek-style"` | `bodyOverrides: { reasoning_format: 'deepseek-style' }` | batch-05 | +| **GLM / Zhipu**（bigmodel/zai/zhipu_v4） | `thinking:{type:"disabled"}` | `provider_options: { zhipuV4: { thinking: { type: 'disabled' } } }` | 跨代稳定（batch-01/06） | +| **Qwen 系**（alibaba/baidu） | `enable_thinking: false`；预算 `thinking_budget` [100,16384] | `provider_options: { alibaba: { enable_thinking: false } }`；开思考 `{ enable_thinking: true, thinking_budget: 16384 }` | qwen3 混合（batch-01）；**纯思考版不可关**；新版混合用消息级 `/no_think` | +| **Kimi / Moonshot** | k3: `reasoning_effort`（low/high/max）；k2.5/k2.6: `thinking:{type:"disabled"}`；k2.7-code: **不可关** | `provider_options: { moonshotai: { thinking: { type: 'disabled' } } }`（k2.5/k2.6 系） | by-model 三套（batch-03） | +| **MiniMax** | M3: `thinking:{type:"disabled"}`；M2.x: **不可关**（传 disabled 仍思考） | `provider_options: { minimax: { thinking: { type: 'disabled' } } }` | 开启值 M3 用 `"adaptive"`（batch-04） | +| **方舟 / 豆包**（bytedance/byteplus） | `thinking:{type:"disabled"}`；带预算 `{type:"enabled",budget_tokens:N}` | `provider_options: { bytedance: { thinking: { type: 'disabled' } } }` | batch-01/02 | +| **DeepInfra** | `reasoning:{enabled:false}` | `provider_options: { deepinfra: { reasoning: { enabled: false } } }` | 非 thinking 对象（batch-02） | +| **SiliconFlow** | `thinking_budget`（思维链 token 上限，Qwen3 系强制截断；**无 0=关 语义**，关思考走 qwen 系 `enable_thinking`） | `provider_options: { siliconflow: { thinking_budget: 1024 } }`（调低预算） | batch-05 | +| **Perplexity** | `reasoning_effort` 四档 + `stream_mode`；推理 token **不可强制关闭** | `provider_options: { perplexity: { stream_mode: 'concise' } }` | batch-05 | +| **Groq** | `reasoning_format`（如 raw）；effort 按包映射 | `provider_options: { groq: { reasoningFormat: 'raw' } }` | `reasoningEffort` 接受 none/default/low/medium/high | +| **Heroku** | `extended_thinking:{enabled,budget_tokens,include_reasoning}`；未知参数需 `allow_ignored_params` | `provider_options: { heroku: { extended_thinking: { enabled: true, budget_tokens: 2000 } } }`；非标准参数一并 `provider_options: { heroku: { allow_ignored_params: true, ... } }` | batch-03 | +| **Hetzner** | `chat_template_kwargs:{enable_thinking:false}` | `provider_options: { hetzner: { chat_template_kwargs: { enable_thinking: false } } }` | 社区实测（batch-03） | +| **Venice** | `venice_parameters:{disable_thinking:true}` | `provider_options: { venice: { venice_parameters: { disable_thinking: true } } }` | 封闭字段（batch-06） | +| **腾讯 hy3**（TokenHub） | `thinking:{type:"enabled"}` + `reasoning_effort`（默认 low） | `provider_options: { tencentTokenhub: { thinking: { type: 'enabled' }, reasoning_effort: 'low' } }` | batch-06（C 级官方文档） | +| **StepFun** | `reasoning_format: "general"/"deepseek-style"` | `provider_options: { stepfun: { reasoning_format: 'deepseek-style' } }` | batch-05 | ## 4. max_tokens_key（内置修复，用户无感） -aimux 内置了 8 家厂商的 max tokens 字段名差异（**纯内部数据，用户不需要配置**）： +registry 内置了 7 家兼容厂商的 max tokens 字段名差异（**纯内部数据，用户不需要配置**）： - 只认 `max_tokens`：stepfun / siliconflow / sarvam / reka_ai / publicai / perplexity -- 只认 `max_completion_tokens`：groq（max_tokens 已弃用）/ heroku（官方要求） +- 只认 `max_completion_tokens`：heroku（官方要求） +- 原生 Groq 和 DeepSeek 使用 `max_tokens`，不走 registry 推断。 ```ts // 用户始终写 max_output_tokens，aimux 按厂商自动选字段名 @@ -82,31 +87,44 @@ await generateText(model, prompt, { max_output_tokens: 4096 }) // stepfun → {"max_tokens":4096} ``` -> ⚠️ **不要用 `provider_options.maxCompletionTokens` 显式指定**：对只认 `max_tokens` -> 的厂商（stepfun / siliconflow / sarvam / reka_ai / publicai / perplexity），该 -> 选项会被**静默丢弃**（不报错、无 warning，backlog B10）。请改用顶层 -> `max_output_tokens`——它会被自动映射到正确字段名。 +> 使用顶层 `max_output_tokens`，由包选择 wire 字段名。`provider_options` +> 必须是 namespace → object 的结构；`provider_options.maxCompletionTokens` +> 直接放数字不符合结构，不能作为 max tokens 的配置入口。 > -> 未内置 `max_tokens_key` 的厂商（含 **DeepSeek**）走默认推断：推理模型发 -> `max_completion_tokens`，非推理发 `max_tokens`。DeepSeek 对 -> `max_completion_tokens` 的接受性**尚未实测**（backlog B6）——当前行为 -> 由默认推断路径决定，实测结论落地前请勿依赖其行为。 +> 未内置 `max_tokens_key` 的兼容厂商走默认推断：推理模型发 +> `max_completion_tokens`，非推理发 `max_tokens`。原生 DeepSeek 始终将 +> `max_output_tokens` 写为 `max_tokens`。 -## 5. bodyOverrides 用法速查 +## 5. `provider_options` 透传用法速查 ```ts -// provider 级（创建时，每次请求生效）—— 适合固定字段 -const model = await openai(key, 'qwen3-coder', { - bodyOverrides: { enable_thinking: false }, +// OpenAI-compatible per-call：厂商命名空间下 schema 不认识的字段原样进入请求体 +await generateText(model, prompt, { + provider_options: { alibaba: { enable_thinking: false } }, }) -// per-call 级（单次请求，覆盖 provider 级） -await generateText(model, prompt, { bodyOverrides: { enable_thinking: false } }) +// OpenAI-compatible 已知字段（user / reasoningEffort / textVerbosity / strictJsonSchema）按 schema 处理， +// 其余字段（thinking、enable_thinking、extended_thinking …）直接透传。 +// 通用命名空间 openaiCompatible 只认 schema 字段，未知字段会被丢弃——厂商专属字段请写在厂商自己的命名空间下。 +``` -// null = 删除字段（删除通用路径注入的字段） -await generateText(model, prompt, { bodyOverrides: { 'reasoning_effort': null } }) +Rust 的 OpenAI-compatible 包支持对**最终**请求体做整体改写（provider 级、每次请求生效、可删字段），用该包 settings 的 `transform_request_body`： + +```rust +let provider = create_openai_compatible(OpenAICompatibleProviderSettings { + name: "acme".into(), + base_url: "https://api.acme.example/v1".into(), + transform_request_body: Some(Arc::new(|mut body| { + body["enable_thinking"] = json!(false); + body.as_object_mut().unwrap().remove("reasoning_effort"); + body + })), + ..Default::default() +})?; ``` +JS / Python / Go 等绑定没有闭包入口，provider 级的固定字段请在每次调用的 `provider_options` 里带上。 + ## 6. 核实日期与来源 - 手册条目核实日期：2026-08-02 @@ -117,7 +135,7 @@ await generateText(model, prompt, { bodyOverrides: { 'reasoning_effort': null } ## 7. Model List API 与模型配置补充（RFC-0027） aimux 提供两个**独立原语**,host 各自调用并按需合并(aimux 是库、不持有状态): -- `Provider::list_models()` — 运行时从 provider `/models` 发现可用模型,返回 `RuntimeModel[]`(provider 官方数据,通常只有 id)。 +- Rust: `Provider::discovery()` → `ProviderDiscovery::list_models()` — 运行时从 provider `/models` 发现可用模型,返回 `RuntimeModel[]`(provider 官方数据,通常只有 id)。 - `get_model_specs()` — 独立拉取 `models.anya2a.com` 社区聚合数据,返回 `ModelSpec` 配置/能力。 aimux **不自动合并**二者:`list_models` 不带 anya2a 配置,`get_model_specs` 不带可用性。host 自行按 modelId 把 spec 合并进 `list_models` 的结果。 @@ -139,7 +157,7 @@ const spec = catalogue.specs?.['deepseek']?.['deepseek-v4'] // 4. 用户读 spec,按业务自己定 options const model = await p.model('deepseek-v4') -await generateText(model, prompt, { max_output_tokens: 8000, bodyOverrides: { thinking: { type: 'enabled' } } }) +await generateText(model, prompt, { max_output_tokens: 8000, provider_options: { deepseek: { thinking: { type: 'enabled' } } } }) ``` ### 7.2 config 是咨询性的 @@ -157,7 +175,7 @@ anya2a 只补 provider 列表里出现的 modelId,不作可用性依据。缺字 ### 7.4 覆盖范围 -292/325 家真 LLM provider 已覆盖(89.8%):251 registry + 23 standalone + 10 vertex MaaS + 8 native。33 家 modality-only(speech/image/embed/search)返回 `Unsupported`。 +Rust 通过 `Provider::discovery()` 检查厂商是否提供模型发现,返回 `None` 表示未提供。预设表有 282 行;Groq 和 DeepSeek 由各自厂商包提供发现。绑定的 provider 句柄保留 `listModels()` / `list_models()` API。 ### 7.5 无内置缓存(不存在缓存/离线环境变量) @@ -172,6 +190,6 @@ aimux 是库、不持有状态:`list_models` 与 `get_model_specs` 都是**实 | `limits.context` / `limits.output` | 用户自行做上下文截断;`max_output_tokens` 设多少 | | `capabilities.tool_call` | 决定是否传 `tools` | | `capabilities.structured_output` | 决定是否用 `response_format: Json` | -| `reasoning.effort_default` / `mode` | `bodyOverrides` 里填 `thinking:{enabled}` / `reasoning_effort` | +| `reasoning.effort_default` / `mode` | `provider_options.` 里填 `thinking:{enabled}` / `reasoning_effort` | | `cost` | 发送前预估成本(用户自行算) | | `modalities` | 决定能否传 image/audio 内容 | diff --git a/docs/session-affinity-guide.md b/docs/session-affinity-guide.md index 83002d681..7ea15317f 100644 --- a/docs/session-affinity-guide.md +++ b/docs/session-affinity-guide.md @@ -106,7 +106,7 @@ A: 大概率不会(HTTP 语义下未知头被忽略,RFC 9110;OpenCode 对全部 A: 会话头是"你的网关"的属性,不是 provider 的属性。同一个 OpenAI 兼容 provider,直连官方、走 OpenRouter、走自建网关,要传的头完全不同;而大多数自建网关根本不在 registry 里。预置会误导,且与集成方的部署强耦合——所以 aimux 交付文档,不交付预置数据。 **Q: 我的上游是 OpenRouter,`session_id` 传 body 还是 header?** -A: OpenRouter 两者都接受,body 顶层 `session_id` 优先于 `x-session-id` 头。aimux 的 `body_overrides` 可传 body 字段,`headers` 可传头;二选一即可。 +A: OpenRouter 两者都接受,body 顶层 `session_id` 优先于 `x-session-id` 头。aimux 的 `provider_options.`(厂商命名空间下的未知字段原样进 body)可传 body 字段,`headers` 可传头;二选一即可。 **Q: 会话 id 谁来生成?** A: 你。aimux 不做会话状态管理(生命周期由你持有);同一会话内保持值稳定即可获得粘性/缓存收益。 diff --git a/fixtures/aisdk/anthropic/messages-basic.json b/fixtures/aisdk/anthropic/messages-basic.json new file mode 100644 index 000000000..51d4e4c90 --- /dev/null +++ b/fixtures/aisdk/anthropic/messages-basic.json @@ -0,0 +1,108 @@ +{ + "case": "messages-basic", + "model": { + "modelId": "claude-sonnet-4-5", + "provider": "anthropic.messages" + }, + "package": "@ai-sdk/anthropic", + "request": { + "body": { + "max_tokens": 64000, + "messages": [ + { + "content": [ + { + "text": "Say hello.", + "type": "text" + } + ], + "role": "user" + } + ], + "model": "claude-sonnet-4-5" + }, + "headers": { + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-api-key": "" + }, + "method": "POST", + "url": "https://api.anthropic.com/v1/messages" + }, + "response": { + "body": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "id": "msg_fixture_1", + "model": "claude-sonnet-4-5", + "role": "assistant", + "stop_reason": "end_turn", + "stop_sequence": null, + "type": "message", + "usage": { + "input_tokens": 12, + "output_tokens": 9 + } + }, + "headers": { + "content-type": "application/json", + "request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "anthropic": { + "container": null, + "contextManagement": null, + "iterations": null, + "stopSequence": null, + "usage": { + "input_tokens": 12, + "output_tokens": 9 + } + } + }, + "rawFinishReason": "end_turn", + "response": { + "headers": { + "content-type": "application/json", + "request-id": "req_fixture_1" + }, + "id": "msg_fixture_1", + "modelId": "claude-sonnet-4-5" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "cacheWriteTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": {}, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "4.0.71" +} diff --git a/fixtures/aisdk/anthropic/messages-custom-name.json b/fixtures/aisdk/anthropic/messages-custom-name.json new file mode 100644 index 000000000..579a40f11 --- /dev/null +++ b/fixtures/aisdk/anthropic/messages-custom-name.json @@ -0,0 +1,141 @@ +{ + "case": "messages-custom-name", + "model": { + "modelId": "claude-sonnet-4-5", + "provider": "myproxy" + }, + "observations": { + "bodyMetadataUserId": "from-myproxy-key", + "modelProvider": "myproxy", + "providerMetadataKeys": [ + "anthropic", + "myproxy" + ] + }, + "package": "@ai-sdk/anthropic", + "request": { + "body": { + "max_tokens": 64000, + "messages": [ + { + "content": [ + { + "text": "Say hello.", + "type": "text" + } + ], + "role": "user" + } + ], + "metadata": { + "user_id": "from-myproxy-key" + }, + "model": "claude-sonnet-4-5" + }, + "headers": { + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-api-key": "" + }, + "method": "POST", + "url": "https://api.anthropic.com/v1/messages" + }, + "response": { + "body": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "id": "msg_fixture_1", + "model": "claude-sonnet-4-5", + "role": "assistant", + "stop_reason": "end_turn", + "stop_sequence": null, + "type": "message", + "usage": { + "input_tokens": 12, + "output_tokens": 9 + } + }, + "headers": { + "content-type": "application/json", + "request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "anthropic": { + "container": null, + "contextManagement": null, + "iterations": null, + "stopSequence": null, + "usage": { + "input_tokens": 12, + "output_tokens": 9 + } + }, + "myproxy": { + "container": null, + "contextManagement": null, + "iterations": null, + "stopSequence": null, + "usage": { + "input_tokens": 12, + "output_tokens": 9 + } + } + }, + "rawFinishReason": "end_turn", + "response": { + "headers": { + "content-type": "application/json", + "request-id": "req_fixture_1" + }, + "id": "msg_fixture_1", + "modelId": "claude-sonnet-4-5" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "cacheWriteTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": {}, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello.", + "providerOptions": { + "anthropic": { + "metadata": { + "userId": "from-anthropic-key" + } + }, + "myproxy": { + "metadata": { + "userId": "from-myproxy-key" + } + } + } + }, + "operation": "generateText" + }, + "version": "4.0.71" +} diff --git a/fixtures/aisdk/anthropic/messages-provider-options.json b/fixtures/aisdk/anthropic/messages-provider-options.json new file mode 100644 index 000000000..5750690b6 --- /dev/null +++ b/fixtures/aisdk/anthropic/messages-provider-options.json @@ -0,0 +1,135 @@ +{ + "case": "messages-provider-options", + "model": { + "modelId": "claude-sonnet-4-5", + "provider": "anthropic.messages" + }, + "package": "@ai-sdk/anthropic", + "request": { + "body": { + "max_tokens": 6096, + "messages": [ + { + "content": [ + { + "text": "What is 17 * 23?", + "type": "text" + } + ], + "role": "user" + } + ], + "model": "claude-sonnet-4-5", + "thinking": { + "budget_tokens": 2000, + "type": "enabled" + } + }, + "headers": { + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-api-key": "" + }, + "method": "POST", + "url": "https://api.anthropic.com/v1/messages" + }, + "response": { + "body": { + "content": [ + { + "signature": "sig_fixture", + "thinking": "17 * 23 = 391.", + "type": "thinking" + }, + { + "text": "391", + "type": "text" + } + ], + "id": "msg_fixture_1", + "model": "claude-sonnet-4-5", + "role": "assistant", + "stop_reason": "end_turn", + "stop_sequence": null, + "type": "message", + "usage": { + "input_tokens": 12, + "output_tokens": 9 + } + }, + "headers": { + "content-type": "application/json", + "request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "providerMetadata": { + "anthropic": { + "signature": "sig_fixture" + } + }, + "text": "17 * 23 = 391.", + "type": "reasoning" + }, + { + "text": "391", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "anthropic": { + "container": null, + "contextManagement": null, + "iterations": null, + "stopSequence": null, + "usage": { + "input_tokens": 12, + "output_tokens": 9 + } + } + }, + "rawFinishReason": "end_turn", + "response": { + "headers": { + "content-type": "application/json", + "request-id": "req_fixture_1" + }, + "id": "msg_fixture_1", + "modelId": "claude-sonnet-4-5" + }, + "text": "391", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "cacheWriteTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": {}, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "maxOutputTokens": 4096, + "prompt": "What is 17 * 23?", + "providerOptions": { + "anthropic": { + "thinking": { + "budgetTokens": 2000, + "type": "enabled" + } + } + } + }, + "operation": "generateText" + }, + "version": "4.0.71" +} diff --git a/fixtures/aisdk/anthropic/messages-stream-basic.json b/fixtures/aisdk/anthropic/messages-stream-basic.json new file mode 100644 index 000000000..25f9040f7 --- /dev/null +++ b/fixtures/aisdk/anthropic/messages-stream-basic.json @@ -0,0 +1,166 @@ +{ + "case": "messages-stream-basic", + "model": { + "modelId": "claude-sonnet-4-5", + "provider": "anthropic.messages" + }, + "package": "@ai-sdk/anthropic", + "request": { + "body": { + "max_tokens": 64000, + "messages": [ + { + "content": [ + { + "text": "Say hello.", + "type": "text" + } + ], + "role": "user" + } + ], + "model": "claude-sonnet-4-5", + "stream": true + }, + "headers": { + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "user-agent": "ai-sdk-anthropic/4.0.71 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-api-key": "" + }, + "method": "POST", + "url": "https://api.anthropic.com/v1/messages" + }, + "response": { + "body": "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_fixture_2\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-4-5\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":12,\"output_tokens\":1}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\nevent: ping\ndata: {\"type\":\"ping\"}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\" world\"}}\n\nevent: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":2}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + "headers": { + "cache-control": "no-cache", + "content-type": "text/event-stream" + }, + "status": 200 + }, + "result": { + "finishReason": "stop", + "parts": [ + { + "type": "start" + }, + { + "request": {}, + "type": "start-step", + "warnings": [] + }, + { + "id": "0", + "type": "text-start" + }, + { + "id": "0", + "text": "Hello", + "type": "text-delta" + }, + { + "id": "0", + "text": " world", + "type": "text-delta" + }, + { + "id": "0", + "type": "text-end" + }, + { + "finishReason": "stop", + "performance": { + "effectiveOutputTokensPerSecond": 0, + "effectiveTotalTokensPerSecond": 0, + "inputTokensPerSecond": 0, + "outputTokensPerSecond": 0, + "responseTimeMs": 0, + "stepTimeMs": 0, + "timeBetweenOutputChunksMs": { + "avg": 0, + "max": 0, + "median": 0, + "min": 0, + "p10": 0, + "p90": 0 + }, + "timeToFirstOutputMs": 0, + "toolExecutionMs": {} + }, + "providerMetadata": { + "anthropic": { + "container": null, + "contextManagement": null, + "iterations": null, + "stopSequence": null, + "usage": { + "input_tokens": 12, + "output_tokens": 2 + } + } + }, + "rawFinishReason": "end_turn", + "response": { + "headers": { + "cache-control": "no-cache", + "content-type": "text/event-stream" + }, + "id": "msg_fixture_2", + "modelId": "claude-sonnet-4-5", + "timestamp": "2025-01-01T00:00:00.000Z" + }, + "type": "finish-step", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "cacheWriteTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": {}, + "outputTokens": 2, + "raw": { + "input_tokens": 12, + "output_tokens": 2 + }, + "totalTokens": 14 + } + }, + { + "finishReason": "stop", + "rawFinishReason": "end_turn", + "totalUsage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "cacheWriteTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": {}, + "outputTokens": 2, + "totalTokens": 14 + }, + "type": "finish" + } + ], + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "cacheWriteTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": {}, + "outputTokens": 2, + "totalTokens": 14 + } + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "streamText" + }, + "version": "4.0.71" +} diff --git a/fixtures/aisdk/anthropic/messages-system-tools.json b/fixtures/aisdk/anthropic/messages-system-tools.json new file mode 100644 index 000000000..45e0b13d5 --- /dev/null +++ b/fixtures/aisdk/anthropic/messages-system-tools.json @@ -0,0 +1,166 @@ +{ + "case": "messages-system-tools", + "model": { + "modelId": "claude-sonnet-4-5", + "provider": "anthropic.messages" + }, + "package": "@ai-sdk/anthropic", + "request": { + "body": { + "max_tokens": 256, + "messages": [ + { + "content": [ + { + "text": "What is the weather in Paris?", + "type": "text" + } + ], + "role": "user" + } + ], + "model": "claude-sonnet-4-5", + "system": [ + { + "text": "You are a weather assistant.", + "type": "text" + } + ], + "tool_choice": { + "type": "any" + }, + "tools": [ + { + "description": "Get the current weather for a city.", + "input_schema": { + "additionalProperties": false, + "properties": { + "city": { + "description": "City name", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + }, + "name": "get_weather" + } + ] + }, + "headers": { + "anthropic-beta": "structured-outputs-2025-11-13", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-api-key": "" + }, + "method": "POST", + "url": "https://api.anthropic.com/v1/messages" + }, + "response": { + "body": { + "content": [ + { + "id": "toolu_fixture_1", + "input": { + "city": "Paris" + }, + "name": "get_weather", + "type": "tool_use" + } + ], + "id": "msg_fixture_1", + "model": "claude-sonnet-4-5", + "role": "assistant", + "stop_reason": "tool_use", + "stop_sequence": null, + "type": "message", + "usage": { + "input_tokens": 12, + "output_tokens": 9 + } + }, + "headers": { + "content-type": "application/json", + "request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "input": { + "city": "Paris" + }, + "toolCallId": "toolu_fixture_1", + "toolName": "get_weather", + "type": "tool-call" + } + ], + "finishReason": "tool-calls", + "providerMetadata": { + "anthropic": { + "container": null, + "contextManagement": null, + "iterations": null, + "stopSequence": null, + "usage": { + "input_tokens": 12, + "output_tokens": 9 + } + } + }, + "rawFinishReason": "tool_use", + "response": { + "headers": { + "content-type": "application/json", + "request-id": "req_fixture_1" + }, + "id": "msg_fixture_1", + "modelId": "claude-sonnet-4-5" + }, + "text": "", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "cacheWriteTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": {}, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "maxOutputTokens": 256, + "prompt": "What is the weather in Paris?", + "system": "You are a weather assistant.", + "toolChoice": "required", + "tools": { + "get_weather": { + "description": "Get the current weather for a city.", + "inputSchema": { + "additionalProperties": false, + "properties": { + "city": { + "description": "City name", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + } + } + }, + "operation": "generateText" + }, + "version": "4.0.71" +} diff --git a/fixtures/aisdk/google/embedding-basic.json b/fixtures/aisdk/google/embedding-basic.json new file mode 100644 index 000000000..91658315e --- /dev/null +++ b/fixtures/aisdk/google/embedding-basic.json @@ -0,0 +1,94 @@ +{ + "case": "embedding-basic", + "model": { + "modelId": "gemini-embedding-001", + "provider": "google.generative-ai" + }, + "package": "@ai-sdk/google", + "request": { + "body": { + "requests": [ + { + "content": { + "parts": [ + { + "text": "sunny day at the beach" + } + ], + "role": "user" + }, + "model": "models/gemini-embedding-001" + }, + { + "content": { + "parts": [ + { + "text": "rainy day in the city" + } + ], + "role": "user" + }, + "model": "models/gemini-embedding-001" + } + ] + }, + "headers": { + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-goog-api-key": "" + }, + "method": "POST", + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents" + }, + "response": { + "body": { + "embeddings": [ + { + "values": [ + 0.125, + -0.25, + 0.5 + ] + }, + { + "values": [ + 0.0625, + 0.375, + -0.5 + ] + } + ] + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "embeddings": [ + [ + 0.125, + -0.25, + 0.5 + ], + [ + 0.0625, + 0.375, + -0.5 + ] + ], + "usage": { + "tokens": null + } + }, + "sdk": { + "input": { + "values": [ + "sunny day at the beach", + "rainy day in the city" + ] + }, + "operation": "embedMany" + }, + "version": "4.0.87" +} diff --git a/fixtures/aisdk/google/generate-basic.json b/fixtures/aisdk/google/generate-basic.json new file mode 100644 index 000000000..4a61501eb --- /dev/null +++ b/fixtures/aisdk/google/generate-basic.json @@ -0,0 +1,113 @@ +{ + "case": "generate-basic", + "model": { + "modelId": "gemini-2.5-flash", + "provider": "google.generative-ai" + }, + "package": "@ai-sdk/google", + "request": { + "body": { + "contents": [ + { + "parts": [ + { + "text": "Say hello." + } + ], + "role": "user" + } + ], + "generationConfig": {} + }, + "headers": { + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-goog-api-key": "" + }, + "method": "POST", + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent" + }, + "response": { + "body": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Hello! How can I help you today?" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "modelVersion": "gemini-2.5-flash", + "responseId": "resp-fixture-1", + "usageMetadata": { + "candidatesTokenCount": 3, + "promptTokenCount": 4, + "totalTokenCount": 7 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "google": { + "finishMessage": null, + "groundingMetadata": null, + "promptFeedback": null, + "safetyRatings": null, + "serviceTier": null, + "urlContextMetadata": null, + "usageMetadata": { + "candidatesTokenCount": 3, + "promptTokenCount": 4, + "totalTokenCount": 7 + } + } + }, + "rawFinishReason": "STOP", + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "resp-fixture-1", + "modelId": "gemini-2.5-flash" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 4 + }, + "inputTokens": 4, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 3 + }, + "outputTokens": 3, + "totalTokens": 7 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "4.0.87" +} diff --git a/fixtures/aisdk/google/generate-provider-options.json b/fixtures/aisdk/google/generate-provider-options.json new file mode 100644 index 000000000..0f6ce01da --- /dev/null +++ b/fixtures/aisdk/google/generate-provider-options.json @@ -0,0 +1,138 @@ +{ + "case": "generate-provider-options", + "model": { + "modelId": "gemini-2.5-flash", + "provider": "google.generative-ai" + }, + "package": "@ai-sdk/google", + "request": { + "body": { + "contents": [ + { + "parts": [ + { + "text": "Say hello." + } + ], + "role": "user" + } + ], + "generationConfig": { + "thinkingConfig": { + "includeThoughts": true, + "thinkingBudget": 1024 + } + }, + "safetySettings": [ + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "threshold": "BLOCK_ONLY_HIGH" + } + ] + }, + "headers": { + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-goog-api-key": "" + }, + "method": "POST", + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent" + }, + "response": { + "body": { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Hello! How can I help you today?" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "modelVersion": "gemini-2.5-flash", + "responseId": "resp-fixture-1", + "usageMetadata": { + "candidatesTokenCount": 3, + "promptTokenCount": 4, + "totalTokenCount": 7 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "google": { + "finishMessage": null, + "groundingMetadata": null, + "promptFeedback": null, + "safetyRatings": null, + "serviceTier": null, + "urlContextMetadata": null, + "usageMetadata": { + "candidatesTokenCount": 3, + "promptTokenCount": 4, + "totalTokenCount": 7 + } + } + }, + "rawFinishReason": "STOP", + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "resp-fixture-1", + "modelId": "gemini-2.5-flash" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 4 + }, + "inputTokens": 4, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 3 + }, + "outputTokens": 3, + "totalTokens": 7 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello.", + "providerOptions": { + "google": { + "safetySettings": [ + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "threshold": "BLOCK_ONLY_HIGH" + } + ], + "thinkingConfig": { + "includeThoughts": true, + "thinkingBudget": 1024 + } + } + } + }, + "operation": "generateText" + }, + "version": "4.0.87" +} diff --git a/fixtures/aisdk/google/generate-stream-basic.json b/fixtures/aisdk/google/generate-stream-basic.json new file mode 100644 index 000000000..cb0113eae --- /dev/null +++ b/fixtures/aisdk/google/generate-stream-basic.json @@ -0,0 +1,172 @@ +{ + "case": "generate-stream-basic", + "model": { + "modelId": "gemini-2.5-flash", + "provider": "google.generative-ai" + }, + "package": "@ai-sdk/google", + "request": { + "body": { + "contents": [ + { + "parts": [ + { + "text": "Say hello." + } + ], + "role": "user" + } + ], + "generationConfig": {} + }, + "headers": { + "content-type": "application/json", + "user-agent": "ai-sdk-google/4.0.87 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-goog-api-key": "" + }, + "method": "POST", + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse" + }, + "response": { + "body": "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"Hello\"}]},\"finishReason\":\"STOP\",\"index\":0}],\"modelVersion\":\"gemini-2.5-flash\",\"responseId\":\"resp-fixture-2\"}\n\ndata: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\" world\"}]},\"finishReason\":\"STOP\",\"index\":0}],\"usageMetadata\":{\"promptTokenCount\":4,\"candidatesTokenCount\":2,\"totalTokenCount\":6},\"modelVersion\":\"gemini-2.5-flash\",\"responseId\":\"resp-fixture-2\"}\n\n", + "headers": { + "cache-control": "no-cache", + "content-type": "text/event-stream" + }, + "status": 200 + }, + "result": { + "finishReason": "stop", + "parts": [ + { + "type": "start" + }, + { + "request": {}, + "type": "start-step", + "warnings": [] + }, + { + "id": "0", + "type": "text-start" + }, + { + "id": "0", + "text": "Hello", + "type": "text-delta" + }, + { + "id": "0", + "text": " world", + "type": "text-delta" + }, + { + "id": "0", + "type": "text-end" + }, + { + "finishReason": "stop", + "performance": { + "effectiveOutputTokensPerSecond": 0, + "effectiveTotalTokensPerSecond": 0, + "inputTokensPerSecond": 0, + "outputTokensPerSecond": 0, + "responseTimeMs": 0, + "stepTimeMs": 0, + "timeBetweenOutputChunksMs": { + "avg": 0, + "max": 0, + "median": 0, + "min": 0, + "p10": 0, + "p90": 0 + }, + "timeToFirstOutputMs": 0, + "toolExecutionMs": {} + }, + "providerMetadata": { + "google": { + "finishMessage": null, + "groundingMetadata": null, + "promptFeedback": null, + "safetyRatings": null, + "serviceTier": null, + "urlContextMetadata": null, + "usageMetadata": { + "candidatesTokenCount": 2, + "promptTokenCount": 4, + "totalTokenCount": 6 + } + } + }, + "rawFinishReason": "STOP", + "response": { + "headers": { + "cache-control": "no-cache", + "content-type": "text/event-stream" + }, + "id": "resp-fixture-2", + "modelId": "gemini-2.5-flash", + "timestamp": "2025-01-01T00:00:00.000Z" + }, + "type": "finish-step", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 4 + }, + "inputTokens": 4, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 2 + }, + "outputTokens": 2, + "raw": { + "candidatesTokenCount": 2, + "promptTokenCount": 4, + "totalTokenCount": 6 + }, + "totalTokens": 6 + } + }, + { + "finishReason": "stop", + "rawFinishReason": "STOP", + "totalUsage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 4 + }, + "inputTokens": 4, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 2 + }, + "outputTokens": 2, + "totalTokens": 6 + }, + "type": "finish" + } + ], + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 4 + }, + "inputTokens": 4, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 2 + }, + "outputTokens": 2, + "totalTokens": 6 + } + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "streamText" + }, + "version": "4.0.87" +} diff --git a/fixtures/aisdk/google/generate-tools.json b/fixtures/aisdk/google/generate-tools.json new file mode 100644 index 000000000..a00a5a0dd --- /dev/null +++ b/fixtures/aisdk/google/generate-tools.json @@ -0,0 +1,169 @@ +{ + "case": "generate-tools", + "model": { + "modelId": "gemini-2.5-flash", + "provider": "google.generative-ai" + }, + "package": "@ai-sdk/google", + "request": { + "body": { + "contents": [ + { + "parts": [ + { + "text": "What is the weather in Paris?" + } + ], + "role": "user" + } + ], + "generationConfig": {}, + "toolConfig": { + "functionCallingConfig": { + "mode": "ANY" + } + }, + "tools": [ + { + "functionDeclarations": [ + { + "description": "Get the current weather for a city.", + "name": "get_weather", + "parametersJsonSchema": { + "additionalProperties": false, + "properties": { + "city": { + "description": "City name", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + } + ] + } + ] + }, + "headers": { + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-goog-api-key": "" + }, + "method": "POST", + "url": "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent" + }, + "response": { + "body": { + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "args": { + "city": "Paris" + }, + "name": "get_weather" + } + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "modelVersion": "gemini-2.5-flash", + "responseId": "resp-fixture-1", + "usageMetadata": { + "candidatesTokenCount": 3, + "promptTokenCount": 4, + "totalTokenCount": 7 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "input": { + "city": "Paris" + }, + "toolCallId": "5WtKk8A3IuO3vm8Q", + "toolName": "get_weather", + "type": "tool-call" + } + ], + "finishReason": "tool-calls", + "providerMetadata": { + "google": { + "finishMessage": null, + "groundingMetadata": null, + "promptFeedback": null, + "safetyRatings": null, + "serviceTier": null, + "urlContextMetadata": null, + "usageMetadata": { + "candidatesTokenCount": 3, + "promptTokenCount": 4, + "totalTokenCount": 7 + } + } + }, + "rawFinishReason": "STOP", + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "resp-fixture-1", + "modelId": "gemini-2.5-flash" + }, + "text": "", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 4 + }, + "inputTokens": 4, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 3 + }, + "outputTokens": 3, + "totalTokens": 7 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "What is the weather in Paris?", + "toolChoice": "required", + "tools": { + "get_weather": { + "description": "Get the current weather for a city.", + "inputSchema": { + "additionalProperties": false, + "properties": { + "city": { + "description": "City name", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + } + } + }, + "operation": "generateText" + }, + "version": "4.0.87" +} diff --git a/fixtures/aisdk/openai-compatible/chat-basic.json b/fixtures/aisdk/openai-compatible/chat-basic.json new file mode 100644 index 000000000..7e4abaf00 --- /dev/null +++ b/fixtures/aisdk/openai-compatible/chat-basic.json @@ -0,0 +1,95 @@ +{ + "case": "chat-basic", + "model": { + "modelId": "llama-3.3-70b-versatile", + "provider": "groq.chat" + }, + "package": "@ai-sdk/openai-compatible", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "llama-3.3-70b-versatile" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.groq.com/openai/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "llama-3.3-70b-versatile", + "object": "chat.completion", + "usage": { + "completion_tokens": 9, + "prompt_tokens": 12, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "groq": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "chatcmpl-fixture-1", + "modelId": "llama-3.3-70b-versatile" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "3.0.62" +} diff --git a/fixtures/aisdk/openai-compatible/chat-no-api-key.json b/fixtures/aisdk/openai-compatible/chat-no-api-key.json new file mode 100644 index 000000000..d747d9337 --- /dev/null +++ b/fixtures/aisdk/openai-compatible/chat-no-api-key.json @@ -0,0 +1,97 @@ +{ + "case": "chat-no-api-key", + "model": { + "modelId": "local-model", + "provider": "local.chat" + }, + "observations": { + "hasAuthorizationHeader": false + }, + "package": "@ai-sdk/openai-compatible", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "local-model" + }, + "headers": { + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "http://localhost:1234/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "llama-3.3-70b-versatile", + "object": "chat.completion", + "usage": { + "completion_tokens": 9, + "prompt_tokens": 12, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "local": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "chatcmpl-fixture-1", + "modelId": "llama-3.3-70b-versatile" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "3.0.62" +} diff --git a/fixtures/aisdk/openai-compatible/chat-provider-options-generic-key-unknown.json b/fixtures/aisdk/openai-compatible/chat-provider-options-generic-key-unknown.json new file mode 100644 index 000000000..63aa1513c --- /dev/null +++ b/fixtures/aisdk/openai-compatible/chat-provider-options-generic-key-unknown.json @@ -0,0 +1,105 @@ +{ + "case": "chat-provider-options-generic-key-unknown", + "model": { + "modelId": "llama-3.3-70b-versatile", + "provider": "groq.chat" + }, + "observations": { + "bodyHasExtraField": false + }, + "package": "@ai-sdk/openai-compatible", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "llama-3.3-70b-versatile", + "user": "u1" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.groq.com/openai/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "llama-3.3-70b-versatile", + "object": "chat.completion", + "usage": { + "completion_tokens": 9, + "prompt_tokens": 12, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "groq": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "chatcmpl-fixture-1", + "modelId": "llama-3.3-70b-versatile" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello.", + "providerOptions": { + "openaiCompatible": { + "extra_field": "x", + "user": "u1" + } + } + }, + "operation": "generateText" + }, + "version": "3.0.62" +} diff --git a/fixtures/aisdk/openai-compatible/chat-provider-options-namespace.json b/fixtures/aisdk/openai-compatible/chat-provider-options-namespace.json new file mode 100644 index 000000000..565908a7a --- /dev/null +++ b/fixtures/aisdk/openai-compatible/chat-provider-options-namespace.json @@ -0,0 +1,110 @@ +{ + "case": "chat-provider-options-namespace", + "model": { + "modelId": "llama-3.3-70b-versatile", + "provider": "groq.chat" + }, + "observations": { + "bodyHasFooAtTopLevel": true, + "bodyUser": "u1", + "providerIsGroqChat": true + }, + "package": "@ai-sdk/openai-compatible", + "request": { + "body": { + "foo": "bar", + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "llama-3.3-70b-versatile", + "user": "u1" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.groq.com/openai/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "llama-3.3-70b-versatile", + "object": "chat.completion", + "usage": { + "completion_tokens": 9, + "prompt_tokens": 12, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "groq": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "chatcmpl-fixture-1", + "modelId": "llama-3.3-70b-versatile" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello.", + "providerOptions": { + "groq": { + "foo": "bar" + }, + "openaiCompatible": { + "user": "u1" + } + } + }, + "operation": "generateText" + }, + "version": "3.0.62" +} diff --git a/fixtures/aisdk/openai-compatible/chat-query-params-and-headers.json b/fixtures/aisdk/openai-compatible/chat-query-params-and-headers.json new file mode 100644 index 000000000..3b2707b21 --- /dev/null +++ b/fixtures/aisdk/openai-compatible/chat-query-params-and-headers.json @@ -0,0 +1,109 @@ +{ + "case": "chat-query-params-and-headers", + "model": { + "modelId": "llama-3.3-70b-versatile", + "provider": "groq.chat" + }, + "observations": { + "headerNamesAsSent": [ + "x-a", + "x-b", + "x-c" + ] + }, + "package": "@ai-sdk/openai-compatible", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "llama-3.3-70b-versatile" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-a": "1", + "x-b": "3", + "x-c": "4" + }, + "method": "POST", + "url": "https://api.groq.com/openai/v1/chat/completions?api-version=1" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "llama-3.3-70b-versatile", + "object": "chat.completion", + "usage": { + "completion_tokens": 9, + "prompt_tokens": 12, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "groq": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "chatcmpl-fixture-1", + "modelId": "llama-3.3-70b-versatile" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "headers": { + "x-b": "3", + "x-c": "4" + }, + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "3.0.62" +} diff --git a/fixtures/aisdk/openai-compatible/chat-stream-basic.json b/fixtures/aisdk/openai-compatible/chat-stream-basic.json new file mode 100644 index 000000000..9973beed4 --- /dev/null +++ b/fixtures/aisdk/openai-compatible/chat-stream-basic.json @@ -0,0 +1,157 @@ +{ + "case": "chat-stream-basic", + "model": { + "modelId": "llama-3.3-70b-versatile", + "provider": "groq.chat" + }, + "package": "@ai-sdk/openai-compatible", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "llama-3.3-70b-versatile", + "stream": true + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai-sdk-openai-compatible/3.0.62 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.groq.com/openai/v1/chat/completions" + }, + "response": { + "body": "data: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"llama-3.3-70b-versatile\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"llama-3.3-70b-versatile\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"llama-3.3-70b-versatile\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"llama-3.3-70b-versatile\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"llama-3.3-70b-versatile\",\"choices\":[],\"usage\":{\"prompt_tokens\":12,\"completion_tokens\":2,\"total_tokens\":14}}\n\ndata: [DONE]\n\n", + "headers": { + "cache-control": "no-cache", + "content-type": "text/event-stream" + }, + "status": 200 + }, + "result": { + "finishReason": "stop", + "parts": [ + { + "type": "start" + }, + { + "request": {}, + "type": "start-step", + "warnings": [] + }, + { + "id": "txt-0", + "type": "text-start" + }, + { + "id": "txt-0", + "text": "Hello", + "type": "text-delta" + }, + { + "id": "txt-0", + "text": " world", + "type": "text-delta" + }, + { + "id": "txt-0", + "type": "text-end" + }, + { + "finishReason": "stop", + "performance": { + "effectiveOutputTokensPerSecond": 0, + "effectiveTotalTokensPerSecond": 0, + "inputTokensPerSecond": 0, + "outputTokensPerSecond": 0, + "responseTimeMs": 0, + "stepTimeMs": 0, + "timeBetweenOutputChunksMs": { + "avg": 0, + "max": 0, + "median": 0, + "min": 0, + "p10": 0, + "p90": 0 + }, + "timeToFirstOutputMs": 0, + "toolExecutionMs": {} + }, + "providerMetadata": { + "groq": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "cache-control": "no-cache", + "content-type": "text/event-stream" + }, + "id": "chatcmpl-fixture-2", + "modelId": "llama-3.3-70b-versatile", + "timestamp": "2025-01-01T00:00:00.000Z" + }, + "type": "finish-step", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 2 + }, + "outputTokens": 2, + "raw": { + "completion_tokens": 2, + "prompt_tokens": 12, + "total_tokens": 14 + }, + "totalTokens": 14 + } + }, + { + "finishReason": "stop", + "rawFinishReason": "stop", + "totalUsage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 2 + }, + "outputTokens": 2, + "totalTokens": 14 + }, + "type": "finish" + } + ], + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 2 + }, + "outputTokens": 2, + "totalTokens": 14 + } + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "streamText" + }, + "version": "3.0.62" +} diff --git a/fixtures/aisdk/openai-compatible/embedding-basic.json b/fixtures/aisdk/openai-compatible/embedding-basic.json new file mode 100644 index 000000000..9c54480c4 --- /dev/null +++ b/fixtures/aisdk/openai-compatible/embedding-basic.json @@ -0,0 +1,86 @@ +{ + "case": "embedding-basic", + "model": { + "modelId": "text-embedding-fixture", + "provider": "groq.embedding" + }, + "package": "@ai-sdk/openai-compatible", + "request": { + "body": { + "encoding_format": "float", + "input": [ + "sunny day at the beach", + "rainy day in the city" + ], + "model": "text-embedding-fixture" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.groq.com/openai/v1/embeddings" + }, + "response": { + "body": { + "data": [ + { + "embedding": [ + 0.125, + -0.25, + 0.5 + ], + "index": 0, + "object": "embedding" + }, + { + "embedding": [ + 0.0625, + 0.375, + -0.5 + ], + "index": 1, + "object": "embedding" + } + ], + "model": "text-embedding-fixture", + "object": "list", + "usage": { + "prompt_tokens": 10, + "total_tokens": 10 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "embeddings": [ + [ + 0.125, + -0.25, + 0.5 + ], + [ + 0.0625, + 0.375, + -0.5 + ] + ], + "usage": { + "tokens": 10 + } + }, + "sdk": { + "input": { + "values": [ + "sunny day at the beach", + "rainy day in the city" + ] + }, + "operation": "embedMany" + }, + "version": "3.0.62" +} diff --git a/fixtures/aisdk/openai/chat-basic.json b/fixtures/aisdk/openai/chat-basic.json new file mode 100644 index 000000000..bea6c44fe --- /dev/null +++ b/fixtures/aisdk/openai/chat-basic.json @@ -0,0 +1,104 @@ +{ + "case": "chat-basic", + "model": { + "modelId": "gpt-4o", + "provider": "openai.chat" + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "gpt-4o" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.openai.com/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "gpt-4o-2024-08-06", + "object": "chat.completion", + "system_fingerprint": "fp_fixture", + "usage": { + "completion_tokens": 9, + "completion_tokens_details": { + "reasoning_tokens": 0 + }, + "prompt_tokens": 12, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "openai": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "id": "chatcmpl-fixture-1", + "modelId": "gpt-4o-2024-08-06" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/chat-custom-name.json b/fixtures/aisdk/openai/chat-custom-name.json new file mode 100644 index 000000000..f1fea9733 --- /dev/null +++ b/fixtures/aisdk/openai/chat-custom-name.json @@ -0,0 +1,107 @@ +{ + "case": "chat-custom-name", + "model": { + "modelId": "gpt-4o", + "provider": "proxy.chat" + }, + "observations": { + "providerIsProxyChat": true + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "gpt-4o" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.openai.com/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "gpt-4o-2024-08-06", + "object": "chat.completion", + "system_fingerprint": "fp_fixture", + "usage": { + "completion_tokens": 9, + "completion_tokens_details": { + "reasoning_tokens": 0 + }, + "prompt_tokens": 12, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "openai": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "id": "chatcmpl-fixture-1", + "modelId": "gpt-4o-2024-08-06" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/chat-headers-merge.json b/fixtures/aisdk/openai/chat-headers-merge.json new file mode 100644 index 000000000..d24d734f6 --- /dev/null +++ b/fixtures/aisdk/openai/chat-headers-merge.json @@ -0,0 +1,119 @@ +{ + "case": "chat-headers-merge", + "model": { + "modelId": "gpt-4o", + "provider": "openai.chat" + }, + "observations": { + "headerNamesAsSent": [ + "x-a", + "x-b", + "x-c" + ], + "xB": "3" + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "gpt-4o" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22", + "x-a": "1", + "x-b": "3", + "x-c": "4" + }, + "method": "POST", + "url": "https://api.openai.com/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "gpt-4o-2024-08-06", + "object": "chat.completion", + "system_fingerprint": "fp_fixture", + "usage": { + "completion_tokens": 9, + "completion_tokens_details": { + "reasoning_tokens": 0 + }, + "prompt_tokens": 12, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "openai": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "id": "chatcmpl-fixture-1", + "modelId": "gpt-4o-2024-08-06" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "headers": { + "x-b": "3", + "x-c": "4" + }, + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/chat-provider-options.json b/fixtures/aisdk/openai/chat-provider-options.json new file mode 100644 index 000000000..719052387 --- /dev/null +++ b/fixtures/aisdk/openai/chat-provider-options.json @@ -0,0 +1,153 @@ +{ + "case": "chat-provider-options", + "model": { + "modelId": "gpt-4o", + "provider": "openai.chat" + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "messages": [ + { + "content": "What is the weather in Paris?", + "role": "user" + } + ], + "model": "gpt-4o", + "parallel_tool_calls": false, + "tool_choice": "auto", + "tools": [ + { + "function": { + "description": "Get the current weather for a city.", + "name": "get_weather", + "parameters": { + "additionalProperties": false, + "properties": { + "city": { + "description": "City name", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + }, + "type": "function" + } + ], + "user": "u1" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.openai.com/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "gpt-4o-2024-08-06", + "object": "chat.completion", + "system_fingerprint": "fp_fixture", + "usage": { + "completion_tokens": 9, + "completion_tokens_details": { + "reasoning_tokens": 0 + }, + "prompt_tokens": 12, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "openai": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "id": "chatcmpl-fixture-1", + "modelId": "gpt-4o-2024-08-06" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "What is the weather in Paris?", + "providerOptions": { + "openai": { + "parallelToolCalls": false, + "user": "u1" + } + }, + "tools": { + "get_weather": { + "description": "Get the current weather for a city.", + "inputSchema": { + "additionalProperties": false, + "properties": { + "city": { + "description": "City name", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + } + } + }, + "operation": "generateText" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/chat-stream-basic.json b/fixtures/aisdk/openai/chat-stream-basic.json new file mode 100644 index 000000000..2c927f6e3 --- /dev/null +++ b/fixtures/aisdk/openai/chat-stream-basic.json @@ -0,0 +1,160 @@ +{ + "case": "chat-stream-basic", + "model": { + "modelId": "gpt-4o", + "provider": "openai.chat" + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "gpt-4o", + "stream": true, + "stream_options": { + "include_usage": true + } + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai-sdk-openai/4.0.83 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.openai.com/v1/chat/completions" + }, + "response": { + "body": "data: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"gpt-4o-2024-08-06\",\"system_fingerprint\":\"fp_fixture\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"gpt-4o-2024-08-06\",\"system_fingerprint\":\"fp_fixture\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"gpt-4o-2024-08-06\",\"system_fingerprint\":\"fp_fixture\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n\ndata: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"gpt-4o-2024-08-06\",\"system_fingerprint\":\"fp_fixture\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: {\"id\":\"chatcmpl-fixture-2\",\"object\":\"chat.completion.chunk\",\"created\":1735689600,\"model\":\"gpt-4o-2024-08-06\",\"choices\":[],\"usage\":{\"prompt_tokens\":12,\"completion_tokens\":2,\"total_tokens\":14}}\n\ndata: [DONE]\n\n", + "headers": { + "cache-control": "no-cache", + "content-type": "text/event-stream" + }, + "status": 200 + }, + "result": { + "finishReason": "stop", + "parts": [ + { + "type": "start" + }, + { + "request": {}, + "type": "start-step", + "warnings": [] + }, + { + "id": "0", + "type": "text-start" + }, + { + "id": "0", + "text": "Hello", + "type": "text-delta" + }, + { + "id": "0", + "text": " world", + "type": "text-delta" + }, + { + "id": "0", + "type": "text-end" + }, + { + "finishReason": "stop", + "performance": { + "effectiveOutputTokensPerSecond": 0, + "effectiveTotalTokensPerSecond": 0, + "inputTokensPerSecond": 0, + "outputTokensPerSecond": 0, + "responseTimeMs": 0, + "stepTimeMs": 0, + "timeBetweenOutputChunksMs": { + "avg": 0, + "max": 0, + "median": 0, + "min": 0, + "p10": 0, + "p90": 0 + }, + "timeToFirstOutputMs": 0, + "toolExecutionMs": {} + }, + "providerMetadata": { + "openai": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "cache-control": "no-cache", + "content-type": "text/event-stream" + }, + "id": "chatcmpl-fixture-2", + "modelId": "gpt-4o-2024-08-06", + "timestamp": "2025-01-01T00:00:00.000Z" + }, + "type": "finish-step", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 2 + }, + "outputTokens": 2, + "raw": { + "completion_tokens": 2, + "prompt_tokens": 12, + "total_tokens": 14 + }, + "totalTokens": 14 + } + }, + { + "finishReason": "stop", + "rawFinishReason": "stop", + "totalUsage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 2 + }, + "outputTokens": 2, + "totalTokens": 14 + }, + "type": "finish" + } + ], + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 2 + }, + "outputTokens": 2, + "totalTokens": 14 + } + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "streamText" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/chat-system-and-settings.json b/fixtures/aisdk/openai/chat-system-and-settings.json new file mode 100644 index 000000000..1031365da --- /dev/null +++ b/fixtures/aisdk/openai/chat-system-and-settings.json @@ -0,0 +1,130 @@ +{ + "case": "chat-system-and-settings", + "model": { + "modelId": "gpt-4o", + "provider": "openai.chat" + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "max_tokens": 64, + "messages": [ + { + "content": "You are a terse assistant.", + "role": "system" + }, + { + "content": "Name a color.", + "role": "user" + } + ], + "model": "gpt-4o", + "seed": 42, + "stop": [ + "END", + "###" + ], + "temperature": 0.3, + "top_p": 0.9 + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.openai.com/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Blue.", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "gpt-4o-2024-08-06", + "object": "chat.completion", + "system_fingerprint": "fp_fixture", + "usage": { + "completion_tokens": 9, + "completion_tokens_details": { + "reasoning_tokens": 0 + }, + "prompt_tokens": 12, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Blue.", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "openai": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "id": "chatcmpl-fixture-1", + "modelId": "gpt-4o-2024-08-06" + }, + "text": "Blue.", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "maxOutputTokens": 64, + "messages": [ + { + "content": "Name a color.", + "role": "user" + } + ], + "seed": 42, + "stopSequences": [ + "END", + "###" + ], + "system": "You are a terse assistant.", + "temperature": 0.3, + "topP": 0.9 + }, + "operation": "generateText" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/chat-tools.json b/fixtures/aisdk/openai/chat-tools.json new file mode 100644 index 000000000..9a4c0cd4a --- /dev/null +++ b/fixtures/aisdk/openai/chat-tools.json @@ -0,0 +1,158 @@ +{ + "case": "chat-tools", + "model": { + "modelId": "gpt-4o", + "provider": "openai.chat" + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "messages": [ + { + "content": "What is the weather in Paris?", + "role": "user" + } + ], + "model": "gpt-4o", + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get the current weather for a city.", + "name": "get_weather", + "parameters": { + "additionalProperties": false, + "properties": { + "city": { + "description": "City name", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + }, + "type": "function" + } + ] + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.openai.com/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": null, + "role": "assistant", + "tool_calls": [ + { + "function": { + "arguments": "{\"city\":\"Paris\"}", + "name": "get_weather" + }, + "id": "call_fixture_1", + "type": "function" + } + ] + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "gpt-4o-2024-08-06", + "object": "chat.completion", + "system_fingerprint": "fp_fixture", + "usage": { + "completion_tokens": 9, + "completion_tokens_details": { + "reasoning_tokens": 0 + }, + "prompt_tokens": 12, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "input": { + "city": "Paris" + }, + "toolCallId": "call_fixture_1", + "toolName": "get_weather", + "type": "tool-call" + } + ], + "finishReason": "tool-calls", + "providerMetadata": { + "openai": {} + }, + "rawFinishReason": "tool_calls", + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "chatcmpl-fixture-1", + "modelId": "gpt-4o-2024-08-06" + }, + "text": "", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "What is the weather in Paris?", + "toolChoice": "required", + "tools": { + "get_weather": { + "description": "Get the current weather for a city.", + "inputSchema": { + "additionalProperties": false, + "properties": { + "city": { + "description": "City name", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + } + } + }, + "operation": "generateText" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/embedding-basic.json b/fixtures/aisdk/openai/embedding-basic.json new file mode 100644 index 000000000..344b54a68 --- /dev/null +++ b/fixtures/aisdk/openai/embedding-basic.json @@ -0,0 +1,86 @@ +{ + "case": "embedding-basic", + "model": { + "modelId": "text-embedding-3-small", + "provider": "openai.embedding" + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "encoding_format": "float", + "input": [ + "sunny day at the beach", + "rainy day in the city" + ], + "model": "text-embedding-3-small" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.openai.com/v1/embeddings" + }, + "response": { + "body": { + "data": [ + { + "embedding": [ + 0.125, + -0.25, + 0.5 + ], + "index": 0, + "object": "embedding" + }, + { + "embedding": [ + 0.0625, + 0.375, + -0.5 + ], + "index": 1, + "object": "embedding" + } + ], + "model": "text-embedding-3-small", + "object": "list", + "usage": { + "prompt_tokens": 10, + "total_tokens": 10 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "embeddings": [ + [ + 0.125, + -0.25, + 0.5 + ], + [ + 0.0625, + 0.375, + -0.5 + ] + ], + "usage": { + "tokens": 10 + } + }, + "sdk": { + "input": { + "values": [ + "sunny day at the beach", + "rainy day in the city" + ] + }, + "operation": "embedMany" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/empty-api-key.json b/fixtures/aisdk/openai/empty-api-key.json new file mode 100644 index 000000000..0dac40642 --- /dev/null +++ b/fixtures/aisdk/openai/empty-api-key.json @@ -0,0 +1,109 @@ +{ + "case": "empty-api-key", + "model": { + "modelId": "gpt-4o", + "provider": "openai.chat" + }, + "observations": { + "authorizationIsEmptyBearer": true, + "authorizationRawLength": 6, + "usedEnvKey": false + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "messages": [ + { + "content": "Say hello.", + "role": "user" + } + ], + "model": "gpt-4o" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.openai.com/v1/chat/completions" + }, + "response": { + "body": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "Hello! How can I help you today?", + "role": "assistant" + } + } + ], + "created": 1735689600, + "id": "chatcmpl-fixture-1", + "model": "gpt-4o-2024-08-06", + "object": "chat.completion", + "system_fingerprint": "fp_fixture", + "usage": { + "completion_tokens": 9, + "completion_tokens_details": { + "reasoning_tokens": 0 + }, + "prompt_tokens": 12, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "status": 200 + }, + "result": { + "content": [ + { + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "openai": {} + }, + "rawFinishReason": "stop", + "response": { + "headers": { + "content-type": "application/json", + "x-request-id": "req_fixture_1" + }, + "id": "chatcmpl-fixture-1", + "modelId": "gpt-4o-2024-08-06" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/missing-api-key.json b/fixtures/aisdk/openai/missing-api-key.json new file mode 100644 index 000000000..b4b9dcece --- /dev/null +++ b/fixtures/aisdk/openai/missing-api-key.json @@ -0,0 +1,24 @@ +{ + "case": "missing-api-key", + "error": { + "message": "OpenAI API key is missing. Pass it using the 'apiKey' parameter or the OPENAI_API_KEY environment variable.", + "name": "AI_LoadAPIKeyError" + }, + "model": { + "modelId": "gpt-4o", + "provider": "openai.chat" + }, + "observations": { + "requestsSent": 0 + }, + "package": "@ai-sdk/openai", + "request": null, + "response": null, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "4.0.83" +} diff --git a/fixtures/aisdk/openai/responses-basic.json b/fixtures/aisdk/openai/responses-basic.json new file mode 100644 index 000000000..242183ce7 --- /dev/null +++ b/fixtures/aisdk/openai/responses-basic.json @@ -0,0 +1,120 @@ +{ + "case": "responses-basic", + "model": { + "modelId": "gpt-4o", + "provider": "openai.responses" + }, + "package": "@ai-sdk/openai", + "request": { + "body": { + "input": [ + { + "content": [ + { + "text": "Say hello.", + "type": "input_text" + } + ], + "role": "user" + } + ], + "model": "gpt-4o" + }, + "headers": { + "authorization": "", + "content-type": "application/json", + "user-agent": "ai/7.0.127 ai-sdk-provider-utils/5.0.53 node.js/22" + }, + "method": "POST", + "url": "https://api.openai.com/v1/responses" + }, + "response": { + "body": { + "created_at": 1735689600, + "error": null, + "id": "resp_fixture_1", + "incomplete_details": null, + "model": "gpt-4o-2024-08-06", + "object": "response", + "output": [ + { + "content": [ + { + "annotations": [], + "text": "Hello! How can I help you today?", + "type": "output_text" + } + ], + "id": "msg_fixture_1", + "role": "assistant", + "status": "completed", + "type": "message" + } + ], + "status": "completed", + "usage": { + "input_tokens": 12, + "input_tokens_details": { + "cached_tokens": 0 + }, + "output_tokens": 9, + "output_tokens_details": { + "reasoning_tokens": 0 + }, + "total_tokens": 21 + } + }, + "headers": { + "content-type": "application/json" + }, + "status": 200 + }, + "result": { + "content": [ + { + "providerMetadata": { + "openai": { + "itemId": "msg_fixture_1" + } + }, + "text": "Hello! How can I help you today?", + "type": "text" + } + ], + "finishReason": "stop", + "providerMetadata": { + "openai": { + "responseId": "resp_fixture_1" + } + }, + "response": { + "headers": { + "content-type": "application/json" + }, + "id": "resp_fixture_1", + "modelId": "gpt-4o-2024-08-06" + }, + "text": "Hello! How can I help you today?", + "usage": { + "inputTokenDetails": { + "cacheReadTokens": 0, + "noCacheTokens": 12 + }, + "inputTokens": 12, + "outputTokenDetails": { + "reasoningTokens": 0, + "textTokens": 9 + }, + "outputTokens": 9, + "totalTokens": 21 + }, + "warnings": [] + }, + "sdk": { + "input": { + "prompt": "Say hello." + }, + "operation": "generateText" + }, + "version": "4.0.83" +} diff --git a/rfc/0002-provider-improvements.md b/rfc/0002-provider-improvements.md index defb75dcb..eabfc0b04 100644 --- a/rfc/0002-provider-improvements.md +++ b/rfc/0002-provider-improvements.md @@ -1,6 +1,7 @@ # Provider Adapter Layer Improvements > **Status**: IMPLEMENTED (2026-07-28 — OpenAICompatProfile 差异描述 + 145 个 thin wrapper 集成,见下方 Implementation Progress;"配置即数据"的后续演进为 [RFC-0017](0017-provider-config-dx.md) 的 provider registry;cassette 测试独立为 [RFC-0003](0003-test-cassette.md)) +> **修订**: 目标设计已由 [aisdk-architecture-alignment.md](../docs/aisdk-architecture-alignment.md)(D12、D20、§3.3)改写:`OpenAICompatProfile` 与 thin wrapper 由 preset 描述数据加 openai-compatible 包的 `create_*` 工厂取代。对应的 V4 分层 PR 合入主线之前,本文仍描述主线现有代码。 ## Current Problems diff --git a/rfc/0006-provider-development.md b/rfc/0006-provider-development.md index c0e19d55b..73ca33809 100644 --- a/rfc/0006-provider-development.md +++ b/rfc/0006-provider-development.md @@ -1,6 +1,7 @@ # RFC-0006: Provider Development Specification > **Status**: ACCEPTED (in force — required reading in CONTRIBUTING.md and applied by provider-research batches; 2026-07-28) +> **Revised by**: [aisdk-architecture-alignment.md](../docs/aisdk-architecture-alignment.md) (D12, D20, §3.3, §4.1). In the target design each vendor is a package with a `create_*` factory and presets are descriptor rows reusing the openai-compatible package; the thin-wrapper and profile steps below describe the code on master until the corresponding V4 layer merges. > **Date**: 2026-07-28 > **Scope**: Adding or redoing provider adapters in `aimux-providers` > **Related**: [Archived Provider Inventory](../provider-inventory/README.md), [Provider Adapter Layer Improvements](0002-provider-improvements.md), [Cassette Test Plan](0003-test-cassette.md), [Protocol Conversion and Adapter Layer Design](0005-protocol-conversion.md) diff --git a/rfc/0015-cache-trace-audit.md b/rfc/0015-cache-trace-audit.md index af48a66e3..3bcdc2321 100644 --- a/rfc/0015-cache-trace-audit.md +++ b/rfc/0015-cache-trace-audit.md @@ -1,6 +1,7 @@ # RFC-0015: 缓存命中探测(cache-hit probe) > **Status**: IMPLEMENTED (2026-08-05 — 探测本身进 core:TraceLayer + 判定引擎 + RingTraceStore + 查询 API + FFI/Node/Python/Go 透传;探测业务归 CLI(RFC-0025),告警外部消费) +> **修订**: V4 层的 model 不再提供 `config_snapshot()`(base_url、key 指纹),作用域键改为 `H(scope 名 ‖ model)`,scope 名默认取 `model.provider()`:同一 provider 名下的不同端点不隔离,由宿主用 `TraceLayer::with_scope_name` 给每个端点各起一个名字(aimux-web 对自定义 `base_url` 即如此)。对应的 V4 分层 PR 合入主线之前,§6 的 `ScopeKey` 公式仍描述主线现有代码。 > **Date**: 2026-08-01(初稿)/ 2026-08-05(重写) > **Scope**: aimux 统一 LLM 访问层上的可选缓存命中探测能力——采集 request_body 指纹 + usage 快照,客户端对连续 agent 调用的原始请求体做前缀对比,判定各 provider 服务端上报的 cache 命中率是否掺水,暴露 verdict 数据与查询接口供消费。**探测本身进 core,探测业务(CLI)独立,告警外部消费。** > **Related**: [RFC-0014](0014-logging.md) 统一日志(挂载其 span 树)、[RFC-0023](0023-runtime-request-recording.md) 录制与回放(明文兄弟子系统,call_id 关联)、[RFC-0024](0024-session-aggregation.md) 会话聚合(session_id 统一来源)、[RFC-0009](0009-request-resilience.md) retry/超时(重试语义影响判定规则)、研究存档 [cache-tracing](../docs/internal/cache-tracing/00-research-plan.md) diff --git a/rfc/0016-align-with-aisdk.md b/rfc/0016-align-with-aisdk.md index 89acb6c7c..9764fe490 100644 --- a/rfc/0016-align-with-aisdk.md +++ b/rfc/0016-align-with-aisdk.md @@ -1,6 +1,7 @@ # RFC-0016: 对齐 Vercel AI SDK 能力缺口 > **Status**: IMPLEMENTED (2026-09-24 核定) — §2 缺口绝大多数落地:H1/H2/H3、M1-M7、M9-M12 经 #99(2026-08-13,M6 proxy/M7 聚合/M11 流聚合/M12 generateObject)、#164(2026-09-12,请求管线对齐:重试分层/错误域/超时)、#165(工具输入解析边界)落地;M13 经核不做、H4 明确不做;工具调用修复由 [RFC-0035](0035-host-side-tool-call-repair.md) 扩展。未落地余项(M8 StreamPart 变体等)逐条追踪见 [§7.2](#72-未落地清单逐条追踪) +> **修订**: L2(tool-call input 维持 `Value`)由 [aisdk-architecture-alignment.md](../docs/aisdk-architecture-alignment.md) 第二部分取代:V4 层 input 为原文 String,解析在 core 层。L7 中“metadataExtractor 已被 usage.raw 超出”不成立:上游 metadataExtractor 从整个响应和流中提取 providerMetadata,不限于 usage(openai-compatible 3.0.62 `chat/openai-compatible-chat-language-model.ts:408`、`:458-459`);该文 D15 要求 openai-compatible 保留上游设置。 > **Date**: 2026-08-01 > **Scope**: 系统对比 aimux 0.1.2 与 Vercel AI SDK (`@ai-sdk/openai` / `@ai-sdk/openai-compatible` / `@ai-sdk/provider` V4) 的接口与实现,识别能力缺口,并按优先级规划补齐路径 > **Related**: [RFC-0009](0009-request-resilience.md) request resilience(retry/timeout,本 RFC 的 abortSignal/timeout 缺口与之相关),[RFC-0014](0014-logging.md) 统一日志体系(可观测性缺口依赖本 RFC 的 span 树) diff --git a/rfc/0017-provider-config-dx.md b/rfc/0017-provider-config-dx.md index 2d5e499a5..b79be3724 100644 --- a/rfc/0017-provider-config-dx.md +++ b/rfc/0017-provider-config-dx.md @@ -1,6 +1,7 @@ # RFC-0017: 统一 provider 配置与 request override(DX 提升) > **Status**: IMPLEMENTED (2026-08-02 — 阶段 1-4 全部完成,详见 [§3](#3-实现计划) 实施状态与 [§4](#4-向后兼容);全量测试 2769 绿) +> **修订**: 通用 body override 路线已由 [aisdk-architecture-alignment.md](../docs/aisdk-architecture-alignment.md)(D15,§9.4)撤销:openai-compatible 保留 providerOptions 透传与 `transformRequestBody`,原生包不提供统一的请求体改写。对应的 V4 分层 PR 合入主线之前,本文仍描述主线现有代码。 > **Date**: 2026-08-01 > **Scope**: 统一 aimux 的 provider 配置层——把 Rust core 已有但 Node 未透出的能力(headers/org/project/retry)暴露出来,引入通用 `bodyOverrides`(JSON deep merge)替代闭包式 transform,并以**配置数据(而非代码枚举)**驱动 Qwen/Kimi/GLM 等厂商的思考开关翻译 > **Related**: [RFC-0016](0016-align-with-aisdk.md) 对齐 Vercel AI SDK 能力缺口(本 RFC 补齐 M1/M5 + 思考开关翻译层),[RFC-0009](0009-request-resilience.md) retry(本 RFC 透出 maxRetries) diff --git a/rfc/0018-codex-subscription.md b/rfc/0018-codex-subscription.md index a30703656..f894e377d 100644 --- a/rfc/0018-codex-subscription.md +++ b/rfc/0018-codex-subscription.md @@ -1,6 +1,7 @@ # RFC-0018: Codex 订阅通道 provider 集成评估 > **Status**: IMPLEMENTED (2026-08-03 核验→定案→实现完成:Path A + Path B + codex_refresh + C ABI,测试全绿) +> **修订**: 无状态刷新边界保留,见 [aisdk-architecture-alignment.md](../docs/aisdk-architecture-alignment.md) §4.1 的 Codex 行与 §9.3「状态操作」行:`codex_refresh` 仍是一次 `/oauth/token` 调用的纯函数,401 映射 `TokenExpired`,库内不自动重试旋转 refresh token;刷新后由宿主持久化新 token,并用它重建 provider 再重试。 > **Date**: 2026-08-02 > **Scope**: 评估并设计 Codex(OpenAI 编程模型/订阅通道)在 aimux 的集成路径。 > 本 RFC 基于 104 项目参考审计的发现,先给出动机与现状,再列出**必须核验的协议事实**, diff --git a/rfc/0020-external-provider-config.md b/rfc/0020-external-provider-config.md index 5a2397a97..da436bffd 100644 --- a/rfc/0020-external-provider-config.md +++ b/rfc/0020-external-provider-config.md @@ -1,6 +1,7 @@ # RFC-0020: 外部 Provider 配置 > **Status**: IMPLEMENTED (运行时 overlay 落地于 #97(2026-08-12)— `provider::register_provider` 运行时注册/覆盖条目,查找顺序 overlay → 内置 registry;FFI `aimux_register_providers` 及 Node/Python 透传) +> **修订**: 全局 overlay 已由 [aisdk-architecture-alignment.md](../docs/aisdk-architecture-alignment.md)(§3.3、§3.4 `providers_from_config`)撤销:外部配置改为纯组装函数,不修改全局 registry。删除 overlay 的那层 PR 合入时,本文状态改为 SUPERSEDED。 > **Date**: 2026-08-05 > **Scope**: `aimux-providers` 新增运行时覆盖层,允许外部(配置文件 / 编程式 API)注册、覆盖 OpenAI 兼容 provider 条目;各 binding 薄透传 > **Related**: [RFC-0017](0017-provider-config-dx.md) 配置 DX(本 RFC 落地其 §363 预留的"用户覆盖内置条目"语义)、[RFC-0019](0019-session-affinity.md) 会话亲和、[调研报告](../docs/external-provider-config-research.md) diff --git a/rfc/0023-runtime-request-recording.md b/rfc/0023-runtime-request-recording.md index 8bd7336dd..e3eed3be6 100644 --- a/rfc/0023-runtime-request-recording.md +++ b/rfc/0023-runtime-request-recording.md @@ -1,6 +1,7 @@ # RFC-0023: 调用上下文录制与回放 > **Status**: IMPLEMENTED (P1-P6 + config_snapshot 已实施,2026-08-10;实施进度见 [docs/plan/rfc0023-recording.md](../docs/plan/rfc0023-recording.md)) +> **修订**: 录制格式与回放由 [aisdk-architecture-alignment.md](../docs/aisdk-architecture-alignment.md)(D24、§4.3、§4.4)改写:迁移各层仍写 `schema: 2`(完整的 schema 3 排在最后几层),读取 schema 1 与 2,schema 3 及以上给出明确错误;`ProviderRecord` 只存身份 `{provider_id, provider, model_id}`,不再有 `config_snapshot`。读主线(aimux 0.5 及更早)写的录制时,只有主线写的字段(`base_url`、`api_key_source`、`profile`、`provider_options`)被忽略,只有新层写的字段(`provider_id`、`input.decision_capabilities`)可缺省,本版本读不了的 prompt 读成空;回放不了的部分(无 `provider_id` 不能从 registry 重建、旧格式 prompt 与 options、decision 能力只在 `provider.profile` 中)在回放时报错说明,不在读取时报错。live replay 用宿主 registry(`default_providers()`,给了 key 时用 `create_provider` 重建录制的 provider)重建,decision 录制同样适用于任何有 decision model 的 provider,不只 Jev。另有两条规则:调用丢了录制事件(JSONL 队列已满或 writer 已退出)或被标为不一致(重复 attempt、`ExchangeUpdate` 找不到对应 exchange)时写出 `complete: false`,JSONL 与 ring 两种 recorder 相同,不论由该调用的最后一个事件还是 flush 写出;§4 的敏感头还包括 `x-prodia-key` 与以 `subscription-key` 结尾的头(如 Azure 的 `Ocp-Apim-Subscription-Key`),不区分大小写与分隔符。对应的 V4 分层 PR 合入主线之前,本文仍描述主线现有代码。 > **Date**: 2026-08-05(重写:从"仅 HTTP 录制"扩展为"三层完整上下文录制 + 两种回放模式") > **Scope**: `aimux-core` + `aimux-provider-utils` 新增可选的完整调用上下文录制(输入侧 + 配置侧 + HTTP 侧)与两种回放模式(请求回放 / mock 响应回放),opt-in,默认关闭 > **Related**: [RFC-0003](0003-test-cassette.md) 测试 cassette(格式与匹配基础)、[RFC-0014](0014-logging.md) 统一日志(零成本门控范本)、[RFC-0015](0015-cache-trace-audit.md) 缓存审计(sink 抽象兄弟)、[RFC-0020](0020-external-provider-config.md) 外部 provider 配置(配置侧录制关联) diff --git a/rfc/0029-web-console.md b/rfc/0029-web-console.md index 1f47827a3..d7f229c92 100644 --- a/rfc/0029-web-console.md +++ b/rfc/0029-web-console.md @@ -227,7 +227,7 @@ POST /api/calls 2. **Settings 已存该 provider 的 key**(KeyStore,含启动时从磁盘加载的); 3. `None` —— provider 读自己注册的默认 env var。 - 请求自带 `base_url`(Replay 时为录制里的地址)且不是该 provider 的默认地址时,只用请求内直接给的 key(明文,仅回环):第 2、3 级不适用,`env:` 引用被拒绝,不给 key 则以空 key 发出(与 AI SDK 显式传 `apiKey: ''` 相同,不回落到 Settings key 或环境变量;OpenAI 等 provider 因此发一个空凭据头,如 `Authorization: Bearer`)。否则任何客户端都能把服务端的 key 发往自己指定的主机。 + 请求自带 `base_url` 且不是该 provider 的默认地址时,只用请求内直接给的 key(明文,仅回环):第 2、3 级不适用,`env:` 引用被拒绝,不给 key 则以空 key 发出(与 AI SDK 显式传 `apiKey: ''` 相同,不回落到 Settings key 或环境变量;部分 provider 因此发一个空凭据头,如 OpenAI 的 `Authorization: Bearer`)。除 key 外还从服务端环境读凭据的 provider(AWS 凭据、Vertex access token、Codex 登录、DataForSEO 账号)不接受自定义 `base_url`。Replay 按录制的 provider 身份重建,总是发往该 provider 的默认地址(录制不存地址)。否则任何客户端都能把服务端的 key 发往自己指定的主机。 **威胁模型**:信任边界 = 本机回环客户端。回环下浏览器 ↔ 后端信道视为可信,故允许网页内明文录入与请求内明文直传;非回环(如 `--host 0.0.0.0` LAN 共享)下任何局域网客户端都能 POST,故全部明文入口关闭。明文 key 永不进日志、永不回传前端、永不进 Recording(录制层脱敏复用 RFC-0023 规则)。 diff --git a/rfc/0037-decision-api.md b/rfc/0037-decision-api.md index a15461d79..3c6d0e118 100644 --- a/rfc/0037-decision-api.md +++ b/rfc/0037-decision-api.md @@ -117,12 +117,12 @@ server 验证；另于 2026-10-06 使用 Python binding 对官方 endpoint 完 CI 无需线上 key。 `decide` 复用统一 JSONL/ring recorder，输入标记为 decision（旧记录默认 -language model），state/questions 存入 input.options，能力快照存入 provider.profile， +language model），state/questions 存入 input.options，能力声明存入 input.decision_capabilities，provider 仅保留身份， 规范化结果存入 outcome.decision_result。每次 retry 分配独立 attempt， 录制器快照跨 HTTP、取消和收尾保持一致。覆盖成功、错误、超时、取消、 脱敏、无重复占位记录、严格匹配和 recorder 替换。 `aimux-replay --mock` 支持标准 JSONL 的离线 decision 回放；不带 --mock -时可重建官方 Jev 执行请求回放。 +时通过调用方 registry 配置重建官方 Jev 执行请求回放。 官方字段覆盖、Provider 精度规则、全语言 capability 查询和宿主包装均已实现。 其他官方决策接口按 [RFC-0038](0038-decision-provider-expansion.md) 接入。 diff --git a/rfc/0040-provider-auth.md b/rfc/0040-provider-auth.md index d14f819ee..675733554 100644 --- a/rfc/0040-provider-auth.md +++ b/rfc/0040-provider-auth.md @@ -2,6 +2,8 @@ > **状态：Draft，待评审；本文不表示实现完成或维护者已接受。** > +> **修订（2026-10-15，V4 [5/8] #253）：** §3.2 第三段关于 redirect 的规则已按实现改写，见该段后的修订说明；“默认不跟随 redirect”不再适用于模型 API 请求。 +> > **范围：** ROADMAP 的 B1–B9、#174 / #175、L0 native passthrough；区分架构切换必需工作与新增公开能力。 > > **核对基线：** master `72a37b5`；目标架构参考尚未合并的 [PR #200](https://github.com/arcships/aimux/pull/200) head `d8c3d15`。#200 的 §0.4 已记录四项决定，不能说没有记录；PR 仍开放且文档一致性 review 尚待处理，不能据此宣称已在 master 生效。 @@ -74,6 +76,8 @@ header 名大小写不敏感，以覆盖插入避免重复认证头，删除项 禁用携密认证请求的自动跨 origin redirect；默认不跟随 redirect。批准的同 origin 重定向也必须重新做目标校验与签名，不能复用旧签名。媒体下载/外部 URL 不自动继承 provider headers。TLS/代理属于 transport 实例，不能由单次 model 参数偷偷改写。 +> **修订（2026-10-15，V4 [5/8] #253）：** 上段前两句按实现改为：模型 API 请求与上游 `postToApi` 一致，默认跟随 redirect（`@ai-sdk/provider-utils@5.0.55` `post-to-api.ts:98-107` 调 fetch 时不传 `redirect`，取 fetch 默认的 `follow`）。凭证锁定在请求的 origin：某一跳离开该 origin 后只保留 `User-Agent`，其余请求头全部丢弃，transport 收到 `RequestCredentials::Omit`，签名与 token 装饰器不再添加凭证；之后即使跳回原 origin 也不恢复。每一跳都重新经过 transport，签名按实际发送的 URL 重新计算，不复用旧签名。跳转目标只接受 http(s)，最多 10 跳；303，以及 POST 上的 301/302，改为不带 body 的 GET。模型 API 请求的跳转目标不做下载那样的地址校验；响应给出的 URL 与媒体下载仍逐跳校验，不继承 provider 头。这比上游严：上游 API 请求由 fetch 跟随，fetch 规范跨 origin 只去掉 `Authorization`、最多 20 跳（`fetch-with-validated-redirects.ts:89-92` 的说明；`:196-206` 的逐跳剥离只用于带校验的请求与下载）。实现与测试：`aimux-provider-utils/src/http.rs` `send_following_redirects`、`tests/api_redirect_test.rs`；对齐文档 §3.2、§5 已登记。§8 A03 的“跨 origin redirect”按本条验收。 + secret 值不可进入 descriptor、Debug、错误、trace、metrics label、settings snapshot 或录制；URL query 中的 token/签名也要按字段脱敏。日志只能记录 credential source 类别、namespace 的非敏感标识、generation 和结果码。调试不得 dump 原始 headers。录制 body 仍可能包含用户敏感数据，必须遵守录制单独的启用与保留策略，不能用“认证已脱敏”宣称整个文件无敏感信息。 ## 4. 刷新所有权与并发（G3，明确后置） diff --git a/scripts/aisdk-fixtures/README.md b/scripts/aisdk-fixtures/README.md new file mode 100644 index 000000000..520cadd67 --- /dev/null +++ b/scripts/aisdk-fixtures/README.md @@ -0,0 +1,65 @@ +# AI SDK protocol fixtures + +Records what the official Vercel AI SDK (`ai` + `@ai-sdk/*`, versions pinned in +`package.json` and `fixtures/aisdk/VERSIONS.json`) sends and returns for a fixed input over +a mocked transport. These are the AI SDK half of the RFC-0036 section 9.3 protocol +differential: Rust tests replay the same inputs against aimux's providers and compare URL, +headers, body, normalized result, stream parts and error classification. + +## Run + +```sh +cd scripts/aisdk-fixtures +npm ci +node record.mjs # all cases +node record.mjs openai # one package +node record.mjs openai/chat-basic # one case (VERSIONS.json / SKIPPED.md untouched) +``` + +Each `fixtures/aisdk/VERSIONS.json` entry is `{version, integrity}`: `node record.mjs` keeps the +integrity of an unchanged version and takes a new one from `package-lock.json`, and +`bash scripts/fetch-aisdk.sh` fetches exactly those sources into `target/aisdk`. + +No network is used. The mock `fetch` (`mock-fetch.mjs`) throws on any URL other than the +case's expected one and on a second request. + +## Layout + +- `record.mjs` - runner. `mock-fetch.mjs` - `createMockFetch(canned)`. `common.mjs` - fake keys, SSE helper. +- `cases/.mjs` - default-exports an array of cases: + `{ name, build(mock), call(model, run), response, env?, expectError?, expectRequests?, observe?, skip? }`. + `build` creates the provider/model with `fetch: mock`; `call` invokes `run('generateText' | 'streamText' | 'embed' | 'embedMany' | 'generateImage', input)`; + `response` is `{ url, status, headers, body }` (JSON object, or a string with `content-type: text/event-stream` for SSE). + `observe` adds an `observations` object (e.g. which header names were sent, `usedEnvKey`). + `skip: ''` skips the case and records it in `fixtures/aisdk/SKIPPED.md`. +- Output: `fixtures/aisdk//.json`, where `` is `openai`, `openai-compatible`, `anthropic`, `google`. + +## What a fixture captures + +`package`, `version`, `case`; `sdk.{operation,input}` (the options passed to the SDK, minus the model; tools +are serialized as `{ description, inputSchema }`); `model.{provider,modelId}`; `request.{method,url,headers,body}` +(headers lower-cased and sorted, JSON body parsed); `response.{status,headers,body}` (the canned reply); +and either `result` (normalized SDK output) or `error.{name,message}` (then `request` is `null` if nothing was sent); +optionally `observations`. All object keys are sorted so diffs are stable. + +`result` per operation: `generateText` -> `content, finishReason, rawFinishReason, providerMetadata, response{id,modelId,headers}, text, usage, warnings`; +`streamText` -> `parts` (full `fullStream`), `finishReason`, `usage`; `embed`/`embedMany` -> `embeddings, usage, providerMetadata`; +`generateImage` -> `images[{base64Length, mediaType}], warnings`. + +## Redaction and determinism + +- Headers named `authorization`, `x-api-key`, `api-key`, `x-goog-api-key`, or whose value contains one of the fake keys, + become `""` (name kept). Only fake keys are used (`sk-test-fixture`, `sk-env-should-not-be-used`); + the recorder aborts if either string appears anywhere in a fixture. +- `Date` is frozen (2025-01-01T00:00:00Z), `Math.random` is seeded, and `_internal.generateId/generateCallId/now` + are injected into `generateText`/`streamText` (these `_internal` options are not recorded in `sdk.input`). +- `navigator.userAgent` is pinned to `Node.js/22` so the `user-agent` header (`ai/ ai-sdk-provider-utils/ node.js/22`) + does not depend on the local Node version. The user-agent string is otherwise kept verbatim. +- Provider env vars (`OPENAI_*`, `ANTHROPIC_*`, `GOOGLE_*`, `GROQ_*`, `AI_GATEWAY_*`, `VERCEL_*`) are cleared per case; a case sets what it needs via `env`. +- Canned responses always carry explicit ids, so provider-side generated ids never enter the fixtures. + +## Policy + +Fixtures are committed and are the reference that Rust tests replay. Regenerate them only on purpose: when bumping the +pinned versions (update `package.json` + `package-lock.json` together) or adding/changing a case. Review the diff; a changed +fixture means the SDK's wire behavior changed. Running `node record.mjs` twice must yield no diff. diff --git a/scripts/aisdk-fixtures/cases/anthropic.mjs b/scripts/aisdk-fixtures/cases/anthropic.mjs new file mode 100644 index 000000000..bf73ebfb4 --- /dev/null +++ b/scripts/aisdk-fixtures/cases/anthropic.mjs @@ -0,0 +1,111 @@ +import { jsonSchema, tool } from 'ai'; +import { createAnthropic } from '@ai-sdk/anthropic'; +import { API_KEY, JSON_HEADERS, SSE_HEADERS, sse } from '../common.mjs'; + +const URL_ = 'https://api.anthropic.com/v1/messages'; +const MODEL = 'claude-sonnet-4-5'; +const provider = (mock, extra = {}) => createAnthropic({ apiKey: API_KEY, fetch: mock, ...extra }); + +const message = (content, stop_reason = 'end_turn') => ({ + id: 'msg_fixture_1', + type: 'message', + role: 'assistant', + model: MODEL, + content, + stop_reason, + stop_sequence: null, + usage: { input_tokens: 12, output_tokens: 9 }, +}); + +const ok = (content = [{ type: 'text', text: 'Hello! How can I help you today?' }], stop_reason) => ({ + url: URL_, + status: 200, + headers: { ...JSON_HEADERS, 'request-id': 'req_fixture_1' }, + body: message(content, stop_reason), +}); + +const weatherTool = tool({ + description: 'Get the current weather for a city.', + inputSchema: jsonSchema({ + type: 'object', + properties: { city: { type: 'string', description: 'City name' } }, + required: ['city'], + additionalProperties: false, + }), +}); + +export default [ + { + name: 'messages-basic', + build: (mock) => provider(mock)(MODEL), + call: (model, run) => run('generateText', { prompt: 'Say hello.' }), + response: ok(), + }, + { + name: 'messages-system-tools', + build: (mock) => provider(mock)(MODEL), + call: (model, run) => + run('generateText', { + system: 'You are a weather assistant.', + prompt: 'What is the weather in Paris?', + tools: { get_weather: weatherTool }, + toolChoice: 'required', + maxOutputTokens: 256, + }), + response: ok([{ type: 'tool_use', id: 'toolu_fixture_1', name: 'get_weather', input: { city: 'Paris' } }], 'tool_use'), + }, + { + name: 'messages-stream-basic', + build: (mock) => provider(mock)(MODEL), + call: (model, run) => run('streamText', { prompt: 'Say hello.' }), + response: { + url: URL_, + status: 200, + headers: SSE_HEADERS, + body: sse([ + { event: 'message_start', data: { type: 'message_start', message: { id: 'msg_fixture_2', type: 'message', role: 'assistant', model: MODEL, content: [], stop_reason: null, stop_sequence: null, usage: { input_tokens: 12, output_tokens: 1 } } } }, + { event: 'content_block_start', data: { type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } } }, + { event: 'ping', data: { type: 'ping' } }, + { event: 'content_block_delta', data: { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'Hello' } } }, + { event: 'content_block_delta', data: { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: ' world' } } }, + { event: 'content_block_stop', data: { type: 'content_block_stop', index: 0 } }, + { event: 'message_delta', data: { type: 'message_delta', delta: { stop_reason: 'end_turn', stop_sequence: null }, usage: { output_tokens: 2 } } }, + { event: 'message_stop', data: { type: 'message_stop' } }, + ]), + }, + }, + { + // Documented option: extended thinking via providerOptions.anthropic.thinking. + name: 'messages-provider-options', + build: (mock) => provider(mock)(MODEL), + call: (model, run) => + run('generateText', { + prompt: 'What is 17 * 23?', + maxOutputTokens: 4096, + providerOptions: { anthropic: { thinking: { type: 'enabled', budgetTokens: 2000 } } }, + }), + response: ok([ + { type: 'thinking', thinking: '17 * 23 = 391.', signature: 'sig_fixture' }, + { type: 'text', text: '391' }, + ]), + }, + { + // Which key wins when the same option is given under `anthropic` and under the custom provider name? + name: 'messages-custom-name', + build: (mock) => provider(mock, { name: 'myproxy' })(MODEL), + call: (model, run) => + run('generateText', { + prompt: 'Say hello.', + providerOptions: { + anthropic: { metadata: { userId: 'from-anthropic-key' } }, + myproxy: { metadata: { userId: 'from-myproxy-key' } }, + }, + }), + response: ok(), + observe: ({ model, request, result }) => ({ + bodyMetadataUserId: request.body.metadata?.user_id, + modelProvider: model.provider, + providerMetadataKeys: Object.keys(result.providerMetadata ?? {}).sort(), + }), + }, +]; diff --git a/scripts/aisdk-fixtures/cases/google.mjs b/scripts/aisdk-fixtures/cases/google.mjs new file mode 100644 index 000000000..b2be8e920 --- /dev/null +++ b/scripts/aisdk-fixtures/cases/google.mjs @@ -0,0 +1,87 @@ +import { jsonSchema, tool } from 'ai'; +import { createGoogleGenerativeAI } from '@ai-sdk/google'; +import { API_KEY, JSON_HEADERS, SSE_HEADERS, sse } from '../common.mjs'; + +const BASE = 'https://generativelanguage.googleapis.com/v1beta'; +const MODEL = 'gemini-2.5-flash'; +const provider = (mock, extra = {}) => createGoogleGenerativeAI({ apiKey: API_KEY, fetch: mock, ...extra }); + +const candidates = (parts, finishReason = 'STOP') => [{ content: { role: 'model', parts }, finishReason, index: 0 }]; +const usageMetadata = { promptTokenCount: 4, candidatesTokenCount: 3, totalTokenCount: 7 }; + +const ok = (parts = [{ text: 'Hello! How can I help you today?' }]) => ({ + url: `${BASE}/models/${MODEL}:generateContent`, + status: 200, + headers: JSON_HEADERS, + body: { candidates: candidates(parts), usageMetadata, modelVersion: MODEL, responseId: 'resp-fixture-1' }, +}); + +const weatherTool = tool({ + description: 'Get the current weather for a city.', + inputSchema: jsonSchema({ + type: 'object', + properties: { city: { type: 'string', description: 'City name' } }, + required: ['city'], + additionalProperties: false, + }), +}); + +export default [ + { + name: 'generate-basic', + build: (mock) => provider(mock)(MODEL), + call: (model, run) => run('generateText', { prompt: 'Say hello.' }), + response: ok(), + }, + { + name: 'generate-stream-basic', + build: (mock) => provider(mock)(MODEL), + call: (model, run) => run('streamText', { prompt: 'Say hello.' }), + response: { + url: `${BASE}/models/${MODEL}:streamGenerateContent?alt=sse`, + status: 200, + headers: SSE_HEADERS, + body: sse([ + { data: { candidates: candidates([{ text: 'Hello' }], undefined), modelVersion: MODEL, responseId: 'resp-fixture-2' } }, + { data: { candidates: candidates([{ text: ' world' }]), usageMetadata: { promptTokenCount: 4, candidatesTokenCount: 2, totalTokenCount: 6 }, modelVersion: MODEL, responseId: 'resp-fixture-2' } }, + ]), + }, + }, + { + name: 'generate-tools', + build: (mock) => provider(mock)(MODEL), + call: (model, run) => + run('generateText', { + prompt: 'What is the weather in Paris?', + tools: { get_weather: weatherTool }, + toolChoice: 'required', + }), + response: ok([{ functionCall: { name: 'get_weather', args: { city: 'Paris' } } }]), + }, + { + name: 'generate-provider-options', + build: (mock) => provider(mock)(MODEL), + call: (model, run) => + run('generateText', { + prompt: 'Say hello.', + providerOptions: { + google: { + thinkingConfig: { thinkingBudget: 1024, includeThoughts: true }, + safetySettings: [{ category: 'HARM_CATEGORY_HATE_SPEECH', threshold: 'BLOCK_ONLY_HIGH' }], + }, + }, + }), + response: ok(), + }, + { + name: 'embedding-basic', + build: (mock) => provider(mock).embedding('gemini-embedding-001'), + call: (model, run) => run('embedMany', { values: ['sunny day at the beach', 'rainy day in the city'] }), + response: { + url: `${BASE}/models/gemini-embedding-001:batchEmbedContents`, + status: 200, + headers: JSON_HEADERS, + body: { embeddings: [{ values: [0.125, -0.25, 0.5] }, { values: [0.0625, 0.375, -0.5] }] }, + }, + }, +]; diff --git a/scripts/aisdk-fixtures/cases/openai-compatible.mjs b/scripts/aisdk-fixtures/cases/openai-compatible.mjs new file mode 100644 index 000000000..55e3baa8f --- /dev/null +++ b/scripts/aisdk-fixtures/cases/openai-compatible.mjs @@ -0,0 +1,116 @@ +import { createOpenAICompatible } from '@ai-sdk/openai-compatible'; +import { API_KEY, JSON_HEADERS, SSE_HEADERS, sse } from '../common.mjs'; + +const BASE = 'https://api.groq.com/openai/v1'; + +const completion = { + id: 'chatcmpl-fixture-1', + object: 'chat.completion', + created: 1735689600, + model: 'llama-3.3-70b-versatile', + choices: [{ index: 0, message: { role: 'assistant', content: 'Hello! How can I help you today?' }, finish_reason: 'stop' }], + usage: { prompt_tokens: 12, completion_tokens: 9, total_tokens: 21 }, +}; + +const ok = (url = `${BASE}/chat/completions`, headers = JSON_HEADERS) => ({ url, status: 200, headers, body: completion }); + +const groq = (mock, extra = {}) => createOpenAICompatible({ name: 'groq', baseURL: BASE, apiKey: API_KEY, fetch: mock, ...extra }); +const MODEL = 'llama-3.3-70b-versatile'; + +const chunk = (delta, finish_reason = null) => ({ + id: 'chatcmpl-fixture-2', + object: 'chat.completion.chunk', + created: 1735689600, + model: MODEL, + choices: [{ index: 0, delta, finish_reason }], +}); + +export default [ + { + name: 'chat-basic', + build: (mock) => groq(mock).chatModel(MODEL), + call: (model, run) => run('generateText', { prompt: 'Say hello.' }), + response: ok(), + }, + { + name: 'chat-stream-basic', + build: (mock) => groq(mock).chatModel(MODEL), + call: (model, run) => run('streamText', { prompt: 'Say hello.' }), + response: { + url: `${BASE}/chat/completions`, + status: 200, + headers: SSE_HEADERS, + body: sse([ + { data: chunk({ role: 'assistant', content: '' }) }, + { data: chunk({ content: 'Hello' }) }, + { data: chunk({ content: ' world' }) }, + { data: chunk({}, 'stop') }, + { data: { id: 'chatcmpl-fixture-2', object: 'chat.completion.chunk', created: 1735689600, model: MODEL, choices: [], usage: { prompt_tokens: 12, completion_tokens: 2, total_tokens: 14 } } }, + 'data: [DONE]', + ]), + }, + }, + { + name: 'embedding-basic', + build: (mock) => groq(mock).embeddingModel('text-embedding-fixture'), + call: (model, run) => run('embedMany', { values: ['sunny day at the beach', 'rainy day in the city'] }), + response: { + url: `${BASE}/embeddings`, + status: 200, + headers: JSON_HEADERS, + body: { + object: 'list', + data: [ + { object: 'embedding', index: 0, embedding: [0.125, -0.25, 0.5] }, + { object: 'embedding', index: 1, embedding: [0.0625, 0.375, -0.5] }, + ], + model: 'text-embedding-fixture', + usage: { prompt_tokens: 10, total_tokens: 10 }, + }, + }, + }, + { + name: 'chat-provider-options-namespace', + build: (mock) => groq(mock).chatModel(MODEL), + call: (model, run) => + run('generateText', { + prompt: 'Say hello.', + providerOptions: { groq: { foo: 'bar' }, openaiCompatible: { user: 'u1' } }, + }), + response: ok(), + observe: ({ model, request }) => { + if (model.provider !== 'groq.chat') throw new Error(`expected provider "groq.chat", got "${model.provider}"`); + return { bodyHasFooAtTopLevel: request.body.foo === 'bar', bodyUser: request.body.user, providerIsGroqChat: true }; + }, + }, + { + // Extra: unknown fields under the generic `openaiCompatible` key (as opposed to the provider's + // own name) are NOT forwarded; only the schema-known ones (user, reasoningEffort, ...) are used. + name: 'chat-provider-options-generic-key-unknown', + build: (mock) => groq(mock).chatModel(MODEL), + call: (model, run) => + run('generateText', { prompt: 'Say hello.', providerOptions: { openaiCompatible: { user: 'u1', extra_field: 'x' } } }), + response: ok(), + observe: ({ request }) => ({ bodyHasExtraField: 'extra_field' in request.body }), + }, + { + name: 'chat-no-api-key', + env: {}, + build: (mock) => createOpenAICompatible({ name: 'local', baseURL: 'http://localhost:1234/v1', fetch: mock }).chatModel('local-model'), + call: (model, run) => run('generateText', { prompt: 'Say hello.' }), + response: ok('http://localhost:1234/v1/chat/completions'), + observe: ({ request }) => { + const has = Object.keys(request.headers).some((k) => k === 'authorization'); + if (has) throw new Error('unexpected authorization header'); + return { hasAuthorizationHeader: has }; + }, + }, + { + name: 'chat-query-params-and-headers', + build: (mock) => + groq(mock, { queryParams: { 'api-version': '1' }, headers: { 'x-a': '1', 'X-B': '2' } }).chatModel(MODEL), + call: (model, run) => run('generateText', { prompt: 'Say hello.', headers: { 'x-b': '3', 'x-c': '4' } }), + response: ok(`${BASE}/chat/completions?api-version=1`), + observe: ({ request }) => ({ headerNamesAsSent: request.rawHeaderNames.filter((n) => /^x-/i.test(n)).sort() }), + }, +]; diff --git a/scripts/aisdk-fixtures/cases/openai.mjs b/scripts/aisdk-fixtures/cases/openai.mjs new file mode 100644 index 000000000..c882e0679 --- /dev/null +++ b/scripts/aisdk-fixtures/cases/openai.mjs @@ -0,0 +1,233 @@ +import { jsonSchema, tool } from 'ai'; +import { createOpenAI } from '@ai-sdk/openai'; +import { API_KEY, ENV_KEY, JSON_HEADERS, SSE_HEADERS, sse } from '../common.mjs'; + +const CHAT_URL = 'https://api.openai.com/v1/chat/completions'; + +const chatCompletion = (message, finish_reason = 'stop', extra = {}) => ({ + id: 'chatcmpl-fixture-1', + object: 'chat.completion', + created: 1735689600, + model: 'gpt-4o-2024-08-06', + choices: [{ index: 0, message, finish_reason }], + usage: { + prompt_tokens: 12, + completion_tokens: 9, + total_tokens: 21, + prompt_tokens_details: { cached_tokens: 0 }, + completion_tokens_details: { reasoning_tokens: 0 }, + }, + system_fingerprint: 'fp_fixture', + ...extra, +}); + +const chatOk = (content = 'Hello! How can I help you today?') => ({ + url: CHAT_URL, + status: 200, + headers: { ...JSON_HEADERS, 'x-request-id': 'req_fixture_1' }, + body: chatCompletion({ role: 'assistant', content }), +}); + +const provider = (mock, extra = {}) => createOpenAI({ apiKey: API_KEY, fetch: mock, ...extra }); + +const weatherTool = tool({ + description: 'Get the current weather for a city.', + inputSchema: jsonSchema({ + type: 'object', + properties: { city: { type: 'string', description: 'City name' } }, + required: ['city'], + additionalProperties: false, + }), +}); + +const chunk = (delta, finish_reason = null, extra = {}) => ({ + id: 'chatcmpl-fixture-2', + object: 'chat.completion.chunk', + created: 1735689600, + model: 'gpt-4o-2024-08-06', + system_fingerprint: 'fp_fixture', + choices: [{ index: 0, delta, finish_reason }], + ...extra, +}); + +export default [ + { + name: 'chat-basic', + build: (mock) => provider(mock).chat('gpt-4o'), + call: (model, run) => run('generateText', { prompt: 'Say hello.' }), + response: chatOk(), + }, + { + name: 'chat-system-and-settings', + build: (mock) => provider(mock).chat('gpt-4o'), + call: (model, run) => + run('generateText', { + system: 'You are a terse assistant.', + messages: [{ role: 'user', content: 'Name a color.' }], + temperature: 0.3, + maxOutputTokens: 64, + topP: 0.9, + stopSequences: ['END', '###'], + seed: 42, + }), + response: chatOk('Blue.'), + }, + { + name: 'chat-tools', + build: (mock) => provider(mock).chat('gpt-4o'), + call: (model, run) => + run('generateText', { + prompt: 'What is the weather in Paris?', + tools: { get_weather: weatherTool }, + toolChoice: 'required', + }), + response: { + url: CHAT_URL, + status: 200, + headers: JSON_HEADERS, + body: chatCompletion( + { + role: 'assistant', + content: null, + tool_calls: [{ id: 'call_fixture_1', type: 'function', function: { name: 'get_weather', arguments: '{"city":"Paris"}' } }], + }, + 'tool_calls', + ), + }, + }, + { + name: 'chat-provider-options', + build: (mock) => provider(mock).chat('gpt-4o'), + call: (model, run) => + run('generateText', { + prompt: 'What is the weather in Paris?', + tools: { get_weather: weatherTool }, + providerOptions: { openai: { parallelToolCalls: false, user: 'u1' } }, + }), + response: chatOk(), + }, + { + name: 'chat-headers-merge', + build: (mock) => provider(mock, { headers: { 'x-a': '1', 'X-B': '2' } }).chat('gpt-4o'), + call: (model, run) => run('generateText', { prompt: 'Say hello.', headers: { 'x-b': '3', 'x-c': '4' } }), + response: chatOk(), + observe: ({ request }) => ({ + headerNamesAsSent: request.rawHeaderNames.filter((n) => /^x-/i.test(n)).sort(), + xB: request.headers['x-b'], + }), + }, + { + name: 'chat-custom-name', + build: (mock) => provider(mock, { name: 'proxy' }).chat('gpt-4o'), + call: (model, run) => run('generateText', { prompt: 'Say hello.' }), + response: chatOk(), + observe: ({ model }) => { + if (model.provider !== 'proxy.chat') throw new Error(`expected provider "proxy.chat", got "${model.provider}"`); + return { providerIsProxyChat: true }; + }, + }, + { + name: 'chat-stream-basic', + build: (mock) => provider(mock).chat('gpt-4o'), + call: (model, run) => run('streamText', { prompt: 'Say hello.' }), + response: { + url: CHAT_URL, + status: 200, + headers: SSE_HEADERS, + body: sse([ + { data: chunk({ role: 'assistant', content: '' }) }, + { data: chunk({ content: 'Hello' }) }, + { data: chunk({ content: ' world' }) }, + { data: chunk({}, 'stop') }, + { + data: { + id: 'chatcmpl-fixture-2', + object: 'chat.completion.chunk', + created: 1735689600, + model: 'gpt-4o-2024-08-06', + choices: [], + usage: { prompt_tokens: 12, completion_tokens: 2, total_tokens: 14 }, + }, + }, + 'data: [DONE]', + ]), + }, + }, + { + name: 'responses-basic', + build: (mock) => provider(mock).responses('gpt-4o'), + call: (model, run) => run('generateText', { prompt: 'Say hello.' }), + response: { + url: 'https://api.openai.com/v1/responses', + status: 200, + headers: JSON_HEADERS, + body: { + id: 'resp_fixture_1', + object: 'response', + created_at: 1735689600, + status: 'completed', + error: null, + incomplete_details: null, + model: 'gpt-4o-2024-08-06', + output: [ + { + type: 'message', + id: 'msg_fixture_1', + status: 'completed', + role: 'assistant', + content: [{ type: 'output_text', text: 'Hello! How can I help you today?', annotations: [] }], + }, + ], + usage: { + input_tokens: 12, + output_tokens: 9, + total_tokens: 21, + input_tokens_details: { cached_tokens: 0 }, + output_tokens_details: { reasoning_tokens: 0 }, + }, + }, + }, + }, + { + name: 'embedding-basic', + build: (mock) => provider(mock).embedding('text-embedding-3-small'), + call: (model, run) => run('embedMany', { values: ['sunny day at the beach', 'rainy day in the city'] }), + response: { + url: 'https://api.openai.com/v1/embeddings', + status: 200, + headers: JSON_HEADERS, + body: { + object: 'list', + data: [ + { object: 'embedding', index: 0, embedding: [0.125, -0.25, 0.5] }, + { object: 'embedding', index: 1, embedding: [0.0625, 0.375, -0.5] }, + ], + model: 'text-embedding-3-small', + usage: { prompt_tokens: 10, total_tokens: 10 }, + }, + }, + }, + { + name: 'missing-api-key', + env: { OPENAI_API_KEY: undefined }, + build: (mock) => createOpenAI({ fetch: mock }).chat('gpt-4o'), + call: (model, run) => run('generateText', { prompt: 'Say hello.' }), + expectError: true, + expectRequests: 0, + response: chatOk(), + observe: ({ requests }) => ({ requestsSent: requests.length }), + }, + { + name: 'empty-api-key', + env: { OPENAI_API_KEY: ENV_KEY }, + // Pins the "empty string does not fall back to env" rule: whatever the SDK does is recorded. + build: (mock) => createOpenAI({ apiKey: '', fetch: mock }).chat('gpt-4o'), + call: (model, run) => run('generateText', { prompt: 'Say hello.' }), + response: chatOk(), + observe: ({ request }) => ({ + authorizationIsEmptyBearer: request.rawHeaders.authorization?.trim() === 'Bearer', + authorizationRawLength: request.rawHeaders.authorization?.length, + usedEnvKey: (request.rawHeaders.authorization ?? '').includes(ENV_KEY), + }), + }, +]; diff --git a/scripts/aisdk-fixtures/common.mjs b/scripts/aisdk-fixtures/common.mjs new file mode 100644 index 000000000..d6c43da13 --- /dev/null +++ b/scripts/aisdk-fixtures/common.mjs @@ -0,0 +1,27 @@ +// Shared constants for the fixture recorder. Keys are obviously fake; they are +// redacted from every fixture and the recorder fails if one ever leaks. +export const API_KEY = 'sk-test-fixture'; +export const ENV_KEY = 'sk-env-should-not-be-used'; +export const FAKE_KEYS = [API_KEY, ENV_KEY]; + +// Deterministic clock: every `new Date()` / `Date.now()` returns this instant. +export const FIXED_NOW_MS = Date.UTC(2025, 0, 1, 0, 0, 0); + +/** Build an SSE text body from event objects (JSON `data:` lines) or raw strings. */ +export function sse(events) { + return ( + events + .map((e) => { + if (typeof e === 'string') return e; + const { event, data } = e; + return `${event ? `event: ${event}\n` : ''}data: ${typeof data === 'string' ? data : JSON.stringify(data)}`; + }) + .join('\n\n') + '\n\n' + ); +} + +export const JSON_HEADERS = { 'content-type': 'application/json' }; +export const SSE_HEADERS = { + 'cache-control': 'no-cache', + 'content-type': 'text/event-stream', +}; diff --git a/scripts/aisdk-fixtures/mock-fetch.mjs b/scripts/aisdk-fixtures/mock-fetch.mjs new file mode 100644 index 000000000..585cc7d7a --- /dev/null +++ b/scripts/aisdk-fixtures/mock-fetch.mjs @@ -0,0 +1,81 @@ +// Mock transport. `createMockFetch(canned)` returns a fetch-compatible function that +// records every outgoing request and answers with the canned response. It never touches +// the network: a request to any URL other than `canned.url`, or a second request, throws. + +function lowerSortedHeaders(entries) { + const out = {}; + for (const [k, v] of entries.map(([k, v]) => [String(k).toLowerCase(), String(v)]).sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0))) { + out[k] = v; + } + return out; +} + +function headerEntries(headers) { + if (headers == null) return []; + if (typeof Headers !== 'undefined' && headers instanceof Headers) return [...headers.entries()]; + if (Array.isArray(headers)) return headers; + return Object.entries(headers); +} + +export function createMockFetch(canned) { + const requests = []; + + async function mockFetch(input, init = {}) { + const url = typeof input === 'string' ? input : input instanceof URL ? input.href : input.url; + const method = (init.method ?? (typeof input === 'object' && input.method) ?? 'GET').toUpperCase(); + const rawEntries = headerEntries(init.headers); + const headers = lowerSortedHeaders(rawEntries); + + let text; + const b = init.body; + if (b == null) text = undefined; + else if (typeof b === 'string') text = b; + else if (b instanceof Uint8Array || b instanceof ArrayBuffer) text = new TextDecoder().decode(b); + else throw new Error(`mock fetch: unsupported request body type ${Object.prototype.toString.call(b)}`); + + let body = text; + if (text !== undefined && (headers['content-type'] ?? '').includes('json')) { + try { + body = JSON.parse(text); + } catch { + body = text; + } + } + requests.push({ method, url, headers, body, rawHeaderNames: rawEntries.map(([k]) => String(k)), rawHeaders: Object.fromEntries(rawEntries.map(([k, v]) => [String(k), String(v)])) }); + + if (requests.length > 1 || !canned || url !== canned.url) { + throw new Error(`mock fetch: unmocked request #${requests.length} ${method} ${url} (expected ${canned?.url ?? ''})`); + } + + const status = canned.status ?? 200; + const respHeaders = { ...(canned.headers ?? {}) }; + const isSse = Object.entries(respHeaders).some(([k, v]) => k.toLowerCase() === 'content-type' && /text\/event-stream/i.test(v)); + let payload; + if (typeof canned.body === 'string') { + if (isSse) { + const enc = new TextEncoder(); + const events = canned.body.split('\n\n').filter((e) => e.length > 0).map((e) => enc.encode(e + '\n\n')); + payload = new ReadableStream({ + start(controller) { + for (const e of events) controller.enqueue(e); + controller.close(); + }, + }); + } else payload = canned.body; + } else { + payload = JSON.stringify(canned.body); + if (!Object.keys(respHeaders).some((k) => k.toLowerCase() === 'content-type')) respHeaders['content-type'] = 'application/json'; + } + const response = new Response(payload, { status, headers: respHeaders }); + const last = requests[requests.length - 1]; + last.responseStatus = response.status; + last.responseHeaders = Object.fromEntries(response.headers.entries()); + return response; + } + + mockFetch.requests = requests; + mockFetch.canned = canned; + return mockFetch; +} + +export { lowerSortedHeaders }; diff --git a/scripts/aisdk-fixtures/package-lock.json b/scripts/aisdk-fixtures/package-lock.json new file mode 100644 index 000000000..45e09a595 --- /dev/null +++ b/scripts/aisdk-fixtures/package-lock.json @@ -0,0 +1,331 @@ +{ + "name": "aimux-aisdk-fixtures", + "version": "0.0.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "aimux-aisdk-fixtures", + "version": "0.0.0", + "dependencies": { + "@ai-sdk/anthropic": "4.0.71", + "@ai-sdk/google": "4.0.87", + "@ai-sdk/openai": "4.0.83", + "@ai-sdk/openai-compatible": "3.0.62", + "@ai-sdk/provider": "4.0.21", + "@ai-sdk/provider-utils": "5.0.55", + "ai": "7.0.127" + } + }, + "node_modules/@ai-sdk/anthropic": { + "version": "4.0.71", + "resolved": "https://registry.npmjs.org/@ai-sdk/anthropic/-/anthropic-4.0.71.tgz", + "integrity": "sha512-wiv3jhUH0RrvGzSolJMZfZ2cI9As45LhwYmcAoWvNivXjC7wzJ4W/Udq9/u7EGnGsT8lpmFKEKBrMXVeUO97Dg==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@ai-sdk/provider-utils": "5.0.53" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/anthropic/node_modules/@ai-sdk/provider-utils": { + "version": "5.0.53", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-5.0.53.tgz", + "integrity": "sha512-VVe6UDd0y0/B4TfauuKKzcppqvhTmSluwwUVT4WG2/8h7gFbyFm7PM4UGM+ZXe3/RI6bv+Vl1cw/7aiBoQqdxg==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@standard-schema/spec": "^1.1.0", + "@workflow/serde": "4.1.0", + "eventsource-parser": "^3.0.8", + "undici": "^7.29.0" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/gateway": { + "version": "4.0.103", + "resolved": "https://registry.npmjs.org/@ai-sdk/gateway/-/gateway-4.0.103.tgz", + "integrity": "sha512-nnGTUHzPRQ14jhv9CL4OQrzcRBIFTBMIPKGIi9Z9u/4GtwSRxa/nxUX6np2Iwp5phOzO0a0njj/j9Jolf5AEaA==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@ai-sdk/provider-utils": "5.0.53", + "@vercel/oidc": "3.2.0" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/gateway/node_modules/@ai-sdk/provider-utils": { + "version": "5.0.53", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-5.0.53.tgz", + "integrity": "sha512-VVe6UDd0y0/B4TfauuKKzcppqvhTmSluwwUVT4WG2/8h7gFbyFm7PM4UGM+ZXe3/RI6bv+Vl1cw/7aiBoQqdxg==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@standard-schema/spec": "^1.1.0", + "@workflow/serde": "4.1.0", + "eventsource-parser": "^3.0.8", + "undici": "^7.29.0" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/google": { + "version": "4.0.87", + "resolved": "https://registry.npmjs.org/@ai-sdk/google/-/google-4.0.87.tgz", + "integrity": "sha512-s0tWc+QeSce7O2ZLm+VWSRQy16dGfqEcMIUxEWh+LFvXX3h0TLH8jxCQ2H1qhiBgTta+6A97eHPhBlRFu1xy+g==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@ai-sdk/provider-utils": "5.0.53" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/google/node_modules/@ai-sdk/provider-utils": { + "version": "5.0.53", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-5.0.53.tgz", + "integrity": "sha512-VVe6UDd0y0/B4TfauuKKzcppqvhTmSluwwUVT4WG2/8h7gFbyFm7PM4UGM+ZXe3/RI6bv+Vl1cw/7aiBoQqdxg==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@standard-schema/spec": "^1.1.0", + "@workflow/serde": "4.1.0", + "eventsource-parser": "^3.0.8", + "undici": "^7.29.0" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/openai": { + "version": "4.0.83", + "resolved": "https://registry.npmjs.org/@ai-sdk/openai/-/openai-4.0.83.tgz", + "integrity": "sha512-NgqZVWoya7hfmtLkWyIjX9THbyYMLfxFqpShdAmvMnkoT+R9AZbdULGan8+4BiBWIZ1aHiAsc90l3SL8QU/0+Q==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@ai-sdk/provider-utils": "5.0.53" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/openai-compatible": { + "version": "3.0.62", + "resolved": "https://registry.npmjs.org/@ai-sdk/openai-compatible/-/openai-compatible-3.0.62.tgz", + "integrity": "sha512-IAz1xv8zot4KLGa8kN7UfAA6H5+ARzAgAKLsoKxHPBzZuOc2E3ohaELg0x/JR9HVTbMbiftb0dPyCb8cT2KBRg==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@ai-sdk/provider-utils": "5.0.53" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/openai-compatible/node_modules/@ai-sdk/provider-utils": { + "version": "5.0.53", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-5.0.53.tgz", + "integrity": "sha512-VVe6UDd0y0/B4TfauuKKzcppqvhTmSluwwUVT4WG2/8h7gFbyFm7PM4UGM+ZXe3/RI6bv+Vl1cw/7aiBoQqdxg==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@standard-schema/spec": "^1.1.0", + "@workflow/serde": "4.1.0", + "eventsource-parser": "^3.0.8", + "undici": "^7.29.0" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/openai/node_modules/@ai-sdk/provider-utils": { + "version": "5.0.53", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-5.0.53.tgz", + "integrity": "sha512-VVe6UDd0y0/B4TfauuKKzcppqvhTmSluwwUVT4WG2/8h7gFbyFm7PM4UGM+ZXe3/RI6bv+Vl1cw/7aiBoQqdxg==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@standard-schema/spec": "^1.1.0", + "@workflow/serde": "4.1.0", + "eventsource-parser": "^3.0.8", + "undici": "^7.29.0" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/provider": { + "version": "4.0.21", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider/-/provider-4.0.21.tgz", + "integrity": "sha512-UpbC9C1oht8dhfKPbVXSLRZS3DI8uk8n5v2uMjBPNIYGsC2kL045ywq7D8Hh9KgnH9rP5W/GxQdYtWScRouqlA==", + "license": "Apache-2.0", + "dependencies": { + "json-schema": "^0.4.0" + }, + "engines": { + "node": ">=22" + } + }, + "node_modules/@ai-sdk/provider-utils": { + "version": "5.0.55", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-5.0.55.tgz", + "integrity": "sha512-VeY9IqwJGcHSmezQTZknaYUT3sRcaXpuIt5VWyGOlJtoc5wkpDPstXZtfa41bGUsbQwCijvN6ASbggUCeTcMkw==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.23", + "@standard-schema/spec": "^1.1.0", + "@workflow/serde": "4.1.0", + "eventsource-parser": "^3.0.8", + "undici": "^7.29.0" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/@ai-sdk/provider-utils/node_modules/@ai-sdk/provider": { + "version": "4.0.23", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider/-/provider-4.0.23.tgz", + "integrity": "sha512-TdEzgE1ndBGhDDwZcg/WZG3IeUj8Yz7/DkUG4ctMyK5HdNXTEKpjkCdwX8ayvkgluTv7yjnqCdOkB8/q1nfaNQ==", + "license": "Apache-2.0", + "dependencies": { + "json-schema": "^0.4.0" + }, + "engines": { + "node": ">=22" + } + }, + "node_modules/@standard-schema/spec": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@standard-schema/spec/-/spec-1.1.0.tgz", + "integrity": "sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w==", + "license": "MIT" + }, + "node_modules/@vercel/oidc": { + "version": "3.2.0", + "resolved": "https://registry.npmjs.org/@vercel/oidc/-/oidc-3.2.0.tgz", + "integrity": "sha512-UycprH3T6n3jH0k44NHMa7pnFHGu/N05MjojYr+Mc6I7obkoLIJujSWwin1pCvdy/eOxrI/l3uDLQsmcrOb4ug==", + "license": "Apache-2.0", + "engines": { + "node": ">= 20" + } + }, + "node_modules/@workflow/serde": { + "version": "4.1.0", + "resolved": "https://registry.npmjs.org/@workflow/serde/-/serde-4.1.0.tgz", + "integrity": "sha512-pav4F2BoirECWR7Nf1TKt+2eETcBj7jj4cBefQ8VXQCA6NPkaKeLfj/zMgi+3zYV5ZIBT4GuUiphsj0/b9hPQQ==", + "license": "Apache-2.0" + }, + "node_modules/ai": { + "version": "7.0.127", + "resolved": "https://registry.npmjs.org/ai/-/ai-7.0.127.tgz", + "integrity": "sha512-JNsNPk4ZvRGEJuqdVo2lYFkm7uAuf3p/RwKFWn1Wh/+IK3U1RyFWw2GXoFnhzqNfxEuj7XJB4WpIUhl1sHToWg==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/gateway": "4.0.103", + "@ai-sdk/provider": "4.0.21", + "@ai-sdk/provider-utils": "5.0.53" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/ai/node_modules/@ai-sdk/provider-utils": { + "version": "5.0.53", + "resolved": "https://registry.npmjs.org/@ai-sdk/provider-utils/-/provider-utils-5.0.53.tgz", + "integrity": "sha512-VVe6UDd0y0/B4TfauuKKzcppqvhTmSluwwUVT4WG2/8h7gFbyFm7PM4UGM+ZXe3/RI6bv+Vl1cw/7aiBoQqdxg==", + "license": "Apache-2.0", + "dependencies": { + "@ai-sdk/provider": "4.0.21", + "@standard-schema/spec": "^1.1.0", + "@workflow/serde": "4.1.0", + "eventsource-parser": "^3.0.8", + "undici": "^7.29.0" + }, + "engines": { + "node": ">=22" + }, + "peerDependencies": { + "zod": "^3.25.76 || ^4.1.8" + } + }, + "node_modules/eventsource-parser": { + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/eventsource-parser/-/eventsource-parser-3.1.1.tgz", + "integrity": "sha512-EKN1vKAMcZ8MlYMpaNuxN6R9yakzH6uajHcHVTqWJzvu5pWw9DyhbP35HH8MVBQ+dZjAfDxk+A8NiR9KWaXiyQ==", + "license": "MIT", + "engines": { + "node": ">=18.0.0" + } + }, + "node_modules/json-schema": { + "version": "0.4.0", + "resolved": "https://registry.npmjs.org/json-schema/-/json-schema-0.4.0.tgz", + "integrity": "sha512-es94M3nTIfsEPisRafak+HDLfHXnKBhV3vU5eqPcS3flIWqcxJWgXHXiey3YrpaNsanY5ei1VoYEbOzijuq9BA==", + "license": "(AFL-2.1 OR BSD-3-Clause)" + }, + "node_modules/undici": { + "version": "7.30.0", + "resolved": "https://registry.npmjs.org/undici/-/undici-7.30.0.tgz", + "integrity": "sha512-dkrQXeHSaoamnItlYbmzG0wFYrM0ZwDxCIg0A7aKjTyyhh9svRzCNFEzV+Vm05/yehjCzjDZ31KXfGEjYSztDQ==", + "license": "MIT", + "engines": { + "node": ">=20.18.1" + } + }, + "node_modules/zod": { + "version": "4.6.5", + "resolved": "https://registry.npmjs.org/zod/-/zod-4.6.5.tgz", + "integrity": "sha512-v5l/aFXZQeai4awLbOpSoHecE9UiMrnfx75tEXLjNonXVARxQ5mOeipTjROUchszUNCqnE+hqAMujRsRHsut2Q==", + "license": "MIT", + "peer": true, + "funding": { + "url": "https://github.com/sponsors/colinhacks" + } + } + } +} diff --git a/scripts/aisdk-fixtures/package.json b/scripts/aisdk-fixtures/package.json new file mode 100644 index 000000000..43fa6b36b --- /dev/null +++ b/scripts/aisdk-fixtures/package.json @@ -0,0 +1,19 @@ +{ + "name": "aimux-aisdk-fixtures", + "version": "0.0.0", + "private": true, + "type": "module", + "description": "Records Vercel AI SDK request/response fixtures under a mocked fetch transport (RFC-0036 section 9.3).", + "scripts": { + "record": "node record.mjs" + }, + "dependencies": { + "@ai-sdk/anthropic": "4.0.71", + "@ai-sdk/google": "4.0.87", + "@ai-sdk/openai": "4.0.83", + "@ai-sdk/openai-compatible": "3.0.62", + "@ai-sdk/provider": "4.0.21", + "@ai-sdk/provider-utils": "5.0.55", + "ai": "7.0.127" + } +} diff --git a/scripts/aisdk-fixtures/record.mjs b/scripts/aisdk-fixtures/record.mjs new file mode 100644 index 000000000..c6a9f0250 --- /dev/null +++ b/scripts/aisdk-fixtures/record.mjs @@ -0,0 +1,260 @@ +#!/usr/bin/env node +// Protocol fixture recorder: runs the official Vercel AI SDK packages against a mock +// `fetch` and writes what was sent / returned to fixtures/aisdk//.json. +// See README.md. Usage: `npm ci && node record.mjs [package[/case] ...]`. +import { readFileSync, writeFileSync, mkdirSync, readdirSync, rmSync, existsSync } from 'node:fs'; +import { dirname, join, resolve } from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { generateText, streamText, embed, embedMany, generateImage } from 'ai'; +import { createMockFetch, lowerSortedHeaders } from './mock-fetch.mjs'; +import { FAKE_KEYS, FIXED_NOW_MS } from './common.mjs'; +import { lockIntegrity, readJson, readVersions, versionsDocument } from './versions.mjs'; + +const here = dirname(fileURLToPath(import.meta.url)); +const OUT = resolve(here, '../../fixtures/aisdk'); + +const PACKAGES = [ + { dir: 'openai', pkg: '@ai-sdk/openai' }, + { dir: 'openai-compatible', pkg: '@ai-sdk/openai-compatible' }, + { dir: 'anthropic', pkg: '@ai-sdk/anthropic' }, + { dir: 'google', pkg: '@ai-sdk/google' }, +]; +const VERSIONED = ['ai', '@ai-sdk/provider', '@ai-sdk/provider-utils', ...PACKAGES.map((p) => p.pkg)]; + +function installedVersion(pkg) { + return JSON.parse(readFileSync(join(here, 'node_modules', pkg, 'package.json'), 'utf8')).version; +} + +// ---------------------------------------------------------------- determinism +const RealDate = Date; +class FixedDate extends RealDate { + constructor(...args) { + if (args.length === 0) super(FIXED_NOW_MS); + else super(...args); + } + static now() { + return FIXED_NOW_MS; + } +} +globalThis.Date = FixedDate; + +// The provider-utils user-agent embeds `runtime/`; pin it so the +// fixtures do not depend on the Node major version of whoever regenerates them. +Object.defineProperty(globalThis, 'navigator', { value: { userAgent: 'Node.js/22' }, configurable: true, writable: true }); + +function seedRandom(seed = 0x2f6e2b1) { + let a = seed >>> 0; + Math.random = () => { + a = (a + 0x6d2b79f5) >>> 0; + let t = a; + t = Math.imul(t ^ (t >>> 15), t | 1); + t ^= t + Math.imul(t ^ (t >>> 7), t | 61); + return ((t ^ (t >>> 14)) >>> 0) / 4294967296; + }; +} + +function makeIdGenerator() { + let n = 0; + return () => `id-${++n}`; +} + +// ---------------------------------------------------------------- serialization +function sortKeys(v) { + if (Array.isArray(v)) return v.map(sortKeys); + if (v && typeof v === 'object') { + return Object.fromEntries(Object.keys(v).sort().map((k) => [k, sortKeys(v[k])])); + } + return v; +} + +function replacer(_key, value) { + if (value instanceof Error) return { name: value.name, message: value.message }; + if (value instanceof Uint8Array) return { uint8ArrayLength: value.length }; + if (typeof value === 'bigint') return value.toString(); + // AI SDK tool objects: keep description + the JSON schema, drop functions/symbols. + if (value && typeof value === 'object' && 'inputSchema' in value && !Array.isArray(value)) { + const s = value.inputSchema; + return { description: value.description, inputSchema: s && typeof s === 'object' && 'jsonSchema' in s ? s.jsonSchema : s }; + } + return value; +} + +/** JSON-roundtrip with SDK-object handling, then key-sorted. */ +function toPlain(v) { + const s = JSON.stringify(v, replacer); + return s === undefined ? undefined : sortKeys(JSON.parse(s)); +} + +function redactHeaders(headers) { + const secretNames = new Set(['authorization', 'x-api-key', 'api-key', 'x-goog-api-key']); + const out = {}; + for (const [k, v] of Object.entries(headers)) { + const name = k.toLowerCase(); + const leaks = FAKE_KEYS.some((key) => key !== '' && v.includes(key)); + out[name] = secretNames.has(name) || leaks ? '' : v; + } + return sortKeys(out); +} + +// ---------------------------------------------------------------- SDK operations +function generateTextResult(r) { + return { + content: r.content, + finishReason: r.finishReason, + providerMetadata: r.providerMetadata, + rawFinishReason: r.rawFinishReason, + response: { headers: r.response?.headers, id: r.response?.id, modelId: r.response?.modelId }, + text: r.text, + usage: r.usage, + warnings: r.warnings, + }; +} + +async function runOperation(operation, model, input, internal) { + switch (operation) { + case 'generateText': + return generateTextResult(await generateText({ model, ...input, _internal: internal })); + case 'streamText': { + const r = streamText({ model, ...input, _internal: internal, onError: () => {} }); + const parts = []; + for await (const p of r.fullStream) parts.push(p); + return { finishReason: await r.finishReason, parts, usage: await r.usage }; + } + case 'embed': { + const r = await embed({ model, ...input }); + return { embeddings: [r.embedding], providerMetadata: r.providerMetadata, usage: r.usage }; + } + case 'embedMany': { + const r = await embedMany({ model, ...input }); + return { embeddings: r.embeddings, providerMetadata: r.providerMetadata, usage: r.usage }; + } + case 'generateImage': { + const r = await generateImage({ model, ...input }); + return { + images: r.images.map((i) => ({ base64Length: i.base64?.length ?? 0, mediaType: i.mediaType })), + warnings: r.warnings, + }; + } + default: + throw new Error(`unknown operation ${operation}`); + } +} + +// ---------------------------------------------------------------- one case +const SCRUB_ENV = /^(OPENAI|ANTHROPIC|GOOGLE|GROQ|AI_GATEWAY|VERCEL)_/; + +async function recordCase({ pkg, version, c }) { + seedRandom(); + const generateId = makeIdGenerator(); + const internal = { generateId, generateCallId: generateId, now: () => FIXED_NOW_MS }; + + // Hermetic environment: no ambient provider env vars, plus whatever the case sets. + const savedEnv = { ...process.env }; + for (const k of Object.keys(process.env)) if (SCRUB_ENV.test(k)) delete process.env[k]; + for (const [k, v] of Object.entries(c.env ?? {})) { + if (v === undefined) delete process.env[k]; + else process.env[k] = v; + } + + const mock = createMockFetch(c.response); + const ctx = { operation: undefined, input: undefined }; + let model; + let result; + let error; + try { + model = c.build(mock); + const run = (operation, input) => { + ctx.operation = operation; + ctx.input = toPlain(input); + return runOperation(operation, model, input, internal); + }; + try { + result = toPlain(await c.call(model, run)); + } catch (e) { + error = e; + } + } finally { + for (const k of Object.keys(process.env)) if (!(k in savedEnv)) delete process.env[k]; + Object.assign(process.env, savedEnv); + } + + if (c.expectError && !error) throw new Error(`${pkg}/${c.name}: expected an error, got a result`); + if (!c.expectError && error) throw new Error(`${pkg}/${c.name}: unexpected error: ${error.stack ?? error}`); + const expected = c.expectRequests ?? 1; + if (mock.requests.length !== expected) { + throw new Error(`${pkg}/${c.name}: expected ${expected} request(s), mock saw ${mock.requests.length}`); + } + + const raw = mock.requests[0]; + const observations = c.observe ? toPlain(c.observe({ request: raw, requests: mock.requests, result, error, model })) : undefined; + + const fixture = { + package: pkg, + version, + case: c.name, + sdk: { operation: ctx.operation, input: ctx.input }, + model: { modelId: model?.modelId, provider: model?.provider }, + request: raw ? { body: raw.body, headers: redactHeaders(raw.headers), method: raw.method, url: raw.url } : null, + response: raw ? { body: c.response.body, headers: lowerSortedHeaders(Object.entries(raw.responseHeaders)), status: raw.responseStatus } : null, + ...(error ? { error: { message: error.message, name: error.name } } : { result }), + ...(observations ? { observations } : {}), + }; + const text = JSON.stringify(sortKeys(JSON.parse(JSON.stringify(fixture, replacer))), null, 2) + '\n'; + for (const key of FAKE_KEYS) { + if (key && text.includes(key)) throw new Error(`${pkg}/${c.name}: API key "${key}" leaked into fixture`); + } + return text; +} + +// ---------------------------------------------------------------- main +const filters = process.argv.slice(2); +const wanted = (dir, name) => filters.length === 0 || filters.some((f) => f === dir || f === `${dir}/${name}`); + +const skipped = []; +let recorded = 0; +for (const { dir, pkg } of PACKAGES) { + const cases = (await import(`./cases/${dir}.mjs`)).default; + const version = installedVersion(pkg); + const pkgOut = join(OUT, dir); + mkdirSync(pkgOut, { recursive: true }); + const names = new Set(); + for (const c of cases) { + if (names.has(c.name)) throw new Error(`duplicate case ${dir}/${c.name}`); + names.add(c.name); + if (!wanted(dir, c.name)) continue; + if (c.skip) { + skipped.push({ dir, pkg, name: c.name, reason: c.skip }); + console.log(`skip ${dir}/${c.name}: ${c.skip}`); + continue; + } + const text = await recordCase({ pkg, version, c }); + writeFileSync(join(pkgOut, `${c.name}.json`), text); + recorded++; + console.log(`record ${dir}/${c.name}`); + } + if (filters.length === 0) { + // Drop fixtures of cases that no longer exist. + for (const f of readdirSync(pkgOut)) { + if (f.endsWith('.json') && !names.has(f.slice(0, -5))) rmSync(join(pkgOut, f)); + } + } +} + +if (filters.length === 0) { + // Each entry is {version, integrity}; the integrity is kept from the + // previous file or taken from package-lock.json (see versions.mjs). + const versionsPath = join(OUT, 'VERSIONS.json'); + const versions = versionsDocument( + Object.fromEntries(VERSIONED.map((p) => [p, installedVersion(p)])), + readVersions(versionsPath), + lockIntegrity(readJson(join(here, 'package-lock.json'), {})), + ); + writeFileSync(versionsPath, JSON.stringify(versions, null, 2) + '\n'); + const skippedPath = join(OUT, 'SKIPPED.md'); + if (skipped.length) { + const lines = ['# Skipped cases', '', 'Generated by `scripts/aisdk-fixtures/record.mjs`. These cases could not be recorded against the pinned SDK versions.', '']; + for (const s of skipped) lines.push(`- \`${s.dir}/${s.name}\` (${s.pkg}): ${s.reason}`); + writeFileSync(skippedPath, lines.join('\n') + '\n'); + } else if (existsSync(skippedPath)) rmSync(skippedPath); +} +console.log(`done: ${recorded} recorded, ${skipped.length} skipped`); diff --git a/scripts/check_provider_boundaries.sh b/scripts/check_provider_boundaries.sh new file mode 100755 index 000000000..add64bb2e --- /dev/null +++ b/scripts/check_provider_boundaries.sh @@ -0,0 +1,295 @@ +#!/usr/bin/env bash +# Provider boundary grep gates (RFC-0036, AI SDK architecture alignment). +# +# Keeps the provider layer on the same side of the line as the AI SDK: +# 1. Providers never read call-level plumbing (retry budget, timeouts, +# session / call ids) from `CallOptions`. Core owns retry and recording; +# a provider hands the options to `HttpRequest::new(url, headers, options)` +# and lets the transport read what it needs. +# 2. Names removed by the provider-factory rewrite stay removed. +# 3. The OpenAI package keeps no builder-era `OpenAIConfig` reads. +# 4. The shared protocol packages never compare a provider name: vendor +# behavior is injected (flags, hooks, a dialect), not branched on. +# 5. The builder-era compat types and the placeholder key stay removed. +# 6. The Anthropic family (anthropic, anthropic_aws, anthropic on Vertex) +# keeps no builder-era `AnthropicConfig` state, and the canonical +# providerOptions key is spelled once, in `anthropic/options.rs`. +# 7. Google, Vertex and Bedrock keep no builder-era config state, the SigV4 +# shim stays deleted, and their providerOptions / providerMetadata keys +# are spelled once, in `google/options.rs` and `bedrock/options.rs`. +# 8. Azure, xAI, Mistral, Cohere, Hugging Face, Codex, Open Responses, Voyage +# and ElevenLabs keep no builder-era config state (and `StaticBearerConfig` +# stays deleted), and their providerOptions / providerMetadata keys are +# spelled once, in each package's `options.rs`. +# +# 9. No provider code calls the operation-level retry: `prepare_retries` and +# `aimux_core::retry` are not used anywhere in aimux-providers/src (the +# exception list is empty). A provider's own stages (poll, download) are +# bounded by package constants through `shared::poll`; the call-level +# retry belongs to Core. +# 10. The single-modality vendor packages (search, speech, transcription, +# image, video, reranking) keep no builder-era config state, and their +# providerOptions namespace keys are spelled once, in each package's +# `options.rs`. +# +# Usage: bash scripts/check_provider_boundaries.sh (exit 1 on any violation) + +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")/.." + +status=0 + +# violation : print and remember a failure when the grep +# output is non-empty. +violation() { + local rule="$1" hits="$2" + if [[ -n "$hits" ]]; then + printf 'provider boundary violation: %s\n%s\n\n' "$rule" "$hits" >&2 + status=1 + fi +} + +# ── Rule 1: no call-level reads inside aimux-providers/src ─────────────────── +# +# Allowed: lines that build `HttpRequest::new(`, and aimux-providers/src/replay.rs. +# WebSocket realtime sessions are the one transport that the provider itself +# opens, so they carry `options.timeout` into `WebSocketRequest` (no core +# operation owns that timeout); the two sites are listed explicitly. +rule1_pattern='options\.(max_retries|timeout|session_id|call_id)\b' +rule1_hits="$( + { grep -rnE "$rule1_pattern" aimux-providers/src --include='*.rs' || true; } | + grep -v 'HttpRequest::new(' | + grep -v '^aimux-providers/src/replay\.rs:' | + grep -vE '^aimux-providers/src/(elevenlabs\.rs|openai/transcription\.rs):[0-9]+: *timeout: options\.timeout,$' || + true +)" +violation 'aimux-providers/src reads options.max_retries / timeout / session_id / call_id' "$rule1_hits" + +# ── Rule 2: removed names stay removed (whole repository) ──────────────────── +# +# Source files only; docs/, rfc/, fixtures/, build output and vendored +# node_modules are history or data, not code. +search_removed() { + local pattern="$1" + grep -rnE "$pattern" . \ + --include='*.rs' --include='*.py' --include='*.pyi' --include='*.ts' \ + --include='*.js' --include='*.go' --include='*.java' --include='*.kt' \ + --include='*.swift' --include='*.dart' --include='*.c' --include='*.h' \ + --exclude-dir=target --exclude-dir=docs --exclude-dir=rfc \ + --exclude-dir=fixtures --exclude-dir=node_modules --exclude-dir=.git || + true +} + +for pattern in \ + 'fn config_snapshot' \ + 'fn retry_config' \ + 'RetryConfig' \ + 'Provider::name' \ + 'StaticBearerConfig' \ + 'api_key_source' \ + 'body_merge|apply_body_overrides|deep_merge_json'; do + violation "removed name \`$pattern\` is back" "$(search_removed "$pattern")" +done + +# `shared_client()` lives only in the injectable transport. +shared_client_hits="$( + search_removed 'shared_client\(\)' | + grep -v '^\./aimux-provider-utils/src/fetch\.rs:' || + true +)" +violation 'shared_client() used outside aimux-provider-utils/src/fetch.rs' "$shared_client_hits" + +# ── Rule 3: the OpenAI package reads its settings only through the model config ─ +# +# `openai/` has no `OpenAIConfig` field reads, builder-era names or `from_env`: +# the credential is loaded in the request headers, the URL comes from the +# config's `url` closure and provider-level body rewrites are +# `transform_request_body`. +rule3_hits="$( + grep -rnE 'config\.api_key\b|config\.base_url\b|body_overrides|api_key_source|from_env' \ + aimux-providers/src/openai --include='*.rs' || true +)" +violation 'aimux-providers/src/openai reads builder-era OpenAIConfig state' "$rule3_hits" + +# ── Rule 4: no provider-name comparisons in the shared protocol packages ───── +# +# `openai/`, `openai_compatible/`, `anthropic/`, `google/`, `bedrock/` and +# `cohere/` implement a protocol for every vendor that speaks it. Which vendor +# is in play is data the caller injects (`ChatDialect`, hooks, settings), so a +# `== "groq"` / `!= "openai"` comparison in there is the branch the factory +# rewrite removed. The names checked are every registry row plus the native +# packages; comparisons against media types, tool names and wire tags are fine. +# (Unit tests live next to the code and compare no provider names.) +rule4_names="$( + python3 - <<'PY' +import json + +with open("aimux-providers/src/provider_registry.json", encoding="utf-8") as f: + names = {row["name"] for row in json.load(f)} +names |= { + "openai", "anthropic", "google", "vertex", "bedrock", "cohere", "mistral", "xai", + "azure", "groq", "deepseek", "codex", "huggingface", "openai_compatible", +} +print("|".join(sorted(names, key=lambda n: (-len(n), n)))) +PY +)" +rule4_hits="$( + grep -rnE "(==|!=)[[:space:]]*\"($rule4_names)\"" \ + aimux-providers/src/openai aimux-providers/src/openai_compatible \ + aimux-providers/src/anthropic aimux-providers/src/google \ + aimux-providers/src/bedrock aimux-providers/src/cohere --include='*.rs' || true +)" +violation 'provider-name comparison in a shared protocol package (aimux-providers/src/{openai,openai_compatible,anthropic,google,bedrock,cohere})' "$rule4_hits" + +# `"groq"` / `"deepseek"` literals belong to their own packages and the registry. +rule4b_hits="$( + grep -rnE '"(groq|deepseek)"' \ + aimux-providers/src/openai aimux-providers/src/openai_compatible --include='*.rs' || true +)" +violation 'groq / deepseek literal in the shared OpenAI or compatible package' "$rule4b_hits" + +# ── Rule 5: the builder-era compat names stay removed in the providers crate ── +rule5_hits="$( + grep -rnE 'OpenAIConfig\b|OpenAIConfigProvider|OpenAICompatProfile|PLACEHOLDER_API_KEY|openai_legacy' \ + aimux-providers/src aimux-providers/tests --include='*.rs' || true +)" +violation 'OpenAIConfig / OpenAIConfigProvider / OpenAICompatProfile / PLACEHOLDER_API_KEY / openai_legacy is back in aimux-providers' "$rule5_hits" + +# ── Rule 6: the Anthropic family reads its settings only through the model config ─ +# +# No `AnthropicConfig` / `body_overrides` / `api_key_source` / `from_env`, and +# the providerOptions namespace goes through `options::anthropic_options` / +# `CANONICAL`: a `"anthropic"` literal anywhere else would read or write a key +# that ignores the provider's custom name. +anthropic_family='aimux-providers/src/anthropic aimux-providers/src/anthropic_aws aimux-providers/src/vertex/anthropic_model.rs' +rule6_hits="$( + # shellcheck disable=SC2086 + grep -rnE 'AnthropicConfig\b|body_overrides|api_key_source|from_env' \ + $anthropic_family --include='*.rs' || true +)" +violation 'Anthropic family reads builder-era AnthropicConfig state' "$rule6_hits" +rule6b_hits="$( + # shellcheck disable=SC2086 + { grep -rnE '"anthropic"' $anthropic_family --include='*.rs' || true; } | + grep -v '^aimux-providers/src/anthropic/options\.rs:' || + true +)" +violation '"anthropic" literal outside aimux-providers/src/anthropic/options.rs (use options::CANONICAL / anthropic_options)' "$rule6b_hits" + +# ── Rule 7: Google, Vertex and Bedrock read their settings through the model config ─ +# +# No `GoogleConfig` / `VertexProviderConfig` / `BedrockProviderConfig` / +# `body_overrides` / `api_key_source` / `from_env`; `bedrock::sigv4` is gone +# (signing is the `SigV4Fetch` decorator); and the canonical namespace keys +# (`google`, `googleVertex`, `amazonBedrock`) are quoted only in the two +# `options.rs` helpers (comments, unit tests and the SigV4 service name +# aside). The historical `vertex` / `bedrock` aliases are not ported at all, so +# they are flagged too: no model may read or write them. +gvb_family='aimux-providers/src/google aimux-providers/src/vertex aimux-providers/src/bedrock' +rule7_hits="$( + # shellcheck disable=SC2086 + grep -rnE 'GoogleConfig\b|VertexProviderConfig|BedrockProviderConfig|VertexAuth\b|BedrockAuth\b|api_key_source|from_env|body_overrides' \ + $gvb_family --include='*.rs' || true +)" +violation 'Google / Vertex / Bedrock reads builder-era config state' "$rule7_hits" +rule7b_hits="$( + search_removed 'bedrock::sigv4|bedrock/sigv4' | grep -v '^\./scripts/check_provider_boundaries\.sh:' || true +)" +violation 'bedrock::sigv4 is back (use SigV4Fetch)' "$rule7b_hits" +rule7c_hits="$( + # Non-test, non-comment lines of every file but the two helpers. + find $gvb_family -name '*.rs' ! -name options.rs -print0 | + xargs -0 awk ' + FNR == 1 { in_tests = 0 } + /^#\[cfg\(test\)\]/ { in_tests = 1 } + in_tests { next } + /^[[:space:]]*\/\// { next } + /"(google|googleVertex|vertex|amazonBedrock|bedrock)"/ && !/SIGV4_SERVICE/ { + printf "%s:%d:%s\n", FILENAME, FNR, $0 + } + ' || true +)" +violation 'providerOptions namespace key spelled outside google/options.rs or bedrock/options.rs' "$rule7c_hits" + +# ── Rule 8: the vendor packages read their settings through the model config ── +# +# No `XxxConfig` / `AzureAuth` / `TokenProvider` / `StaticBearerConfig` / +# `body_overrides` / `api_key_source` / `from_env`, and the namespace keys +# (`azure`, `xai`, `mistral`, `cohere`, `huggingface`, `voyage`, `elevenlabs`) +# are quoted only in the package's `options.rs` (comments, unit tests and the +# `DEFAULT_NAME` provider-name constants aside). Codex reads the OpenAI +# Responses namespace through the shared model, and Open Responses derives its +# key from the provider name, so neither spells a key. +vendor_family='aimux-providers/src/azure aimux-providers/src/xai aimux-providers/src/mistral aimux-providers/src/cohere aimux-providers/src/voyage aimux-providers/src/huggingface aimux-providers/src/huggingface.rs aimux-providers/src/codex.rs aimux-providers/src/open_responses.rs aimux-providers/src/elevenlabs aimux-providers/src/elevenlabs.rs' +rule8_hits="$( + # shellcheck disable=SC2086 + grep -rnE '\b(AzureConfig|AzureAuth|TokenProvider|XAIConfig|MistralConfig|CohereConfig|VoyageConfig|HuggingFaceConfig|CodexConfig|OpenResponsesConfig|ElevenLabsConfig|StaticBearerConfig)\b|body_overrides|api_key_source|from_env' \ + $vendor_family --include='*.rs' || true +)" +violation 'Azure / xAI / Mistral / Cohere / Hugging Face / Codex / Open Responses / Voyage / ElevenLabs reads builder-era config state' "$rule8_hits" +rule8b_hits="$( + # Non-test, non-comment lines of every file but the options helpers. + find $vendor_family -name '*.rs' ! -name options.rs -print0 | + xargs -0 awk ' + FNR == 1 { in_tests = 0 } + /^#\[cfg\(test\)\]/ { in_tests = 1 } + in_tests { next } + /^[[:space:]]*\/\// { next } + /"(azure|xai|mistral|cohere|huggingface|voyage|elevenlabs)"/ && !/DEFAULT_NAME/ { + printf "%s:%d:%s\n", FILENAME, FNR, $0 + } + ' || true +)" +violation 'vendor providerOptions namespace key spelled outside the package options.rs' "$rule8b_hits" + +# ── Rule 9: no operation-level retry inside aimux-providers/src ────────────── +# +# Core owns `prepare_retries`. Providers run their poll / download stages with +# the bounded helpers of `shared::poll`, and never re-submit a job. +rule9_hits="$( + grep -rnE 'prepare_retries|aimux_core::retry|use aimux_core::\{[^}]*\bretry\b' \ + aimux-providers/src --include='*.rs' || true +)" +violation 'prepare_retries / aimux_core::retry used in aimux-providers/src (exceptions: none)' "$rule9_hits" + +# ── Rule 10: the single-modality vendor packages read their settings through the model config ─ +# +# No `XxxConfig` / `api_key_source` / `from_env` / `with_base_url` / `body_overrides`, +# and the namespace keys (`lmnt`, `hume`, `deepgram`, `cartesia`, `stability`, +# `aws_polly`, `prodia`, `luma`, `klingai`, `replicate`, `fal`, `recraft`, +# `gladia`, `blackForestLabs`, `assemblyai`, `revai`, `linkup`, `jina`, +# `google_pse`) are quoted only in the package's `options.rs` (comments, unit +# tests and the `DEFAULT_NAME` provider-name constants aside). +single_family='serper prodia deepgram you_com luma lmnt klingai dataforseo replicate fal recraft aws_polly jina_ai gladia tavily linkup tinyfish black_forest_labs assemblyai revai hume cartesia searxng parallel_ai firecrawl runwayml exa_ai stability google_pse' +single_paths='' +for name in $single_family; do + single_paths="$single_paths aimux-providers/src/$name.rs" + [[ -d "aimux-providers/src/$name" ]] && single_paths="$single_paths aimux-providers/src/$name" +done +rule10_hits="$( + # shellcheck disable=SC2086 + grep -rnE '\b(SerperConfig|ProdiaConfig|DeepgramConfig|YouComConfig|LumaConfig|LMNTConfig|KlingAIConfig|DataforseoConfig|ReplicateConfig|FalConfig|RecraftConfig|AwsPollyConfig|JinaAiConfig|GladiaConfig|TavilyConfig|LinkupConfig|TinyfishConfig|BlackForestLabsConfig|AssemblyAIConfig|RevaiConfig|HumeConfig|CartesiaConfig|SearxngConfig|ParallelAiConfig|FirecrawlConfig|RunwaymlConfig|ExaAiConfig|StabilityConfig|GooglePseConfig)\b|body_overrides|api_key_source|from_env|with_base_url' \ + $single_paths --include='*.rs' || true +)" +violation 'single-modality vendor packages read builder-era config state' "$rule10_hits" +rule10b_hits="$( + # Non-test, non-comment lines of every file but the options helpers. + # shellcheck disable=SC2086 + find $single_paths -name '*.rs' ! -name options.rs -print0 | + xargs -0 awk ' + FNR == 1 { in_tests = 0 } + /^#\[cfg\(test\)\]/ { in_tests = 1 } + in_tests { next } + /^[[:space:]]*\/\// { next } + /"(lmnt|hume|deepgram|cartesia|stability|aws_polly|prodia|luma|klingai|replicate|fal|recraft|gladia|blackForestLabs|assemblyai|revai|linkup|jina|google_pse)"/ && !/DEFAULT_NAME/ { + printf "%s:%d:%s\n", FILENAME, FNR, $0 + } + ' || true +)" +violation 'single-modality providerOptions namespace key spelled outside the package options.rs' "$rule10b_hits" + +if [[ "$status" -eq 0 ]]; then + echo 'provider boundaries: ok' +fi +exit "$status" diff --git a/scripts/gen_providers_doc.py b/scripts/gen_providers_doc.py index e62626cad..c02044bc8 100644 --- a/scripts/gen_providers_doc.py +++ b/scripts/gen_providers_doc.py @@ -2,7 +2,9 @@ """Generate docs/api/providers.md — the full provider lookup reference. Single source of truth: -- aimux-providers/src/provider_registry.json (registry of OpenAI-compatible entries) +- aimux-providers/src/provider_registry.json (registry of OpenAI-compatible presets: + name, display, base_url, env_var, auth, + base_url_env, params, profile) - aimux-providers/src/lib.rs (non-registry modules + categories) A `pub mod` in lib.rs counts as a provider module iff its `pub use` re-exports @@ -121,7 +123,7 @@ def module_exports(module): j = text.find(";", start) body = text[start:j] pos = j + 1 - for name in re.findall(r"\b([A-Z][A-Za-z0-9_]+)", body): + for name in re.findall(r"\b([A-Z][A-Za-z0-9_]+|create_[a-z0-9_]+)", body): if name not in exports: exports.append(name) return exports @@ -190,17 +192,28 @@ def build_page(): w("") w( f"**{len(registry)} registry-backed OpenAI-compatible providers** " - f"(construct via `provider(name, ...)`) + " + f"(create presets by name via `provider(name, ...)`) + " f"**{non_registry} non-registry providers** " "(construct via the typed factories listed below)." ) w("") w(f"## Registry-backed (OpenAI-compatible) — {len(registry)}") w("") - w("| name | display | env var | base_url |") - w("|------|---------|---------|----------|") + w("| name | display | auth | env var | base_url |") + w("|------|---------|------|---------|----------|") for e in sorted(registry, key=lambda x: x["name"]): - w(f"| `{e['name']}` | {e.get('display', '')} | `{e.get('env_var', '')}` | `{e.get('base_url', '')}` |") + auth = e.get("auth", "api_key") + env_var = f"`{e['env_var']}`" if e.get("env_var") else "—" + base = f"`{e.get('base_url', '')}`" + notes = [] + if e.get("base_url_env"): + notes.append(f"URL from `{e['base_url_env']}`") + params = [p["name"] for p in e.get("params", []) if "derive" not in p] + if params: + notes.append("params: " + ", ".join(f"`{p}`" for p in params)) + if notes: + base += " (" + "; ".join(notes) + ")" + w(f"| `{e['name']}` | {e.get('display', '')} | {auth} | {env_var} | {base} |") w("") w("## Typed factories (non-registry)") w("") @@ -218,9 +231,14 @@ def build_page(): w("|--------|--------------------|") for mod in mods: exports = module_exports(mod) - display = display_from_exports(exports) - if display: - cols = f"`{display}Config` / `{display}Provider`" if f"{display}Config" in exports or f"{display}Provider" in exports else ", ".join(f"`{e}`" for e in exports) + entry_points = [ + e + for e in exports + if e.endswith(("Config", "ProviderConfig", "ProviderSettings", "Provider")) + or e.startswith("create_") + ] + if entry_points: + cols = " / ".join(f"`{e}`" for e in entry_points) else: cols = "`" + "`, `".join(exports) + "`" if exports else "—" w(f"| `{mod}` | {cols} |") diff --git a/scripts/generate_thin_wrapper_cassettes.py b/scripts/generate_thin_wrapper_cassettes.py index 999514290..bc8394605 100644 --- a/scripts/generate_thin_wrapper_cassettes.py +++ b/scripts/generate_thin_wrapper_cassettes.py @@ -2,7 +2,7 @@ """Generate cassettes for OpenAI-compatible thin-wrapper providers. These providers (alibaba, baseten, bytedance, deepinfra, fireworks, -moonshotai, togetherai, vercel) are thin wrappers around OpenAIProvider. +moonshotai, togetherai, vercel) are thin wrappers around OpenAIConfigProvider. Their responses are byte-for-byte OpenAI Chat Completions format, so we can derive valid cassettes from existing OpenAI recordings by rewriting the request path and model field. diff --git a/tools/aimux-cli/src/probe/provider.rs b/tools/aimux-cli/src/probe/provider.rs index e8b59785c..113e890b1 100644 --- a/tools/aimux-cli/src/probe/provider.rs +++ b/tools/aimux-cli/src/probe/provider.rs @@ -7,7 +7,6 @@ use std::sync::Arc; use aimux_core::trace::{RingTraceStore, TraceFilter, TraceLayer, VerdictKind}; -use aimux_providers::ProviderOptions; use crate::report; use crate::{Format, ProviderArgs}; @@ -81,44 +80,22 @@ condition under which provider-side prompt caching should activate; observe the usage fields to confirm the cache read tokens; compare with the client LCP bound; report any discrepancy as a suspect overclaim; the test sequence is complete."#; -/// 构造 provider 模型:原生单 key provider(openai/anthropic/mistral/xai/cohere/ -/// google)直接构造;其余走注册表(compat,如 deepseek/groq/moonshotai 等)。 -/// azure/bedrock/vertex 需额外参数(资源名/region/凭证),CLI 第一版不直接 -/// 支持,可用其 OpenAI-compat 注册表镜像或后续扩展。 +/// 构造 provider 模型:名字可以是任一内置厂商(厂商包或注册表行)。 fn build_model( provider: &str, api_key: String, model_id: &str, base_url: Option<&str>, ) -> anyhow::Result> { - macro_rules! native { - ($provider_mod:ident, $config:ident, $provider_type:ident) => {{ - let mut cfg = aimux_providers::$provider_mod::$config::new(api_key.clone()); - if let Some(url) = base_url { - cfg = cfg.with_base_url(url); - } - let p = aimux_providers::$provider_mod::$provider_type::new(cfg); - Arc::from(p.model(model_id)) - }}; - } - - match provider { - "openai" => Ok(native!(openai, OpenAIConfig, OpenAIProvider)), - "anthropic" => Ok(native!(anthropic, AnthropicConfig, AnthropicProvider)), - "mistral" => Ok(native!(mistral, MistralConfig, MistralProvider)), - "xai" => Ok(native!(xai, XAIConfig, XAIProvider)), - "cohere" => Ok(native!(cohere, CohereConfig, CohereProvider)), - "google" => Ok(native!(google, GoogleConfig, GoogleProvider)), - _ => { - let mut options = ProviderOptions::default(); - if let Some(url) = base_url { - options.base_url = Some(url.to_string()); - } - let model = aimux_providers::provider(provider, Some(api_key), model_id, Some(options)) - .map_err(|e| anyhow::anyhow!("provider '{provider}': {e}"))?; - Ok(Arc::from(model)) - } - } + let provider = aimux_providers::create_provider( + provider, + aimux_providers::PresetSettings { + api_key: Some(api_key.into()), + base_url: base_url.map(str::to_string), + ..Default::default() + }, + )?; + Ok(provider.language_model(model_id)?) } /// 解析 `api_key` 参数:`env:VAR` 引用环境变量,否则按字面 key 使用。 @@ -224,8 +201,10 @@ pub async fn run(args: &ProviderArgs) -> anyhow::Result anyhow::Result ( std::net::SocketAddr, @@ -334,14 +313,13 @@ mod tests { let cached = if n == 0 { 0 } else { 512 }; let body = format!( r#"{{ - "id": "chatcmpl-mock", - "model": "gpt-4o", - "choices": [{{"message": {{"role": "assistant", "content": "ok"}}, "finish_reason": "stop"}}], + "id": "resp-mock", + "output": [{{"type": "message", "content": [{{"type": "output_text", "text": "ok"}}]}}], "usage": {{ - "prompt_tokens": 2048, - "completion_tokens": 5, + "input_tokens": 2048, + "output_tokens": 5, "total_tokens": 2053, - "prompt_tokens_details": {{"cached_tokens": {cached}}} + "input_tokens_details": {{"cached_tokens": {cached}}} }} }}"# ); diff --git a/tools/aimux-replay/src/main.rs b/tools/aimux-replay/src/main.rs index 3e4622c45..ad7d6511b 100644 --- a/tools/aimux-replay/src/main.rs +++ b/tools/aimux-replay/src/main.rs @@ -1,16 +1,17 @@ //! aimux-replay — RFC-0023 请求回放 CLI(层 2 消费端)。 //! -//! 读录制 jsonl(每行一个 `Recording`),按 `ProviderRecord` 自动重建 -//! provider(语言模型 OpenAI 兼容族、官方 TypeSafe Jev decision),再用录制输入 **重发真实 -//! API**。用途:离线重跑线上流量、改 prompt 重发(A/B)、回归对比、CI 集成。 +//! 读录制 jsonl(每行一个 `Recording`,schema 1 或 2),按 `ProviderRecord` 的 +//! `provider_id` + `model_id` 重建 model,再用录制输入经 `replay_with_model` +//! **重发真实 API**。用途:离线重跑线上流量、改 prompt 重发(A/B)、回归对比、CI 集成。 //! //! 安全: -//! - `--dry-run` 只打印重建后的 provider/prompt/目标 URL,**不发请求、 -//! 不输出任何凭据**(api_key 来源按 `api_key_source` 打印)。 +//! - `--dry-run` 只打印录制的 provider/model/prompt,**不发请求、 +//! 不输出任何凭据**。 //! - 重发会消耗真实 token/费用,文档(§3.6.1)已有警示。 -//! - `--mock` 使用统一录制的响应离线回放；decision 使用规范化结果。 -//! - 其他原生语言协议 provider(anthropic/google/...)`rebuild_provider` 明确 -//! `Unsupported`——CLI 打印错误并跳过,需自行传 model 实例走库 API。 +//! - `--mock` 离线回放语言或 decision 录制；真实回放只使用调用方配置， +//! 不从录制恢复连接设置。语言与 decision 录制都从同一个 registry 重建 +//! ([`replay_registry`]:全部内置 provider,`--api-key` 替换录制的那个); +//! Jev 不在内置集合里,录制指向它时加入,使用显式 key 或 TYPESAFE_API_KEY。 use std::path::PathBuf; @@ -19,6 +20,7 @@ use clap::Parser; use tokio::runtime::Runtime; use aimux_core::language_model_message::{AssistantPart, LanguageModelMessage, UserPart}; +use aimux_core::provider_registry::ProviderRegistry; use aimux_core::recording::Recording; use aimux_core::replay::{ReplayOverrides, replay_with_model}; use aimux_providers::rebuild_provider; @@ -35,16 +37,16 @@ struct Cli { /// 仅回放该 call_id(默认全部)。 #[arg(long)] call_id: Option, - /// 显式 api key。api_key_source 为 explicit/unknown 时需要。 + /// 显式 api key;缺省时读 registry 条目的环境变量。 #[arg(long)] api_key: Option, /// 覆盖 prompt(替换所有录制的 prompt;A/B 重发)。 #[arg(long)] prompt: Option, - /// dry-run:打印重建后的 provider/prompt/目标 URL,不发请求、不输出凭据。 + /// dry-run:打印录制的 provider/model/prompt,不发请求、不输出凭据。 #[arg(long)] dry_run: bool, - /// Replay recorded responses offline instead of sending a live request. + /// Replay recorded responses offline. #[arg(long)] mock: bool, } @@ -104,14 +106,9 @@ fn load_recordings(path: &PathBuf) -> Result> { fn run_dry(rec: &Recording) -> Result<()> { println!("call_id: {}", rec.call_id); println!( - " provider: {} / model: {}", - rec.provider.provider, rec.provider.model_id + " provider_id: {} / provider: {} / model: {}", + rec.provider.provider_id, rec.provider.provider, rec.provider.model_id ); - println!( - " base_url: {}", - rec.provider.base_url.as_deref().unwrap_or("(default)") - ); - println!(" api_key_source: {}", rec.provider.api_key_source); if rec.input.operation == aimux_core::recording::RecordingOperation::Decision { println!( " decision: {} question(s)", @@ -119,11 +116,10 @@ fn run_dry(rec: &Recording) -> Result<()> { .as_array() .map_or(0, Vec::len) ); - } else { - println!(" prompt: {} message(s)", rec.input.prompt.len()); - if let Some(text) = rec.input.prompt.first().and_then(prompt_text) { - println!(" first message: {text:?}"); - } + } + println!(" prompt: {} message(s)", rec.input.prompt.len()); + if let Some(text) = rec.input.prompt.first().and_then(prompt_text) { + println!(" first message: {text:?}"); } println!(" (dry-run — no request sent)"); Ok(()) @@ -158,14 +154,22 @@ fn run_replay( "--prompt applies to language recordings; decision replay preserves state/questions" ); } - let model: Box = if mock { + let model: Box = if mock { Box::new(aimux_core::replay::MockDecisionReplayModel::new( &rec.provider.provider, &rec.provider.model_id, vec![rec.clone()], )?) } else { - aimux_providers::rebuild_decision_provider(&rec.provider, api_key)? + let registry = replay_registry(rec.provider.registry_id()?, api_key)?; + aimux_providers::rebuild_decision_provider(&rec.provider, ®istry).with_context( + || { + format!( + "rebuild decision provider for '{}'", + rec.provider.provider_id + ) + }, + )? }; let result = Runtime::new()?.block_on(aimux_core::replay::replay_decision_with_model( rec, @@ -180,15 +184,16 @@ fn run_replay( ); return Ok(()); } - let model: Box = if mock { - Box::new(aimux_core::replay::MockReplayModel::new( + let model: std::sync::Arc = if mock { + std::sync::Arc::new(aimux_core::replay::MockReplayModel::new( &rec.provider.provider, &rec.provider.model_id, vec![rec.clone()], )) } else { - rebuild_provider(&rec.provider, api_key) - .with_context(|| format!("rebuild provider for '{}'", rec.provider.provider))? + let registry = replay_registry(rec.provider.registry_id()?, api_key)?; + rebuild_provider(&rec.provider, ®istry) + .with_context(|| format!("rebuild provider for '{}'", rec.provider.provider_id))? }; let overrides = prompt.map(|text| ReplayOverrides { @@ -214,6 +219,40 @@ fn run_replay( Ok(()) } +/// The providers live replay rebuilds a recorded model from, for language and +/// decision recordings alike: every built-in provider with default settings +/// (`default_providers()`), the recorded one created again with `api_key` +/// when given. Jev, the decision provider outside the built-in set, is added +/// when the recording names it, with `api_key` or `TYPESAFE_API_KEY`. +fn replay_registry(provider_id: &str, api_key: Option<&str>) -> Result { + let mut providers = aimux_providers::default_providers(); + if provider_id == "jev" { + let config = match api_key { + Some(key) => aimux_providers::JevConfig::new(key), + None => aimux_providers::JevConfig::from_env()?, + }; + providers.insert( + provider_id.to_string(), + std::sync::Arc::new(aimux_providers::JevProvider::new(config)), + ); + } else if let Some(key) = api_key { + providers.insert( + provider_id.to_string(), + aimux_providers::create_provider( + provider_id, + aimux_providers::PresetSettings { + api_key: Some(key.to_string().into()), + ..Default::default() + }, + )?, + ); + } + Ok(aimux_core::create_provider_registry( + providers, + Default::default(), + )) +} + fn truncate(s: &str, max: usize) -> String { if s.chars().count() <= max { s.to_string() @@ -222,3 +261,158 @@ fn truncate(s: &str, max: usize) -> String { format!("{cut}…(truncated)") } } + +#[cfg(test)] +mod tests { + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::sync::{Arc, Mutex}; + + use aimux_core::Provider; + use aimux_core::decision_model::{DecisionCallOptions, decide}; + use aimux_core::recording::{self, RecordingOperation, RingRecorder}; + use aimux_providers::{OpenAIProviderSettings, create_openai}; + use serde_json::Value; + + use super::*; + + /// The OpenAI Decisions contract fixture of the provider tests. + const OPENAI_DECISIONS: &str = + include_str!("../../../aimux-providers/tests/fixtures/openai_decisions.json"); + + /// An HTTP server on a thread that answers every request with `body` and + /// keeps each raw request. + fn serve(body: String) -> (String, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}/v1", listener.local_addr().unwrap()); + let seen = Arc::new(Mutex::new(Vec::new())); + let sink = seen.clone(); + std::thread::spawn(move || { + for mut socket in listener.incoming().flatten() { + let mut raw = Vec::new(); + let mut buf = [0u8; 4096]; + loop { + let n = socket.read(&mut buf).unwrap_or(0); + raw.extend_from_slice(&buf[..n]); + let text = String::from_utf8_lossy(&raw).to_string(); + let complete = text.split_once("\r\n\r\n").is_some_and(|(head, rest)| { + let length = head + .lines() + .find_map(|l| { + l.to_ascii_lowercase() + .strip_prefix("content-length:") + .map(|v| v.trim().parse::().unwrap_or(0)) + }) + .unwrap_or(0); + rest.len() >= length + }); + if n == 0 || complete { + break; + } + } + sink.lock() + .unwrap() + .push(String::from_utf8_lossy(&raw).to_string()); + let reply = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + let _ = socket.write_all(reply.as_bytes()); + } + }); + (url, seen) + } + + /// Sets an environment variable for the test and restores it on drop. + struct EnvVar(&'static str, Option); + + impl EnvVar { + fn set(name: &'static str, value: &str) -> Self { + let previous = std::env::var_os(name); + // SAFETY: the only test in this crate that touches the variable. + unsafe { std::env::set_var(name, value) }; + Self(name, previous) + } + } + + impl Drop for EnvVar { + fn drop(&mut self) { + // SAFETY: as in `set`. + unsafe { + match &self.1 { + Some(value) => std::env::set_var(self.0, value), + None => std::env::remove_var(self.0), + } + } + } + } + + /// B17: a decision recording of a provider other than Jev replays live + /// through the registry the language path uses. The recorded OpenAI + /// decision is rebuilt from `default_providers()` with `--api-key` and + /// sent again (`OPENAI_BASE_URL` points the default provider at the + /// local server). + #[test] + fn an_openai_decision_recording_replays_live() { + let fixture: Value = serde_json::from_str(OPENAI_DECISIONS).unwrap(); + let (base_url, seen) = serve(fixture["response"].to_string()); + + let recorder = Arc::new(RingRecorder::new()); + recording::init_recording(Some(recorder.clone())); + let model = create_openai(OpenAIProviderSettings { + api_key: Some("record-key".into()), + base_url: Some(base_url.clone()), + ..Default::default() + }) + .unwrap() + .decision_model("gpt-6-luna") + .unwrap(); + let options: DecisionCallOptions = + serde_json::from_value(fixture["options"].clone()).unwrap(); + Runtime::new() + .unwrap() + .block_on(decide(model.as_ref(), options)) + .unwrap(); + recording::init_recording(None); + let rec = recorder.completed().pop().expect("one recording"); + assert_eq!(rec.input.operation, RecordingOperation::Decision); + assert_eq!(rec.provider.provider_id, "openai"); + + let _base_url = EnvVar::set("OPENAI_BASE_URL", &base_url); + run_replay(&rec, Some("replay-key"), None, false).unwrap(); + + let seen = seen.lock().unwrap(); + assert_eq!(seen.len(), 2, "the recorded call and its replay"); + assert!(seen[1].starts_with("POST /v1/decisions"), "{}", seen[1]); + assert!( + seen[1].to_ascii_lowercase().contains("bearer replay-key"), + "the replay uses the caller's key" + ); + } + + /// Every decision provider resolves in the replay registry: a built-in + /// one by its name, Jev when the recording names it; an unknown name is + /// reported. + #[test] + fn the_replay_registry_holds_every_decision_provider() { + let registry = replay_registry("ollama", None).unwrap(); + assert!( + registry + .provider("ollama") + .unwrap() + .decision_model("m") + .is_ok() + ); + let registry = replay_registry("jev", Some("k")).unwrap(); + assert!( + registry + .provider("jev") + .unwrap() + .decision_model("m") + .is_ok() + ); + let record = aimux_core::recording::ProviderRecord::new("no-such", "no-such", "m"); + let registry = replay_registry("no-such", None).unwrap(); + assert!(aimux_providers::rebuild_decision_provider(&record, ®istry).is_err()); + } +} diff --git a/tools/aimux-web/README.md b/tools/aimux-web/README.md index afd90e759..0ecee5fe4 100644 --- a/tools/aimux-web/README.md +++ b/tools/aimux-web/README.md @@ -74,7 +74,8 @@ env var。非回环绑定(`--host 0.0.0.0` 等)时网页内 key 管理自动禁 自填 `base_url`(不是该 provider 的默认地址)时只用请求里直接给的 key(仅回环 绑定可填明文):Settings 已存 key 和环境变量都不会发往该地址,`env:` 引用被拒绝; key 留空则以空 key 调用(与 AI SDK 显式传 `apiKey: ''` 相同,不回落到任何已存或环境 key; -OpenAI 等 provider 因此发一个空凭据头,如 `Authorization: Bearer`)。Replay 一条指向非默认地址的录制时同样如此。 +部分 provider 因此发一个空凭据头,如 OpenAI 的 `Authorization: Bearer`)。还会从环境 +读取其他凭据的 provider(Bedrock、Vertex、Codex 等)不能自填 `base_url`。Replay 总是发往 provider 的默认地址。 打开后使用 Playground / Agent 页跑一次调用,Traces 页即可看到完整录制。 diff --git a/tools/aimux-web/examples/gen_demo_fixtures.rs b/tools/aimux-web/examples/gen_demo_fixtures.rs index 071b57af9..242628865 100644 --- a/tools/aimux-web/examples/gen_demo_fixtures.rs +++ b/tools/aimux-web/examples/gen_demo_fixtures.rs @@ -45,18 +45,12 @@ fn recording( call_id: call_id.into(), recorded_at: recorded_at.into(), input: InputRecord { + decision_capabilities: None, operation: Default::default(), prompt: prompt.clone(), options: serde_json::to_value(&options).unwrap(), }, - provider: ProviderRecord { - provider: provider.into(), - model_id: model_id.into(), - base_url: None, - api_key_source: "none".into(), - profile: None, - provider_options: None, - }, + provider: ProviderRecord::from_model(provider, model_id), exchanges: vec![HttpExchange { step: None, attempt: 0, diff --git a/tools/aimux-web/examples/gen_fixture.rs b/tools/aimux-web/examples/gen_fixture.rs index edf52366e..ac1e394b0 100644 --- a/tools/aimux-web/examples/gen_fixture.rs +++ b/tools/aimux-web/examples/gen_fixture.rs @@ -35,18 +35,12 @@ fn main() { call_id: "call-fixture-1".into(), recorded_at: "2026-08-14T00:00:00.000Z".into(), input: InputRecord { + decision_capabilities: None, operation: Default::default(), prompt: prompt.clone(), options: serde_json::to_value(&options).unwrap(), }, - provider: ProviderRecord { - provider: "openai".into(), - model_id: "gpt-4o-mini".into(), - base_url: None, - api_key_source: "none".into(), - profile: None, - provider_options: None, - }, + provider: ProviderRecord::new("openai", "openai.chat", "gpt-4o-mini"), exchanges: vec![HttpExchange { step: None, attempt: 0, diff --git a/tools/aimux-web/src/api/cache_probe.rs b/tools/aimux-web/src/api/cache_probe.rs index a2492541e..ad14a60ac 100644 --- a/tools/aimux-web/src/api/cache_probe.rs +++ b/tools/aimux-web/src/api/cache_probe.rs @@ -133,8 +133,10 @@ async fn run_inner(state: AppState, req: ProbeRequest) -> Result Result); + + impl EnvVar { + fn set(name: &'static str, value: &str) -> Self { + let previous = std::env::var_os(name); + // SAFETY: no other test reads or writes these providers' variables. + unsafe { std::env::set_var(name, value) }; + Self(name, previous) + } + } + + impl Drop for EnvVar { + fn drop(&mut self) { + // SAFETY: as in `set`. + unsafe { + match &self.1 { + Some(value) => std::env::set_var(self.0, value), + None => std::env::remove_var(self.0), + } + } + } + } + /// A caller-supplied base_url never receives the Settings key or the /// provider's environment key: with no key in the request, the call goes /// out without one. #[tokio::test] async fn custom_base_url_gets_no_stored_or_env_key() { - // A registry-style provider whose key variable this test owns. - aimux_providers::load_providers_from_json( - r#"{"providers":[{"name":"web-test-custom-base","base_url":"https://example.invalid/v1","env_var":"AIMUX_WEB_TEST_CUSTOM_BASE_KEY"}]}"#, - ) - .unwrap(); - // SAFETY: This test owns a unique process-wide env key and removes it below. - unsafe { std::env::set_var("AIMUX_WEB_TEST_CUSTOM_BASE_KEY", "env-secret-7") }; + // A vendor package and a registry row, with their own key variables. + let _package_key = EnvVar::set("DEEPSEEK_API_KEY", "env-secret-7"); + let _row_key = EnvVar::set("API2GPT_API_KEY", "env-secret-8"); let keys = KeyStore::from_path(None); - keys.set("openai", "sk-stored-secret-1", false).unwrap(); - keys.set("web-test-custom-base", "sk-stored-secret-2", false) + keys.set("deepseek", "sk-stored-secret-1", false).unwrap(); + keys.set("api2gpt", "sk-stored-secret-2", false).unwrap(); + keys.set("amazon_bedrock", "sk-stored-secret-3", false) .unwrap(); let app = router(AppState::with_bind_host_and_store("127.0.0.1", keys)); let (url, seen) = capture_server().await; - for provider in ["openai", "web-test-custom-base"] { + for provider in ["deepseek", "api2gpt"] { let (status, body) = call(app.clone(), request(provider, &url, None)).await; assert_eq!(status, StatusCode::OK, "{provider}: {body}"); } // An env: reference is refused outright for a custom base_url. let (status, body) = call( app.clone(), - request( - "web-test-custom-base", - &url, - Some("env:AIMUX_WEB_TEST_CUSTOM_BASE_KEY"), - ), + request("api2gpt", &url, Some("env:API2GPT_API_KEY")), ) .await; assert_eq!(status, StatusCode::BAD_REQUEST, "{body}"); + // A provider that reads other credentials from the environment (AWS) + // is not called at a custom base_url. + let (status, body) = call(app.clone(), request("amazon_bedrock", &url, None)).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{body}"); // A key carried in the request (loopback) is what the endpoint gets. - let (status, body) = call(app, request("openai", &url, Some("sk-for-this-endpoint"))).await; + let (status, body) = + call(app, request("deepseek", &url, Some("sk-for-this-endpoint"))).await; assert_eq!(status, StatusCode::OK, "{body}"); - // SAFETY: Cleanup of the unique test key set above. - unsafe { std::env::remove_var("AIMUX_WEB_TEST_CUSTOM_BASE_KEY") }; - let seen = seen.lock().unwrap(); - assert_eq!(seen.len(), 3, "the refused call must not reach the server"); + assert_eq!(seen.len(), 3, "the refused calls must not reach the server"); for raw in seen.iter() { - for secret in ["sk-stored-secret-1", "sk-stored-secret-2", "env-secret-7"] { + for secret in [ + "sk-stored-secret-1", + "sk-stored-secret-2", + "sk-stored-secret-3", + "env-secret-7", + "env-secret-8", + ] { assert!(!raw.contains(secret), "{secret} leaked: {raw}"); } } assert!(seen[2].contains("sk-for-this-endpoint"), "{}", seen[2]); } + + /// B18: calls of one provider and model at two endpoints get two trace + /// scopes; calls at one endpoint share a scope. + #[tokio::test] + async fn each_endpoint_gets_its_own_trace_scope() { + let state = AppState::with_bind_host_and_store("127.0.0.1", KeyStore::from_path(None)); + let app = router(state.clone()); + let (first, _) = capture_server().await; + let (second, _) = capture_server().await; + + for url in [&first, &first, &second] { + let (status, body) = + call(app.clone(), request("deepseek", url, Some("sk-endpoint"))).await; + assert_eq!(status, StatusCode::OK, "{body}"); + } + + let scopes: Vec = state + .trace_sink + .records() + .iter() + .map(|record| record.scope_key) + .collect(); + assert_eq!(scopes.len(), 3, "{scopes:?}"); + assert_eq!(scopes[0], scopes[1], "one endpoint, one scope"); + assert_ne!(scopes[0], scopes[2], "two endpoints, two scopes"); + // Records still report the model's provider, not the scope name. + let records = state.trace_sink.records(); + assert!( + records + .iter() + .all(|record| record.provider == records[0].provider + && !record.provider.contains(&first)) + ); + } } diff --git a/tools/aimux-web/src/api/mod.rs b/tools/aimux-web/src/api/mod.rs index b8beed577..cb74c3ca0 100644 --- a/tools/aimux-web/src/api/mod.rs +++ b/tools/aimux-web/src/api/mod.rs @@ -59,6 +59,8 @@ pub fn err_response(e: AiMuxError) -> Response { let status = match &e { AiMuxError::InvalidArgument(_) | AiMuxError::InvalidPrompt(_) + | AiMuxError::LoadApiKey { .. } + | AiMuxError::LoadSetting { .. } | AiMuxError::NoSuchProvider { .. } | AiMuxError::NoSuchModel { .. } => StatusCode::BAD_REQUEST, AiMuxError::UnsupportedFunctionality(_) => StatusCode::NOT_IMPLEMENTED, diff --git a/tools/aimux-web/src/api/providers.rs b/tools/aimux-web/src/api/providers.rs index 5393794c7..544dfa2ce 100644 --- a/tools/aimux-web/src/api/providers.rs +++ b/tools/aimux-web/src/api/providers.rs @@ -4,16 +4,9 @@ use axum::Json; use axum::response::{IntoResponse, Response}; use serde_json::{Value, json}; -/// Provider names: the six native single-key protocols plus every -/// registry-backed OpenAI-compatible provider. +/// Provider names: every built-in provider. pub fn provider_names() -> Vec { - const NATIVE: [&str; 6] = ["openai", "anthropic", "google", "mistral", "xai", "cohere"]; - let mut names: Vec = NATIVE - .iter() - .map(std::string::ToString::to_string) - .collect(); - names.extend(aimux_providers::provider_names().map(str::to_owned)); - names + aimux_providers::default_providers().into_keys().collect() } /// `GET /api/providers` — the provider list and suggested models for common ones. diff --git a/tools/aimux-web/src/api/replay.rs b/tools/aimux-web/src/api/replay.rs index 2c6fddb03..d5f7ac676 100644 --- a/tools/aimux-web/src/api/replay.rs +++ b/tools/aimux-web/src/api/replay.rs @@ -12,11 +12,10 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::convert_to_language_model_prompt; use aimux_core::message::ModelPrompt; -use aimux_core::recording::{ProviderRecord, Recording}; +use aimux_core::recording::Recording; use aimux_core::replay::{MockReplayModel, ReplayOverrides, replay_with_model}; use crate::api::err_response; -use crate::model_builder; use crate::state::AppState; use crate::wire::{self, WireMessage}; @@ -61,39 +60,37 @@ async fn run_inner(state: AppState, req: ReplayRequest) -> Result { - wire::resolve_api_key(Some(s), &provider_name, &state.keys, state.loopback, custom)? - } - "none" => wire::resolve_api_key(None, &provider_name, &state.keys, state.loopback, custom)?, - _ => wire::resolve_api_key( - req.api_key.as_deref(), - &provider_name, - &state.keys, - state.loopback, - custom, - ) - .map_err(|_| { - AiMuxError::InvalidArgument( - "this recording has no env key source — pass api_key=\"env:VAR\" or save the \ - provider key in Settings to replay" - .into(), - ) - })?, - }; - - let model = aimux_providers::rebuild_provider(&recording.provider, key.as_deref())?; + // Resolve the key for provider rebuild: the request spec first, then the + // Settings key store; otherwise the provider's registered env var is read + // (RFC-0029 §5.5). The recording stores no credential source and no base + // URL (only the provider identity), so the replay goes to the provider's + // default endpoint: never a custom base URL. + // A recording written by aimux 0.5 or earlier has no provider_id to + // rebuild from; refuse it before any key is resolved. + let provider_id = recording.provider.registry_id()?; + let key = wire::resolve_api_key( + req.api_key.as_deref(), + provider_id, + &state.keys, + state.loopback, + false, + )?; + + let mut providers = aimux_providers::default_providers(); + if let Some(key) = key { + providers.insert( + provider_id.to_string(), + aimux_providers::create_provider( + provider_id, + aimux_providers::PresetSettings { + api_key: Some(key.into()), + ..Default::default() + }, + )?, + ); + } + let registry = aimux_core::create_provider_registry(providers, Default::default()); + let model = aimux_providers::rebuild_provider(&recording.provider, ®istry)?; let mut overrides = ReplayOverrides { prompt: None, @@ -129,17 +126,6 @@ async fn run_inner(state: AppState, req: ReplayRequest) -> Result Option { - p.provider_options - .as_ref() - .and_then(|options| options.get("base_url")) - .and_then(serde_json::Value::as_str) - .or(p.base_url.as_deref()) - .map(str::to_string) -} - /// `POST /api/mock/load` — load NDJSON recordings into a `MockReplayModel` /// and switch the console to offline mock mode (no real API calls). pub async fn mock_load( @@ -159,8 +145,14 @@ fn mock_load_inner(state: AppState, req: MockLoadRequest) -> Result, body: String) -> Respon count += 1; } Err(e) => { - return err_response(AiMuxError::JsonParse(format!( - "import line {}: {e}", - idx + 1 - ))); + let msg = format!("import line {}: {e}", idx + 1); + return err_response(if e.classify() == serde_json::error::Category::Data { + AiMuxError::InvalidArgument(msg) + } else { + AiMuxError::JsonParse(msg) + }); } } } diff --git a/tools/aimux-web/src/model_builder.rs b/tools/aimux-web/src/model_builder.rs index 9fa6f4640..e5699f879 100644 --- a/tools/aimux-web/src/model_builder.rs +++ b/tools/aimux-web/src/model_builder.rs @@ -1,104 +1,71 @@ -//! Provider construction for the console (RFC-0029 §8.2): native single-key -//! providers and the registry-backed OpenAI-compatible family. -//! -//! Adapted from `aimux-cli` `probe::provider::build_model` so the console and -//! the CLI share the same provider construction semantics. +//! Provider construction for the console (RFC-0029 §8.2). use std::sync::Arc; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; -use aimux_providers::provider::ProviderOptions; +use aimux_providers::PresetSettings; /// Build a model from a provider name + model id + optional key / base URL. /// -/// - Native providers (openai/anthropic/google/mistral/xai/cohere) need a key: -/// explicit `api_key` or the provider's standard env var. -/// - Everything else goes through the registry (`provider(name, …)`), where a -/// `None` key reads the provider's registered env var. +/// The name is any built-in provider (a vendor package or a registry row); a +/// `None` key is read from that provider's environment variable when the +/// request is made. pub fn build_model( provider: &str, api_key: Option, model_id: &str, base_url: Option<&str>, ) -> Result, AiMuxError> { - macro_rules! native { - ($provider_mod:ident, $config:ident, $provider_type:ident) => {{ - let env = native_key_env(provider).unwrap_or_default(); - let key = native_key(provider, env, api_key.clone())?; - let mut cfg = aimux_providers::$provider_mod::$config::new(key); - if let Some(url) = base_url { - cfg = cfg.with_base_url(url); - } - let p = aimux_providers::$provider_mod::$provider_type::new(cfg); - Ok(Arc::from(p.model(model_id))) - }}; - } - - match provider { - "openai" => native!(openai, OpenAIConfig, OpenAIProvider), - "anthropic" => native!(anthropic, AnthropicConfig, AnthropicProvider), - "google" => native!(google, GoogleConfig, GoogleProvider), - "mistral" => native!(mistral, MistralConfig, MistralProvider), - "xai" => native!(xai, XAIConfig, XAIProvider), - "cohere" => native!(cohere, CohereConfig, CohereProvider), - _ => { - let mut options = ProviderOptions::default(); - if let Some(url) = base_url { - options.base_url = Some(url.to_string()); - } - let model = aimux_providers::provider(provider, api_key, model_id, Some(options)) - .map_err(|e| AiMuxError::InvalidArgument(format!("provider '{provider}': {e}")))?; - Ok(Arc::from(model)) - } - } + aimux_providers::create_provider( + provider, + PresetSettings { + api_key: api_key.map(Into::into), + base_url: base_url.map(str::to_string), + ..Default::default() + }, + )? + .language_model(model_id) } /// Whether `base_url` points somewhere other than `provider`'s own default -/// endpoint. Stored and environment keys are never sent to such a URL: it -/// came from the caller (a request field, or an imported recording). +/// endpoint ([`aimux_providers::provider_default_base_url`]). Stored and +/// environment keys are never sent to such a URL: it came from the caller. /// -/// The default is what [`build_model`] uses when no base URL is given; when -/// it cannot be determined (an unknown provider, a templated registry URL) -/// every given URL counts as custom. +/// When the default cannot be determined (an unknown provider, an endpoint +/// built from settings such as a region or a template), every given URL +/// counts as custom. pub fn is_custom_base_url(provider: &str, base_url: Option<&str>) -> bool { let Some(url) = base_url else { return false; }; - let default = build_model(provider, Some(String::new()), "default-base-url", None) - .ok() - .and_then(|model| model.config_snapshot().base_url); - default.is_none_or(|default| default.trim_end_matches('/') != url.trim_end_matches('/')) + aimux_providers::provider_default_base_url(provider) + .is_none_or(|default| default.trim_end_matches('/') != url.trim_end_matches('/')) } -/// The env var a native provider (see [`build_model`]) reads its key from. -pub fn native_key_env(provider: &str) -> Option<&'static str> { - Some(match provider { - "openai" => "OPENAI_API_KEY", - "anthropic" => "ANTHROPIC_API_KEY", - "google" => "GOOGLE_GENERATIVE_AI_API_KEY", - "mistral" => "MISTRAL_API_KEY", - "xai" => "XAI_API_KEY", - "cohere" => "COHERE_API_KEY", - _ => return None, - }) -} +/// Providers that read credentials from the server environment besides the +/// `api_key` setting: AWS credentials (SigV4 signing, session token), +/// the Vertex access token, the Codex sign-in, the DataForSEO login. An empty +/// `api_key` does not stop them, so the console does not call them at a +/// caller-chosen base URL. +const AMBIENT_CREDENTIALS: [&str; 6] = [ + "amazon_bedrock", + "anthropic_aws", + "aws_polly", + "codex", + "dataforseo", + "google_vertex", +]; -/// Resolve the key for a native provider: explicit key, else the standard env var. -fn native_key(provider: &str, env: &str, api_key: Option) -> Result { - match api_key { - Some(k) => Ok(k), - None => std::env::var(env).map_err(|_| { - AiMuxError::InvalidArgument(format!( - "provider '{provider}' needs an API key — set `{env}` or pass api_key=\"env:{env}\"" - )) - }), - } +/// Whether `provider` may read credentials other than the `api_key` setting +/// from the server environment ([`AMBIENT_CREDENTIALS`]). +pub fn reads_ambient_credentials(provider: &str) -> bool { + AMBIENT_CREDENTIALS.contains(&provider) } #[cfg(test)] mod tests { - use super::is_custom_base_url; + use super::{AMBIENT_CREDENTIALS, is_custom_base_url}; #[test] fn only_the_providers_own_default_is_not_custom() { @@ -115,6 +82,10 @@ mod tests { "groq", Some("https://api.groq.com/openai/v1") )); + assert!(!is_custom_base_url( + "openrouter", + Some("https://openrouter.ai/api/v1") + )); assert!(is_custom_base_url( "openai", Some("https://attacker.example/v1") @@ -128,5 +99,17 @@ mod tests { "no-such-provider", Some("https://x.example") )); + assert!(is_custom_base_url( + "amazon_bedrock", + Some("https://bedrock-runtime.us-east-1.amazonaws.com") + )); + } + + #[test] + fn ambient_credential_providers_are_built_in() { + let names = aimux_providers::provider_names(); + for name in AMBIENT_CREDENTIALS { + assert!(names.contains(&name), "{name}"); + } } } diff --git a/tools/aimux-web/src/state.rs b/tools/aimux-web/src/state.rs index 034ab3fea..248d4b307 100644 --- a/tools/aimux-web/src/state.rs +++ b/tools/aimux-web/src/state.rs @@ -129,8 +129,25 @@ impl AppState { /// Wrap a model in a `TraceLayer` bound to the shared sink, so every call /// is both recorded (RFC-0023) and probed (RFC-0015). - pub fn traced(&self, model: Arc) -> Arc { - Arc::new(TraceLayer::new(model, self.trace_sink.clone())) + /// + /// The trace scope isolates by provider name and model, and every call of + /// one provider has the same name, so a call at a custom `base_url` gets + /// a scope named after that endpoint: two endpoints never share cache + /// history. + pub fn traced( + &self, + model: Arc, + provider: &str, + base_url: Option<&str>, + ) -> Arc { + let layer = match base_url { + Some(url) if crate::model_builder::is_custom_base_url(provider, Some(url)) => { + let scope = format!("{} @ {}", model.provider(), url.trim_end_matches('/')); + TraceLayer::new(model, self.trace_sink.clone()).with_scope_name(scope) + } + _ => TraceLayer::new(model, self.trace_sink.clone()), + }; + Arc::new(layer) } /// All completed recordings: in-ring plus imported (imported appended). diff --git a/tools/aimux-web/src/wire.rs b/tools/aimux-web/src/wire.rs index 39ec89779..a1af8a821 100644 --- a/tools/aimux-web/src/wire.rs +++ b/tools/aimux-web/src/wire.rs @@ -85,8 +85,6 @@ pub struct WireOptions { #[serde(default, skip_serializing_if = "Option::is_none")] pub headers: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] - pub body_overrides: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] pub max_retries: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub include_raw_chunks: Option, @@ -312,7 +310,6 @@ pub fn to_generate_options( .map(|ts| ts.iter().map(to_function_tool).collect()), output, headers: o.headers.clone(), - body_overrides: o.body_overrides.clone(), max_retries: o.max_retries, session_id: session_id.map(str::to_string), include_raw_chunks: o.include_raw_chunks, @@ -338,7 +335,9 @@ pub fn to_generate_options( /// goes to a base URL the caller chose. Then only a key carried in the request /// itself is used: no Settings key, no environment variable (`env:` is /// refused), and with no key the call is sent without one (`Some("")`), so -/// the provider does not fall back to its env var either. +/// the provider does not fall back to its env var either. Providers that read +/// other credentials from the environment +/// ([`crate::model_builder::reads_ambient_credentials`]) are refused. pub fn resolve_api_key( spec: Option<&str>, provider: &str, @@ -346,6 +345,12 @@ pub fn resolve_api_key( allow_plaintext: bool, custom_base_url: bool, ) -> Result, AiMuxError> { + if custom_base_url && crate::model_builder::reads_ambient_credentials(provider) { + return Err(AiMuxError::InvalidArgument(format!( + "provider '{provider}' reads credentials from the server environment, so the \ + console does not call it at a custom base_url — leave base_url empty" + ))); + } let blank = spec.is_none_or(|s| s.trim().is_empty()); if custom_base_url && blank { return Ok(Some(String::new())); @@ -413,12 +418,11 @@ pub fn resolve_api_key( } /// The environment variable a provider's key is registered under: the -/// native providers' standard variable, or the registry / runtime overlay -/// entry's. The only name an `env:` reference may use for that provider. -pub fn registered_key_env(provider: &str) -> Option { - crate::model_builder::native_key_env(provider) - .map(str::to_string) - .or_else(|| aimux_providers::provider_api_key_env(provider)) +/// vendor package's key variable or the registry row's `env_var` +/// ([`aimux_providers::provider_api_key_env`]). The only name an `env:` +/// reference may use for that provider. +pub fn registered_key_env(provider: &str) -> Option<&'static str> { + aimux_providers::provider_api_key_env(provider) } #[cfg(test)] @@ -426,6 +430,38 @@ mod tests { use super::*; use crate::settings::KeyStore; + /// Sets (or, with `None`, removes) a registry provider's key variable for + /// one test and restores the previous value on drop. Each test uses its + /// own provider, so the variables never overlap. + struct KeyVar(&'static str, Option); + + impl KeyVar { + fn set(provider: &str, value: Option<&str>) -> Self { + let name = registered_key_env(provider).expect("a registry row with a key variable"); + let previous = std::env::var_os(name); + // SAFETY: no other test reads or writes this provider's variable. + unsafe { + match value { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + Self(name, previous) + } + } + + impl Drop for KeyVar { + fn drop(&mut self) { + // SAFETY: as in `set`. + unsafe { + match &self.1 { + Some(value) => std::env::set_var(self.0, value), + None => std::env::remove_var(self.0), + } + } + } + } + #[test] fn empty_spec_resolves_as_unset() { // Empty/whitespace spec behaves exactly like None: fall through to @@ -442,27 +478,15 @@ mod tests { ); } - /// Register a console-test provider whose key variable is `env_var`. - fn test_provider(name: &str, env_var: &str) { - aimux_providers::load_providers_from_json(&format!( - r#"{{"providers":[{{"name":"{name}","base_url":"https://example.invalid/v1","env_var":"{env_var}"}}]}}"# - )) - .unwrap(); - } - #[test] fn resolve_key_priority_explicit_env_over_store() { - test_provider("web-test-priority", "AIMUX_WEB_TEST_KEY"); let store = KeyStore::from_path(None); - store - .set("web-test-priority", "sk-from-store", false) - .unwrap(); - // SAFETY: This test owns a unique process-wide env key and removes it below. - unsafe { std::env::set_var("AIMUX_WEB_TEST_KEY", "k-123") }; + store.set("xpersona", "sk-from-store", false).unwrap(); + let _var = KeyVar::set("xpersona", Some("k-123")); assert_eq!( resolve_api_key( - Some("env:AIMUX_WEB_TEST_KEY"), - "web-test-priority", + Some("env:XPERSONA_API_KEY"), + "xpersona", &store, true, false @@ -470,8 +494,6 @@ mod tests { .unwrap(), Some("k-123".to_string()) ); - // SAFETY: Cleanup of the unique test key set above. - unsafe { std::env::remove_var("AIMUX_WEB_TEST_KEY") }; } #[test] @@ -509,13 +531,11 @@ mod tests { let err = resolve_api_key(Some("sk-literal"), "openai", &store, false, false).unwrap_err(); assert!(err.to_string().contains("loopback"), "err: {err}"); // env: references still work on non-loopback. - test_provider("web-test-lan", "AIMUX_WEB_TEST_KEY_LAN"); - // SAFETY: This test owns a unique process-wide env key and removes it below. - unsafe { std::env::set_var("AIMUX_WEB_TEST_KEY_LAN", "k-lan") }; + let _var = KeyVar::set("the_grid_ai", Some("k-lan")); assert_eq!( resolve_api_key( - Some("env:AIMUX_WEB_TEST_KEY_LAN"), - "web-test-lan", + Some("env:THEGRIDAI_API_KEY"), + "the_grid_ai", &store, false, false @@ -523,8 +543,6 @@ mod tests { .unwrap(), Some("k-lan".to_string()) ); - // SAFETY: Cleanup of the unique test key set above. - unsafe { std::env::remove_var("AIMUX_WEB_TEST_KEY_LAN") }; } #[test] @@ -545,11 +563,12 @@ mod tests { "{err}" ); - assert_eq!(registered_key_env("groq").as_deref(), Some("GROQ_API_KEY")); + assert_eq!(registered_key_env("groq"), Some("GROQ_API_KEY")); assert_eq!( - registered_key_env("google").as_deref(), + registered_key_env("google"), Some("GOOGLE_GENERATIVE_AI_API_KEY") ); + assert_eq!(registered_key_env("openrouter"), Some("OPENROUTER_API_KEY")); } #[test] @@ -574,22 +593,30 @@ mod tests { Some("sk-literal".to_string()) ); assert!(resolve_api_key(Some("sk-literal"), "openai", &store, false, true).is_err()); + // A provider that also reads other credentials from the environment + // is not called at a custom base_url at all. + for provider in ["amazon_bedrock", "google_vertex", "codex"] { + let err = resolve_api_key(Some("k"), provider, &store, true, true).unwrap_err(); + assert!(err.to_string().contains("server environment"), "{err}"); + assert!(resolve_api_key(None, provider, &store, true, false).is_ok()); + } } #[test] fn resolve_key_rejects_plaintext_and_missing_env() { let store = KeyStore::from_path(None); assert!(resolve_api_key(Some("sk-literal"), "openai", &store, false, false).is_err()); - assert!( - resolve_api_key( - Some("env:AIMUX_WEB_MISSING_VAR"), - "openai", - &store, - true, - false - ) - .is_err() - ); + // The provider's own variable, unset. + let _var = KeyVar::set("siliconflow", None); + let err = resolve_api_key( + Some("env:SILICONFLOW_API_KEY"), + "siliconflow", + &store, + true, + false, + ) + .unwrap_err(); + assert!(err.to_string().contains("is not set"), "{err}"); assert!(resolve_api_key(Some("env:"), "openai", &store, true, false).is_err()); } diff --git a/tools/aimux-web/web/src/types/generated/WireOptions.ts b/tools/aimux-web/web/src/types/generated/WireOptions.ts index 3100d9b47..7d0ae849c 100644 --- a/tools/aimux-web/web/src/types/generated/WireOptions.ts +++ b/tools/aimux-web/web/src/types/generated/WireOptions.ts @@ -6,4 +6,4 @@ export type WireOptions = { temperature?: number | null, max_output_tokens?: num /** * `"text"` or `{"json": …}` — passed through as `ResponseFormat`. */ -response_format?: JsonValue | null, headers?: { [key in string]: string } | null, body_overrides?: JsonValue | null, max_retries?: number | null, include_raw_chunks?: boolean | null, }; +response_format?: JsonValue | null, headers?: { [key in string]: string } | null, max_retries?: number | null, include_raw_chunks?: boolean | null, }; diff --git a/tools/aimux-web/web/src/views/Playground.vue b/tools/aimux-web/web/src/views/Playground.vue index 11402ac80..c93a851fe 100644 --- a/tools/aimux-web/web/src/views/Playground.vue +++ b/tools/aimux-web/web/src/views/Playground.vue @@ -34,7 +34,6 @@ const maxTokens = ref(1024) const selectedTools = ref([]) const responseFormat = ref('text') const headersJson = ref('{}') -const overridesJson = ref('{}') const suggested = computed(() => store.suggestedModels[provider.value] ?? []) const modelOptions = computed(() => @@ -134,7 +133,6 @@ async function send() { tools, response_format: responseFormat.value === 'json' ? { json: {} } : null, headers: safeJson(headersJson.value), - body_overrides: safeJson(overridesJson.value), }, session_id: sessionId.value, messages: wireMessages.value, @@ -327,10 +325,6 @@ function openTrace(item: ChatItem) {