diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d17dfa22c..eaf7a4d97 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -398,6 +398,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 - name: ts-rs export drift (checked-in types must match generator output) 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 746dfdf78..3dfb36e36 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,14 +9,30 @@ 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, `DEFAULT_MAX_JSON_RESPONSE_SIZE` and `TungsteniteConnector`. + - 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 @@ -48,11 +64,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 @@ -74,8 +90,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()`. @@ -141,6 +157,248 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 name ends the stream with `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. + AI SDK model constructors return an `Arc`. Removed: `Provider::name()`, + `specification_version()` (from `Provider` and the nine model traits), + `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 = 3`. `ProviderRecord` holds identity only + (`provider_id`, `provider`, `model_id`); base URL, key source, profile, + options and retry settings are no longer recorded and schema-2 files are + rejected on read. `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 are canonical only: `googleVertex` and + `amazonBedrock` (the legacy `vertex` / `bedrock` keys are neither read nor + written; Bedrock still 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 281 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`; 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`; `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. +- 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. `VideoModel::poll_config()` stays as + the package-constant source. +- 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 = 18` and `AIMUX_E_LOAD_SETTING = 19` + (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 18 / 19 (`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. + ## [0.5.0] - 2026-09-27 **Breaking release.** 13 PRs since 0.3.0: the cross-language error model diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 240a759ab..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 stream 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 3508b67e3..4e5da174f 100644 --- a/README.md +++ b/README.md @@ -4,12 +4,12 @@ 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 [326 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) [![Rust](https://img.shields.io/badge/rust-stable-orange.svg)](https://www.rust-lang.org/) -[![Providers](https://img.shields.io/badge/providers-327-green.svg)](docs/api/providers.md) +[![Providers](https://img.shields.io/badge/providers-326-green.svg)](docs/api/providers.md) [![Bindings](https://img.shields.io/badge/bindings-8-9cf.svg)](bindings/) [![crates.io](https://img.shields.io/crates/v/aimux-core)](https://crates.io/crates/aimux-core) [![npm](https://img.shields.io/npm/v/@arcships/aimux)](https://www.npmjs.com/package/@arcships/aimux) @@ -30,13 +30,10 @@ difference: aimux is an access layer, those are orchestration layers. ## 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). +- **326 provider names** — 281 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, @@ -49,10 +46,12 @@ difference: aimux is an access layer, those are orchestration layers. - **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, @@ -148,14 +147,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, @@ -182,8 +179,8 @@ let result = stream_text( let mut stream = result.stream; while let Some(part) = stream.next().await { match part? { - StreamPart::TextDelta { delta, .. } => print!("{}", delta), - StreamPart::Finish { .. } => println!("\n[done]"), + TextStreamPart::TextDelta { delta, .. } => print!("{}", delta), + TextStreamPart::Finish { .. } => println!("\n[done]"), _ => {} } } @@ -220,9 +217,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(), @@ -234,7 +232,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(), ); @@ -273,35 +271,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 281 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 bcd4c3d8e..4ea40c032 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -26,7 +26,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）+ 文档三层重组 │ @@ -66,30 +66,30 @@ 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 周） | 项 | 内容 | |---|---| -| C1 | 9 个导出 + `dispatch`，**同一 dispatch 同时服务 FFI 与 stdio CLI**；旧导出保留共存 | +| C1 | 9 个导出 + `dispatch`，**同一 dispatch 同时服务 FFI 与 stdio CLI**；旧符号按切换门一次替换，不共存（调整：见 [docs/aisdk-architecture-alignment.md §0.7](docs/aisdk-architecture-alignment.md#07-与-roadmap--rfc-0036-既有承诺的关系)） | | C3 | `aimux_error_*` 访问器 → 错误 JSON 信封（含派生 `retry_after_ms`） | | D1–D7 | 8 种绑定迁到 ops 薄封装。D1 Kotlin（只依赖 Java artifact）与 D5 Flutter ffigen 过渡版不依赖 C1，可在 0.6/0.7 提前做 | | 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 × 能力，每格由探测或录制支撑 | | CLI | `aimux probe / replay / diff`（#179 replay 子命令并入）；#181 debug CLI 并入 | @@ -98,7 +98,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 c48aaee7c..70ab4aeec 100644 --- a/aimux-core/Cargo.toml +++ b/aimux-core/Cargo.toml @@ -26,6 +26,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 d4f3b5505..e737471cb 100644 --- a/aimux-core/src/composite.rs +++ b/aimux-core/src/composite.rs @@ -8,9 +8,9 @@ //! //! This module holds the pieces both share: //! - [`ChildModel`] — the child-handle shape (`Arc`). -//! - [`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 @@ -34,7 +34,7 @@ pub type ChildModel = Arc; /// 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), @@ -64,7 +64,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 { @@ -83,7 +83,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 6166b146f..af3646461 100644 --- a/aimux-core/src/decision_model.rs +++ b/aimux-core/src/decision_model.rs @@ -418,21 +418,9 @@ impl DecisionResult { #[async_trait] pub trait DecisionModel: Send + Sync { - 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. @@ -496,11 +484,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); @@ -510,9 +494,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 9e35f362f..e588c77c9 100644 --- a/aimux-core/src/decision_replay.rs +++ b/aimux-core/src/decision_replay.rs @@ -49,16 +49,9 @@ 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".into()) + })?; Ok(Self { provider, model_id, diff --git a/aimux-core/src/embedding_model.rs b/aimux-core/src/embedding_model.rs index 7f1807048..e97ff1dab 100644 --- a/aimux-core/src/embedding_model.rs +++ b/aimux-core/src/embedding_model.rs @@ -117,21 +117,12 @@ pub struct EmbeddingResponse { /// `true`; otherwise they should return an error. #[async_trait] pub trait EmbeddingModel: Send + Sync { - /// Specification version (always `"v4"`). - fn specification_version(&self) -> &'static str { - "v4" - } - /// Provider name, e.g. `"openai"`. fn provider(&self) -> &str; /// 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 @@ -159,11 +150,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 e39bf2be8..d9c7bfd9b 100644 --- a/aimux-core/src/error.rs +++ b/aimux-core/src/error.rs @@ -224,6 +224,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 access token has expired (or was invalidated) and must be refreshed /// by the caller. RFC-0018 subscription mode: the library maps a 401 from /// the Codex subscription endpoint to this variant; the integrator @@ -252,8 +280,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), @@ -296,6 +329,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/files_model.rs b/aimux-core/src/files_model.rs index d2fcb1120..091488838 100644 --- a/aimux-core/src/files_model.rs +++ b/aimux-core/src/files_model.rs @@ -98,11 +98,6 @@ pub struct UploadFileResult { /// et al. #[async_trait] pub trait Files: Send + Sync { - /// Specification version (always `"v4"`). - fn specification_version(&self) -> &'static str { - "v4" - } - /// Provider ID. fn provider(&self) -> &str; diff --git a/aimux-core/src/generate.rs b/aimux-core/src/generate.rs index 622219230..3a71543df 100644 --- a/aimux-core/src/generate.rs +++ b/aimux-core/src/generate.rs @@ -102,8 +102,6 @@ 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, /// Per-call retry count override. `None` = provider default, `Some(0)` = disable. pub max_retries: Option, /// Per-call timeout configuration (total / first-chunk / chunk idle). @@ -159,7 +157,6 @@ impl GenerateTextOptions { headers: self.headers, provider_options: self.provider_options, reasoning: self.reasoning, - body_overrides: self.body_overrides, max_retries: self.max_retries, timeout: self.timeout, session_id: self.session_id, @@ -579,11 +576,7 @@ pub(crate) async fn generate_text_from_language_model_prompt( let operation_timeout = timeout::OperationTimeout::new(call_options.timeout.unwrap_or_default())?; let abort_signal = call_options.abort_signal.clone(); - let retries = retry::prepare_retries( - call_options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(call_options.max_retries, abort_signal.clone()); // 2a. RFC-0023: 关闭时零成本(M2 评审)——仅在录制开启时生成 call_id // 并绑定 recorder 快照;传输封闭由层 A 收尾声明(P1 无层 B,barrier @@ -599,8 +592,10 @@ pub(crate) async fn generate_text_from_language_model_prompt( 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 完成。双发幂等。 @@ -941,11 +936,7 @@ pub async fn stream_text( .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( - call_options.max_retries, - model.retry_config(), - abort_signal.clone(), - ); + let retries = retry::prepare_retries(call_options.max_retries, abort_signal.clone()); // 2a. RFC-0023: 关闭时零成本(M2 评审)——仅在录制开启时生成 call_id // 并绑定 recorder 快照;传输封闭由层 A 收尾声明(P1 无层 B,barrier @@ -961,8 +952,10 @@ pub async fn stream_text( 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/image_model.rs b/aimux-core/src/image_model.rs index 2a1ae9418..159de19df 100644 --- a/aimux-core/src/image_model.rs +++ b/aimux-core/src/image_model.rs @@ -181,21 +181,12 @@ pub struct ImageResponse { /// Aligned with V4 `ImageModelV4`. #[async_trait] pub trait ImageModel: Send + Sync { - /// Specification version (always `"v4"`). - fn specification_version(&self) -> &'static str { - "v4" - } - /// Provider name, e.g. `"openai"`. fn provider(&self) -> &str; /// 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 @@ -203,6 +194,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. @@ -220,11 +221,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..8cdcbbcbe 100644 --- a/aimux-core/src/language_model.rs +++ b/aimux-core/src/language_model.rs @@ -3,12 +3,31 @@ //! 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) + } +} + /// The unified language model trait (provider-facing). /// /// Aligned with Vercel AI SDK `LanguageModelV4`. @@ -23,33 +42,25 @@ use crate::result::{GenerateResult, StreamResult}; /// rather than being silently dropped. #[async_trait] pub trait LanguageModel: Send + Sync { - /// Specification version (always `"v4"`). - fn specification_version(&self) -> &'static str { - "v4" - } - /// Provider name, e.g. `"openai"`. fn provider(&self) -> &str; /// 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`. + /// + /// Default: empty, meaning every file part is sent as inline data. + /// Transparent decorators must forward the inner model's value. Core does + /// not consult this yet. + fn supported_urls(&self) -> SupportedUrls { + SupportedUrls::default() } } diff --git a/aimux-core/src/language_model_middleware.rs b/aimux-core/src/language_model_middleware.rs new file mode 100644 index 000000000..058566092 --- /dev/null +++ b/aimux-core/src/language_model_middleware.rs @@ -0,0 +1,150 @@ +//! 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 + } + + 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, + supported_urls: SupportedUrls, +} + +/// 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()); + let supported_urls = middleware + .override_supported_urls(model.as_ref()) + .unwrap_or_else(|| model.supported_urls()); + Arc::new(WrappedLanguageModel { + model, + middleware: Arc::clone(middleware), + provider, + model_id, + supported_urls, + }) + }) +} + +#[async_trait] +impl LanguageModel for WrappedLanguageModel { + fn provider(&self) -> &str { + &self.provider + } + + fn model_id(&self) -> &str { + &self.model_id + } + + fn supported_urls(&self) -> SupportedUrls { + self.supported_urls.clone() + } + + 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 + } +} diff --git a/aimux-core/src/lib.rs b/aimux-core/src/lib.rs index 638e63a0b..cd62b8dba 100644 --- a/aimux-core/src/lib.rs +++ b/aimux-core/src/lib.rs @@ -30,9 +30,11 @@ 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; @@ -42,6 +44,7 @@ pub mod openai_output; pub mod options; pub mod parse_tool_call; pub mod provider; +pub mod provider_registry; pub mod recording; pub mod replay; pub mod reranking_model; @@ -52,6 +55,7 @@ pub mod router; pub mod search_model; pub mod session; pub mod shared; +pub mod skills_model; pub mod speech_model; pub mod stream_part; mod timeout; @@ -84,8 +88,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::{MessageContent, ModelMessage, ModelPrompt, Role}; pub use crate::moa::{MoaConfig, MoaFailMode, MoaModel}; pub use crate::model_catalogue::{ @@ -125,6 +133,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::stream_part::{StreamPart, TextStreamPart}; pub use crate::tool::{ @@ -155,8 +164,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 129371048..6123c1924 100644 --- a/aimux-core/src/moa.rs +++ b/aimux-core/src/moa.rs @@ -125,11 +125,8 @@ impl MoaModel { let child_opts = ref_opts.for_step(format!("moa.ref[{i}]:{}/{}", m.provider(), m.model_id())); async move { - let retries = retry::prepare_retries( - child_opts.max_retries, - m.retry_config(), - child_opts.abort_signal.clone(), - ); + let retries = + retry::prepare_retries(child_opts.max_retries, child_opts.abort_signal.clone()); retries .retry(|| { if let Some(context) = &child_opts.recording_context { @@ -183,15 +180,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?; @@ -206,11 +194,7 @@ impl LanguageModel for MoaModel { agg_opts.prompt = agg_prompt; // 3. Run the aggregator. - let retries = retry::prepare_retries( - agg_opts.max_retries, - self.aggregator.retry_config(), - agg_opts.abort_signal.clone(), - ); + let retries = retry::prepare_retries(agg_opts.max_retries, agg_opts.abort_signal.clone()); let agg_opts = agg_opts.for_step(format!( "moa.aggregator:{}/{}", self.aggregator.provider(), @@ -248,11 +232,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( - agg_opts.max_retries, - self.aggregator.retry_config(), - agg_opts.abort_signal.clone(), - ); + let retries = retry::prepare_retries(agg_opts.max_retries, agg_opts.abort_signal.clone()); let agg_opts = agg_opts.for_step(format!( "moa.aggregator:{}/{}", self.aggregator.provider(), diff --git a/aimux-core/src/model_catalogue.rs b/aimux-core/src/model_catalogue.rs index fc045ed32..a7ef0d8e9 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)] #[ts(export)] pub struct ModelSpec { diff --git a/aimux-core/src/options.rs b/aimux-core/src/options.rs index e9a6a302a..d43fbb159 100644 --- a/aimux-core/src/options.rs +++ b/aimux-core/src/options.rs @@ -115,11 +115,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, - /// Per-call retry count override. `None` uses the provider's configured /// Core operation retry. `Some(0)` disables retries. pub max_retries: Option, @@ -208,7 +203,6 @@ impl CallOptions { headers: None, provider_options: None, reasoning: None, - body_overrides: None, max_retries: None, timeout: None, session_id: None, diff --git a/aimux-core/src/provider.rs b/aimux-core/src/provider.rs index 70658318c..755833137 100644 --- a/aimux-core/src/provider.rs +++ b/aimux-core/src/provider.rs @@ -1,69 +1,198 @@ -//! The `Provider` trait — a factory that creates `LanguageModel` instances. +//! 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 API keys / config and creates `LanguageModel` instances by model name. +/// 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) and +/// [`search_model`](Self::search_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; - - /// 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"`). + /// + /// # 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 /// - /// 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). + /// 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 /// - /// 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: "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 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 204ecabfb..c7282c597 100644 --- a/aimux-core/src/recording.rs +++ b/aimux-core/src/recording.rs @@ -29,7 +29,23 @@ use ts_rs::TS; // ── 数据模型(三层 + call_id 关联;schema 版本)──────────────────────────── /// 录制格式版本(用于未来字段迁移与绑定层兼容)。 -pub const RECORDING_SCHEMA: u32 = 2; +pub const RECORDING_SCHEMA: u32 = 3; + +/// 只接受当前 [`RECORDING_SCHEMA`]:反序列化旧版本录制时给出明确错误。 +fn deserialize_current_schema<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + let schema = u32::deserialize(deserializer)?; + if schema == RECORDING_SCHEMA { + Ok(schema) + } else { + Err(serde::de::Error::custom(format!( + "unsupported recording schema {schema}: only schema {RECORDING_SCHEMA} is \ + readable, recordings written by older versions are no longer read" + ))) + } +} /// Operation stored in a unified recording. Older recordings are language calls. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, TS)] @@ -51,7 +67,9 @@ impl RecordingOperation { #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[ts(export)] pub struct Recording { - /// 格式版本。 + /// 格式版本。读取时只接受当前 [`RECORDING_SCHEMA`];旧 schema 的录制带有 + /// 配置快照(base_url / profile 等),不再读取。 + #[serde(deserialize_with = "deserialize_current_schema")] pub schema: u32, /// 暴露逻辑调用 ID,关联三层。 pub call_id: String, @@ -134,12 +152,16 @@ impl Recording { #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[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 等多模态)。 pub prompt: LanguageModelPrompt, /// 序列化的 CallOptions(abort_signal/call_id 已 serde skip); - /// headers/provider_options/body_overrides 已递归脱敏。 + /// headers/provider_options 已递归脱敏。 pub options: serde_json::Value, } @@ -149,6 +171,7 @@ impl InputRecord { pub fn from_call_options(options: &CallOptions) -> Self { let value = serde_json::to_value(options).unwrap_or(serde_json::Value::Null); Self { + decision_capabilities: None, operation: RecordingOperation::LanguageModel, prompt: options.prompt.clone(), options: redact_json(value), @@ -156,8 +179,12 @@ 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()), @@ -166,36 +193,39 @@ impl InputRecord { } /// ② 配置侧:provider 身份(请求回放重建用)。 +/// +/// 只记身份,不记配置:回放时按 `provider_id` + `model_id` 重新构建 model, +/// 凭证与 base URL 只来自环境/registry(RFC-0036 §3.3)。 #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[ts(export)] pub struct ProviderRecord { - /// `model.provider()`,如 "openai"。 + /// provider 名(registry 名),回放重建的查找键,如 "openai"、"deepseek"。 + 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, } } + + /// 由 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 交换。 @@ -385,6 +415,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 + 会话内步号。 @@ -392,7 +423,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); @@ -517,18 +548,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 { @@ -548,7 +567,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()) @@ -589,7 +608,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) => { @@ -823,18 +842,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), }); } @@ -847,14 +867,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(), }); } @@ -915,11 +930,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("", ""), ) }) } @@ -1387,11 +1403,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("", ""), ), ); } @@ -1464,12 +1481,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) { @@ -1479,7 +1497,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); } } @@ -1493,13 +1511,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) { @@ -1979,11 +1993,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": { @@ -1999,9 +2015,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); @@ -2014,7 +2033,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(); @@ -2063,11 +2082,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 { @@ -2142,7 +2162,7 @@ mod tests { 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 { @@ -2218,7 +2238,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}"); @@ -2760,14 +2780,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] @@ -2782,12 +2795,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"); } @@ -2803,11 +2815,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); } @@ -2848,11 +2857,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 8d641d7a0..ebc162e09 100644 --- a/aimux-core/src/replay.rs +++ b/aimux-core/src/replay.rs @@ -51,10 +51,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/recording_context)不参与比较。 /// -/// headers/provider_options/body_overrides 用**脱敏感知比较**:录制侧这些 +/// headers/provider_options 用**脱敏感知比较**:录制侧这些 /// 字段经 `recording::redact_json` 脱敏 /// (敏感键值→`"[REDACTED]"`),脱敏值视为通配,匹配任意显式请求值;非脱敏 /// 部分仍精确比较。规则见 `redaction_aware_eq`。 @@ -96,9 +96,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/recording_context)。序列化为 `Value` 后比较, -/// 等价于"稳定规范哈希"但无碰撞风险。headers/provider_options/body_overrides +/// 等价于"稳定规范哈希"但无碰撞风险。headers/provider_options /// 在比较时走脱敏感知语义(见 [`canonical_keys_match`])。 fn canonical_call_key(opts: &CallOptions) -> serde_json::Value { serde_json::json!({ @@ -111,7 +111,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(), }) } @@ -128,11 +127,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/...)的**值** @@ -174,7 +172,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" 子串的字段在录制侧也会被 @@ -182,7 +180,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, @@ -441,8 +439,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}")))?; @@ -451,8 +450,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() { @@ -945,13 +952,11 @@ 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`. +/// deserialized, and propagates errors from the generation pipeline. pub async fn replay_with_model( recording: &Recording, model: &dyn LanguageModel, @@ -1007,7 +1012,6 @@ fn generate_options_from_call_options(o: CallOptions) -> GenerateTextOptions { provider_options: o.provider_options, reasoning: o.reasoning, instructions: None, - body_overrides: o.body_overrides, max_retries: o.max_retries, timeout: o.timeout, session_id: o.session_id, @@ -1057,18 +1061,12 @@ 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, options: serde_json::json!({ "temperature": 0.7 }), }, - 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 a07c20514..cf677ebc6 100644 --- a/aimux-core/src/reranking_model.rs +++ b/aimux-core/src/reranking_model.rs @@ -124,21 +124,12 @@ pub struct RerankingResponse { /// Aligned with V4 `RerankingModelV4`. #[async_trait] pub trait RerankingModel: Send + Sync { - /// Specification version (always `"v4"`). - fn specification_version(&self) -> &'static str { - "v4" - } - /// Provider name, e.g. `"cohere"`. fn provider(&self) -> &str; /// 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. @@ -159,11 +150,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 dd97a1470..da3f99c06 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 41b5858db..4151dbff1 100644 --- a/aimux-core/src/router.rs +++ b/aimux-core/src/router.rs @@ -177,11 +177,7 @@ impl RouterModel { if i == exclude { continue; } - let retries = retry::prepare_retries( - options.max_retries, - m.retry_config(), - options.abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, options.abort_signal.clone()); let child_options = options.for_step(self.step_label(i)); match retries .retry(|| { @@ -229,24 +225,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( - options.max_retries, - model.retry_config(), - options.abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, options.abort_signal.clone()); let child_options = options.for_step(self.step_label(idx)); match retries .retry(|| { @@ -272,11 +255,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( - options.max_retries, - model.retry_config(), - options.abort_signal.clone(), - ); + let retries = retry::prepare_retries(options.max_retries, options.abort_signal.clone()); let child_options = options.for_step(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. diff --git a/aimux-core/src/search_model.rs b/aimux-core/src/search_model.rs index cd03f7b10..4c94c65dd 100644 --- a/aimux-core/src/search_model.rs +++ b/aimux-core/src/search_model.rs @@ -138,11 +138,6 @@ pub struct SearchResponse { /// implement `do_search`; users never call it directly. #[async_trait] pub trait SearchModel: Send + Sync { - /// Specification version (always `"v4"`). - fn specification_version(&self) -> &'static str { - "v4" - } - /// Provider name, e.g. `"tavily"`. fn provider(&self) -> &str; @@ -150,10 +145,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. @@ -171,11 +162,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 315c00c0e..38a3869a2 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 db673e627..24cddc49d 100644 --- a/aimux-core/src/shared.rs +++ b/aimux-core/src/shared.rs @@ -311,3 +311,13 @@ 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, + } + } +} diff --git a/aimux-core/src/skills_model.rs b/aimux-core/src/skills_model.rs new file mode 100644 index 000000000..deb5171c0 --- /dev/null +++ b/aimux-core/src/skills_model.rs @@ -0,0 +1,47 @@ +//! 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)] +#[ts(export)] +pub struct SkillFile { + pub path: String, + pub data: UploadFileData, +} + +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct UploadSkillCallOptions { + pub files: Vec, + pub display_title: Option, + pub provider_options: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[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 { + 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 5add06e97..d83d38c8f 100644 --- a/aimux-core/src/speech_model.rs +++ b/aimux-core/src/speech_model.rs @@ -137,21 +137,12 @@ pub struct SpeechResponse { /// Aligned with V4 `SpeechModelV4`. #[async_trait] pub trait SpeechModel: Send + Sync { - /// Specification version (always `"v4"`). - fn specification_version(&self) -> &'static str { - "v4" - } - /// Provider name, e.g. `"openai"`. fn provider(&self) -> &str; /// 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. @@ -169,11 +160,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 c5a73b349..588e6dc7c 100644 --- a/aimux-core/src/trace/layer.rs +++ b/aimux-core/src/trace/layer.rs @@ -163,24 +163,18 @@ impl TraceLayer { .session_id .clone() .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 (provider, model_id), so two TraceLayers + // over the same model_id but different providers sharing one + // RingTraceStore never cross-match their LCP history. A model's + // `provider()` string already carries the settings-level name, so two + // endpoints of one vendor stay apart when they are given different + // provider names. 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), - ); + let mut scope_bytes = Vec::with_capacity(provider.len() + model_id.len()); scope_bytes.extend_from_slice(provider.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 +499,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 +507,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() + fn supported_urls(&self) -> crate::language_model::SupportedUrls { + self.inner.supported_urls() } async fn do_generate(&self, options: &CallOptions) -> Result { 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 abca26a15..d23939e7a 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 efeb2c75b..516d345cf 100644 --- a/aimux-core/src/transcription_model.rs +++ b/aimux-core/src/transcription_model.rs @@ -284,21 +284,12 @@ impl std::fmt::Debug for TranscriptionStreamResult { /// Aligned with V4 `TranscriptionModelV4`. #[async_trait] pub trait TranscriptionModel: Send + Sync { - /// Specification version (always `"v4"`). - fn specification_version(&self) -> &'static str { - "v4" - } - /// Provider name, e.g. `"openai"`. fn provider(&self) -> &str; /// 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. @@ -334,11 +325,7 @@ pub async fn transcribe( ) -> 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/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 459086e03..647f3625a 100644 --- a/aimux-core/src/video_model.rs +++ b/aimux-core/src/video_model.rs @@ -287,23 +287,17 @@ impl Default for VideoPollConfig { /// same operation reference. #[async_trait] pub trait VideoModel: Send + Sync { - /// Specification version (always `"v4"`). - fn specification_version(&self) -> &'static str { - "v4" - } - /// Provider name, e.g. `"fal"`. fn provider(&self) -> &str; /// Provider-specific model ID, e.g. `"kling-video"`. fn model_id(&self) -> &str; - fn retry_config(&self) -> retry::RetryConfig { - retry::RetryConfig::default() - } - - /// Poll pacing for this model. Defaults to the AI SDK values; providers - /// with configurable polling should surface their configuration here. + /// Poll pacing for this model: the package's own constants. Providers have + /// no poll settings (the AI SDK's `pollIntervalMs` / `pollTimeoutMs` are + /// per-call options, here [`VideoCallOptions::poll`], which overrides this + /// per field). Defaults to the AI SDK values; a package whose job runs + /// longer or shorter returns its own constants. fn poll_config(&self) -> VideoPollConfig { VideoPollConfig::default() } @@ -381,11 +375,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 mut poll = model.poll_config(); if let Some(overrides) = options.poll { @@ -596,13 +586,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) } @@ -717,7 +700,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(); @@ -898,7 +882,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/tests/error_value_golden_test.rs b/aimux-core/tests/error_value_golden_test.rs index 7426f4758..fa62d5598 100644 --- a/aimux-core/tests/error_value_golden_test.rs +++ b/aimux-core/tests/error_value_golden_test.rs @@ -158,6 +158,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(), @@ -168,8 +182,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()), diff --git a/aimux-ffi/aimux-error.h b/aimux-ffi/aimux-error.h index bfe50be7c..abb1dda4c 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, @@ -60,7 +60,13 @@ typedef enum aimux_error_code { AIMUX_E_NO_SUCH_TOOL = 15, 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 = 18, + /* A required provider setting was not passed and its fallback environment + * variable is unset (the AI SDK's LoadSettingError). */ + AIMUX_E_LOAD_SETTING = 19, + AIMUX_E_NO_OUTPUT_GENERATED = 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 16db8b2d7..a781cacca 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..19 * [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. */ diff --git a/aimux-ffi/src/lib.rs b/aimux-ffi/src/lib.rs index ebb4d01f9..05094752a 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. //! @@ -71,18 +71,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; @@ -96,7 +98,7 @@ use tokio::runtime::Runtime; #[derive(Clone)] enum HandleEntry { Language(Arc), - Provider(Arc), + Provider(ProviderEntry), Embedding(Arc), Speech(Arc), Image(Arc), @@ -111,6 +113,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(); @@ -439,7 +449,12 @@ pub const AIMUX_E_RETRY: i32 = 14; 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 17 (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 = 18; +pub const AIMUX_E_LOAD_SETTING: i32 = 19; +pub const AIMUX_E_NO_OUTPUT_GENERATED: i32 = 20; // 100..105 preserve `RecordingError` as a separate high-level type while C // uses one code space for every returned error. @@ -479,6 +494,8 @@ fn aimux_error_code_of(err: &AiMuxError) -> i32 { AiMuxError::ToolCallRepair { .. } => AIMUX_E_TOOL_CALL_REPAIR, AiMuxError::InvalidArgument(_) => AIMUX_E_INVALID_ARGUMENT, AiMuxError::InvalidPrompt(_) => 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, @@ -601,19 +618,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()), ) } @@ -748,7 +781,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(), @@ -1044,6 +1077,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( @@ -1053,8 +1100,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)?, + )) }) } @@ -1070,12 +1118,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() }) } @@ -1088,8 +1162,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)) }) } @@ -1105,12 +1179,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)) }) } @@ -1126,10 +1197,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)) }) } @@ -1146,12 +1215,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)) }) } @@ -1171,14 +1237,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)) }) } @@ -1196,17 +1258,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( @@ -1227,13 +1401,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)) }) } @@ -1258,12 +1428,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)) }) } @@ -1287,9 +1459,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)) }) } @@ -1314,12 +1486,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)) }) } @@ -1332,8 +1501,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)) }) } @@ -1347,12 +1516,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)) }) } @@ -1365,8 +1530,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)) }) } @@ -1380,12 +1545,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)) }) } @@ -1398,8 +1559,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)) }) } @@ -1413,12 +1574,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)) }) } @@ -1429,11 +1586,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, @@ -1449,12 +1606,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, @@ -1464,8 +1621,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)) }) } @@ -1493,8 +1650,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, + }))) }) } @@ -1516,7 +1675,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()) }) } @@ -1539,8 +1701,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)) }) } @@ -2170,7 +2332,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)))) }) } @@ -2184,11 +2346,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)))) }) } @@ -2201,7 +2360,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)))) }) } @@ -2215,11 +2374,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)))) }) } @@ -2232,7 +2387,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)))) }) } @@ -2246,11 +2401,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)))) }) } @@ -2296,7 +2447,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)))) }) } @@ -2310,11 +2461,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)))) }) } @@ -2349,7 +2497,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)))) }) } @@ -2363,11 +2511,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)))) }) } @@ -2380,7 +2525,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)))) }) } @@ -2394,11 +2539,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)))) }) } @@ -2431,7 +2572,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)))) }) } @@ -2445,11 +2586,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)))) }) } @@ -2689,7 +2827,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)))) }) } @@ -2702,11 +2840,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)))) }) } @@ -2748,7 +2883,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)))) }) } @@ -2762,11 +2897,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)))) }) } @@ -2805,7 +2936,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)))) }) } @@ -2819,11 +2950,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)))) }) } @@ -2947,6 +3074,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)] @@ -2957,7 +3097,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)))) }) } @@ -2971,11 +3111,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)))) }) } @@ -3333,7 +3469,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. @@ -3744,8 +3880,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); @@ -3924,7 +4109,7 @@ mod tests { ); } - /// Pin the full 16-variant → code mapping. + /// Pin the full 18-variant → code mapping. #[test] fn error_code_mapping_covers_all_variants() { let s = |t: &str| t.to_string(); @@ -3981,6 +4166,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")), @@ -3996,6 +4195,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 5a7cf609d..2b515f81e 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( @@ -569,23 +569,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) ────────────────────────────── @@ -900,12 +894,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..3538b3942 --- /dev/null +++ b/aimux-provider-utils/src/fetch.rs @@ -0,0 +1,429 @@ +//! 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`] — a mock, a signing decorator +//! such as [`crate::sigv4_fetch::SigV4Fetch`], a proxy-aware client — and +//! when none is set the exchange uses [`default_fetch`], resolved *per +//! request* (never frozen when a provider is created), so a process-wide +//! default can change after a provider exists. +//! +//! This module owns the real network leaf: [`ReqwestFetch`] (one shared +//! connection pool per tokio runtime) 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}; +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 url::Url; + +use aimux_core::AbortSignal; + +use crate::http::{ProxyConfig, global_proxy}; + +/// 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. + #[must_use] + 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; + +/// The process-wide default transport, resolved at request time. +#[must_use] +pub fn default_fetch() -> FetchFunction { + static DEFAULT: OnceLock = OnceLock::new(); + DEFAULT.get_or_init(|| Arc::new(ReqwestFetch)).clone() +} + +/// The default transport: reqwest over rustls, with a connection pool shared +/// by every `ReqwestFetch` in the process (one pool per tokio runtime). +/// +/// 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. The process-wide [`ProxyConfig`] is snapshotted when a runtime's +/// client is first built. +#[derive(Debug, Clone, Copy, Default)] +pub struct ReqwestFetch; + +#[async_trait] +impl Fetch for ReqwestFetch { + async fn fetch(&self, request: FetchRequest) -> Result { + let client = shared_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 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, +} + +impl PinnedFetch { + /// Pin the connection to these validated addresses. + #[must_use] + pub fn new(addresses: Vec) -> Self { + Self { addresses } + } + + 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, &global_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()) + } +} + +// 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; + +/// Return the shared client for the current runtime, building it on first use. +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| { + FetchError::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 = 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::*; + + #[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/generate_id.rs b/aimux-provider-utils/src/generate_id.rs new file mode 100644 index 000000000..70c0211d8 --- /dev/null +++ b/aimux-provider-utils/src/generate_id.rs @@ -0,0 +1,17 @@ +use std::hash::{BuildHasher, RandomState}; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Generate a non-cryptographic, 16-character alphanumeric ID, like `generateId`. +#[must_use] +pub 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() +} 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 9d8723414..554a6e88f 100644 --- a/aimux-provider-utils/src/http.rs +++ b/aimux-provider-utils/src/http.rs @@ -1,18 +1,17 @@ -//! 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::sync::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::{ @@ -20,9 +19,14 @@ use aimux_core::recording::{ }; 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. +/// Process-wide proxy configuration, fixed before the first HTTP client is built. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct ProxyConfig { pub http_url: Option, @@ -32,23 +36,6 @@ pub struct ProxyConfig { } 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. /// @@ -68,125 +55,9 @@ 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 +96,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 +118,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 every model's call options carry into an @@ -353,6 +228,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 { @@ -361,10 +240,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) { crate::download_guard::retain_user_agent(&mut headers); + credentials = RequestCredentials::Omit; } PreparedRequest { method, @@ -380,6 +261,8 @@ impl HttpRequest { trusted_origin: self.trusted_origin, credentialed_origin: self.credentialed_origin, }), + fetch: self.fetch, + credentials, } } } @@ -444,7 +327,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 @@ -472,14 +355,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) => { @@ -491,7 +367,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" ); @@ -506,51 +382,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()), @@ -558,53 +431,61 @@ 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 a redirect only while it stays on the origin +/// of the hop that issued it. A cross-origin `3xx` is returned as it is, so +/// neither the request headers 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. + crate::download_guard::sanitize_download_headers(&mut current.headers); + Some( + crate::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. + Some(pins) => send_one_request(&PinnedFetch::new(pins.clone()), ¤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 @@ -621,17 +502,33 @@ 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 = + crate::download_guard::hop_trusted_origin(trusted_origin, ¤t.url); + pinned = + Some(crate::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. + // 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 !crate::download_guard::same_origin(&next, credential_anchor) { crate::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; @@ -642,28 +539,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); - } - 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, + // 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); + } + 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, }) } @@ -691,16 +631,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 { @@ -742,23 +681,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, @@ -766,7 +704,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 { @@ -967,7 +905,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 81f1606dd..556753782 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -2,25 +2,31 @@ //! //! 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; pub mod streaming_tool_call_argument_state; pub mod streaming_tool_call_tracker; pub mod url; @@ -29,16 +35,26 @@ pub mod url; #[cfg(feature = "ws")] pub mod ws; -pub use api_key::load_api_key; +pub use api_key::{load_api_key, load_optional_setting, load_setting}; pub use download_guard::same_origin; +pub use fetch::{ + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, PinnedFetch, RequestCredentials, + ReqwestFetch, default_fetch, +}; +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, ProxyConfig, init_proxy, 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, @@ -46,7 +62,7 @@ pub use response_handler::{ create_standard_json_error_response_handler, create_status_code_error_response_handler, stream_error_api_call, }; -pub use retry::RetryConfig; +pub use sigv4_fetch::{AwsCredentials, SigV4Fetch}; pub use streaming_tool_call_argument_state::{ StreamingToolCallArgumentState, starts_with_structured_value, }; 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 0d8917bd3..ad0c09139 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. 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..4dc43e76b --- /dev/null +++ b/aimux-provider-utils/src/sigv4_fetch.rs @@ -0,0 +1,712 @@ +//! 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. +#[derive(Debug, 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, +} + +/// 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(), + } + } + + 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..c45458c6a 100644 --- a/aimux-provider-utils/src/ws.rs +++ b/aimux-provider-utils/src/ws.rs @@ -22,8 +22,10 @@ //! 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 +36,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 +51,47 @@ 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 the built-in + /// tungstenite connector (with the process-wide proxy tunnel). Injection + /// point for tests 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 +103,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 @@ -398,8 +442,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 +452,24 @@ 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 => TungsteniteConnector.connect(req).await, + } +} + +/// The built-in connector: tungstenite, with the process-wide proxy tunnel. +#[derive(Debug, Clone, Copy, Default)] +struct TungsteniteConnector; + +#[async_trait] +impl WsConnector for TungsteniteConnector { + async fn connect(&self, request: &WebSocketRequest) -> Result { + connect_tungstenite(request).await + } +} + +async fn connect_tungstenite(req: &WebSocketRequest) -> Result { let mut http_req = req .url .as_str() @@ -517,27 +580,52 @@ pub async fn ws_connect(req: &WebSocketRequest) -> 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 +728,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..566c99f41 --- /dev/null +++ b/aimux-provider-utils/tests/api_redirect_test.rs @@ -0,0 +1,179 @@ +//! 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_provider_utils::{ + AwsCredentials, HttpRequest, ProviderErrorParts, Resolvable, SigV4Fetch, + create_json_error_response_handler, create_json_response_handler, default_fetch, + post_json_to_api, +}; +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"); +} 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..abd85c325 100644 --- a/aimux-provider-utils/tests/ws_proxy_test.rs +++ b/aimux-provider-utils/tests/ws_proxy_test.rs @@ -205,6 +205,7 @@ fn ws_request(url: String) -> WebSocketRequest { step_ms: None, total_ms: Some(10_000), }), + connector: None, } } @@ -563,6 +564,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 af679e5f5..0a63c7ed5 100644 --- a/aimux-providers/src/anthropic/convert.rs +++ b/aimux-providers/src/anthropic/convert.rs @@ -28,10 +28,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. @@ -98,18 +102,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 +155,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 +170,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 +196,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 +207,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 +239,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); @@ -316,11 +348,6 @@ pub fn convert_prompt_to_anthropic( } } -/// 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 +357,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 +374,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 +383,7 @@ fn convert_file_part( is_last_part: bool, part_context: &str, message_context: &str, + options_name: &str, ) -> Result { let FilePart { data, @@ -374,7 +395,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 +412,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 +499,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 +515,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 +578,7 @@ fn convert_assistant_part( is_last_part, "assistant message part", "assistant message", + options_name, )?, AssistantPart::Reasoning(ReasoningPart { @@ -567,6 +591,7 @@ fn convert_assistant_part( send_reasoning, warnings, validator, + options_name, )); } @@ -582,16 +607,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 +692,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 +718,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 +731,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 +746,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 +808,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 +1154,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 +1202,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 +1244,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 +1259,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 +1272,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 +1283,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 +1298,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 +1306,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 +1320,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: &std::collections::HashMap, ) -> 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 +1365,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 +1379,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 +1635,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 +1686,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 +1850,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 +1979,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 +2195,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 +2225,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 +2233,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 +2262,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 +2318,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 +2339,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 +2412,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.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 +2469,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 +2483,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 +2498,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 +2535,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 4bc7b46cc..59ee58cb1 100644 --- a/aimux-providers/src/anthropic/files.rs +++ b/aimux-providers/src/anthropic/files.rs @@ -14,9 +14,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 +59,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 +94,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 +104,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 +136,7 @@ impl Files for AnthropicFiles { } let mut provider_ref = HashMap::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,7 +148,9 @@ impl Files for AnthropicFiles { provider_reference: provider_ref, media_type: Some(result_media_type), filename: result_filename, - provider_metadata: Some(HashMap::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..d3a96fc5f 100644 --- a/aimux-providers/src/anthropic/mod.rs +++ b/aimux-providers/src/anthropic/mod.rs @@ -1,25 +1,55 @@ -//! 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, load_optional_setting, 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 +65,303 @@ 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"; +const DEFAULT_BASE_URL: &str = "https://api.anthropic.com/v1"; +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. Default + /// `https://api.anthropic.com/v1`; a trailing slash is removed. The bare + /// `https://api.anthropic.com` means the default. + pub base_url: Option, + /// The API key, sent as `x-api-key`. `None` (with no `auth_token`) loads + /// `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 `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, not here. +pub fn create_anthropic( + settings: AnthropicProviderSettings, +) -> Result { + let base_url = match load_optional_setting(settings.base_url.as_deref(), "ANTHROPIC_BASE_URL") { + Some(url) => normalize_base_url(&url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + 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. The base URL is loaded at 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("invalid ANTHROPIC_BASE_URL") + }) } -/// 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, + base_url: String, + 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 = self.base_url.clone(); + AnthropicModelConfig { + provider, + url: Arc::new(move |path| format!("{base}{path}")), + headers: self.headers.clone(), + fetch: self.fetch.clone(), + supported_urls: self.supported_urls.clone(), + transform_request_body: None, + base_url: self.base_url.clone(), + provider_options_name: options::options_name_of(&self.name), + hooks: AnthropicModelHooks::default(), + resolve: None, + } } -} -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 ef9155db7..efd4edac3 100644 --- a/aimux-providers/src/anthropic/model.rs +++ b/aimux-providers/src/anthropic/model.rs @@ -1,146 +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, - })), - } + 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(), - options.recording_context.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(), ) @@ -148,59 +117,19 @@ 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(), - options.recording_context.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.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 c7b4e606a..c7ff89a33 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::ToolChoice; 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<&ToolChoice>, + 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, )); @@ -249,13 +277,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..b10fffb00 --- /dev/null +++ b/aimux-providers/src/anthropic/skills.rs @@ -0,0 +1,183 @@ +//! Skill uploads and version metadata from the Anthropic Skills API. + +use std::collections::{BTreeSet, HashMap}; + +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: HashMap::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 48c185ec5..4ba960f05 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,19 @@ 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::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,78 +65,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, - recording_context: 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()), - }; - - 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!({ @@ -354,6 +282,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!({ @@ -367,7 +296,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"), ) } @@ -404,6 +333,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, @@ -460,7 +390,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), }), @@ -535,7 +465,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") @@ -556,6 +486,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), @@ -602,7 +533,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 { @@ -612,7 +543,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"), ), })); @@ -647,7 +578,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"), @@ -659,7 +590,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"), ), })); @@ -697,7 +628,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), @@ -826,7 +757,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") @@ -841,36 +772,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, - recording_context: 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, - recording_context, - )?; - 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?; @@ -879,7 +800,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 @@ -895,12 +821,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, @@ -942,38 +877,26 @@ 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, - recording_context: Option, + config: &AnthropicModelConfig, + used_custom_provider_key: bool, tool_names: ToolNameMapping, generate_id: std::sync::Arc String + Send + Sync>, ) -> Result { - let (request, request_body) = build_anthropic_request( - endpoint, - &body, - &build_headers, - body_encoding, - abort_signal, - recording_context, - )?; - let resp = aimux_provider_utils::post_to_api( + 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, + body.clone(), aimux_provider_utils::create_event_source_response_handler::(), - super::anthropic_failed_response_handler(), + config.failed_response_handler(), ) .await?; @@ -986,12 +909,12 @@ pub(crate) async fn anthropic_stream_core( 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(); @@ -1003,6 +926,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 @@ -1022,6 +952,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 { @@ -1073,7 +1004,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 @@ -1129,10 +1060,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. @@ -1140,7 +1071,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, @@ -1157,6 +1088,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, @@ -1306,7 +1238,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") }), }); } @@ -1328,7 +1260,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(), @@ -1345,11 +1278,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() @@ -1417,7 +1359,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 a1831e853..524870370 100644 --- a/aimux-providers/src/anthropic_aws/model.rs +++ b/aimux-providers/src/anthropic_aws/model.rs @@ -1,198 +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(), - options.recording_context.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(), - options.recording_context.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 8c0683aa0..5b825cb59 100644 --- a/aimux-providers/src/assemblyai.rs +++ b/aimux-providers/src/assemblyai.rs @@ -7,8 +7,15 @@ //! 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 +23,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 +69,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; + +const DEFAULT_BASE_URL: &str = "https://api.assemblyai.com"; +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)] @@ -177,48 +258,19 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { 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 = HashMap::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 = HashMap::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..5ab024c63 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; + +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 a911d34c7..8d0b4b20e 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` @@ -693,10 +676,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..221ced35a 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"; +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 cc9a7852e..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,98 +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 { - url, - headers, - - abort_signal: options.abort_signal.clone(), - call_id: options.call_id.clone(), - recording_context: options.recording_context.clone(), - ..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(), @@ -267,19 +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 { - url, - headers, - - abort_signal: options.abort_signal.clone(), - call_id: options.call_id.clone(), - recording_context: options.recording_context.clone(), - ..Default::default() - }, + exchange.request(url, options), HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), bedrock_event_stream_response_handler(), super::bedrock_failed_response_handler(), @@ -586,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(), @@ -597,10 +507,8 @@ impl LanguageModel for BedrockModel { let provider_metadata = if finish_meta.is_empty() { None } else { - Some(HashMap::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 { @@ -624,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. @@ -644,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 { @@ -672,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 e915d95ca..000000000 --- a/aimux-providers/src/bedrock/sigv4.rs +++ /dev/null @@ -1,223 +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. -#[derive(Debug, 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, -} - -/// 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 { - 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 = parsed.path(); - let query = parsed.query().unwrap_or(""); - - // Timestamp:yyyyMMddTHHmmssZ and date:yyyyMMdd - let now = chrono::Utc::now(); - 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, v.trim().to_string())); - } - } - - 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 - let k_date = hmac_sha256( - date_stamp.as_bytes(), - credentials.secret_access_key.as_bytes(), - ); - let k_region = hmac_sha256(credentials.region.as_bytes(), &k_date); - let k_service = hmac_sha256(service.as_bytes(), &k_region); - let k_signing = hmac_sha256(b"aws4_request", &k_service); - - let signature = hex::encode(hmac_sha256(string_to_sign.as_bytes(), &k_signing)); - - 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 } -} - -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(), - } - } - - 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" - ); - } -} 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 1482fc9e8..be3b50f6e 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; + +const DEFAULT_BASE_URL: &str = "https://api.bfl.ai"; +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..b760f4eaf 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 - } +const DEFAULT_BASE_URL: &str = "https://api.cartesia.ai"; +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> { @@ -700,7 +734,7 @@ fn audio_input_to_bytes_stt(audio: &AudioInput) -> Result, AiMuxError> { #[async_trait] impl TranscriptionModel for CartesiaTranscriptionModel { fn provider(&self) -> &str { - "cartesia" + &self.config.provider } fn model_id(&self) -> &str { @@ -723,9 +757,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 +801,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 1621a8ff8..31b08c679 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,26 +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 { - url: endpoint.clone(), - headers: build_header_list(&headers), - abort_signal: opts.abort_signal.clone(), - call_id: options.call_id.clone(), - recording_context: options.recording_context.clone(), - ..Default::default() - }, + 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; @@ -415,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, @@ -436,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, } } } @@ -482,7 +513,7 @@ impl LanguageModel for CodexModel { /// Tokens returned by the Codex OAuth token endpoint (RFC-0018 §2 V3). #[derive(Debug, 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). @@ -563,22 +594,3 @@ pub async fn codex_refresh_at( expires_in_secs: data.get("expires_in").and_then(serde_json::Value::as_u64), }) } - -#[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"); - } -} diff --git a/aimux-providers/src/cohere/convert.rs b/aimux-providers/src/cohere/convert.rs index 0c2d866a2..97c96443f 100644 --- a/aimux-providers/src/cohere/convert.rs +++ b/aimux-providers/src/cohere/convert.rs @@ -12,6 +12,7 @@ use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; 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 ────────────────────────────────────────────────────────── @@ -123,7 +124,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 { @@ -148,7 +149,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 { @@ -158,9 +159,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 { .. } => { @@ -170,7 +173,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(), )); } }; @@ -319,7 +322,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, @@ -432,15 +435,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() @@ -449,7 +453,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..aec4b1d61 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, -} +const DEFAULT_BASE_URL: &str = "https://api.cohere.com/v2"; +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 374140fa7..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,11 +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(HashMap::from([("cohere".to_string(), cohere_meta)])), + provider_metadata: Some(super::options::cohere_metadata(Value::Object( + cohere_meta, + ))), })); } } @@ -265,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(), @@ -571,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!({}))), }); }; @@ -614,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..01672c1d8 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 - } +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/deepgram.rs b/aimux-providers/src/deepgram.rs index 1fcfa51c9..29edd2701 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; + +const DEFAULT_BASE_URL: &str = "https://api.deepgram.com"; +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); @@ -222,39 +288,19 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { 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 +364,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..2597c283e --- /dev/null +++ b/aimux-providers/src/deepseek/mod.rs @@ -0,0 +1,193 @@ +//! 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; + +const DEFAULT_BASE_URL: &str = "https://api.deepseek.com"; +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, + aimux_provider_utils::create_standard_json_error_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..fecea4fc6 --- /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 + } + + 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::(), + aimux_provider_utils::create_standard_json_error_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::(), + aimux_provider_utils::create_standard_json_error_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..d5f35206f --- /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::ToolChoice; +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<&ToolChoice>, + 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 { + ToolChoice::Auto => json!("auto"), + ToolChoice::None => json!("none"), + ToolChoice::Required => json!("required"), + ToolChoice::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..839a729d4 --- /dev/null +++ b/aimux-providers/src/deepseek/types.rs @@ -0,0 +1,128 @@ +//! 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; + +#[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..1edebc13f --- /dev/null +++ b/aimux-providers/src/default_providers.rs @@ -0,0 +1,403 @@ +//! 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"), + } +} + +/// 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. +/// +/// Nothing is read from the environment here and no entry can fail to be +/// created: keys and environment-sourced base URLs are evaluated per request, +/// so a missing key fails the call that needs it. +#[must_use] +pub fn default_providers() -> BTreeMap> { + provider_names() + .into_iter() + .map(|name| { + let provider = create_provider(name, PresetSettings::default()) + .expect("a built-in provider with default settings is always valid"); + (name.to_string(), provider) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 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")); + } +} 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..399695fc1 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, - } - } +const DEFAULT_BASE_URL: &str = "https://api.elevenlabs.io"; +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) } } @@ -630,7 +706,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 +737,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 +763,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 +851,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 +880,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 +909,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 +919,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..9f5472049 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 4bd3daa9c..d18a45069 100644 --- a/aimux-providers/src/fal.rs +++ b/aimux-providers/src/fal.rs @@ -5,21 +5,33 @@ //! //! 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 +72,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 - } +const DEFAULT_BASE_URL: &str = "https://queue.fal.run"; +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)))) } } @@ -167,50 +279,27 @@ fn audio_input_to_base64(audio: &AudioInput) -> Result { 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..04ed8affa 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, +const DEFAULT_BASE_URL: &str = "https://api.firecrawl.dev"; +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..d3fd38c16 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; + +const DEFAULT_BASE_URL: &str = "https://api.gladia.io"; +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)] @@ -162,47 +238,19 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { 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 +267,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 43fbca977..43762c3b2 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, ResponseFormat, ToolChoice}; 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(); @@ -1042,9 +991,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)); @@ -1052,9 +1002,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) } @@ -1272,10 +1223,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. @@ -1295,12 +1247,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 @@ -1313,24 +1261,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))) }); @@ -1397,23 +1338,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 ecab1f955..783776217 100644 --- a/aimux-providers/src/google/files.rs +++ b/aimux-providers/src/google/files.rs @@ -27,7 +27,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 +45,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 +113,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 +133,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 +161,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 +177,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 +233,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 +288,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 +327,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 +350,7 @@ impl Files for GoogleFiles { } let mut provider_ref = HashMap::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()) @@ -402,7 +362,7 @@ impl Files for GoogleFiles { provider_reference: provider_ref, media_type: result_media_type, filename: None, - provider_metadata: Some(HashMap::from([("google".to_string(), metadata)])), + provider_metadata: Some(HashMap::from([(GOOGLE.to_string(), metadata)])), warnings, }) } diff --git a/aimux-providers/src/google/image.rs b/aimux-providers/src/google/image.rs index 098864a03..e113c32cb 100644 --- a/aimux-providers/src/google/image.rs +++ b/aimux-providers/src/google/image.rs @@ -19,9 +19,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 +42,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 +58,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_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 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 +152,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 +269,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 +297,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 +329,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 +351,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 +359,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 +399,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 +465,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..ed4347317 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, +const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta"; +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..3a77425c1 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, - } + 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..d61f04a52 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(), - } - } +const DEFAULT_BASE_URL: &str = "https://www.googleapis.com/customsearch/v1"; +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..d69f4bcfb --- /dev/null +++ b/aimux-providers/src/groq/error.rs @@ -0,0 +1,23 @@ +//! Mirrors `groq-error.ts`: the `{ "error": { "message", "type" } }` shape of +//! a failed response. + +use serde::Deserialize; + +use aimux_core::error::AiMuxError; +use aimux_provider_utils::{ResponseHandler, create_standard_json_error_response_handler}; + +pub(crate) fn groq_failed_response_handler() -> ResponseHandler { + create_standard_json_error_response_handler() +} + +/// `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..14e72bdf3 --- /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}; + +const DEFAULT_BASE_URL: &str = "https://api.groq.com/openai/v1"; +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..a466543fa --- /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 + } + + 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..0087f028c --- /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::{Tool, ToolChoice}; +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 { + ToolChoice::Auto => json!("auto"), + ToolChoice::None => json!("none"), + ToolChoice::Required => json!("required"), + ToolChoice::Tool { tool_name } => { + json!({ "type": "function", "function": { "name": tool_name } }) + } + }); + prepared +} diff --git a/aimux-providers/src/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..72d49a97f 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; +pub use crate::shared::TransformRequestBody; +pub use responses::HuggingFaceResponsesModel; -use crate::openai::{OpenAIConfig, OpenAIModel}; +use std::sync::{Arc, OnceLock}; -const DEFAULT_BASE_URL: &str = "https://router.huggingface.co/v1"; -const ENV_VAR: &str = "HUGGINGFACE_API_KEY"; -const PROVIDER_NAME: &str = "huggingface"; +use futures::future::BoxFuture; -/// Configuration for the Hugging Face provider (wraps [`OpenAIConfig`]). -#[derive(Debug, Clone)] -pub struct HuggingFaceConfig(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}; -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)) - } +use crate::openai::config::OpenAIModelConfig; +use crate::shared::{Credential, EndpointConfig, provider_headers}; - /// 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"))) - } +const DEFAULT_BASE_URL: &str = "https://router.huggingface.co/v1"; +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 e5b89df5b..818ba840f 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, ResponseFormat, ToolChoice}; 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..9300b9b9a 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, - } - } +const DEFAULT_BASE_URL: &str = "https://api.hume.ai"; +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 8c5e3162e..78a1ba2f2 100644 --- a/aimux-providers/src/jev.rs +++ b/aimux-providers/src/jev.rs @@ -22,7 +22,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 +32,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 +62,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() { @@ -304,23 +302,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 { rounding: JEV_ROUNDING, diff --git a/aimux-providers/src/jina_ai.rs b/aimux-providers/src/jina_ai.rs index 51ab3dd7e..c042176b6 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, - } - } +const DEFAULT_BASE_URL: &str = "https://api.jina.ai"; +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..71aa4aa3f 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; + +const DEFAULT_BASE_URL: &str = "https://api.klingai.com"; +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 5c60bc5b8..44ba8b352 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_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,25 @@ 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_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 +207,138 @@ 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, OpenAIEmbeddingModel, OpenAIImageModel, OpenAIProvider, OpenAIResponsesModel, - OpenAISpeechModel, OpenAITranscriptionModel, + 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,14 +353,37 @@ 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, +}; diff --git a/aimux-providers/src/linkup.rs b/aimux-providers/src/linkup.rs index 697931187..0b2bc93a4 100644 --- a/aimux-providers/src/linkup.rs +++ b/aimux-providers/src/linkup.rs @@ -7,20 +7,26 @@ //! 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. const MODEL_ID: &str = "linkup-search"; @@ -37,9 +43,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 +54,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(), - } - } +const DEFAULT_BASE_URL: &str = "https://api.linkup.so"; +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 +232,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,21 +255,10 @@ 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(), 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 c621eea51..000000000 --- a/aimux-providers/src/llamacpp.rs +++ /dev/null @@ -1,79 +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()) - .with_api_key_source(Some("none")), - ) - } - - /// 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 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 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))) - } - - 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..2937cfb91 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, - } - } +const DEFAULT_BASE_URL: &str = "https://api.lmnt.com"; +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 ff6b55d38..000000000 --- a/aimux-providers/src/localai.rs +++ /dev/null @@ -1,79 +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()) - .with_api_key_source(Some("none")), - ) - } - - /// 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 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 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))) - } - - crate::delegate_list_models!(); -} diff --git a/aimux-providers/src/luma.rs b/aimux-providers/src/luma.rs index f99e163a1..d98875565 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>, +const DEFAULT_BASE_URL: &str = "https://api.lumalabs.ai"; +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 683392aa2..221909126 100644 --- a/aimux-providers/src/mistral/convert.rs +++ b/aimux-providers/src/mistral/convert.rs @@ -12,6 +12,7 @@ use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; use aimux_core::shared::{FileBytes, FileData}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified}; +use aimux_provider_utils::{get_top_level_media_type, resolve_full_media_type}; use serde_json::{Value, json}; // ── Prepared tools ────────────────────────────────────────────────────────── @@ -106,7 +107,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> { @@ -229,21 +230,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, .. } => { @@ -256,12 +257,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(), )); } }; @@ -279,7 +280,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, @@ -288,21 +289,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..be1690333 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, -} +const DEFAULT_BASE_URL: &str = "https://api.mistral.ai/v1"; +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..125009937 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 + 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 6c78d5ebe..000000000 --- a/aimux-providers/src/ollama.rs +++ /dev/null @@ -1,103 +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()) - .with_api_key_source(Some("none")), - ) - } - - /// 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 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 - } -} - -/// 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))) - } - - 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 bc6f46dbe..f84a4027b 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, ToolChoice}; -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 f79d608f0..ccecb76bd 100644 --- a/aimux-providers/src/openai/convert.rs +++ b/aimux-providers/src/openai/convert.rs @@ -9,10 +9,10 @@ use aimux_core::options::{CallOptions, ResponseFormat, ToolChoice}; use aimux_core::shared::{FileBytes, FileData, JsonObject, SharedProviderOptions}; use aimux_core::tool::{FunctionTool, Tool}; use aimux_core::types::{FinishReason, FinishReasonUnified, ReasoningEffort, 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,106 +97,6 @@ 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<&ToolChoice>, - 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 { - ToolChoice::Auto => Some(json!("auto")), - ToolChoice::None => Some(json!("none")), - ToolChoice::Required => Some(json!("required")), - ToolChoice::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. @@ -231,7 +131,11 @@ pub fn convert_prompt_to_openai_messages_with_mode_fallible( prompt: &LanguageModelPrompt, system_message_mode: SystemMessageMode, ) -> Result, AiMuxError> { - convert_prompt_to_openai_messages_with_provider(prompt, system_message_mode, "openai") + let mut result = Vec::new(); + for msg in prompt { + result.extend(convert_message_to_openai(msg, system_message_mode)?); + } + Ok(result) } /// Convert a `LanguageModelPrompt` to OpenAI `messages` array with a system @@ -258,40 +162,6 @@ pub fn convert_prompt_to_openai_messages_with_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, - )?); - } - Ok(result) -} - /// Get the prompt cache breakpoint from provider options. fn get_prompt_cache_breakpoint(provider_options: &Option) -> Option { provider_options @@ -310,38 +180,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: &std::collections::HashMap, @@ -359,36 +197,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 +240,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 +266,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 +293,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 +328,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 +339,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 +353,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 +366,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 +387,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 +397,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 +440,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 +545,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 +586,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(std::collections::HashMap::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 +619,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 +627,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 +646,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 +672,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 +797,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 +810,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 +834,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 +855,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 +910,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 +922,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 +945,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| { + std::collections::HashMap::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 +994,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 +1009,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/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 ba9e875df..9c88bdb67 100644 --- a/aimux-providers/src/openai/files.rs +++ b/aimux-providers/src/openai/files.rs @@ -14,9 +14,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 +89,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 +131,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 +140,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 df80a6133..ea93697ab 100644 --- a/aimux-providers/src/openai/image.rs +++ b/aimux-providers/src/openai/image.rs @@ -20,9 +20,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 +65,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 +84,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 +108,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 +120,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 +133,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 +174,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 220758c83..50416b015 100644 --- a/aimux-providers/src/openai/mod.rs +++ b/aimux-providers/src/openai/mod.rs @@ -1,34 +1,60 @@ -//! 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 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 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, load_optional_setting, 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 { @@ -87,334 +113,238 @@ 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>, -} - -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 - } +const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; +const API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; - /// 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. Default `https://api.openai.com/v1`; a + /// trailing slash is removed. + pub base_url: Option, + /// The API key. `None` loads `OPENAI_API_KEY` when a request is made and + /// fails that request with `AiMuxError::LoadApiKey` if it is unset. An + /// 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 `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_openai(settings: OpenAIProviderSettings) -> Result { + let base_url = match load_optional_setting(settings.base_url.as_deref(), "OPENAI_BASE_URL") { + Some(url) => validate_base_url(&url)?, + None => DEFAULT_BASE_URL.to_string(), + }; + 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. The base URL is loaded at 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("invalid OPENAI_BASE_URL") + }) } -/// 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, + base_url: String, + 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 { + OpenAIModelConfig::fixed( + format!("{}.{method}", self.name), + self.base_url.clone(), + self.headers.clone(), + self.fetch.clone(), + None, + ) } - /// Use a custom base URL (for Azure, Groq, etc.). + /// A chat-completions model; `provider()` is `"{name}.chat"`. #[must_use] - pub fn with_base_url(mut self, url: impl Into) -> Self { - self.base_url = without_trailing_slash(&url.into()); - self + pub fn chat(&self, model_id: &str) -> OpenAIModel { + OpenAIModel::from_config(model_id.to_string(), self.model_config("chat")) } + /// A text-completion model; `provider()` is `"{name}.completion"`. #[must_use] - pub fn with_org_id(mut self, org_id: impl Into) -> Self { - self.org_id = Some(org_id.into()); - self + pub fn completion(&self, model_id: &str) -> OpenAICompletionModel { + OpenAICompletionModel::from_native_config( + model_id.to_string(), + self.model_config("completion"), + ) } - /// Set the OpenAI project ID (sent via the `OpenAI-Project` header). + /// A Responses API model; `provider()` is `"{name}.responses"`. #[must_use] - pub fn with_project(mut self, project: impl Into) -> Self { - self.project = Some(project.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) } - /// Attach extra headers merged into every request. + /// An embedding model; `provider()` is `"{name}.embedding"`. #[must_use] - pub fn with_headers(mut self, headers: HashMap) -> Self { - self.headers = Some(headers); - self + pub fn embedding(&self, model_id: &str) -> OpenAIEmbeddingModel { + OpenAIEmbeddingModel::from_config(model_id.to_string(), self.model_config("embedding")) } - /// Set the provider name (e.g. "groq") for provider-specific behaviour. + /// An image model; `provider()` is `"{name}.image"`. #[must_use] - pub fn with_provider(mut self, provider: impl Into) -> Self { - self.provider = provider.into(); - self + pub fn image(&self, model_id: &str) -> OpenAIImageModel { + OpenAIImageModel::from_config(model_id.to_string(), self.model_config("image")) } - /// 设置厂商能力差异描述。 + /// A speech (TTS) model; `provider()` is `"{name}.speech"`. #[must_use] - pub fn with_profile(mut self, profile: OpenAICompatProfile) -> Self { - self.profile = profile; - self + pub fn speech(&self, model_id: &str) -> OpenAISpeechModel { + OpenAISpeechModel::from_config(model_id.to_string(), self.model_config("speech")) } - /// Set the retry configuration. Pass `max_retries: 0` to disable retries. + /// A transcription (STT) model; `provider()` is `"{name}.transcription"`. #[must_use] - pub fn with_retry_config(mut self, config: RetryConfig) -> Self { - self.retry_config = config; - self + 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 } - /// 设置 provider 级请求体覆盖（RFC-0017）。 + /// The files interface; `provider()` is `"{name}.files"`. #[must_use] - pub fn with_body_overrides(mut self, overrides: Value) -> Self { - self.body_overrides = Some(overrides); - self + pub fn files(&self) -> OpenAIFiles { + OpenAIFiles::from_config(self.model_config("files")) } - /// Create from environment variable `OPENAI_API_KEY`. - /// - /// # 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"))) - } - - /// 标注 api_key 来源(RFC-0023 回放重建用)。 + /// 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 with_api_key_source(mut self, source: Option<&str>) -> Self { - self.api_key_source = source.map(std::string::ToString::to_string); - self + pub fn call(&self, model_id: &str) -> Arc { + Arc::new(self.responses(model_id)) } } -/// 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 { - #[must_use] - pub fn new(config: OpenAIConfig) -> Self { - Self { config } - } - - /// Create a model instance for the given model name (e.g. `"gpt-4o"`). - #[must_use] - pub fn model(&self, model_id: &str) -> model::OpenAIModel { - model::OpenAIModel::new(model_id.to_string(), self.config.clone()) +impl Provider for OpenAIProvider { + fn discovery(&self) -> Option<&dyn ProviderDiscovery> { + Some(self) } - /// 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`. - #[must_use] - pub fn responses_model(&self, model_id: &str) -> responses::OpenAIResponsesModel { - responses::OpenAIResponsesModel::new(model_id.to_string(), self.config.clone()) + fn language_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(self.call(model_id)) } - /// Create a Files interface for uploading files to OpenAI. - #[must_use] - pub fn files(&self) -> files::OpenAIFiles { - files::OpenAIFiles::new(self.config.clone()) + fn embedding_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.embedding(model_id))) } - /// 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()) + fn image_model(&self, model_id: &str) -> Result, AiMuxError> { + Ok(Arc::new(self.image(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()) + fn transcription_model( + &self, + model_id: &str, + ) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.transcription(model_id)))) } - /// 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 speech_model(&self, model_id: &str) -> Option, AiMuxError>> { + Some(Ok(Arc::new(self.speech(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 files(&self) -> Option> { + Some(Arc::new(self.files())) } } -impl Provider for OpenAIProvider { - fn name(&self) -> &str { - "openai" - } - - fn language_model(&self, model_id: &str) -> Result, AiMuxError> { - Ok(Box::new(self.model(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) - }) +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..bb392243c 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) + 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 2ed3a40dc..882911c42 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, ResponseFormat, ToolChoice}; -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; @@ -40,11 +41,71 @@ pub use crate::openai::convert_common::SystemMessageMode as ResponsesSystemMessa // -- 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 +134,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,7 +171,7 @@ 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 })); } @@ -120,7 +182,7 @@ pub fn convert_to_responses_input( 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 +190,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 +210,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,7 +231,7 @@ pub fn convert_to_responses_input( provider_options, .. }) => { - let reasoning_id = openai_sub_option(provider_options, "itemId"); + let reasoning_id = openai_sub_option(ns, provider_options, "itemId"); if (has_previous_response_id || has_conversation) && reasoning_id.is_some() { @@ -180,6 +242,7 @@ pub fn convert_to_responses_input( input.push(json!({ "type": "item_reference", "id": rid })); } else { let encrypted = openai_sub_option( + ns, provider_options, "reasoningEncryptedContent", ); @@ -200,6 +263,7 @@ pub fn convert_to_responses_input( } } else { let encrypted = openai_sub_option( + ns, provider_options, "reasoningEncryptedContent", ); @@ -229,7 +293,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 +303,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 +311,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 +343,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 +407,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 +420,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 +476,7 @@ fn convert_user_part(part: &UserPart) -> Result { } fn convert_tool_result_output( + ns: ResponsesNamespace, output: &ToolResultOutput, warnings: &mut Vec, ) -> Result { @@ -409,12 +498,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 +517,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 +531,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 +558,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 +711,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 +749,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 +762,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 +779,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 +800,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 +838,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 +857,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 +873,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 +893,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 +909,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 +929,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 +940,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 +960,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 +972,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 +1006,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 +1018,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 +1047,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 +1076,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 +1117,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 +1127,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 +1171,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. diff --git a/aimux-providers/src/openai/responses/mod.rs b/aimux-providers/src/openai/responses/mod.rs index eb703841e..b9934f070 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 + 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..cd302c7ba 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 ──────────────────────────────────────────────────────────── @@ -217,42 +215,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 +249,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 +319,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 +413,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 +446,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 +465,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..89b0e3da2 --- /dev/null +++ b/aimux-providers/src/openai_compatible/chat.rs @@ -0,0 +1,883 @@ +//! 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 + } + + fn supported_urls(&self) -> SupportedUrls { + self.config + .chat + .supported_urls + .as_ref() + .map(|urls| urls()) + .unwrap_or_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 = HashMap::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 = HashMap::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..10bc9c161 --- /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 + } + 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..630faa0fe --- /dev/null +++ b/aimux-providers/src/openai_compatible/config.rs @@ -0,0 +1,235 @@ +//! 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 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. +pub type SupportedUrlsFn = Arc 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..b61494981 --- /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, ResponseFormat, ToolChoice}; +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: &ToolChoice) -> Value { + match choice { + ToolChoice::Auto => json!("auto"), + ToolChoice::None => json!("none"), + ToolChoice::Required => json!("required"), + ToolChoice::Tool { tool_name } => { + json!({ "type": "function", "function": { "name": tool_name } }) + } + } +} + +/// 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..aec4882c3 --- /dev/null +++ b/aimux-providers/src/openai_compatible/mod.rs @@ -0,0 +1,327 @@ +//! 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))) + } +} + +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..12365ba7d 100644 --- a/aimux-providers/src/parallel_ai.rs +++ b/aimux-providers/src/parallel_ai.rs @@ -7,82 +7,142 @@ //! 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. 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, +const DEFAULT_BASE_URL: &str = "https://api.parallel.ai"; +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 +197,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,21 +219,10 @@ 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(), diff --git a/aimux-providers/src/preset.rs b/aimux-providers/src/preset.rs new file mode 100644 index 000000000..a188f7f46 --- /dev/null +++ b/aimux-providers/src/preset.rs @@ -0,0 +1,520 @@ +//! 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) resolves no key, injects +//! no placeholder and sends no `Authorization` header, unless the caller +//! gives an explicit 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, Copy)] +pub struct DeriveSpec { + pub from: &'static str, + pub map: &'static [(&'static str, &'static str)], + pub otherwise: &'static str, +} + +/// A parameter of a templated base URL (`{name}` in the template). +#[derive(Debug, Clone, Copy)] +pub struct ParamSpec { + pub name: &'static str, + /// Environment variables tried in order when the parameter is not given. + pub env: &'static [&'static str], + pub default: Option<&'static str>, + /// Set for a derived parameter, which the caller cannot give. + pub derive: Option, +} + +/// One row of the embedded provider registry. +#[derive(Debug, Clone, Copy)] +pub struct PresetDescriptor { + pub name: &'static str, + pub display: &'static str, + /// Default base URL; a template when `params` is not empty. + pub base_url: &'static str, + /// Environment variable of the API key; empty for [`AuthMode::None`]. + pub env_var: &'static str, + pub auth: AuthMode, + /// Environment variable holding the base URL (local servers). + pub base_url_env: Option<&'static str>, + /// The only max-token key the vendor accepts (`"max_tokens"` or + /// `"max_completion_tokens"`); `None` sends `max_tokens`. + pub max_tokens_key: Option<&'static str>, + pub params: &'static [ParamSpec], +} + +/// 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: &'static str, + display: &'static str, + base_url: &'static str, + #[serde(default, deserialize_with = "registry_optional")] + env_var: Option<&'static str>, + #[serde(default, deserialize_with = "registry_optional")] + auth: Option<&'static str>, + #[serde(default, deserialize_with = "registry_optional")] + base_url_env: Option<&'static str>, + #[serde(default)] + params: Vec, + #[serde(default)] + profile: HashMap<&'static str, &'static str>, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RegistryParam { + name: &'static str, + #[serde(default, deserialize_with = "registry_optional")] + env: Option>, + #[serde(default, deserialize_with = "registry_optional")] + default: Option<&'static str>, + #[serde(default, deserialize_with = "registry_optional")] + derive: Option, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RegistryDerive { + from: &'static str, + map: HashMap<&'static str, &'static str>, + otherwise: &'static str, +} + +// 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) +} + +fn load_registry() -> Vec { + // Leak the parsed JSON once so descriptors can borrow even escaped strings. + let json: Vec = serde_json::from_str(include_str!("provider_registry.json")) + .expect("invalid provider_registry.json: expected an array of rows"); + let json = json.leak(); + 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_key(|row| row.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 entries = Vec::with_capacity(rows.len()); + for row in rows { + let check = + |valid: bool, message: &str| assert!(valid, "registry row '{}': {message}", row.name); + check(identifier.is_match(row.name), "invalid name"); + check(seen.insert(row.name), "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.unwrap_or("api_key") { + "api_key" => AuthMode::ApiKey, + "none" => AuthMode::None, + value => panic!("registry row '{}': unknown auth {value:?}", row.name), + }; + check( + match auth { + AuthMode::ApiKey => row.env_var.is_some_and(|var| !var.is_empty()), + AuthMode::None => row.env_var.is_none(), + }, + "env_var must be nonempty for api_key and absent for none", + ); + check( + row.base_url_env.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").copied(); + check( + max_tokens_key.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(param.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::>().leak(), + otherwise: derive.otherwise, + } + }); + params.push(ParamSpec { + name: param.name, + env: param.env.unwrap_or_default().leak(), + default: param.default, + derive, + }); + } + let names: HashSet<_> = params.iter().map(|param| param.name).collect(); + check(names.len() == params.len(), "duplicate param names"); + for param in ¶ms { + if let Some(derive) = param.derive { + check( + params + .iter() + .any(|source| source.name == derive.from && source.derive.is_none()), + "derive.from must name a plain param", + ); + } + } + let placeholders: HashSet<_> = placeholder + .captures_iter(row.base_url) + .map(|c| c.get(1).unwrap().as_str()) + .collect(); + check(placeholders == names, "URL placeholders differ from params"); + entries.push(PresetEntry { + descriptor: Box::leak(Box::new(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: params.leak(), + })), + }); + } + entries +} + +/// 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. +pub fn entries() -> impl Iterator { + static TABLE: OnceLock> = OnceLock::new(); + TABLE.get_or_init(load_registry).iter() +} + +/// The names of every preset, in order. +pub fn names() -> impl Iterator { + entries().map(|entry| entry.descriptor.name) +} + +/// 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.to_string(), + description: descriptor.display.to_string(), + }, + (AuthMode::None, None) => Credential::None, + }; + + let mut dialect = ChatDialect::baseline(); + dialect.supports_top_k = true; + dialect.max_tokens_key = descriptor.max_tokens_key; + 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.to_string(), + 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<&'static str> { + descriptor + .params + .iter() + .filter(|spec| spec.derive.is_none()) + .map(|spec| spec.name) + .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: &[&str]) -> Option { + names.iter().find_map(|name| { + std::env::var(name) + .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 + && 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.map(str::to_string)) + .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).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).to_string()) + .unwrap_or_else(|| { + derive + .otherwise + .replace(&format!("{{{}}}", derive.from), &from) + }); + values.insert(spec.name, derived); + } + + let mut url = descriptor.base_url.to_string(); + 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) +} diff --git a/aimux-providers/src/prodia.rs b/aimux-providers/src/prodia.rs index 5ef3262bb..f5a58f184 100644 --- a/aimux-providers/src/prodia.rs +++ b/aimux-providers/src/prodia.rs @@ -5,8 +5,14 @@ //! //! 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 +22,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 +55,176 @@ fn prodia_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, +pub(crate) mod options; + +const DEFAULT_BASE_URL: &str = "https://prodia.com/api"; +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 +294,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 +306,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 +344,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 +399,7 @@ impl ImageModel for ProdiaImageModel { let mut metadata = HashMap::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 +430,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 +467,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 +516,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 4cb393af4..bcbe94fe0 100644 --- a/aimux-providers/src/provider.rs +++ b/aimux-providers/src/provider.rs @@ -1,23 +1,15 @@ -//! 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 goes away +//! when the bindings mirror the Rust surface. use std::collections::HashMap; -use std::sync::{OnceLock, RwLock}; +use std::sync::{Arc, OnceLock, RwLock}; use serde::Deserialize; use serde_json::Value; @@ -25,27 +17,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 +38,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 +61,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 +110,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 +198,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 +219,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 +229,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 +260,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 +311,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,149 +370,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, - } - } - - /// 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(), - } - } -} +// ── By-name creation for the bindings ──────────────────────────────────────── -/// 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) → 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`] 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 { - // 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) = 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() @@ -447,598 +513,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()) -} - -/// 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_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..8fa5f295e 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", @@ -880,6 +927,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 +942,22 @@ "env_var": "LLAMA_API_KEY", "profile": {} }, + { + "name": "llamacpp", + "display": "llama.cpp", + "base_url": "http://127.0.0.1:8080/v1", + "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 +979,30 @@ "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", + "auth": "none", + "base_url_env": "LOCALAI_BASE_URL", + "profile": {} + }, { "name": "longcat", "display": "LongCat", @@ -1006,6 +1101,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 +1116,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 +1204,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 +1265,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 +1288,14 @@ "env_var": "OHMYGPT_API_KEY", "profile": {} }, + { + "name": "ollama", + "display": "Ollama", + "base_url": "http://127.0.0.1:11434/v1", + "auth": "none", + "base_url_env": "OLLAMA_BASE_URL", + "profile": {} + }, { "name": "ollama_cloud", "display": "Ollama Cloud", @@ -1174,6 +1303,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 +1362,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 +1588,14 @@ "env_var": "SCX_AI_API_KEY", "profile": {} }, + { + "name": "sglang", + "display": "SGLang", + "base_url": "http://127.0.0.1:30000/v1", + "auth": "none", + "base_url_env": "SGLANG_BASE_URL", + "profile": {} + }, { "name": "siliconflow", "display": "SiliconFlow", @@ -1432,15 +1608,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 +1841,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 +2188,14 @@ "env_var": "VIVGRID_API_KEY", "profile": {} }, + { + "name": "vllm", + "display": "vLLM", + "base_url": "http://127.0.0.1:8000/v1", + "auth": "none", + "base_url_env": "VLLM_BASE_URL", + "profile": {} + }, { "name": "volc_engine", "display": "VolcEngine", @@ -1718,6 +2252,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..e0513f8ca 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>, -} +const API_KEY_ENV_VAR: &str = "RECRAFT_API_TOKEN"; +const DEFAULT_NAME: &str = "recraft"; -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, - } - } +/// 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, +} - /// 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 +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() } +} - /// 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 - } +/// 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, + }) +} - /// 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)) - } +/// 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") + }) } -/// Recraft provider — creates [`RecraftImageModel`] instances. +/// 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 ce41c6d1f..f98452f21 100644 --- a/aimux-providers/src/replay.rs +++ b/aimux-providers/src/replay.rs @@ -1,485 +1,96 @@ -//! 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 the official TypeSafe 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 != "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> { + if p.provider_id.is_empty() || 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(&p.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` 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> { + if p.provider_id.is_empty() { + return Err(AiMuxError::InvalidArgument( + "mock replay: provider record has empty provider_id".into(), + )); } 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}"); } - - #[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}" - ); - } } diff --git a/aimux-providers/src/replicate.rs b/aimux-providers/src/replicate.rs index b0dca270d..1d4cc4ee6 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 - } +const DEFAULT_BASE_URL: &str = "https://api.replicate.com/v1"; +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..5fff1f651 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; + +const DEFAULT_BASE_URL: &str = "https://api.rev.ai"; +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)] @@ -153,54 +229,19 @@ fn audio_input_to_bytes(audio: &AudioInput) -> Result, AiMuxError> { 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 +265,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 +300,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/runwayml.rs b/aimux-providers/src/runwayml.rs index be5e8bc5b..9c5a8c19b 100644 --- a/aimux-providers/src/runwayml.rs +++ b/aimux-providers/src/runwayml.rs @@ -8,8 +8,13 @@ //! //! 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 std::time::Duration; use async_trait::async_trait; @@ -17,16 +22,21 @@ 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, VideoPollConfig, 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"; +/// Milliseconds between two status checks of a task. The Core poll loop reads +/// it from [`VideoModel::poll_config`]; a call overrides it with +/// `VideoCallOptions::poll`. +const POLL_INTERVAL_MS: u64 = 2_000; +/// Milliseconds the Core poll loop waits for a task to finish. +const POLL_TIMEOUT_MS: u64 = 300_000; const RUNWAY_VERSION: &str = "2024-11-06"; /// RunwayML returns errors as a flat `{"error": ""}` object. @@ -45,91 +55,116 @@ fn runwayml_failed_response_handler() -> aimux_provider_utils::ResponseHandler>, - /// Interval between task-status polls. Defaults to 2 seconds. - pub poll_interval: Duration, - /// Maximum total time to wait for a task to finish. Defaults to 300 seconds. - pub timeout: Duration, +const DEFAULT_BASE_URL: &str = "https://api.dev.runwayml.com"; +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, - poll_interval: Duration::from_secs(2), - timeout: Duration::from_secs(300), - } - } - - #[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 - } - - /// Override the interval between status polls. - #[must_use] - pub fn with_poll_interval(mut self, interval: Duration) -> Self { - self.poll_interval = interval; - self - } - - /// Override the maximum time to wait for a task to finish. - #[must_use] - pub fn with_timeout(mut self, timeout: Duration) -> Self { - self.timeout = timeout; - 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 ───────────────────────────────────────────────────────────────── @@ -167,40 +202,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 { @@ -213,8 +227,8 @@ impl VideoModel for RunwaymlVideoModel { fn poll_config(&self) -> VideoPollConfig { VideoPollConfig { - interval: self.config.poll_interval, - timeout: self.config.timeout, + interval: Duration::from_millis(POLL_INTERVAL_MS), + timeout: Duration::from_millis(POLL_TIMEOUT_MS), } } @@ -256,19 +270,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(), @@ -304,12 +317,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..41b45a1da 100644 --- a/aimux-providers/src/serper.rs +++ b/aimux-providers/src/serper.rs @@ -3,77 +3,139 @@ //! 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}; 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, +const DEFAULT_BASE_URL: &str = "https://google.serper.dev"; +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 +177,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,18 +200,10 @@ 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(), diff --git a/aimux-providers/src/sglang.rs b/aimux-providers/src/sglang.rs deleted file mode 100644 index 56d95fabe..000000000 --- a/aimux-providers/src/sglang.rs +++ /dev/null @@ -1,79 +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()) - .with_api_key_source(Some("none")), - ) - } - - /// 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 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 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))) - } - - 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..2011ff031 --- /dev/null +++ b/aimux-providers/src/shared/mod.rs @@ -0,0 +1,152 @@ +//! 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'-' +} + +/// 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 }, +} + +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)?)), + } + } +} + +/// 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..b218cc5ee 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; + +const DEFAULT_BASE_URL: &str = "https://api.stability.ai"; +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..5325da514 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..986967e52 100644 --- a/aimux-providers/src/tinyfish.rs +++ b/aimux-providers/src/tinyfish.rs @@ -6,24 +6,25 @@ //! 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. const MODEL_ID: &str = "tinyfish-search"; @@ -33,79 +34,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, -} - -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(), - } - } +const DEFAULT_BASE_URL: &str = "https://api.search.tinyfish.ai"; +const API_KEY_ENV_VAR: &str = "TINYFISH_API_KEY"; +const DEFAULT_NAME: &str = "tinyfish"; - /// 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. +/// +/// # Errors /// -/// TinyFish 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_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 +197,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,27 +218,16 @@ 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(), ) diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index cd6aef2b7..136ba0c0e 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -1,640 +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 { - url: url.clone(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: options.call_id.clone(), - recording_context: options.recording_context.clone(), - ..Default::default() - }, - 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 { - url: url.clone(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: options.call_id.clone(), - recording_context: options.recording_context.clone(), - ..Default::default() - }, - 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..d4f20223d 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"; +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 84c2be4ff..4d868bd89 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, - })), - } + fn supported_urls(&self) -> SupportedUrls { + (self.config.supported_urls)(&self.model_id) } async fn do_generate(&self, options: &CallOptions) -> Result { @@ -228,17 +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 { - url: self.generate_endpoint(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: options.call_id.clone(), - recording_context: options.recording_context.clone(), - ..Default::default() - }, + exchange.request(self.generate_endpoint(exchange.base_url()), options), body.clone(), aimux_provider_utils::create_json_response_handler(), crate::google::google_failed_response_handler(), @@ -369,18 +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 { - url: endpoint.clone(), - headers, - - abort_signal: options.abort_signal.clone(), - call_id: options.call_id.clone(), - recording_context: options.recording_context.clone(), - ..Default::default() - }, + exchange.request(endpoint.clone(), options), body.clone(), aimux_provider_utils::create_event_source_response_handler::(), crate::google::google_failed_response_handler(), @@ -845,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..b63ab1e45 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 ───────────────────────────────────────────────────────── @@ -101,69 +100,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 +131,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 +202,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 38f8d752b..000000000 --- a/aimux-providers/src/vllm.rs +++ /dev/null @@ -1,79 +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()) - .with_api_key_source(Some("none")), - ) - } - - /// 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 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 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))) - } - - 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..0b05fe018 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, +const DEFAULT_BASE_URL: &str = "https://api.voyageai.com/v1"; +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 3434df4bc..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, ResponseFormat, ToolChoice}; -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,123 +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<&ToolChoice>) -> 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 { - ToolChoice::Auto => Some(json!("auto")), - ToolChoice::None => Some(json!("none")), - ToolChoice::Required => Some(json!("required")), - ToolChoice::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 { @@ -204,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. /// @@ -260,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: &std::collections::HashMap, + 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..d3859e8b0 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}; + +use crate::shared::{Credential, EndpointConfig, provider_headers}; const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1"; -const ENV_VAR: &str = "XAI_API_KEY"; -const PROVIDER_NAME: &str = "xai"; +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 0855f9ce2..e47826042 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, ResponseFormat, ToolChoice}; @@ -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..d58b755fc 100644 --- a/aimux-providers/src/you_com.rs +++ b/aimux-providers/src/you_com.rs @@ -7,24 +7,25 @@ //! 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. const MODEL_ID: &str = "youcom-search"; @@ -34,79 +35,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, -} - -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(), - } - } +const DEFAULT_BASE_URL: &str = "https://ydc-index.io"; +const API_KEY_ENV_VAR: &str = "YDC_API_KEY"; +const DEFAULT_NAME: &str = "you_com"; - /// 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. +/// +/// # Errors /// -/// You.com 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_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 +200,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,27 +221,16 @@ 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(), ) 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..c9675d9d5 --- /dev/null +++ b/aimux-providers/tests/anthropic_factory_test.rs @@ -0,0 +1,537 @@ +//! `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, Tool, ToolChoice}; +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(ToolChoice::Required), + Some("none") => Some(ToolChoice::None), + Some("auto") => Some(ToolChoice::Auto), + None => options + .tools + .as_ref() + .filter(|tools| !tools.is_empty()) + .map(|_| ToolChoice::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") + ); +} + +#[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..1b05323ae 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 { @@ -293,18 +298,16 @@ 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(); - - 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(); + // The project has no specification version. 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 76ffbce9d..270d119ca 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 441d8ca0c..0a02b503c 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 8c2f06c6a..3dd5c5459 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 4feffd37c..04fc32981 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 bfa3b39f8..49750418d 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 c3bb88c49..b880d2524 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; @@ -40,12 +41,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, }; // ═══════════════════════════════════════════════════════════════════════════ @@ -298,7 +303,12 @@ fn sources(content: &[GenerateContent]) -> Vec> { // ═══════════════════════════════════════════════════════════════════════════ 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() } // ═══════════════════════════════════════════════════════════════════════════ @@ -321,7 +331,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"); @@ -362,7 +372,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"); @@ -400,7 +410,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"); @@ -430,7 +440,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", @@ -500,7 +510,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", @@ -545,11 +555,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] @@ -956,11 +969,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"); @@ -1013,11 +1029,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"); @@ -1046,8 +1065,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] @@ -1150,15 +1174,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 @@ -1200,9 +1223,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. @@ -1294,7 +1317,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] @@ -1347,9 +1370,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" ); } @@ -1407,7 +1430,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] @@ -1448,5 +1471,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..995e628b3 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(); + 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,1932 @@ 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") +} + +#[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..7e3130d4c --- /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.total.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 4cffe5f12..e2301cbbd 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::{MessageContent, ModelMessage, Role}; 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 @@ -403,9 +422,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 @@ -452,8 +470,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 @@ -473,9 +490,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 @@ -540,8 +555,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( @@ -661,8 +676,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) @@ -733,8 +748,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( @@ -801,8 +816,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( @@ -858,7 +873,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, @@ -897,8 +912,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 @@ -936,8 +956,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 @@ -978,8 +1003,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 @@ -1020,8 +1050,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 @@ -1064,8 +1099,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 @@ -1114,8 +1154,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 690d4404b..43ae8cbef 100644 --- a/aimux-providers/tests/elevenlabs_speech_test.rs +++ b/aimux-providers/tests/elevenlabs_speech_test.rs @@ -17,7 +17,26 @@ 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 鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€鈹€ @@ -53,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!"); @@ -82,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!"); @@ -110,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!"); @@ -139,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!"); @@ -162,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!"); @@ -193,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!"); @@ -241,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..6d1c373c1 100644 --- a/aimux-providers/tests/fal_video_test.rs +++ b/aimux-providers/tests/fal_video_test.rs @@ -8,7 +8,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; 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 { @@ -42,8 +43,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 +63,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 +88,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 +117,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..b2fff1e5d 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":3,"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":3,"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..092850b56 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. + +Stack integration migrates the recording envelope to schema 3: provider identity only, with decision capabilities under input. Recorded requests, responses and normalized outcomes are unchanged. 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..1d4210d88 --- /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, Tool, ToolChoice}; +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(ToolChoice::Required), + Some("none") => Some(ToolChoice::None), + Some("auto") => Some(ToolChoice::Auto), + None => options + .tools + .as_ref() + .filter(|tools| !tools.is_empty()) + .map(|_| ToolChoice::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..d20253b45 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 { @@ -85,11 +80,15 @@ async fn mount_success_mocks(server: &MockServer) { // -- constructor tests ------------------------------------------------------- #[tokio::test] -async fn should_expose_correct_provider_and_specification_version() { - let provider = GoogleProvider::new(GoogleConfig::new("test-api-key")); - let files = provider.files(); +async fn should_expose_correct_provider() { + let provider = create_google(GoogleProviderSettings { + api_key: Some("test-api-key".to_string()), + ..Default::default() + }) + .unwrap(); + let files = Provider::files(&provider).unwrap(); + // The project has no specification version. assert_eq!(files.provider(), "google.generative-ai"); - assert_eq!(files.specification_version(), "v4"); } // -- upload initiation tests ------------------------------------------------- @@ -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..63236c7dc 100644 --- a/aimux-providers/tests/google_video_test.rs +++ b/aimux-providers/tests/google_video_test.rs @@ -1,12 +1,9 @@ //! 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 aimux_core::AiMuxError; +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}; @@ -24,13 +21,6 @@ fn options(p: &str) -> VideoCallOptions { 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()) -} - async fn mock_predict_and_poll(server: &MockServer, result: &Value) { // predictLongRunning Mock::given(method("POST")) @@ -51,247 +41,52 @@ 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(); + 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("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 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(); - 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 - ); + match &r.videos[0] { + aimux_core::video_model::VideoData::Url { url, media_type } => { + assert_eq!( + url, + &format!("{}/files/video-123.mp4?key=test-api-key", server.uri()) + ); + assert_eq!(media_type, "video/mp4"); + } + _ => panic!("expected URL 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 +95,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 +120,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 a206caf91..ededb90a4 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, ResponseFormat, Tool, ToolChoice}; 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,34 @@ 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".to_string(), + name: "unsupported_tool".to_string(), + args: serde_json::Map::new(), + }); + 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 +857,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 +923,113 @@ 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".to_string(), + name: "browser_search".to_string(), + args: serde_json::Map::new(), + }); + 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".to_string(), + name: "browser_search".to_string(), + args: serde_json::Map::new(), + }); + 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".to_string(), + name: "browser_search".to_string(), + args: serde_json::Map::new(), + }); + 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".to_string(), + name: "browser_search".to_string(), + args: serde_json::Map::new(), + }); + 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 +1039,34 @@ mod prepare_tools { mod do_generate { use super::*; + /// 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 +1173,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 +1188,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 +1206,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 +1411,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 +1540,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 +1799,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 +1821,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 +1901,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 +1942,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 +2117,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 81fda04de..15ff20404 100644 --- a/aimux-providers/tests/huggingface_responses_test.rs +++ b/aimux-providers/tests/huggingface_responses_test.rs @@ -16,7 +16,9 @@ //! 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 std::collections::HashMap; use futures::StreamExt; use serde_json::{Value, json}; @@ -27,16 +29,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, ResponseFormat, Tool, ToolChoice}; +use aimux_core::result::{GenerateContent, ReasoningOutput, Source}; use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; +use aimux_core::tool::FunctionTool; use aimux_core::types::{FinishReasonUnified, ResponseMetadata}; use aimux_providers::huggingface::responses::convert_to_huggingface_responses_messages; -use aimux_providers::{HuggingFaceConfig, HuggingFaceProvider}; +use aimux_providers::{HuggingFaceProvider, HuggingFaceProviderSettings, create_huggingface}; // ───────────────────────────────────────────────────────────────────────────── // Shared helpers @@ -54,8 +58,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 +166,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 +189,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 +246,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 +292,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 +309,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 +359,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 +438,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 +526,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 +546,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 +622,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 +703,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 +743,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 +787,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 +815,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 +877,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 +930,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 +955,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 +1139,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 +1158,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 +1197,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 +1240,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 +1276,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 = HashMap::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(ToolChoice::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(ToolChoice::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(ToolChoice::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 +1900,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 cebae4ffd..e74848fbf 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}; @@ -152,9 +151,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() @@ -388,7 +385,11 @@ async fn official_provider_failures_use_standard_retry_rules_and_preserve_codes( #[test] fn non_decision_provider_returns_unsupported() { - let provider = aimux_providers::OpenAIProvider::new(aimux_providers::OpenAIConfig::new("test")); + let provider = aimux_providers::create_openai(aimux_providers::OpenAIProviderSettings { + api_key: Some("test".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..e1cef7309 100644 --- a/aimux-providers/tests/klingai_video_test.rs +++ b/aimux-providers/tests/klingai_video_test.rs @@ -8,7 +8,8 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; 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 { @@ -67,8 +68,12 @@ 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")) @@ -90,8 +95,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 +121,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 +150,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 +178,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 +197,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 456481a49..d6a2c9d5b 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 8cf842c39..6ffb9a3f4 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,9 +18,10 @@ 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, ResponseFormat, Tool}; +use aimux_core::options::{CallOptions, ResponseFormat, Tool, ToolChoice}; use aimux_core::result::{GenerateContent, ReasoningOutput}; use aimux_core::shared::{FileBytes, FileData}; use aimux_core::stream_part::StreamPart; @@ -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 { @@ -522,8 +607,320 @@ mod convert_tests { // -- 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] @@ -550,6 +947,92 @@ mod convert_tests { ); } + #[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\"}"} + ]) + ); + } + + #[test] + fn convert_tool_roundtrip_with_followup_assistant() { + 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: 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: "It is 25 C and cloudy in Tokyo.".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\"}"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "It is 25 C and cloudy in Tokyo."}]} + ]) + ); + } + // -- Provider reference (warning, not error in Rust) -- #[test] @@ -939,10 +1422,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(ToolChoice::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(ToolChoice::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(ToolChoice::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(ToolChoice::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(ToolChoice::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(ToolChoice::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(ToolChoice::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] @@ -994,6 +1736,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] @@ -1001,7 +1812,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 { @@ -1011,8 +1822,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 @@ -1027,7 +1838,7 @@ provider_options: None assert_eq!( body, json!({ - "model": "gpt-4.1-nano", + "model": "test-pdf-model", "input": [{ "type": "message", "role": "user", @@ -1045,7 +1856,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 { @@ -1055,8 +1866,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 @@ -1081,7 +1892,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 { @@ -1091,8 +1902,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 @@ -1306,7 +2117,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 { @@ -1365,7 +2176,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 { @@ -1375,8 +2186,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 @@ -1437,89 +2248,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..8a712b81b --- /dev/null +++ b/aimux-providers/tests/openai_compatible_factory_test.rs @@ -0,0 +1,722 @@ +//! `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::message::Role; +use aimux_core::options::CallOptions; +use aimux_core::result::{GenerateContent, GenerateResult}; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{FunctionTool, Tool, ToolChoice}; +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: Role, text: &str) -> LanguageModelMessage { + match role { + Role::System => LanguageModelMessage::System { + content: text.to_string(), + provider_options: None, + }, + Role::User => LanguageModelMessage::user_text(text), + Role::Assistant => LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: text.to_string(), + provider_options: None, + })], + provider_options: None, + }, + Role::Tool => unreachable!("fixture has no tool 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(Role::System, system)); + } + if let Some(text) = input.get("prompt").and_then(Value::as_str) { + prompt.push(message(Role::User, text)); + } + for entry in input + .get("messages") + .and_then(Value::as_array) + .into_iter() + .flatten() + { + let role = match entry["role"].as_str().unwrap() { + "user" => Role::User, + "assistant" => Role::Assistant, + other => panic!("unmapped role {other}"), + }; + prompt.push(message(role, 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(ToolChoice::Required), + Some("none") => Some(ToolChoice::None), + Some("auto") => Some(ToolChoice::Auto), + None => options + .tools + .as_ref() + .filter(|tools| !tools.is_empty()) + .map(|_| ToolChoice::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"); +} + +#[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::(); +} 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_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..f2f437578 --- /dev/null +++ b/aimux-providers/tests/openai_factory_test.rs @@ -0,0 +1,779 @@ +//! `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::message::Role; +use aimux_core::options::CallOptions; +use aimux_core::result::{GenerateContent, GenerateResult}; +use aimux_core::stream_part::StreamPart; +use aimux_core::tool::{FunctionTool, Tool, ToolChoice}; +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: Role, text: &str) -> LanguageModelMessage { + match role { + Role::System => LanguageModelMessage::System { + content: text.into(), + provider_options: None, + }, + Role::User => LanguageModelMessage::user_text(text), + Role::Assistant => LanguageModelMessage::Assistant { + content: vec![AssistantPart::Text(TextPart { + text: text.into(), + provider_options: None, + })], + provider_options: None, + }, + Role::Tool => panic!("unmapped tool role"), + } +} + +/// 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(Role::System, system)); + } + if let Some(text) = input.get("prompt").and_then(Value::as_str) { + prompt.push(message(Role::User, text)); + } + for entry in input + .get("messages") + .and_then(Value::as_array) + .into_iter() + .flatten() + { + let role = match entry["role"].as_str().unwrap() { + "user" => Role::User, + "assistant" => Role::Assistant, + other => panic!("unmapped role {other}"), + }; + prompt.push(message(role, 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(ToolChoice::Required), + Some("none") => Some(ToolChoice::None), + Some("auto") => Some(ToolChoice::Auto), + None => options + .tools + .as_ref() + .filter(|tools| !tools.is_empty()) + .map(|_| ToolChoice::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)); + + 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, + .. + } => Some((finish_reason, usage)), + _ => None, + }); + let (reason, usage) = finish.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)); +} + +#[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_speech_test.rs b/aimux-providers/tests/openai_speech_test.rs index 6905d105c..fe0b93768 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(HashMap::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 deleted file mode 100644 index 7ed19a2f4..000000000 --- a/aimux-providers/tests/openai_transcription_stream_test.rs +++ /dev/null @@ -1,771 +0,0 @@ -//! Realtime transcription streaming tests (RFC-0028 Phase 1). -//! -//! Each test spins up a local WebSocket server (tokio-tungstenite `accept_async`), -//! plays the OpenAI realtime endpoint's role (consume `session.update` / -//! `input_audio_buffer.append` / `commit`, emit `transcription.delta` / -//! `completed` events), and asserts on the `TranscriptionStreamPart` sequence -//! produced by `do_stream`. -//! -//! 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). - -use std::time::Duration; - -use futures::{SinkExt, StreamExt}; -use tokio::net::TcpListener; -use tokio_tungstenite::tungstenite::Message; - -use aimux_core::AbortSignal; -use aimux_core::error::AiMuxError; -use aimux_core::transcription_model::{ - AudioChunk, InputAudioFormat, TranscriptionModel, TranscriptionStreamOptions, - TranscriptionStreamPart, -}; -use aimux_providers::{OpenAIConfig, OpenAIProvider}; - -// ── helpers ───────────────────────────────────────────────────────────────── - -/// A local WebSocket "OpenAI realtime" server. Handshake + `session.update` -/// ack, then run the scripted behavior. -struct RealtimeServer { - port: u16, -} - -impl RealtimeServer { - /// Start a listener; the accept loop runs on a spawned task. - async fn start() -> Self { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - tokio::spawn(async move { - // Accept a single connection (tests make one). - let (stream, _) = listener.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); - serve(&mut ws).await; - }); - Self { port } - } - - fn base_url(&self) -> String { - format!("http://127.0.0.1:{}", self.port) - } -} - -/// Server behavior shared by the happy-path tests: read session.update, -/// ack it, collect audio appends until commit, then emit delta + completed, -/// and finally assert the client closes with code 1000. -async fn serve(ws: &mut tokio_tungstenite::WebSocketStream) { - // 1. Read session.update, capture the model + format from it. - let mut session_model = String::new(); - let mut format_type = String::new(); - if let Some(Ok(Message::Text(text))) = ws.next().await { - let v: serde_json::Value = serde_json::from_str(&text).unwrap(); - assert_eq!( - v["type"], "session.update", - "first message must be session.update" - ); - assert_eq!(v["session"]["type"], "transcription"); - session_model = v["session"]["audio"]["input"]["transcription"]["model"] - .as_str() - .unwrap() - .to_string(); - format_type = v["session"]["audio"]["input"]["format"]["type"] - .as_str() - .unwrap() - .to_string(); - assert!( - v["session"]["audio"]["input"]["turn_detection"].is_null(), - "turn_detection must be null (nested under audio.input)" - ); - ws.send(Message::Text(r#"{"type":"session.created"}"#.to_string())) - .await - .unwrap(); - ws.send(Message::Text(r#"{"type":"session.updated"}"#.to_string())) - .await - .unwrap(); - } - let _ = (&session_model, &format_type); - - // 2. Collect audio appends until commit. - let mut appended: Vec = Vec::new(); - loop { - match ws.next().await { - Some(Ok(Message::Text(text))) => { - let v: serde_json::Value = serde_json::from_str(&text).unwrap(); - match v["type"].as_str().unwrap_or("") { - "input_audio_buffer.append" => { - appended.push(v["audio"].as_str().unwrap().to_string()); - } - "input_audio_buffer.commit" => break, - other => panic!("unexpected client message: {other}"), - } - } - other => panic!("expected client text message, got {other:?}"), - } - } - assert!( - !appended.is_empty(), - "client must append at least one audio chunk" - ); - - // 3. Emit transcript deltas + completed (the shape OpenAI sends). - ws.send(Message::Text( - r#"{"type":"conversation.item.input_audio_transcription.delta","item_id":"item-1","delta":"Hello"}"#.to_string(), - )) - .await - .unwrap(); - ws.send(Message::Text( - r#"{"type":"conversation.item.input_audio_transcription.delta","item_id":"item-1","delta":" world"}"#.to_string(), - )) - .await - .unwrap(); - ws.send(Message::Text( - r#"{"type":"conversation.item.input_audio_transcription.completed","item_id":"item-1","transcript":"Hello world"}"#.to_string(), - )) - .await - .unwrap(); - - // 4. Expect the client to close with code 1000 (normal). - match ws.next().await { - Some(Ok(Message::Close(frame))) => { - let code = frame.map(|f| u16::from(f.code)); - assert_eq!(code, Some(1000), "client must close(1000) after completion"); - } - other => panic!("expected close frame, got {other:?}"), - } -} - -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") -} - -fn stream_options( - chunks: Vec, - abort: Option, -) -> TranscriptionStreamOptions { - TranscriptionStreamOptions { - audio: Box::pin(futures::stream::iter(chunks)), - input_audio_format: InputAudioFormat { - format_type: "audio/pcm".to_string(), - rate: Some(24000), - }, - provider_options: None, - abort_signal: abort, - headers: None, - include_raw_chunks: false, - timeout: None, - } -} - -/// Collect all parts from a do_stream result into a Vec (Err passthrough). -async fn collect( - result: aimux_core::transcription_model::TranscriptionStreamResult, -) -> Vec> { - let mut out = Vec::new(); - let mut stream = result.stream; - while let Some(part) = stream.next().await { - out.push(part); - } - out -} - -// ── tests ─────────────────────────────────────────────────────────────────── - -/// Happy path: connect → session.update → appends → commit → deltas → -/// completed → final + finish → client close(1000). -#[tokio::test] -async fn stream_realtime_happy_path() { - let server = RealtimeServer::start().await; - let model = realtime_model(&server.base_url()); - - let options = stream_options( - vec![ - AudioChunk::Binary(vec![1, 2, 3]), - AudioChunk::Binary(vec![4, 5, 6]), - ], - None, - ); - let result = model - .do_stream(options) - .await - .expect("do_stream should connect"); - let parts = collect(result).await; - - // StreamStart + 2 deltas + final + finish. - assert_eq!(parts.len(), 5, "parts: {parts:?}"); - - assert!(matches!( - &parts[0], - Ok(TranscriptionStreamPart::StreamStart { .. }) - )); - - match &parts[1] { - Ok(TranscriptionStreamPart::TranscriptDelta { id, delta, .. }) => { - assert_eq!(id.as_deref(), Some("item-1")); - assert_eq!(delta, "Hello"); - } - other => panic!("expected delta 'Hello', got {other:?}"), - } - match &parts[2] { - Ok(TranscriptionStreamPart::TranscriptDelta { delta, .. }) => { - assert_eq!(delta, " world"); - } - other => panic!("expected delta ' world', got {other:?}"), - } - match &parts[3] { - Ok(TranscriptionStreamPart::TranscriptFinal { id, text, .. }) => { - assert_eq!(id.as_deref(), Some("item-1")); - assert_eq!(text, "Hello world"); - } - other => panic!("expected final, got {other:?}"), - } - match &parts[4] { - Ok(TranscriptionStreamPart::Finish { text, segments, .. }) => { - assert_eq!(text, "Hello world"); - assert!(segments.is_empty()); - } - other => panic!("expected finish, got {other:?}"), - } -} - -/// The request snapshot records the exact session.update wire shape. -#[tokio::test] -async fn stream_request_carries_session_update() { - let server = RealtimeServer::start().await; - let model = realtime_model(&server.base_url()); - - let options = stream_options(vec![AudioChunk::Binary(vec![1])], None); - let result = model.do_stream(options).await.unwrap(); - let body = result - .request - .as_ref() - .and_then(|r| r.body.as_ref()) - .expect("request body snapshot"); - let v: serde_json::Value = serde_json::from_str(body).unwrap(); - // B1 regression guard: model in audio.input.transcription.model (not URL), - // turn_detection nested under audio.input, format carries the rate. - assert_eq!(v["type"], "session.update"); - assert_eq!( - v["session"]["audio"]["input"]["transcription"]["model"], - "gpt-realtime-whisper" - ); - assert_eq!( - v["session"]["audio"]["input"]["format"]["type"], - "audio/pcm" - ); - assert_eq!(v["session"]["audio"]["input"]["format"]["rate"], 24000); - assert!(v["session"]["audio"]["input"]["turn_detection"].is_null()); - drop(result); -} - -/// Binary audio chunks are base64-encoded in the append messages. -#[tokio::test] -async fn stream_audio_chunks_are_base64() { - // Server variant that asserts the base64 payload. - struct AssertBase64Server; - let _ = AssertBase64Server; - - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); - // session.update - let _ = ws.next().await; - // first append — verify base64 of [1,2,3] - if let Some(Ok(Message::Text(text))) = ws.next().await { - let v: serde_json::Value = serde_json::from_str(&text).unwrap(); - assert_eq!(v["type"], "input_audio_buffer.append"); - assert_eq!(v["audio"], "AQID", "binary [1,2,3] must arrive as base64"); - } - // second append — Base64 chunk passes through verbatim - if let Some(Ok(Message::Text(text))) = ws.next().await { - let v: serde_json::Value = serde_json::from_str(&text).unwrap(); - assert_eq!( - v["audio"], "cGFzcw==", - "base64 chunk must pass through as-is" - ); - } - // commit - let _ = ws.next().await; - // complete the session so the client closes cleanly - ws.send(Message::Text( - r#"{"type":"conversation.item.input_audio_transcription.completed","item_id":"i","transcript":"x"}"# - .to_string(), - )) - .await - .unwrap(); - let _ = ws.next().await; // close - }); - - let model = realtime_model(&format!("http://127.0.0.1:{port}")); - let options = stream_options( - vec![ - AudioChunk::Binary(vec![1, 2, 3]), - AudioChunk::Base64("cGFzcw==".to_string()), - ], - None, - ); - let result = model.do_stream(options).await.unwrap(); - let parts = collect(result).await; - // StreamStart + final + finish (no deltas in this server script). - assert_eq!(parts.len(), 3, "parts: {parts:?}"); -} - -/// Abort mid-session surfaces `AiMuxError::Aborted`. -#[tokio::test] -async fn stream_abort_mid_session() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let abort = AbortSignal::new(); - let abort_clone = abort.clone(); - - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); - let _ = ws.next().await; // session.update - // Hold the connection open (drain client messages without acting) so - // the abort — not a server-side disconnect — is what ends the session. - while let Some(_msg) = ws.next().await { - tokio::time::sleep(Duration::from_millis(10)).await; - } - }); - - let model = realtime_model(&format!("http://127.0.0.1:{port}")); - let options = stream_options(vec![AudioChunk::Binary(vec![1])], Some(abort_clone)); - - let result = model.do_stream(options).await.unwrap(); - // Abort shortly after the stream starts. - tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(100)).await; - abort.abort(); - }); - let parts = collect(result).await; - assert!( - parts - .iter() - .any(|p| matches!(p, Err(AiMuxError::Aborted(_)))), - "expected Aborted in parts: {parts:?}" - ); -} - -/// 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 model = provider.transcription("whisper-1"); - let err = model - .do_stream(stream_options(vec![], None)) - .await - .unwrap_err(); - match err { - AiMuxError::UnsupportedFunctionality(msg) => { - assert!(msg.contains("whisper-1"), "got: {msg}"); - } - other => panic!("expected UnsupportedFunctionality, got {other:?}"), - } -} - -/// A closed port (connection refused) — `do_stream` itself must fail with a -/// clear connect error, not an in-stream error part (#118). -#[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 model = provider.transcription("gpt-realtime-whisper"); - - let err = model - .do_stream(stream_options(vec![AudioChunk::Binary(vec![1])], None)) - .await - .unwrap_err(); - match &err { - AiMuxError::ApiCall(d) => { - assert!( - d.message.contains("websocket connect failed"), - "connect failure must name the phase, got: {}", - d.message - ); - } - other => panic!("expected ApiCall connect error, got {other:?}"), - } -} - -/// Abort while the WS handshake is still in flight — `Aborted`, not a -/// timeout and not an IO error (#118). -#[tokio::test] -async fn stream_abort_during_connect() { - // A black-hole server: accepts TCP but never completes the WS handshake, - // so `connect_async` pends and only the abort can end the connect phase. - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - // Hold the socket open without writing the 101 response. - let _hold = stream; - std::future::pending::<()>().await; - }); - - let abort = AbortSignal::new(); - let abort_clone = abort.clone(); - tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(150)).await; - abort.abort(); - }); - - let model = realtime_model(&format!("http://127.0.0.1:{port}")); - let err = model - .do_stream(stream_options( - vec![AudioChunk::Binary(vec![1])], - Some(abort_clone), - )) - .await - .unwrap_err(); - // Abort semantics exactly: distinct from timeout and from IO failure. - assert!( - matches!(err, AiMuxError::Aborted(_)), - "connect-phase abort must surface Aborted, got {err:?}" - ); - assert!(!matches!(&err, AiMuxError::Timeout(_))); - assert!(!matches!(&err, AiMuxError::ApiCall(_))); -} - -/// `first_chunk_ms` doubles as the connect budget: against a black-hole -/// server the CONNECT phase itself must time out with its own message — -/// distinguishable from both the first-EVENT timeout and a connect IO -/// failure (#118). -#[tokio::test] -async fn stream_connect_timeout_fires() { - // Accepts TCP but never answers the WS handshake — connect pends. - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let _hold = stream; - 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 model = provider.transcription("gpt-realtime-whisper"); - - let options = TranscriptionStreamOptions { - audio: Box::pin(futures::stream::iter(vec![AudioChunk::Binary(vec![1])])), - input_audio_format: InputAudioFormat { - format_type: "audio/pcm".to_string(), - rate: None, - }, - provider_options: None, - abort_signal: None, - headers: None, - include_raw_chunks: false, - timeout: Some(aimux_core::options::TimeoutConfiguration { - first_chunk_ms: Some(300), - total_ms: None, - step_ms: None, - chunk_ms: None, - }), - }; - let err = model.do_stream(options).await.unwrap_err(); - match &err { - AiMuxError::Timeout(msg) => { - assert_eq!( - msg, "websocket connect timed out", - "connect timeout must carry its own message" - ); - assert!( - !msg.contains("first websocket event"), - "must not be confused with the first-event timeout" - ); - } - other => panic!("expected connect Timeout, got {other:?}"), - } -} - -/// A silent server (accepts + reads session.update, then goes quiet) — the -/// first-chunk timeout must fire (B1/timer coverage). -#[tokio::test] -async fn stream_first_chunk_timeout_fires() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); - // Read session.update, then send NOTHING but KEEP the socket alive - // (drain client appends) well past the client's 500ms first-chunk - // deadline — the timeout must fire before any disconnect could. - let _ = ws.next().await; - 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 model = provider.transcription("gpt-realtime-whisper"); - - let options = TranscriptionStreamOptions { - audio: Box::pin(futures::stream::iter(vec![AudioChunk::Binary(vec![1])])), - input_audio_format: InputAudioFormat { - format_type: "audio/pcm".to_string(), - rate: None, - }, - provider_options: None, - abort_signal: None, - headers: None, - include_raw_chunks: false, - timeout: Some(aimux_core::options::TimeoutConfiguration { - // first_chunk covers connect + session ack; 500ms is generous on - // loopback but far below the server's 2s hold. - first_chunk_ms: Some(500), - total_ms: None, - step_ms: None, - chunk_ms: None, - }), - }; - let result = model.do_stream(options).await.unwrap(); - let parts = collect(result).await; - let timed_out = parts - .iter() - .any(|p| matches!(p, Err(AiMuxError::Timeout(_)))); - assert!( - timed_out, - "expected a Timeout error against a silent server; got: {parts:?}" - ); -} - -/// The server emits a realtime `error` event — the provider message must -/// propagate verbatim and readable (not wrapped in transport noise), and the -/// session must terminate there (#118). -#[tokio::test] -async fn stream_server_error_event_propagates() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); - let _ = ws.next().await; // session.update - ws.send(Message::Text( - r#"{"type":"error","error":{"message":"insufficient quota for realtime transcription","type":"insufficient_quota"}}"# - .to_string(), - )) - .await - .unwrap(); - // The client closes(1000) after surfacing the error. - let _ = ws.next().await; - }); - - let model = realtime_model(&format!("http://127.0.0.1:{port}")); - let result = model - .do_stream(stream_options(vec![AudioChunk::Binary(vec![1])], None)) - .await - .unwrap(); - let parts = collect(result).await; - - // StreamStart then the terminal error — nothing after it. - assert_eq!(parts.len(), 2, "parts: {parts:?}"); - assert!(matches!( - &parts[0], - Ok(TranscriptionStreamPart::StreamStart { .. }) - )); - match &parts[1] { - Err(AiMuxError::ApiCall(d)) => { - assert_eq!( - d.message, "insufficient quota for realtime transcription", - "the provider's own message must propagate verbatim" - ); - assert!( - !d.message.contains("websocket"), - "must not be wrapped in transport noise: {}", - d.message - ); - } - other => panic!("expected ApiCall from server error event, got {other:?}"), - } -} - -/// The server closes the socket (close frame with code + reason) after the -/// handshake — the client must surface peer-closed semantics carrying the -/// code/reason, not a bare EOF or a hang (#118). -#[tokio::test] -async fn stream_peer_close_before_completion() { - use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; - use tokio_tungstenite::tungstenite::protocol::frame::{CloseFrame as WsCloseFrame, Frame}; - - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); - let _ = ws.next().await; // session.update - // Server-initiated abnormal close: code 1011 + a human reason. - let frame = WsCloseFrame { - code: CloseCode::Error, - reason: std::borrow::Cow::Borrowed("server overloaded"), - }; - ws.send(Message::Frame(Frame::close(Some(frame)))) - .await - .unwrap(); - // Drain until the client side tears down. - let _ = ws.next().await; - }); - - let model = realtime_model(&format!("http://127.0.0.1:{port}")); - let result = model - .do_stream(stream_options(vec![AudioChunk::Binary(vec![1])], None)) - .await - .unwrap(); - let parts = collect(result).await; - - assert_eq!(parts.len(), 2, "parts: {parts:?}"); - match &parts[1] { - Err(AiMuxError::ApiCall(d)) => { - assert!( - d.message.contains("websocket closed by peer"), - "peer close must be identifiable, got: {}", - d.message - ); - assert!( - d.message.contains("1011") && d.message.contains("server overloaded"), - "close code and reason must be carried: {}", - d.message - ); - } - other => panic!("expected ApiCall peer-closed error, got {other:?}"), - } -} - -/// chunk-idle timeout: the server delivers one event, then goes silent while -/// holding the socket open — `chunk_ms` must trip with its own message, -/// distinct from the total timeout (#118). -#[tokio::test] -async fn stream_chunk_idle_timeout_fires() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); - let _ = ws.next().await; // session.update - // One event so the stream is established and alive… - ws.send(Message::Text(r#"{"type":"session.created"}"#.to_string())) - .await - .unwrap(); - // …then total silence while KEEPING the socket open (drain client - // messages for 2s, far past the 400ms idle window) — the idle timer, - // not a disconnect, must end the session. - 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 model = provider.transcription("gpt-realtime-whisper"); - - let options = TranscriptionStreamOptions { - audio: Box::pin(futures::stream::iter(vec![AudioChunk::Binary(vec![1])])), - input_audio_format: InputAudioFormat { - format_type: "audio/pcm".to_string(), - rate: None, - }, - provider_options: None, - abort_signal: None, - headers: None, - include_raw_chunks: false, - timeout: Some(aimux_core::options::TimeoutConfiguration { - first_chunk_ms: None, - total_ms: None, - step_ms: None, - chunk_ms: Some(400), - }), - }; - let result = model.do_stream(options).await.unwrap(); - let parts = collect(result).await; - - assert_eq!(parts.len(), 2, "parts: {parts:?}"); - match &parts[1] { - Err(AiMuxError::Timeout(msg)) => { - assert_eq!( - msg, "websocket chunk idle timeout", - "chunk-idle must fire with its own message" - ); - assert!( - !msg.contains("total"), - "must not be confused with the total timeout" - ); - } - other => panic!("expected chunk-idle Timeout, got {other:?}"), - } -} - -/// total timeout: the server keeps events flowing (so chunk-idle never -/// trips) — the hard connection deadline must still fire, with a message -/// distinct from chunk-idle (#118). -#[tokio::test] -async fn stream_total_timeout_fires() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap(); - let _ = ws.next().await; // session.update - // Keep the stream visibly alive: a delta every 100ms for 2s — well - // inside the 300ms chunk window, well past the 700ms total deadline. - for i in 0..20 { - tokio::time::sleep(Duration::from_millis(100)).await; - let event = format!( - r#"{{"type":"conversation.item.input_audio_transcription.delta","item_id":"i","delta":"{i}"}}"# - ); - if ws.send(Message::Text(event)).await.is_err() { - break; // client tore down after the total timeout - } - } - }); - - let config = OpenAIConfig::new("k").with_base_url(format!("http://127.0.0.1:{port}")); - let provider = OpenAIProvider::new(config); - let model = provider.transcription("gpt-realtime-whisper"); - - let options = TranscriptionStreamOptions { - audio: Box::pin(futures::stream::iter(vec![AudioChunk::Binary(vec![1])])), - input_audio_format: InputAudioFormat { - format_type: "audio/pcm".to_string(), - rate: None, - }, - provider_options: None, - abort_signal: None, - headers: None, - include_raw_chunks: false, - timeout: Some(aimux_core::options::TimeoutConfiguration { - first_chunk_ms: None, - total_ms: Some(700), - step_ms: None, - chunk_ms: Some(300), - }), - }; - let result = model.do_stream(options).await.unwrap(); - let parts = collect(result).await; - - // StreamStart + several deltas (alive stream) + the terminal total error. - let deltas = parts - .iter() - .filter(|p| matches!(p, Ok(TranscriptionStreamPart::TranscriptDelta { .. }))) - .count(); - assert!( - deltas >= 2, - "stream must be alive pre-timeout; parts: {parts:?}" - ); - let last = parts.last().expect("must end with an error part"); - match last { - Err(AiMuxError::Timeout(msg)) => { - assert_eq!( - msg, "websocket total timeout", - "total deadline must fire with its own message" - ); - assert!( - !msg.contains("chunk"), - "must not be confused with the chunk-idle timeout" - ); - } - other => panic!("expected total Timeout as the last part, got {other:?}"), - } -} diff --git a/aimux-providers/tests/openai_transcription_test.rs b/aimux-providers/tests/openai_transcription_test.rs index f533f960e..1d44ed1eb 100644 --- a/aimux-providers/tests/openai_transcription_test.rs +++ b/aimux-providers/tests/openai_transcription_test.rs @@ -21,8 +21,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 +142,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 +172,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 +197,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 +229,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 +277,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 +310,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 +337,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 +368,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 +409,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 +448,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..aba8f4eaa --- /dev/null +++ b/aimux-providers/tests/presets_test.rs @@ -0,0 +1,230 @@ +//! 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 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(), 281); + 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)); + variables.extend(d.base_url_env); + variables.extend(d.params.iter().flat_map(|param| param.env.iter().copied())); + } + 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] +async fn a_keyless_row_sends_no_authorization() { + 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() > 281, + "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:?}" + ); +} 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..efbc07091 100644 --- a/aimux-providers/tests/prodia_video_test.rs +++ b/aimux-providers/tests/prodia_video_test.rs @@ -1,8 +1,9 @@ -//! 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::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 wiremock::matchers::{method, path}; @@ -42,8 +43,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 +62,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 +86,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 +113,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 fad95c2d0..0c917e66c 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` @@ -339,13 +337,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: @@ -473,13 +471,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..6b231b68b 100644 --- a/aimux-providers/tests/replicate_video_test.rs +++ b/aimux-providers/tests/replicate_video_test.rs @@ -1,8 +1,9 @@ -//! 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::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 wiremock::matchers::{method, path}; @@ -41,8 +42,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 +57,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 +77,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 +100,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/runwayml_test.rs b/aimux-providers/tests/runwayml_test.rs index 23e2ab715..90b44fd4c 100644 --- a/aimux-providers/tests/runwayml_test.rs +++ b/aimux-providers/tests/runwayml_test.rs @@ -6,14 +6,14 @@ //! async polling loop without a real delay. use std::collections::HashMap; -use std::time::Duration; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; 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"; @@ -33,11 +33,12 @@ fn options(prompt: &str) -> VideoCallOptions { /// A config pointed at the mock server with a short poll interval so tests run /// fast. Uses a dummy, non-secret API key. -fn config(server_uri: String) -> RunwaymlConfig { - RunwaymlConfig::new("test-api-key") - .with_base_url(server_uri) - .with_poll_interval(Duration::from_millis(10)) - .with_timeout(Duration::from_secs(10)) +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: @@ -86,7 +87,7 @@ 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")) @@ -109,7 +110,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")) @@ -130,7 +131,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(); @@ -164,12 +165,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) - .with_poll_interval(Duration::from_millis(10)) - .with_timeout(Duration::from_secs(10)); - 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"); @@ -203,7 +210,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; @@ -224,7 +231,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; @@ -233,9 +240,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..8d52b2ad9 100644 --- a/aimux-providers/tests/vertex_video_test.rs +++ b/aimux-providers/tests/vertex_video_test.rs @@ -1,10 +1,10 @@ -//! 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 aimux_core::Provider; 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}; @@ -22,9 +22,14 @@ fn options(p: &str) -> VideoCallOptions { } 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 +40,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 +54,21 @@ 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 r = generate_video(model.as_ref(), options("A cat")) + .await + .unwrap(); assert_eq!(r.videos.len(), 1); + match &r.videos[0] { + aimux_core::video_model::VideoData::Url { url, media_type } => { + assert_eq!(url, "gs://bucket/video.mp4"); + assert_eq!(media_type, "video/mp4"); + } + _ => panic!("expected URL video"), + } } #[tokio::test] @@ -61,30 +78,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 +105,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 cb7ce7d49..65db224ea 100644 --- a/bindings/flutter/lib/aimux.dart +++ b/bindings/flutter/lib/aimux.dart @@ -498,7 +498,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); @@ -741,20 +743,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, @@ -766,8 +769,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'); } diff --git a/bindings/flutter/lib/errors.dart b/bindings/flutter/lib/errors.dart index f9f6cab81..470ab8b09 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; 18 / 19 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 @@ -45,7 +46,9 @@ abstract final class AimuxErrorCode { static const int noSuchTool = 15; static const int invalidToolInput = 16; static const int toolCallRepair = 17; - static const int noOutputGenerated = 18; + static const int loadApiKey = 18; + static const int loadSetting = 19; + static const int noOutputGenerated = 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 bf1e15391..7293725f7 100644 --- a/bindings/flutter/lib/types.dart +++ b/bindings/flutter/lib/types.dart @@ -428,8 +428,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` @@ -793,7 +800,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()}; } // ───────────────────────────────────────────────────────────────────────────── @@ -1016,8 +1023,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; @@ -1060,7 +1065,6 @@ class GenerateTextOptions { this.providerOptions, this.reasoning, this.instructions, - this.bodyOverrides, this.maxRetries, this.timeout, this.includeRawChunks, @@ -1091,7 +1095,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) @@ -1116,7 +1119,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 5c69ff170..5de8fc8cb 100644 --- a/bindings/flutter/lib/types.g.dart +++ b/bindings/flutter/lib/types.g.dart @@ -303,7 +303,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 @@ -329,7 +328,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 07fe3833f..bd9546cf3 100644 --- a/bindings/flutter/test/errors_test.dart +++ b/bindings/flutter/test/errors_test.dart @@ -75,6 +75,8 @@ void main() { AimuxErrorCode.noSuchTool: NoSuchToolError, AimuxErrorCode.invalidToolInput: InvalidToolInputError, AimuxErrorCode.toolCallRepair: ToolCallRepairError, + AimuxErrorCode.loadApiKey: LoadAPIKeyError, + AimuxErrorCode.loadSetting: LoadSettingError, AimuxErrorCode.other: OtherError, AimuxErrorCode.retry: RetryError, }; @@ -88,10 +90,10 @@ 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 and 18 is the first unassigned value. + // slot, so it resolves; 4 is retired and 20 is the first unassigned value. expect(() => AimuxException.fromCode(999, 'future'), throwsStateError); expect(() => AimuxException.fromCode(4, 'retired'), throwsStateError); - expect(() => AimuxException.fromCode(18, 'unassigned'), throwsStateError); + expect(() => AimuxException.fromCode(20, 'unassigned'), throwsStateError); }); test('bare retry code synthesizes a single-attempt RetryError', () { @@ -181,7 +183,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–19 (4 retired). expect(AimuxErrorCode.other, 1); expect(AimuxErrorCode.aborted, 13); expect(AimuxErrorCode.noSuchTool, 15); @@ -192,6 +194,11 @@ void main() { expect(AimuxErrorCode.toolCallRepair, 17); expect(AimuxErrorCode.name(AimuxErrorCode.toolCallRepair), 'ToolCallRepair'); + // A missing API key / provider setting have their own codes. + expect(AimuxErrorCode.loadApiKey, 18); + expect(AimuxErrorCode.name(AimuxErrorCode.loadApiKey), 'LoadApiKey'); + expect(AimuxErrorCode.loadSetting, 19); + expect(AimuxErrorCode.name(AimuxErrorCode.loadSetting), 'LoadSetting'); }); }); diff --git a/bindings/go/aimux.go b/bindings/go/aimux.go index 39fa00a80..3fb7f88cd 100644 --- a/bindings/go/aimux.go +++ b/bindings/go/aimux.go @@ -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..19), // 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..19 → // *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..804941e37 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,43 @@ 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) + } +} + +// A registry provider with no key and no env var 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) + } + _, err := Provider("groq", "", "llama-3.3-70b") + 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/error.go b/bindings/go/error.go index a92904e03..048e7b441 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 +// 18 / 19. 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. @@ -41,7 +42,13 @@ const ( CodeNoSuchTool Code = 15 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 = 18 + // CodeLoadSetting: a required provider setting was not passed and its + // fallback environment variable is unset (LoadSettingError). + CodeLoadSetting Code = 19 + CodeNoOutputGenerated 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,7 +100,7 @@ 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 { @@ -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 f562cf574..5a9632d75 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"`)}, @@ -472,7 +466,7 @@ func TestEmbedRejectsRawPassThroughOpts(t *testing.T) { func TestCodeFromCRejectsOutOfRange(t *testing.T) { // 4 is retired; 14 is Retry; 15..17 are tool-call errors. - for _, bad := range []int{0, 4, 18, 999} { + for _, bad := range []int{0, 4, 20, 999} { if _, ok := codeFromC(bad); ok { t.Fatalf("%d is not an AiMuxError variant", bad) } @@ -487,6 +481,13 @@ func TestCodeFromCRejectsOutOfRange(t *testing.T) { if c, ok := codeFromC(14); !ok || c != CodeRetry { t.Fatalf("14 → %v, %v; want CodeRetry", c, ok) } + // 18 / 19 are the missing-key / missing-setting codes. + if c, ok := codeFromC(18); !ok || c != CodeLoadAPIKey || c.String() != "LoadApiKey" { + t.Fatalf("18 → %v, %v; want CodeLoadAPIKey", c, ok) + } + if c, ok := codeFromC(19); !ok || c != CodeLoadSetting || c.String() != "LoadSetting" { + t.Fatalf("19 → %v, %v; want CodeLoadSetting", c, ok) + } } // A CodeRetry error keeps the per-attempt history in order; LastError is the diff --git a/bindings/go/types.go b/bindings/go/types.go index 9c70586fd..0b919cf6d 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 { @@ -387,7 +387,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/AimuxException.java b/bindings/java/src/main/java/ai/arcships/aimux/AimuxException.java index c73ca92ff..204f0b799 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; 18 / 19 = 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. @@ -63,7 +64,9 @@ public class AimuxException extends RuntimeException { public static final int AIMUX_E_NO_SUCH_TOOL = 15; 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 = 18; + public static final int AIMUX_E_LOAD_SETTING = 19; + public static final int AIMUX_E_NO_OUTPUT_GENERATED = 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 2540f8483..fe3b5c1c1 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java @@ -222,7 +222,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. */ 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..7c2079c21 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,7 +79,8 @@ 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 18 / 19. if (code < AimuxException.AIMUX_E_OTHER || code > AimuxException.AIMUX_E_NO_OUTPUT_GENERATED || code == 4) { 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 706119913..4755e3a2f 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 1ed7c68a3..ca6eb20e9 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/Types.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/Types.java @@ -1922,7 +1922,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; @@ -1939,7 +1938,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; @@ -1956,7 +1955,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; } @@ -2025,7 +2022,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; @@ -2047,7 +2043,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; } @@ -2058,7 +2053,7 @@ public static class Builder { public GenerateTextOptions build() { GenerateTextOptions options = new GenerateTextOptions(maxOutputTokens, temperature, stopSequences, topP, topK, presencePenalty, frequencyPenalty, responseFormat, 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. @@ -2087,7 +2082,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) @@ -2098,7 +2092,7 @@ public boolean equals(Object o) { public int hashCode() { return Objects.hash(maxOutputTokens, temperature, stopSequences, topP, topK, presencePenalty, frequencyPenalty, responseFormat, seed, tools, toolChoice, headers, providerOptions, reasoning, - instructions, bodyOverrides, maxRetries, includeRawChunks, timeout, sessionId); + instructions, maxRetries, includeRawChunks, timeout, sessionId); } } @@ -3020,34 +3014,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; @@ -3055,10 +3023,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; @@ -3071,6 +3038,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 { @@ -3079,20 +3047,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); } @@ -3100,6 +3073,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()}. @@ -3113,15 +3120,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; @@ -3136,8 +3143,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() { @@ -3164,16 +3171,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 3822b530a..b3d493d49 100644 --- a/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java +++ b/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java @@ -226,16 +226,34 @@ 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(18); + assertThat(AimuxException.AIMUX_E_LOAD_SETTING).isEqualTo(19); + assertThat(AimuxException.codeName(18)).isEqualTo("LoadApiKey"); + assertThat(AimuxException.codeName(19)).isEqualTo("LoadSetting"); + AimuxException key = AimuxException.of(AimuxException.AIMUX_E_LOAD_API_KEY, "no key"); + assertThat(key.getCode()).isEqualTo(18); + assertThat(((AimuxException.LoadAPIKeyError) key).getEnvVar()).isNull(); + AimuxException setting = AimuxException.of(AimuxException.AIMUX_E_LOAD_SETTING, "no region"); + assertThat(setting.getCode()).isEqualTo(19); + assertThat(((AimuxException.LoadSettingError) setting).getEnvVar()).isNull(); } @Test void codesOutsideTheRustEnumAreRejected() { - // Tool errors occupy 15..17; retired code 4 and values beyond 17 are - // header/library mismatches. + // Tool errors occupy 15..17 and 18 / 19 are the missing key / setting; + // retired code 4 and values beyond 19 are header/library mismatches. assertThat(AimuxException.of(15, "")).isInstanceOf(AimuxException.NoSuchToolError.class); assertThatThrownBy(() -> AimuxException.of(4, "")) .isInstanceOf(IllegalStateException.class); - assertThatThrownBy(() -> AimuxException.of(18, "")) + assertThatThrownBy(() -> AimuxException.of(20, "")) .isInstanceOf(IllegalStateException.class); assertThatThrownBy(() -> AimuxException.of(999, "m")) .isInstanceOf(IllegalStateException.class); 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..5de65650c 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`; 18 / 19 = 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]. */ @@ -30,7 +31,9 @@ const val AIMUX_E_ABORTED: Int = 13 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 = 18 +const val AIMUX_E_LOAD_SETTING: Int = 19 +const val AIMUX_E_NO_OUTPUT_GENERATED: 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 2a253fade..8dbb0f439 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`. @@ -289,7 +289,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 { @@ -297,7 +297,8 @@ 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. + // Code 4 is retired; Retry is 14, tool errors are 15..17 and a missing + // API key / setting are 18 / 19. check(code in AIMUX_E_OTHER..AIMUX_E_NO_OUTPUT_GENERATED && code != 4) { "${prefix}aimux ffi: expected AiMuxError code, got $code" } 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..eeaa2fae1 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..19 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 a8ec46957..1d0e0b3ff 100644 --- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt +++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt @@ -146,12 +146,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, @@ -159,6 +154,9 @@ data class GenerateResponse( val body: JsonElement? = null, ) +@Serializable +data class RequestInfo(val body: JsonElement? = null) + /** * A tool call requested by the model. * @@ -635,8 +633,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). */ @@ -1055,8 +1051,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 37a10f71e..0263ae286 100644 --- a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt +++ b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt @@ -215,13 +215,30 @@ class ErrorsTest { assertThat(noKey.retryable).isFalse() } + /** 18 / 19: 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(18) + assertThat(AIMUX_E_LOAD_SETTING).isEqualTo(19) + assertThat(AimuxException.codeName(18)).isEqualTo("LoadApiKey") + assertThat(AimuxException.codeName(19)).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(18) + 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(19) + 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 the first unassigned value and is rejected. - assertThatThrownBy { AimuxException.createByCode(18, "?") } + // 20 is the first unassigned value and is rejected. + assertThatThrownBy { AimuxException.createByCode(20, "?") } .isInstanceOf(IllegalStateException::class.java) assertThatThrownBy { AimuxException.createByCode(AIMUX_OK, "?") } .isInstanceOf(IllegalStateException::class.java) 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__/error.test.ts b/bindings/node/__test__/error.test.ts index 5540c45e5..0bdac9658 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,51 @@ 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 carrying the environment variable', async (t) => { + const saved = process.env.GROQ_API_KEY + delete process.env.GROQ_API_KEY + try { + const err = (await t.throwsAsync(() => provider('groq', null, 'some-model'))) 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__/provider_config.test.ts b/bindings/node/__test__/provider_config.test.ts new file mode 100644 index 000000000..e8c112183 --- /dev/null +++ b/bindings/node/__test__/provider_config.test.ts @@ -0,0 +1,245 @@ +// 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. + const missing = await t.throwsAsync(() => provider('cloudflare_workers_ai', 'k', 'm')) + 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('groq', '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__/trace.test.ts b/bindings/node/__test__/trace.test.ts index 2739063fe..18cb2a5c4 100644 --- a/bindings/node/__test__/trace.test.ts +++ b/bindings/node/__test__/trace.test.ts @@ -69,7 +69,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, 'openai.chat') t.is(stats[0].requests, 2) t.truthy(stats[0].reported_hit_rate) t.truthy(stats[0].client_upper_bound_hit_rate) diff --git a/bindings/node/index.d.ts b/bindings/node/index.d.ts index 324faed9e..2ea3cced5 100644 --- a/bindings/node/index.d.ts +++ b/bindings/node/index.d.ts @@ -491,13 +491,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 229f1bc00..7435a677e 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", @@ -187,6 +189,8 @@ fn aimux_error_class_name(error: &AiMuxError) -> &'static str { AiMuxError::ToolCallRepair { .. } => "ToolCallRepairError", AiMuxError::InvalidArgument(_) => "InvalidArgumentError", AiMuxError::InvalidPrompt(_) => "InvalidPromptError", + AiMuxError::LoadApiKey { .. } => "LoadAPIKeyError", + AiMuxError::LoadSetting { .. } => "LoadSettingError", AiMuxError::TokenExpired(_) => "TokenExpiredError", AiMuxError::UnsupportedFunctionality(_) => "UnsupportedFunctionalityError", AiMuxError::NoSuchModel { .. } => "NoSuchModelError", @@ -293,8 +297,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, @@ -318,7 +341,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 5bd91ef70..df6e25a3e 100644 --- a/bindings/node/src/index.ts +++ b/bindings/node/src/index.ts @@ -88,6 +88,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 47d0ac56d..b259b9d79 100644 --- a/bindings/node/src/lib.rs +++ b/bindings/node/src/lib.rs @@ -34,6 +34,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; @@ -569,8 +570,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 @@ -578,9 +577,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; @@ -631,42 +637,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）。幂等：多次调用无副作用；宿主已自建 @@ -718,7 +763,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}")), ), @@ -947,19 +992,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. @@ -972,24 +1013,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, }) } @@ -1008,39 +1044,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, }) } @@ -1059,10 +1077,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, }) } @@ -1085,26 +1103,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, }) } @@ -1123,26 +1136,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, }) } @@ -1161,26 +1169,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, }) } @@ -1199,26 +1202,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, }) } @@ -1239,27 +1237,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, }) } @@ -1280,27 +1281,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, }) } @@ -1321,28 +1319,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, }) } @@ -1365,34 +1361,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, }) } @@ -1417,10 +1410,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, }) } @@ -1460,6 +1453,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))?; @@ -1481,7 +1480,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, @@ -1512,11 +1511,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 @@ -1543,15 +1540,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); } @@ -1564,12 +1562,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 d93c4e42a..7d2b1ef79 100644 --- a/bindings/node/src/types.ts +++ b/bindings/node/src/types.ts @@ -127,6 +127,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' @@ -164,6 +165,8 @@ export type * from './types/TranscriptionStreamPart' 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/UserPart' diff --git a/bindings/node/src/types/AiMuxError.ts b/bindings/node/src/types/AiMuxError.ts index e41afd40f..0a44a0a3e 100644 --- a/bindings/node/src/types/AiMuxError.ts +++ b/bindings/node/src/types/AiMuxError.ts @@ -14,9 +14,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 } | { "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, } } | { "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/CallOptions.ts b/bindings/node/src/types/CallOptions.ts index db98aa764..80ffacc3c 100644 --- a/bindings/node/src/types/CallOptions.ts +++ b/bindings/node/src/types/CallOptions.ts @@ -76,12 +76,6 @@ provider_options: { [key in string]: { [key in string]: JsonValue } } | null, * Anthropic `thinking` config. */ 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, /** * Per-call retry count override. `None` uses the provider's configured * Core operation retry. `Some(0)` disables retries. diff --git a/bindings/node/src/types/GenerateTextOptions.ts b/bindings/node/src/types/GenerateTextOptions.ts index 68aec90e8..4bbbcfc83 100644 --- a/bindings/node/src/types/GenerateTextOptions.ts +++ b/bindings/node/src/types/GenerateTextOptions.ts @@ -21,10 +21,6 @@ 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, /** * Per-call retry count override. `None` = provider default, `Some(0)` = disable. */ diff --git a/bindings/node/src/types/InputRecord.ts b/bindings/node/src/types/InputRecord.ts index 12ea93e86..61cb96552 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,13 +7,18 @@ import type { JsonValue } from "./serde_json/JsonValue"; /** * ① 输入侧:完整调用参数,足以重建 generate_text 调用。 */ -export type InputRecord = { operation?: RecordingOperation, +export type InputRecord = { +/** + * Decision contract used for validation and deterministic offline replay. + * Connection settings remain outside the recording. + */ +decision_capabilities?: DecisionCapabilities | null, operation?: RecordingOperation, /** * 完整 prompt(消息数组,含 ContentPart::Image 等多模态)。 */ -prompt: Array, +prompt: Array, /** * 序列化的 CallOptions(abort_signal/call_id 已 serde skip); - * headers/provider_options/body_overrides 已递归脱敏。 + * headers/provider_options 已递归脱敏。 */ options: JsonValue, }; 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..a59cfc113 100644 --- a/bindings/node/src/types/ProviderRecord.ts +++ b/bindings/node/src/types/ProviderRecord.ts @@ -1,31 +1,21 @@ // 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 = { /** - * `model.provider()`,如 "openai"。 + * provider 名(registry 名),回放重建的查找键,如 "openai"、"deepseek"。 */ -provider: 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, +provider_id: 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..8ef090200 100644 --- a/bindings/node/src/types/Recording.ts +++ b/bindings/node/src/types/Recording.ts @@ -9,7 +9,8 @@ import type { ProviderRecord } from "./ProviderRecord"; */ export type Recording = { /** - * 格式版本。 + * 格式版本。读取时只接受当前 [`RECORDING_SCHEMA`];旧 schema 的录制带有 + * 配置快照(base_url / profile 等),不再读取。 */ schema: number, /** 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/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/python/python/aimux/__init__.py b/bindings/python/python/aimux/__init__.py index de446b67c..1709a57ad 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 7b54aac03..e2e42d3d9 100644 --- a/bindings/python/python/aimux/wrapper.py +++ b/bindings/python/python/aimux/wrapper.py @@ -1027,7 +1027,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. @@ -1060,8 +1059,6 @@ class GenerateTextOptions(BaseModel): """ - - class RequestInfo(BaseModel): body: Optional[Any] = None diff --git a/bindings/python/src/error.rs b/bindings/python/src/error.rs index 8e317bbc3..f81504fbf 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, @@ -250,6 +262,8 @@ fn variant_instance<'py>(py: Python<'py>, e: &AiMuxError) -> PyResult py.get_type_bound::(), AiMuxError::InvalidArgument(_) => py.get_type_bound::(), AiMuxError::InvalidPrompt(_) => 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::() @@ -328,8 +342,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, @@ -395,6 +428,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 b06f2a3ce..a22c759a9 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; @@ -412,7 +417,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 } @@ -423,23 +428,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, }) } @@ -449,13 +463,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))?; @@ -469,12 +484,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, @@ -487,13 +507,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))?; @@ -508,13 +529,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))?; @@ -529,13 +551,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))?; @@ -550,13 +573,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))?; @@ -577,13 +601,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))?; @@ -604,13 +637,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))?; @@ -630,13 +668,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))?; @@ -659,21 +703,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))?; @@ -686,8 +733,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( @@ -697,15 +745,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, @@ -728,9 +776,14 @@ impl ProviderHandle { /// List models available on this provider (runtime discovery + anya2a spec). /// Returns a JSON array of RuntimeModel. fn list_models(&self) -> 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 = rt - .block_on(async { self.inner.list_models().await }) + .block_on(discovery.list_models()) .map_err(|e| to_py_err(&e))?; serialize_result(&models) } @@ -760,17 +813,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 @@ -813,7 +865,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}"))) } diff --git a/bindings/python/src/multimodal.rs b/bindings/python/src/multimodal.rs index 4212dbcc9..5b89c3038 100644 --- a/bindings/python/src/multimodal.rs +++ b/bindings/python/src/multimodal.rs @@ -401,13 +401,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), }) @@ -421,12 +416,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), @@ -437,12 +427,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), @@ -457,12 +442,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), @@ -473,12 +453,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), @@ -497,13 +472,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), }) @@ -517,13 +493,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), }) @@ -541,13 +518,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), }) @@ -557,12 +535,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), @@ -573,12 +552,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), @@ -589,12 +569,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), @@ -735,13 +716,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..2374ba7e0 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,18 +104,48 @@ 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(): +def test_provider_missing_env_key_raises(monkeypatch): """provider() with api_key=None reads the env var and fails clearly when unset.""" - from aimux import provider + from aimux import LoadAPIKeyError, provider - with pytest.raises(Exception, match="(?i)api key"): + monkeypatch.delenv("ABACUS_API_KEY", raising=False) + with pytest.raises(LoadAPIKeyError, match="(?i)api key"): provider("abacus", None, "m") +def test_provider_params_fill_a_preset_template(monkeypatch): + """config["params"] reaches the preset's template parameters.""" + from aimux import InvalidArgumentError, provider + + monkeypatch.delenv("CLOUDFLARE_ACCOUNT_ID", raising=False) + # Without the parameter the preset cannot build its base URL. + with pytest.raises(InvalidArgumentError, match="account_id"): + provider("cloudflare_workers_ai", "sk-test-fake-key", "m") + 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( + "groq", + "sk-test-fake-key", + "m", + config={"params": {"account_id": "acct123"}}, + ) + + def test_openai_creates_model(): """openai() creates a model instance.""" from aimux import openai diff --git a/bindings/python/tests/test_error.py b/bindings/python/tests/test_error.py index cc6adc58e..e47da0d7f 100644 --- a/bindings/python/tests/test_error.py +++ b/bindings/python/tests/test_error.py @@ -163,3 +163,48 @@ 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, not InvalidArgument.""" + from aimux import AimuxError, InvalidArgumentError, LoadAPIKeyError, provider + + monkeypatch.delenv("GROQ_API_KEY", raising=False) + with pytest.raises(LoadAPIKeyError) as excinfo: + provider("groq", None, "some-model") + 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_trace.py b/bindings/python/tests/test_trace.py index c4f1b3389..36a1dfcae 100644 --- a/bindings/python/tests/test_trace.py +++ b/bindings/python/tests/test_trace.py @@ -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"] == "openai.chat" assert "verdict_counts" in stats[0] chain = json.loads(traced.trace_session_chain("sess-1")) diff --git a/bindings/swift/Sources/Aimux/Aimux.swift b/bindings/swift/Sources/Aimux/Aimux.swift index aa650af1e..f34e67b49 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 { diff --git a/bindings/swift/Sources/Aimux/Types.swift b/bindings/swift/Sources/Aimux/Types.swift index 8b308b546..2b6b8b59b 100644 --- a/bindings/swift/Sources/Aimux/Types.swift +++ b/bindings/swift/Sources/Aimux/Types.swift @@ -1003,12 +1003,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? @@ -1017,13 +1012,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 (the `raw` field of `GenerateTextResult`). public struct GenerateResult: Codable, Equatable { public var content: [GenerateContent] @@ -1031,21 +1030,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 @@ -1295,7 +1293,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? @@ -1348,7 +1345,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" @@ -1362,7 +1358,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, @@ -1373,7 +1369,7 @@ public struct GenerateTextOptions: Codable, Equatable { self.responseFormat = responseFormat; 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 1753dce86..77d22ed00 100644 --- a/bindings/swift/Tests/AimuxTests/AimuxTests.swift +++ b/bindings/swift/Tests/AimuxTests/AimuxTests.swift @@ -186,6 +186,22 @@ final class AimuxTests: XCTestCase { Int32(AIMUX_E_TIMEOUT.rawValue)) } + /// 18 / 19: a missing API key / provider setting have their own codes and + /// carry the consulted environment variable. + func testLoadErrorCodes() { + XCTAssertEqual(AIMUX_E_LOAD_API_KEY.rawValue, 18) + XCTAssertEqual(AIMUX_E_LOAD_SETTING.rawValue, 19) + 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() { diff --git a/docs/API.md b/docs/API.md index eda312d76..41643758a 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 281 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 @@ -122,6 +131,8 @@ for its own declaration): | `usage` | Token usage | | `warnings` | Warnings | | `raw` | Raw provider result (includes full content) | +| `request` | Request information, including the sent body | +| `response` | Response id, timestamp, model id, headers and body; missing identity metadata is filled by the call layer | > **Note**: `text` and `tool_calls` are convenience fields extracted from `raw.content`. > Extraction parses and validates each tool call per the AI SDK @@ -160,7 +171,7 @@ Returns generated content as a stream, output chunk by chunk. Examples: [Node.js](api/node.md#streaming-generation) · [Python](api/python.md#streaming-generation) · [Rust](api/rust.md#streaming-generation) · [Go](api/go.md#streaming-generation) · [Swift](api/swift.md#streaming-generation) · [Kotlin](api/kotlin.md#streaming-generation) · [Flutter](api/flutter.md#streaming-generation) · [C ABI](api/c.md#language-model) -#### StreamPart Types +#### TextStreamPart Types | Variant | Description | |------|------| @@ -170,12 +181,15 @@ Examples: [Node.js](api/node.md#streaming-generation) · [Python](api/python.md# | `ToolCall` | Complete tool call (same shape as a `tool_calls` entry, incl. `invalid?` / `error?`) | | `ToolResult` | Tool result executed by the provider | | `ReasoningStart` / `ReasoningDelta` / `ReasoningEnd` | Reasoning segment lifecycle | -| `ResponseMetadata` | Response metadata (id, timestamp, model_id) | | `Source` | Reference/source | | `Finish` | Stream end (carries usage + finish_reason) | | `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 @@ -243,7 +257,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. diff --git a/docs/PROJECT-OVERVIEW.md b/docs/PROJECT-OVERVIEW.md index 27dd5a4e6..d3be1ec4e 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 | 326 accepted names (281 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 streaming 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 281 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] 326 provider names accepted by `create_provider` (281 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 3220e1547..2a9d6b1cf 100644 --- a/docs/ai-sdk-request-pipeline.md +++ b/docs/ai-sdk-request-pipeline.md @@ -86,7 +86,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` | +| 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） | @@ -294,7 +294,7 @@ redaction helper。不得依赖每个 Provider 自己记得脱敏。 它们遵循 AI SDK 的单次 fetch attempt 契约： -- 不接受 `RetryConfig`； +- 不接受 retry 配置； - 不接受 `TimeoutConfiguration`； - 不执行 backoff； - 不调用自身或另一个 API helper 进行 retry； @@ -355,17 +355,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**： @@ -393,7 +394,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 类型。 @@ -802,10 +803,10 @@ Composite 外层不重放；重试放在已有语义边界内： `Files::upload_file` 按 AI SDK Provider SPI 本身就是公开 operation，没有对应的 `do_upload_file`；本轮不为它虚构第二层。**修订**（RFC-0031 review）：`upload_file` -自己用 Provider 已配置的 `RetryConfig` 重试（与 `execute_list_models` 同一 -primitive）。这是“创建资源类 POST 不自动重试”一般规则的显式例外，前提是 -upload 不计费、失败的 create-file 请求不会返回可复用的 file id。多阶段的 -Google Files 按 exchange 分别重试，而不是包住整个 `upload_file`。 +自己重试（RFC-0031 review 的修订）。**RFC-0036 撤销该修订**：provider 没有 retry +设置，`upload_file`（OpenAI、Anthropic、Google Files 的多阶段上传同样）与 +`list_models` 一样是单次 exchange，不重试——“创建资源类 POST 不自动重试”的一般规则 +不再有例外；需要重试的调用方自己包一层。 --- @@ -831,9 +832,9 @@ search 所有 binding 调用这些 user operations，不得直接调用 provider-facing `do_*`。`do_*` 继续作为 Provider SPI。 -为保留现有 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 @@ -894,9 +895,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 0d0231ac6..5a4eed5fc 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 裁定缺失等评审问题） +> **Date**: 2026-09-29 · 2026-10-04 修订（撤销切换后的 ABI 共存承诺、xAI/HF chat 扩展登记后撤回、preset 改为运行时 descriptor 表、重定向规则收归 helper、修复附录 §6.4-14 裁定缺失等评审问题） > -> **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 工具 > @@ -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 单一生成链驱动 @@ -110,7 +110,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"(与 #207 §12.2 的 C2 / C4 行一致;切换门由 RFC-0039 定义) | | 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);`ToolInput` 的 Raw / Parsed 区分由第二部分的四层消息协议承担(V4 tool-call `input` 为原文 String,core 层 `parse_tool_call` 产出解析值) | @@ -153,7 +153,7 @@ 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 模拟所有厂商 | +| D12 | preset 按名字从运行时 descriptor 表创建，复用 compat 内部模型实现 | registry JSON 是唯一数据来源，首次使用时解析并校验一次；不为每一行生成 `create_` 源码（同一份数据不存两遍）。preset 的 key 懒加载是明确的产品扩展；不以原生 OpenAI model 加 profile 模拟所有厂商 | | D13 | V4 Provider 与厂商扩展分开 | `chat`、`responses`、`getAvailableModels`、tools 等通过注册条目的扩展能力访问；registry、wrap 不丢失这些能力 | | D14 | 删除所有 model/provider 默认 retry 配置 | retry 属于操作；router/MoA 为每个 child 配置预算，最终组合失败不得触发整组重跑 | | D15 | 删除通用 `body_overrides` | compat 保留 SDK 的 providerOptions 透传与 `transformRequestBody`；原生包不增加统一请求体改写契约 | @@ -161,7 +161,7 @@ FFI / Node / Python / 5 个 C-ABI 绑定 / CLI / Web 由 descriptor + manifest | 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 扩展 | @@ -399,7 +399,6 @@ pub struct FetchRequest { pub url: Url, pub headers: Headers, pub body: Bytes, - pub redirect: RedirectPolicy, pub signal: Option, } @@ -424,6 +423,7 @@ pub type FetchFunction = Arc; 关键规则： - multipart 在 helper 内编码成字节，签名处理最终发送字节。 +- 重定向由 helper 逐跳处理，transport（含注入的 `Fetch`）不跟随重定向、原样返回 `3xx`。普通 API 调用只跟随同源重定向，跨 origin 的 `3xx` 按非 2xx 返回，请求头与 fetch 装饰器注入的凭证都不会到达另一个 origin；每一跳重新经过请求的 `Fetch`，签名装饰器对实际发送的 URL 重新签名。校验下载（D26）沿用逐跳校验、pin 与离开 credentialed origin 后剥离调用方头的规则。 - 默认 HTTP client 按 Tokio runtime 分片，保留现有连接池生命周期约束。 - `settings.fetch == None` 时，每次请求读取当前默认 leaf；不在工厂创建时冻结全局 fetch。 - WebSocket 使用独立 connector，和默认 HTTP leaf 共用显式 `TransportSettings`。 @@ -431,7 +431,7 @@ pub type FetchFunction = Arc; - 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 { @@ -492,7 +492,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。 @@ -768,7 +768,7 @@ Unknown 不冒充 OpenAI，不输出需要厂商规则才能成立的结论。 | 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 注入 | @@ -782,7 +782,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 所有历史行为。 @@ -1118,7 +1118,7 @@ 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` 擦除结果类型，保留所有权 | @@ -1128,7 +1128,7 @@ RFC-0014 的 `tracing` span 保留为统一事件的内建日志适配器。HTTP | 产品选择 | 观测范围 | 为 object 和非 LM 模态增加 operation/attempt 事件 | | 产品选择 | telemetry 关闭 | 关闭观测，保留功能性 scope 与 session affinity | | 产品选择 | 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 引用 | @@ -1141,6 +1141,13 @@ RFC-0014 的 `tracing` span 保留为统一事件的内建日志适配器。HTTP | 不采用 | 历史版本和 deprecated 兼容 | 不实现 V2/V3 适配、旧 namespace 双写、旧方法转发、旧录制读取 | | 本期范围外 | Evaluation/Realtime/SpeechTranslation/Batch 等完整链路 | 不注册虚假能力；后续需完整接通操作与观测 | | ABI 范围外 | C 宿主异步 fetch/model/headers/telemetry | 由独立异步 ABI 设计处理；本期提供 Rust 内建句柄 | +| 语言适配 | `specification_version` / `Provider.name()` | `Provider` 与各模型 trait 不带 `specification_version`，`Provider` 不带 `name()`（D-c）：Rust trait 本身就是版本边界，注册名归持有 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 没有 poll 设置；`VideoModel::poll_config()` 保留在 core，作为包内常量的来源，调用级 `VideoCallOptions.poll` 逐字段覆盖 | +| 产品选择 | 录制重建原生协议 provider | `rebuild_provider` 只走 registry 与 overlay；原生协议包返回 `NoSuchProvider`，回放时改用 `replay_with_model` 传入模型 | 本地没有安装源码的厂商包，不能仅凭命名推测宣称“已对齐 AI SDK”。descriptor 记录 `verified_sdk_source` 或 `aimux_extension`；发布前按实际支持承诺完成协议 fixture 验证。 @@ -1320,7 +1327,6 @@ aimux_error_t* aimux_stream_text( ```text 规范类型/Wire DTO + 包 descriptor + preset 输入 - → aimux-codegen 生成 preset Rust 源码 → aimux-manifest 汇总 → manifest.json → 八种语言类型/工厂、FFI 分发、文档、contract vectors @@ -1597,7 +1603,7 @@ P1 可以拆成多个审阅工作包，但这些工作包不被描述为独立 | 未覆盖项：补回 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 宣称已经解决 | diff --git a/docs/api/c.md b/docs/api/c.md index 2bd71cfa1..491fa7c3c 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..19` | `AiMuxError` (4 retired, 14 = `Retry`) | | `100..105` | `RecordingError` | | `200..206` | failure detected by the C ABI | @@ -139,10 +139,10 @@ 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–19 (`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_LOAD_API_KEY` 18, `AIMUX_E_LOAD_SETTING` 19; +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 +154,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` (18) and `AIMUX_E_LOAD_SETTING` (19) 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 +325,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 | diff --git a/docs/api/decision.md b/docs/api/decision.md index 90d7a9eb9..e9ae649aa 100644 --- a/docs/api/decision.md +++ b/docs/api/decision.md @@ -170,10 +170,11 @@ finally: ``` 标准 `recordings.jsonl` 每行记录一次逻辑调用：`input.operation = "decision"`， -`input.options` 保留 state/questions；provider 快照含 endpoint 和 capabilities； +`input.options` 保留 state/questions，`input.decision_capabilities` 保留能力声明； +provider 只记录 provider_id/provider/model_id，endpoint 与凭证来自调用方配置； exchanges 记录各重试 attempt 的 HTTP 请求、响应和时延； `outcome.decision_result` 保存规范化结果。凭据沿用统一脱敏规则。 -旧语言录制缺少 operation 时仍按 language model 解析，schema 2 保持兼容。 +录制使用 schema 3；语言录制缺少 operation 时仍按 language model 解析。 CLI 离线回放不需要 key： @@ -186,7 +187,7 @@ Rust 可用 `aimux_core::replay::MockDecisionReplayModel::from_jsonl` 加载录 再通过 `decide` 或 `replay_decision_with_model` 离线回放。匹配 provider/model、 state/questions、headers 和 provider_options；timeout/retry 控制不参与匹配。 输入未命中或记录不完整时返回错误。请求回放可通过 -`aimux_providers::rebuild_decision_provider` 重建官方 Jev，再调用 +`aimux_providers::rebuild_decision_provider` 从调用方 registry 重建 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 39d70832b..18c1b2d79 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 281 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 281 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..19 (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..19, 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,8 @@ 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 + ├── LoadAPIKeyError // code 18; no API key, fallback env var unset (envVar) + ├── LoadSettingError // code 19; a required setting is missing (envVar) └── OtherError ``` diff --git a/docs/api/go.md b/docs/api/go.md index da8ef02af..266162d2e 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..19 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); a missing API key / provider setting are +`CodeLoadAPIKey` (18) / `CodeLoadSetting` (19). 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–19), `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..19 | 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 281 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 281 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 diff --git a/docs/api/java.md b/docs/api/java.md index c1314c719..a34942a0b 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 281 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 281 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,8 @@ 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 + ├── LoadAPIKeyError // code 18: no API key, fallback env var unset (getEnvVar()) + ├── LoadSettingError // code 19: a required setting is missing (getEnvVar()) └── OtherError ``` @@ -83,7 +84,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–19 (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 +150,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–19 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 154ce3aa8..a49d46a6e 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 281 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 281 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..19 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,8 @@ 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) + ├── LoadAPIKeyError // code 18: no API key, fallback env var unset (envVar) + ├── LoadSettingError // code 19: a required setting is missing (envVar) └── OtherError ``` @@ -104,7 +105,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..19; 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 974751062..f4e3c223d 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 281 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 281 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 @@ -509,7 +516,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 @@ -524,7 +531,7 @@ them, not a local copy): ```typescript import type { - GenerateTextOptions, GenerateTextResult, StreamPart, ModelMessage, + GenerateTextOptions, GenerateTextResult, TextStreamPart, ModelMessage, Tool, ToolChoice, ToolCall, ToolResult, Usage, FinishReason, Warning, Role, MessageContent, ContentPart, ResponseFormat, ReasoningEffort, GenerateResult, FunctionTool, @@ -547,25 +554,29 @@ export type GenerateTextResult = { response_messages: Array // assistant messages ready for the next turn raw_finish_reason: string | null // provider's own finish-reason string provider_metadata: JsonValue | null // mirrored from raw.provider_metadata - response: ResponseMetadata // mirrored from raw.response (id, timestamp, model_id) + request: RequestInfo // sent request body + response: ResponseInfo // id, timestamp, model_id, headers and body total_usage: Usage // usage across all steps (equals usage in single-step mode) } ``` -`StreamPart` is an external-tagged union of 18 variants (each is a one-key +`streamText` emits `TextStreamPart`, an external-tagged union (each is a one-key object — type narrowing via `part.TextDelta` etc. works out of the box): ```typescript -// bindings/node/src/types/StreamPart.ts (variants, abridged) -export type StreamPart = +// bindings/node/src/types/TextStreamPart.ts (variants, abridged) +export type TextStreamPart = | { StreamStart: ... } | { TextStart: ... } | { TextDelta: ... } | { TextEnd: ... } | { ToolInputStart: ... } | { ToolInputDelta: ... } | { ToolInputEnd: ... } | { ToolCall: ... } | { ToolResult: ... } | { ReasoningStart: ... } | { ReasoningDelta: ... } | { ReasoningEnd: ... } - | { ResponseMetadata: ... } | { Source: ... } | { Finish: ... } + | { Source: ... } | { Finish: ... } | { Error: ... } | { Raw: ... } | { File: ... } ``` +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..7ac9aa641 100644 --- a/docs/api/providers.md +++ b/docs/api/providers.md @@ -9,282 +9,308 @@ | category | count | |----------|-------| -| Registry-backed OpenAI-compatible (`provider(name, ...)`) | 251 | +| Registry-backed OpenAI-compatible (`provider(name, ...)`) | 281 | | 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** | **326** | + +**281 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) — 281 + +| 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` | +| `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 | — | `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 | — | `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 | — | `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 | — | `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 | — | `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 +320,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 9c06abedc..c5c83701f 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 +281 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 281 preset rows; +> typed factories remain available. Custom endpoints use the `base_url` param. > Full list: [providers.md](providers.md). ## Text Generation @@ -391,6 +396,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 @@ -429,7 +435,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 ( @@ -465,9 +474,16 @@ class GenerateTextResult(BaseModel): usage: Usage warnings: list[Warning] raw: GenerateResult + request: RequestInfo + response: ResponseInfo ``` -`StreamPart` is a `RootModel` over the external-tagged union dict, e.g. +`stream_text` emits the Rust call-layer `TextStreamPart` JSON. The Python +wrapper calls its validating `RootModel` `StreamPart`, over the external-tagged +union dict, e.g. `{"TextDelta": {"id": ..., "delta": ...}}`. 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 a1796fd3b..edb5b1ab3 100644 --- a/docs/api/reference.md +++ b/docs/api/reference.md @@ -15,10 +15,10 @@ | `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` | -| `GenerateTextResult` | [generate.rs](../../aimux-core/src/generate.rs) | [GenerateTextResult.ts](../../bindings/node/src/types/GenerateTextResult.ts) | `text`, `tool_calls`, `usage`, `warnings`, `raw` | -| `StreamTextResult` | [generate.rs](../../aimux-core/src/generate.rs) | — | use `StreamPart` while iterating | -| `StreamPart` | [stream_part.rs](../../aimux-core/src/stream_part.rs) | [StreamPart.ts](../../bindings/node/src/types/StreamPart.ts) | TextDelta / ToolCallDelta / Finish / … | +| `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`) | +| `GenerateTextResult` | [generate.rs](../../aimux-core/src/generate.rs) | [GenerateTextResult.ts](../../bindings/node/src/types/GenerateTextResult.ts) | `text`, `tool_calls`, `usage`, `warnings`, `raw`, `request`, `response` | +| `StreamTextResult` | [generate.rs](../../aimux-core/src/generate.rs) | — | use `TextStreamPart` while iterating | +| `StreamPart` / `TextStreamPart` | [stream_part.rs](../../aimux-core/src/stream_part.rs) | [TextStreamPart.ts](../../bindings/node/src/types/TextStreamPart.ts) | Provider / call layer; only the provider layer emits ResponseMetadata | | `Usage` / `TokenUsage` | [types.rs](../../aimux-core/src/types.rs) | [Usage.ts](../../bindings/node/src/types/Usage.ts) | input/output/reasoning tokens | | `FinishReason` / `FinishReasonUnified` | [types.rs](../../aimux-core/src/types.rs) | [FinishReason.ts](../../bindings/node/src/types/FinishReason.ts) | | | `Warning` | [types.rs](../../aimux-core/src/types.rs) | [Warning.ts](../../bindings/node/src/types/Warning.ts) | non-fatal provider notices | @@ -44,24 +44,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. @@ -69,9 +68,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 | @@ -79,13 +80,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 | @@ -97,7 +102,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 | @@ -107,7 +112,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 | @@ -118,14 +123,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 4614f3c25..6bacbd6fc 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 281 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; + +// 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")?; -// 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)?; +// 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() +})?; -// 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() }), -)?; +// 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 @@ -95,7 +156,9 @@ let result = task.await.unwrap()?; ## 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; @@ -104,8 +167,8 @@ let result = stream_text(&model, "Write a haiku.", GenerateTextOptions::default( let mut stream = result.stream; while let Some(part) = stream.next().await { match part? { - StreamPart::TextDelta { delta, .. } => print!("{}", delta), - StreamPart::Finish { .. } => println!("\n[done]"), + TextStreamPart::TextDelta { delta, .. } => print!("{}", delta), + TextStreamPart::Finish { .. } => println!("\n[done]"), _ => {} } } @@ -115,7 +178,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). @@ -174,7 +237,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 @@ -240,7 +303,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 @@ -253,7 +316,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 @@ -278,11 +341,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 | @@ -295,14 +359,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 @@ -313,7 +374,7 @@ Rust types are the canonical definitions — one module per feature in |------|------| | `generate` | `generate_text` / `stream_text` functions, `GenerateResult` | | `language_model` | `LanguageModel`, `GenerateTextResult` | -| `stream_part` | `StreamPart` (18 variants) | +| `stream_part` | `StreamPart` (provider layer), `TextStreamPart` (call layer) | | `options` | `GenerateTextOptions`, `CallOptions`, `ResponseFormat`, `ToolChoice`, `ReasoningEffort` | | `message` / `language_model_message` | `ModelMessage`, `ModelPrompt`, `MessageContent`, `Role` | | `content` | `ContentPart` (Text / Image / File / Reasoning / ToolCall / ToolResult) | diff --git a/docs/api/swift.md b/docs/api/swift.md index e87e054d6..7e6571e14 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 281 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 281 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...19; 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,8 @@ 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) | +| `.loadApiKey` | `AIMUX_E_LOAD_API_KEY` (18) | No API key passed and the fallback environment variable is unset; carries `envVar` | +| `.loadSetting` | `AIMUX_E_LOAD_SETTING` (19) | 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..570501187 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..19` 属于核心，`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..69256d5bd 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` 表示未提供。预设表有 281 行;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/VERSIONS.json b/fixtures/aisdk/VERSIONS.json new file mode 100644 index 000000000..f4f49d7ba --- /dev/null +++ b/fixtures/aisdk/VERSIONS.json @@ -0,0 +1,9 @@ +{ + "@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.53", + "ai": "7.0.127" +} 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/0037-decision-api.md b/rfc/0037-decision-api.md index 21b3761a1..6a3ad59fd 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 查询和宿主包装均已实现。 其他官方决策接口 diff --git a/scripts/aisdk-fixtures/README.md b/scripts/aisdk-fixtures/README.md new file mode 100644 index 000000000..15bb3adc1 --- /dev/null +++ b/scripts/aisdk-fixtures/README.md @@ -0,0 +1,61 @@ +# 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) +``` + +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..144b2c2b7 --- /dev/null +++ b/scripts/aisdk-fixtures/package-lock.json @@ -0,0 +1,205 @@ +{ + "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.53", + "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/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/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/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/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.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/@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/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..5014ec873 --- /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.53", + "ai": "7.0.127" + } +} diff --git a/scripts/aisdk-fixtures/record.mjs b/scripts/aisdk-fixtures/record.mjs new file mode 100644 index 000000000..d1cfd4a0a --- /dev/null +++ b/scripts/aisdk-fixtures/record.mjs @@ -0,0 +1,252 @@ +#!/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'; + +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) { + const versions = Object.fromEntries(VERSIONED.map((p) => [p, installedVersion(p)])); + writeFileSync(join(OUT, 'VERSIONS.json'), JSON.stringify(sortKeys(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..30e874432 --- /dev/null +++ b/scripts/check_provider_boundaries.sh @@ -0,0 +1,296 @@ +#!/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' \ + 'fn specification_version' \ + '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 c76569546..43c728ee5 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 使用。 @@ -223,8 +200,10 @@ pub async fn run(args: &ProviderArgs) -> anyhow::Result anyhow::Result ( std::net::SocketAddr, @@ -333,14 +312,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 dd095e062..a77d2ddef 100644 --- a/tools/aimux-replay/src/main.rs +++ b/tools/aimux-replay/src/main.rs @@ -1,16 +1,15 @@ //! aimux-replay — RFC-0023 请求回放 CLI(层 2 消费端)。 //! -//! 读录制 jsonl(每行一个 `Recording`),按 `ProviderRecord` 自动重建 -//! provider(语言模型 OpenAI 兼容族、官方 TypeSafe Jev decision),再用录制输入 **重发真实 -//! API**。用途:离线重跑线上流量、改 prompt 重发(A/B)、回归对比、CI 集成。 +//! 读录制 jsonl(每行一个 `Recording`,仅 schema 3),按 `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 录制；真实回放只使用调用方配置， +//! 不从录制恢复连接设置。Jev 使用显式 key 或 TYPESAFE_API_KEY。 use std::path::PathBuf; @@ -35,16 +34,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 +103,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 +113,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 +151,30 @@ 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)? + if rec.provider.provider_id != "jev" { + bail!( + "decision replay: no configured provider {}", + rec.provider.provider_id + ); + } + let config = match api_key { + Some(key) => aimux_providers::JevConfig::new(key), + None => aimux_providers::JevConfig::from_env()?, + }; + let providers = std::collections::BTreeMap::from([( + "jev".into(), + std::sync::Arc::new(aimux_providers::JevProvider::new(config)) + as std::sync::Arc, + )]); + let registry = aimux_core::create_provider_registry(providers, Default::default()); + aimux_providers::rebuild_decision_provider(&rec.provider, ®istry)? }; let result = Runtime::new()?.block_on(aimux_core::replay::replay_decision_with_model( rec, @@ -180,15 +189,29 @@ 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 mut providers = aimux_providers::default_providers(); + if let Some(key) = api_key { + providers.insert( + rec.provider.provider_id.clone(), + aimux_providers::create_provider( + &rec.provider.provider_id, + aimux_providers::PresetSettings { + api_key: Some(key.to_string().into()), + ..Default::default() + }, + )?, + ); + } + let registry = aimux_core::create_provider_registry(providers, Default::default()); + rebuild_provider(&rec.provider, ®istry) + .with_context(|| format!("rebuild provider for '{}'", rec.provider.provider_id))? }; let overrides = prompt.map(|text| ReplayOverrides { 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 8f5152024..8201743c1 100644 --- a/tools/aimux-web/src/api/cache_probe.rs +++ b/tools/aimux-web/src/api/cache_probe.rs @@ -132,8 +132,10 @@ async fn run_inner(state: AppState, req: ProbeRequest) -> Result 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 beb2a9d10..19233ba7e 100644 --- a/tools/aimux-web/src/api/replay.rs +++ b/tools/aimux-web/src/api/replay.rs @@ -60,31 +60,31 @@ async fn run_inner(state: AppState, req: ReplayRequest) -> Result { - wire::resolve_api_key(Some(s), &provider_name, &state.keys, state.loopback)? - } - "none" => wire::resolve_api_key(None, &provider_name, &state.keys, state.loopback)?, - _ => wire::resolve_api_key( - req.api_key.as_deref(), - &provider_name, - &state.keys, - state.loopback, - ) - .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. + let key = wire::resolve_api_key( + req.api_key.as_deref(), + &recording.provider.provider_id, + &state.keys, + state.loopback, + )?; + + let mut providers = aimux_providers::default_providers(); + if let Some(key) = key { + providers.insert( + recording.provider.provider_id.clone(), + aimux_providers::create_provider( + &recording.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, @@ -135,8 +135,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 3823e10ff..90b8249c0 100644 --- a/tools/aimux-web/src/model_builder.rs +++ b/tools/aimux-web/src/model_builder.rs @@ -1,76 +1,29 @@ -//! 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, $env:literal) => {{ - 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, "OPENAI_API_KEY"), - "anthropic" => native!( - anthropic, - AnthropicConfig, - AnthropicProvider, - "ANTHROPIC_API_KEY" - ), - "google" => native!( - google, - GoogleConfig, - GoogleProvider, - "GOOGLE_GENERATIVE_AI_API_KEY" - ), - "mistral" => native!(mistral, MistralConfig, MistralProvider, "MISTRAL_API_KEY"), - "xai" => native!(xai, XAIConfig, XAIProvider, "XAI_API_KEY"), - "cohere" => native!(cohere, CohereConfig, CohereProvider, "COHERE_API_KEY"), - _ => { - 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)) - } - } -} - -/// 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}\"" - )) - }), - } + 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) } diff --git a/tools/aimux-web/src/wire.rs b/tools/aimux-web/src/wire.rs index 8787cd023..e9457b0d2 100644 --- a/tools/aimux-web/src/wire.rs +++ b/tools/aimux-web/src/wire.rs @@ -83,8 +83,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, @@ -278,7 +276,6 @@ pub fn to_generate_options( .map(|ts| ts.iter().map(to_function_tool).collect()), response_format, 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, 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 c60e4b0d5..08a95cdb1 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) {