diff --git a/.env.example b/.env.example index 04c320b..3d0aa0b 100644 --- a/.env.example +++ b/.env.example @@ -10,6 +10,7 @@ ANTHROPIC_API_KEY= XAI_API_KEY= ARK_API_KEY= OPENROUTER_API_KEY= +WAVESPEED_API_KEY= # Blockchain heartbeat configuration (optional). # If HEARTBEAT_CONTRACT_ADDRESS and HEARTBEAT_FACILITATOR_URL are set, the enclave diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 64d07d9..41ecb61 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -41,6 +41,7 @@ jobs: # accept either secret name so the job runs whichever the repo still holds. OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY || secrets.NOUS_API_KEY }} ZAI_API_KEY: ${{ secrets.ZAI_API_KEY }} + WAVESPEED_API_KEY: ${{ secrets.WAVESPEED_API_KEY }} steps: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v5 @@ -49,7 +50,7 @@ jobs: run: | for name in \ OPENAI_API_KEY ANTHROPIC_API_KEY GOOGLE_API_KEY XAI_API_KEY \ - ARK_API_KEY OPENROUTER_API_KEY ZAI_API_KEY + ARK_API_KEY OPENROUTER_API_KEY ZAI_API_KEY WAVESPEED_API_KEY do if [[ -z "${!name}" ]]; then echo "Missing required repository secret: $name" >&2 diff --git a/CLAUDE.md b/CLAUDE.md index 3077694..8291445 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -75,6 +75,7 @@ API keys (injected at runtime via POST /v1/keys — do NOT bake into the image): - `ARK_API_KEY` (BytePlus / ByteDance ModelArk; injected as `bytedance_api_key`) - `OPENROUTER_API_KEY` (OpenRouter; injected as `openrouter_api_key`) - `ZAI_API_KEY` (Z.ai Model API; injected as `zai_api_key`) +- `WAVESPEED_API_KEY` (WaveSpeed; injected as `wavespeed_api_key`) - `EXA_API_KEY` (Exa search; injected as `exa_api_key`) — backs the in-enclave `/v1/web_search` endpoint, not an LLM provider. Without it the endpoint returns 503 and `/health` reports `web_search_enabled: false`. @@ -126,13 +127,14 @@ Server configuration: `model_registry.py` is the source of truth; this list mirrors it. Model name prefixes determine routing: -- **OpenAI**: gpt-6.1-sol, gpt-6-sol, gpt-6-luna, gpt-6-astra, gpt-4.1, gpt-4.1-mini, gpt-5, gpt-5-mini, gpt-5.2, gpt-5.4, gpt-5.4-mini, gpt-5.4-nano, gpt-5.5, gpt-5.6-sol/terra/luna, o3; image generation: gpt-image-2.5-flare, gpt-image-2.5-sunburst, gpt-image-2 -- **Anthropic**: claude-sonnet-4-5/4-6, claude-sonnet-5, claude-sonnet-5-5, claude-haiku-4-5, claude-opus-4-5/4-6/4-7/4-8, claude-opus-5, claude-opus-5-5, claude-fable-5, claude-fable-5-1 -- **Google**: gemini-3.8-flash, gemini-3.7-flash, gemini-3.6-flash, gemini-3.5-flash, gemini-3.5-flash-lite, gemini-2.5-flash, gemini-2.5-flash-lite, gemini-2.5-pro, gemini-3-flash-preview, gemini-3.1-pro-preview; image generation: gemini-2.5-flash-image, gemini-3.1-flash-image +- **OpenAI**: gpt-6.1-terra, gpt-6.1-sol, gpt-6-sol, gpt-6-luna, gpt-6-astra, gpt-4.1, gpt-4.1-mini, gpt-5, gpt-5-mini, gpt-5.2, gpt-5.4, gpt-5.4-mini, gpt-5.4-nano, gpt-5.5, gpt-5.6-sol/terra/luna, o3; image generation: gpt-image-2.5-flare, gpt-image-2.5-sunburst, gpt-image-2 +- **Anthropic**: claude-sonnet-4-5/4-6, claude-sonnet-5, claude-sonnet-5-5, claude-haiku-5-5, claude-haiku-4-5, claude-opus-4-5/4-6/4-7/4-8, claude-opus-5, claude-opus-5-5, claude-fable-5, claude-fable-5-1 +- **Google**: gemini-4-argon, gemini-3.8-flash, gemini-3.7-flash, gemini-3.6-flash, gemini-3.5-flash, gemini-3.5-flash-lite, gemini-2.5-flash, gemini-2.5-flash-lite, gemini-2.5-pro, gemini-3-flash-preview, gemini-3.1-pro-preview; image generation: gemini-2.5-flash-image, gemini-3.1-flash-image, gemini-nano-banana-2.1 - **xAI**: grok-4.7, grok-4.6, grok-4.5, grok-4.3, grok-4.20-reasoning, grok-4.20-non-reasoning; image generation: grok-2-image, grok-imagine-image-2.0 - **ByteDance** (BytePlus ModelArk, OpenAI-compatible, ap-southeast): seed-1.6, seed-1.8, seed-2.0-lite, deepseek-v4-flash, deepseek-v4-pro, glm-5.2 (Z.ai's model served via a ModelArk deployment endpoint); image generation: seedream-4.0, seedream-5.0-lite, seedance-4.5, seedance-5.0 - **OpenRouter** (OpenAI-compatible): hy4-preview, hermes-4-405b, hy3 - **Z.ai** (Model API, OpenAI-compatible): image generation: glm-image (glm-5.2 chat is routed through BytePlus ModelArk, see ByteDance above) +- **WaveSpeed** (async prediction API): image generation: krea-3 (Krea's model, always `berry` at `reasoning_level: medium`), qwen-image-3.0-pro (Alibaba's model); each is two WaveSpeed endpoints, text-to-image and edit Models kept registered but no longer offered to new clients — each still resolves so older SDK versions keep working, and each is priced at what the @@ -181,8 +183,30 @@ prices by output pixel count can declare **resolution tiers** `resolution` field (`"1.5K"`, `"2K"`) picks the tier, which supplies the `size` keyword, the ratio→pixels table, and the per-image price — omitted means the default tier, and the field is rejected on single-resolution -models. Seedream 5.0 is the only tiered model today (`1.5K` at $0.045 is the -default; `2K` at $0.09 is what every request used to be pinned to). +models. Seedream 5.0 (`1.5K` at $0.045 is the default; `2K` at $0.09 is what +every request used to be pinned to), Krea 3 (`1K` $0.03 default, `2K` $0.12, +each reference billed at the same price as the output, per WaveSpeed's +`$0.03 x MP x (1 + references)` formula) and Qwen Image 3.0 Pro (`1K` $0.04 default, +`2K` $0.075) are the tiered models today. + +**WaveSpeed** (krea-3, qwen-image-3.0-pro) is the one image provider that does not +answer in a single call. `image_generation._generate_wavespeed` submits an +async prediction to `/{model_id}` — the `…/text-to-image` model, or the +`…/edit` model (`image_edit_model`) when the turn carries references — then +polls `/predictions/{id}/result` (2s, backing off to 5s, 180s deadline → 504, +under the relay's 210s read timeout), +fetches the output URL like any hosted image, and finally deletes the +prediction so WaveSpeed's history no longer holds the prompt or output. +WaveSpeed takes references as URLs only: inline `data:` references are +uploaded through its `/media/uploads` ticket (a keyless PUT to the signed +URL, which is a credential and is never logged), plain URLs are passed +through, and at most `image_max_references` (Qwen 3, Krea 8) are sent. Its +input images are billed too — `per_reference_image_price_usd` (Qwen $0.003 +each; Krea per tier, at the output's price) on top of the per-image price, +counted from the references actually forwarded. One image per prediction; `n` +is ignored. Every Qwen request sends `enable_prompt_expansion: false`, and +every Krea request `model: berry` + `reasoning_level: medium` +(`image_extra_params`). ### Web Search diff --git a/README.md b/README.md index 45c1f2a..f43762b 100644 --- a/README.md +++ b/README.md @@ -33,17 +33,17 @@ The gateway solves this by running inside a hardware-isolated Nitro Enclave wher | Provider | Models | |----------|--------| -| OpenAI | gpt-6.1-sol, gpt-6-sol, gpt-6-luna, gpt-6-astra, gpt-4.1, gpt-4.1-mini, gpt-5, gpt-5-mini, gpt-5.2, gpt-5.4, gpt-5.5, gpt-5.6-sol, gpt-5.6-terra, gpt-5.6-luna, o3 | -| Anthropic | claude-sonnet-5-5, claude-fable-5-1, claude-fable-5, claude-opus-5-5, claude-opus-5, claude-opus-4-8, claude-opus-4-7, claude-opus-4-6, claude-opus-4-5, claude-sonnet-5, claude-sonnet-4-6, claude-sonnet-4-5, claude-haiku-4-5 | -| Google | gemini-3.8-flash, gemini-3.7-flash, gemini-3.6-flash, gemini-3.5-flash, gemini-3.5-flash-lite, gemini-2.5-flash, gemini-2.5-flash-lite, gemini-2.5-pro, gemini-3-flash-preview, gemini-3.1-pro-preview | +| OpenAI | gpt-6.1-terra, gpt-6.1-sol, gpt-6-sol, gpt-6-luna, gpt-6-astra, gpt-4.1, gpt-4.1-mini, gpt-5, gpt-5-mini, gpt-5.2, gpt-5.4, gpt-5.5, gpt-5.6-sol, gpt-5.6-terra, gpt-5.6-luna, o3 | +| Anthropic | claude-haiku-5-5, claude-sonnet-5-5, claude-fable-5-1, claude-fable-5, claude-opus-5-5, claude-opus-5, claude-opus-4-8, claude-opus-4-7, claude-opus-4-6, claude-opus-4-5, claude-sonnet-5, claude-sonnet-4-6, claude-sonnet-4-5, claude-haiku-4-5 | +| Google | gemini-4-argon, gemini-3.8-flash, gemini-3.7-flash, gemini-3.6-flash, gemini-3.5-flash, gemini-3.5-flash-lite, gemini-2.5-flash, gemini-2.5-flash-lite, gemini-2.5-pro, gemini-3-flash-preview, gemini-3.1-pro-preview | | xAI | grok-4.7, grok-4.6, grok-4.5, grok-4.3, grok-4.20-reasoning, grok-4.20-non-reasoning | | ByteDance | seed-1.6, seed-1.8, seed-2.0-lite, dola-seed-2.0-mini, deepseek-v4-flash, deepseek-v4-pro, glm-5.2 | | OpenRouter | hy4-preview, hermes-4-405b, hy3 | -Image generation: gpt-image-2, gpt-image-2.5-flare, gpt-image-2.5-sunburst, -gemini-2.5-flash-image, gemini-3.1-flash-image, grok-2-image, -grok-imagine-image-2.0, seedream-4.0, seedream-5.0-lite, seedance-4.5, -seedance-5.0, glm-image. +Image generation: krea-3 and qwen-image-3.0-pro (via WaveSpeed), gpt-image-2, +gpt-image-2.5-flare, gpt-image-2.5-sunburst, gemini-2.5-flash-image, +gemini-nano-banana-2.1, gemini-3.1-flash-image, grok-2-image, grok-imagine-image-2.0, seedream-4.0, +seedream-5.0-lite, seedance-4.5, seedance-5.0, glm-image. Some retired models stay registered so older SDK versions keep working; see "Supported Providers" in CLAUDE.md. `tee_gateway/model_registry.py` is the @@ -65,6 +65,7 @@ export XAI_API_KEY=... export ARK_API_KEY=... # BytePlus / ByteDance ModelArk export OPENROUTER_API_KEY=... # OpenRouter export ZAI_API_KEY=... # Z.ai Model API +export WAVESPEED_API_KEY=... # WaveSpeed (krea-3, qwen-image-3.0-pro) # Run server (starts the Flask/connexion app on port 8000) make test-local # Werkzeug dev server (development only) @@ -418,6 +419,7 @@ Clients use an x402-compatible client (e.g., the [x402 SDK](https://github.com/o | `ARK_API_KEY` | - | BytePlus / ByteDance ModelArk API key (injected as `bytedance_api_key`) | | `OPENROUTER_API_KEY` | - | OpenRouter API key (injected as `openrouter_api_key`) | | `ZAI_API_KEY` | - | Z.ai Model API key (injected as `zai_api_key`) | +| `WAVESPEED_API_KEY` | - | WaveSpeed API key (injected as `wavespeed_api_key`) | | `EVM_PAYMENT_ADDRESS` | - | Wallet address to receive x402 payments | | `FACILITATOR_URL` | see `tee_gateway/__main__.py` | x402 payment facilitator endpoint | diff --git a/scripts/run-enclave.sh b/scripts/run-enclave.sh index d0d32fc..7e74dda 100755 --- a/scripts/run-enclave.sh +++ b/scripts/run-enclave.sh @@ -92,6 +92,7 @@ if [ -f "$ENV_FILE" ]; then ARK_API_KEY="$(grep -E '^ARK_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" OPENROUTER_API_KEY="$(grep -E '^OPENROUTER_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" ZAI_API_KEY="$(grep -E '^ZAI_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" + WAVESPEED_API_KEY="$(grep -E '^WAVESPEED_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" # Backs the in-enclave `web_search` tool (Exa), not an LLM provider. EXA_API_KEY="$(grep -E '^EXA_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" @@ -112,6 +113,7 @@ if [ -f "$ENV_FILE" ]; then --arg bytedance "$ARK_API_KEY" \ --arg openrouter "$OPENROUTER_API_KEY" \ --arg zai "$ZAI_API_KEY" \ + --arg wavespeed "$WAVESPEED_API_KEY" \ --arg exa "$EXA_API_KEY" \ --arg hb_contract "$HEARTBEAT_CONTRACT_ADDRESS" \ --arg facilitator "$FACILITATOR_URL" \ @@ -124,6 +126,7 @@ if [ -f "$ENV_FILE" ]; then bytedance_api_key: $bytedance, openrouter_api_key: $openrouter, zai_api_key: $zai, + wavespeed_api_key: $wavespeed, exa_api_key: $exa } + if $hb_contract != "" then {heartbeat_contract_address: $hb_contract} else {} end @@ -155,7 +158,7 @@ if [ -f "$ENV_FILE" ]; then # Clear key variables from this shell immediately after use unset OPENAI_API_KEY GOOGLE_API_KEY ANTHROPIC_API_KEY XAI_API_KEY ARK_API_KEY - unset OPENROUTER_API_KEY ZAI_API_KEY EXA_API_KEY + unset OPENROUTER_API_KEY ZAI_API_KEY WAVESPEED_API_KEY EXA_API_KEY unset HEARTBEAT_CONTRACT_ADDRESS FACILITATOR_URL TEE_HEARTBEAT_INTERVAL fi else diff --git a/tee_gateway/__main__.py b/tee_gateway/__main__.py index 09ba256..23ac5bf 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -429,6 +429,7 @@ def set_provider_keys(): bytedance_api_key=body.get("bytedance_api_key") or None, openrouter_api_key=body.get("openrouter_api_key") or None, zai_api_key=body.get("zai_api_key") or None, + wavespeed_api_key=body.get("wavespeed_api_key") or None, exa_api_key=body.get("exa_api_key") or None, ) set_provider_config(provider_config) @@ -496,6 +497,10 @@ def _set(val: str | None) -> str: logger.info( " zai_api_key : %s", _set(provider_config.zai_api_key) ) + logger.info( + " wavespeed_api_key : %s", + _set(provider_config.wavespeed_api_key), + ) logger.info( " exa_api_key (web search) : %s", _set(provider_config.exa_api_key) ) @@ -533,6 +538,7 @@ def _set(val: str | None) -> str: "bytedance": provider_config.bytedance_api_key, "openrouter": provider_config.openrouter_api_key, "zai": provider_config.zai_api_key, + "wavespeed": provider_config.wavespeed_api_key, }.items() if k ] diff --git a/tee_gateway/config.py b/tee_gateway/config.py index 2e19bf0..772f88b 100644 --- a/tee_gateway/config.py +++ b/tee_gateway/config.py @@ -30,6 +30,7 @@ class ProviderConfig: bytedance_api_key: Optional[str] = None openrouter_api_key: Optional[str] = None zai_api_key: Optional[str] = None + wavespeed_api_key: Optional[str] = None # Exa, which backs the `web_search` tool the gateway executes inside the # enclave for every model (see web_search.py). Not an LLM provider, so it is # deliberately absent from initialized_providers() — /health reports it @@ -53,6 +54,8 @@ def initialized_providers(self) -> list[str]: providers.append("openrouter") if self.zai_api_key: providers.append("zai") + if self.wavespeed_api_key: + providers.append("wavespeed") return providers diff --git a/tee_gateway/image_generation.py b/tee_gateway/image_generation.py index 129d4cc..757f980 100644 --- a/tee_gateway/image_generation.py +++ b/tee_gateway/image_generation.py @@ -17,6 +17,11 @@ ride out-of-band under the message ``images`` key), signed and billed flat per generated image. +WaveSpeed is the exception to the single-call shape: it runs generation as an +async prediction the gateway submits and polls (``_generate_wavespeed``), and +takes reference images as URLs only, so inline ones are uploaded to its storage +first. Its result is fetched and inlined like any other hosted URL. + Per-provider request quirks (response format, ``n`` support, reference-image editing, extra params) live in the model registry, keeping this code flat. """ @@ -29,7 +34,7 @@ import time import uuid from typing import Any, List, Optional -from urllib.parse import urlparse +from urllib.parse import quote, urlparse import httpx from flask import Response @@ -59,6 +64,7 @@ "x-ai": "xai_http_client", "bytedance": "bytedance_http_client", "zai": "zai_http_client", + "wavespeed": "wavespeed_http_client", } # Bounds on the URL fetch (egress hardening). Provider images are well under the @@ -83,6 +89,33 @@ # (and from there into logs and the client-facing error payload). _MAX_ERROR_DETAIL_CHARS = 600 +# WaveSpeed polling. Its docs ask for no more than one result query every 2s, +# backing off toward 5-10s for long tasks; Qwen Image 3.0 Pro takes ~80s. +_WAVESPEED_POLL_INITIAL_SECONDS = 2.0 +_WAVESPEED_POLL_MAX_SECONDS = 5.0 +# How long one request waits for its prediction before failing as a provider +# timeout (504). An abandoned prediction may still complete and be billed by +# WaveSpeed; the client is not charged for it. Kept under chat-api's 210s relay +# read timeout: the image stream is silent until the final frame, so waiting +# past it would only finish generations no one is still connected to receive. +_WAVESPEED_DEADLINE_SECONDS = 180.0 +# Consecutive failed result queries tolerated before giving up, so one blip +# does not throw away a generation that is already paid for upstream. +_WAVESPEED_MAX_POLL_ERRORS = 3 +_WAVESPEED_FAILED_STATUSES = frozenset({"failed", "cancelled", "timeout", "deleted"}) + +# Keyless client for PUTting reference images to WaveSpeed's signed upload URL. +# The API key must never be sent there, so this is not the provider client. +_wavespeed_upload_client: Optional[httpx.Client] = None + + +class ProviderTaskError(httpx.HTTPError): + """A provider rejected or failed an async task without an HTTP error status. + + Subclasses ``httpx.HTTPError`` so ``errors.describe_exception`` classifies + it as a provider failure (502) rather than a bug in the enclave. + """ + def _provider_error_detail(resp: httpx.Response) -> str: """Extract a compact human-readable detail from a provider error response. @@ -390,6 +423,196 @@ def validate_resolution(model: str, resolution: Optional[str]) -> None: cfg.image_tier(resolution) +def _forwarded_references(cfg: Any, reference_images: Optional[List[str]]) -> list[str]: + """The reference images a request actually sends to its provider. + + Non-empty strings only, capped at the model's ``image_max_references``, and + empty for models that take no references. WaveSpeed only accepts URLs, so + an inline reference that does not decode (and so cannot be uploaded) is + dropped here. Shared by generation and billing, which charges per + reference forwarded on models that price them. + """ + if not reference_images or not cfg.image_supports_reference: + return [] + refs = [r for r in reference_images if isinstance(r, str) and r] + if cfg.provider == "wavespeed": + refs = [r for r in refs if not r.startswith("data:") or _decode_data_uri(r)] + return refs[: cfg.image_max_references] + + +def _wavespeed_data(resp: httpx.Response) -> dict[str, Any]: + """Unwrap WaveSpeed's ``{"code", "message", "data"}`` envelope. + + WaveSpeed reports some failures in the body ``code`` as well as the HTTP + status, so both are checked. + """ + _raise_for_status_with_detail(resp) + body = resp.json() + if not isinstance(body, dict): + raise ProviderTaskError("WaveSpeed returned a non-object response") + code = body.get("code") + if code is not None and code != 200: + message = str(body.get("message") or "no message")[:_MAX_ERROR_DETAIL_CHARS] + raise ProviderTaskError(f"WaveSpeed error {code}: {message}") + data = body.get("data") + if not isinstance(data, dict): + raise ProviderTaskError("WaveSpeed response has no data object") + return data + + +def _wavespeed_reference_url(client: httpx.Client, ref: str, index: int) -> str: + """Return a URL WaveSpeed can read for one reference image. + + A plain URL is passed through for WaveSpeed to fetch itself (as with + ByteDance; the enclave never dereferences client URLs). An inline ``data:`` + reference is uploaded first: an authenticated request for a short-lived + upload ticket, then a keyless PUT of the bytes to the ticket's signed URL. + That URL is a temporary credential, so it is never put in an error or log. + """ + decoded = _decode_data_uri(ref) + if decoded is None: + return ref + raw, mime = decoded + ext = _IMAGE_MIME_EXT.get(mime, "png") + ticket = _wavespeed_data( + client.post( + "/media/uploads", + json={ + "filename": f"image_{index}.{ext}", + "size": len(raw), + "content_type": mime, + }, + ) + ) + upload = ticket.get("upload") + download_url = ticket.get("download_url") + upload_url = upload.get("url") if isinstance(upload, dict) else None + if not isinstance(upload_url, str) or not isinstance(download_url, str): + raise ProviderTaskError("WaveSpeed upload ticket is missing its URLs") + if urlparse(upload_url).scheme != "https": + raise ProviderTaskError("WaveSpeed upload ticket has a non-https upload URL") + headers = upload.get("headers") if isinstance(upload, dict) else None + method = upload.get("method") if isinstance(upload, dict) else None + + global _wavespeed_upload_client + if _wavespeed_upload_client is None: + _wavespeed_upload_client = httpx.Client( + timeout=httpx.Timeout(timeout=120.0, connect=15.0), + follow_redirects=False, + ) + resp = _wavespeed_upload_client.request( + str(method or "PUT"), + upload_url, + content=raw, + headers={str(k): str(v) for k, v in (headers or {}).items()}, + ) + if not resp.is_success: + raise ProviderTaskError( + f"WaveSpeed reference upload failed with HTTP {resp.status_code}" + ) + return download_url + + +def _wavespeed_wait(client: httpx.Client, prediction_id: str) -> list[Any]: + """Poll a WaveSpeed prediction until it completes; return its outputs.""" + result_path = f"/predictions/{quote(prediction_id, safe='')}/result" + deadline = time.monotonic() + _WAVESPEED_DEADLINE_SECONDS + delay = _WAVESPEED_POLL_INITIAL_SECONDS + poll_errors = 0 + while True: + # Nothing is ready at submit time, so sleep before the first query too. + time.sleep(delay) + try: + data = _wavespeed_data(client.get(result_path)) + poll_errors = 0 + except (httpx.TransportError, httpx.HTTPStatusError) as e: + transient = isinstance(e, httpx.TransportError) or ( + e.response.status_code in _FETCH_RETRY_STATUS_CODES | {429} + ) + poll_errors += 1 + if not transient or poll_errors >= _WAVESPEED_MAX_POLL_ERRORS: + raise + logger.warning("WaveSpeed result query failed, retrying: %s", e) + data = {} + status = data.get("status") + if status == "completed": + outputs = data.get("outputs") + return outputs if isinstance(outputs, list) else [] + if status in _WAVESPEED_FAILED_STATUSES: + error = str(data.get("error") or "no error message") + raise ProviderTaskError( + f"WaveSpeed prediction {status}: {error[:_MAX_ERROR_DETAIL_CHARS]}" + ) + if time.monotonic() + delay > deadline: + raise httpx.TimeoutException( + "WaveSpeed prediction did not complete within " + f"{int(_WAVESPEED_DEADLINE_SECONDS)}s" + ) + delay = min(delay * 1.5, _WAVESPEED_POLL_MAX_SECONDS) + + +def _wavespeed_delete(client: httpx.Client, prediction_id: str) -> None: + """Best-effort delete of a finished prediction. + + WaveSpeed keeps a prediction's request (the prompt) and outputs in the + account's history; deleting it once the image is inside the enclave drops + both, leaving only the id and billing record. A still-running prediction + cannot be deleted and is skipped by WaveSpeed. Never fails the request. + """ + try: + _wavespeed_data( + client.post("/predictions/delete", json={"ids": [prediction_id]}) + ) + except Exception as e: + logger.warning("Could not delete WaveSpeed prediction: %s", e) + + +def _generate_wavespeed( + cfg: Any, + client: httpx.Client, + prompt: str, + refs: list[str], + aspect_ratio: Optional[str], + resolution: Optional[str], +) -> list[str]: + """Run one WaveSpeed image prediction and return its images as data URIs. + + Submitted to ``image_edit_model`` when there are references (WaveSpeed + serves text-to-image and edit as separate model endpoints), else to + ``api_name``. One image per prediction; ``n`` is not a WaveSpeed field. + With references and no ratio, the edit endpoint keeps the first + reference's shape. + """ + model_path = cfg.image_edit_model if refs and cfg.image_edit_model else cfg.api_name + payload: dict[str, Any] = {"prompt": prompt} + if cfg.image_extra_params: + payload.update(cfg.image_extra_params) + tier = cfg.image_tier(resolution) + if tier is not None: + payload.update(tier.extra_params) + payload.update(aspect_ratio_params(cfg, aspect_ratio, resolution=resolution)) + if refs: + payload["images"] = [ + _wavespeed_reference_url(client, ref, i) for i, ref in enumerate(refs) + ] + + prediction = _wavespeed_data(client.post(f"/{model_path}", json=payload)) + prediction_id = prediction.get("id") + if not isinstance(prediction_id, str) or not prediction_id: + raise ProviderTaskError("WaveSpeed did not return a prediction id") + try: + outputs = _wavespeed_wait(client, prediction_id) + images: list[str] = [] + for output in outputs: + if isinstance(output, str) and output.startswith(("https://", "http://")): + images.append(_fetch_url_as_data_uri(output)) + else: + logger.warning("Skipping non-URL WaveSpeed output") + return images + finally: + _wavespeed_delete(client, prediction_id) + + def generate_images( model: str, prompt: str, @@ -434,11 +657,9 @@ def generate_images( # n is clamped to the OpenAI-compatible providers' documented 1..10 range. count = max(1, min(int(n), 10)) - # Filter references to non-empty strings and cap at the providers' 10-image - # limit; only kept for models that accept reference images. - refs: Optional[list[str]] = None - if reference_images and cfg.image_supports_reference: - refs = [r for r in reference_images if isinstance(r, str) and r][:10] or None + # Only kept for models that accept reference images, capped at the model's + # own limit (10 unless the registry says otherwise). + refs: Optional[list[str]] = _forwarded_references(cfg, reference_images) or None logger.info( "Generating %d image(s) - Provider: %s, Model: %s", @@ -447,6 +668,12 @@ def generate_images( cfg.api_name, ) + if provider == "wavespeed": + task_images = _generate_wavespeed( + cfg, client, prompt, refs or [], aspect_ratio, resolution + ) + return task_images, len(task_images) + # Two delivery paths, picked by the model's config: multipart file uploads to # a dedicated edits endpoint, or the JSON generations endpoint (with inline # references in the ``image`` field when the provider carries them there). @@ -605,6 +832,11 @@ def _run_image_generation( usage, image_count=image_count, resolution=chat_request.resolution, + reference_count=len( + _forwarded_references( + get_model_config(chat_request.model), reference_images + ) + ), ) return { diff --git a/tee_gateway/llm_backend.py b/tee_gateway/llm_backend.py index c6c1470..6ea591f 100644 --- a/tee_gateway/llm_backend.py +++ b/tee_gateway/llm_backend.py @@ -62,6 +62,10 @@ # the subscription Coding Plan endpoint at /api/coding/paas/v4. ZAI_BASE_URL = "https://api.z.ai/api/paas/v4" +# WaveSpeed REST API. Image models are submitted to /{model_id} as async +# predictions and polled at /predictions/{id}/result (see image_generation). +WAVESPEED_BASE_URL = "https://api.wavespeed.ai/api/v3" + # Shared synchronous HTTP clients for each provider. # Initialized to None; built by set_provider_config() after key injection. openai_http_client: Optional[httpx.Client] = None @@ -69,6 +73,7 @@ bytedance_http_client: Optional[httpx.Client] = None openrouter_http_client: Optional[httpx.Client] = None zai_http_client: Optional[httpx.Client] = None +wavespeed_http_client: Optional[httpx.Client] = None _provider_config: Optional[ProviderConfig] = None @@ -77,13 +82,14 @@ def set_provider_config(config: ProviderConfig) -> None: """Store the provider config and rebuild HTTP clients. Called once after key injection.""" global _provider_config, openai_http_client, xai_http_client, bytedance_http_client - global openrouter_http_client, zai_http_client + global openrouter_http_client, zai_http_client, wavespeed_http_client old_openai = openai_http_client old_xai = xai_http_client old_bytedance = bytedance_http_client old_openrouter = openrouter_http_client old_zai = zai_http_client + old_wavespeed = wavespeed_http_client openai_http_client = httpx.Client( base_url="https://api.openai.com/v1", @@ -132,6 +138,14 @@ def set_provider_config(config: ProviderConfig) -> None: http2=True, follow_redirects=False, ) + wavespeed_http_client = httpx.Client( + base_url=WAVESPEED_BASE_URL, + headers={"Authorization": f"Bearer {config.wavespeed_api_key or ''}"}, + timeout=_TIMEOUT, + limits=_LIMITS, + http2=True, + follow_redirects=False, + ) # Web search runs inside the enclave against Exa rather than through any # provider's native tool, so its client is built here alongside them. @@ -154,6 +168,8 @@ def set_provider_config(config: ProviderConfig) -> None: old_openrouter.close() if old_zai is not None: old_zai.close() + if old_wavespeed is not None: + old_wavespeed.close() def get_provider_config() -> Optional[ProviderConfig]: @@ -235,14 +251,22 @@ def get_chat_model_cached( response_modalities=[Modality.TEXT, Modality.IMAGE], ) - return ChatGoogleGenerativeAI( + # Unlike the other langchain- packages, langchain-google-genai + # types `temperature` as a plain float (default 0.7) and rejects None at + # construction. Build without it, then clear it: the field is not + # re-validated on assignment, and a None temperature is dropped from the + # outgoing GenerationConfig, so the request carries no temperature. + google_llm = ChatGoogleGenerativeAI( model=api_name, google_api_key=config.google_api_key, - temperature=effective_temp, + **({} if effective_temp is None else {"temperature": effective_temp}), max_output_tokens=max_tokens, thinking_budget=cfg.thinking_budget, include_thoughts=False if cfg.thinking_budget is not None else None, ) + if effective_temp is None: + google_llm.temperature = None # type: ignore[assignment] + return google_llm elif provider == "openai": if not config.openai_api_key: @@ -337,7 +361,13 @@ def get_chat_model_cached( http_client=openrouter_http_client, api_key=SecretStr(config.openrouter_api_key), base_url=OPENROUTER_BASE_URL, - extra_body={"provider": {"sort": "price"}}, + extra_body={ + "provider": ( + {"order": [cfg.openrouter_provider], "allow_fallbacks": False} + if cfg.openrouter_provider + else {"sort": "price"} + ) + }, streaming=True, stream_usage=True, ) # type: ignore [call-arg] diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index 86d1a92..2eeb5b3 100644 --- a/tee_gateway/model_registry.py +++ b/tee_gateway/model_registry.py @@ -33,12 +33,17 @@ class ImageResolutionTier: # ModelArk the bare ``size`` keyword). Merged after ``image_extra_params`` # and before the ratio override, so an explicit ratio still wins. extra_params: Mapping[str, Any] = field(default_factory=dict) + # USD per reference image forwarded at this tier, for providers whose + # input-image charge scales with the output size (WaveSpeed bills Krea 3 + # references at the output's megapixels). ``None`` falls back to + # ``ModelConfig.per_reference_image_price_usd``. + per_reference_image_price_usd: Optional[Decimal] = None @dataclass(frozen=True) class ModelConfig: # "openai" | "anthropic" | "google" | "x-ai" | "bytedance" | - # "openrouter" | "zai" + # "openrouter" | "zai" | "wavespeed" provider: str api_name: str # model name sent to provider API input_price_usd: Decimal # USD per token @@ -84,6 +89,19 @@ class ModelConfig: # Used by OpenAI gpt-image, whose editing/compositing lives on a separate # endpoint from text-to-image generation. image_edit_endpoint: Optional[str] = None + # WaveSpeed splits a model into separate text-to-image and edit model + # endpoints (``…/text-to-image`` vs ``…/edit``). When set, a request that + # carries reference images is submitted to this model path instead of + # ``api_name``. WaveSpeed provider only. + image_edit_model: Optional[str] = None + # Most reference images the endpoint accepts. Providers reject a request + # over their cap rather than truncating it, so the gateway forwards at most + # this many (the first ones, in the order the client sent them). + image_max_references: int = 10 + # Flat USD price per reference image forwarded to the provider, for + # providers that bill input images on edits. Added on top of the per-output + # image price; ``None`` means references are not billed. + per_reference_image_price_usd: Optional[Decimal] = None # Static extra params merged verbatim into the request payload (e.g. size, # watermark). Keyed by field name; values must be JSON-serializable. image_extra_params: Optional[Mapping[str, Any]] = None @@ -122,6 +140,9 @@ class ModelConfig: # function tools together (preserving reasoning rather than disabling it). # OpenAI provider only; ignored elsewhere. responses_api_for_tools: bool = False + # OpenRouter provider slug to pin the model to (fallbacks off). ``None`` + # routes to the cheapest provider. + openrouter_provider: Optional[str] = None def image_tier(self, resolution: Optional[str]) -> Optional[ImageResolutionTier]: """The tier a request runs at, or ``None`` for a single-resolution model. @@ -251,6 +272,60 @@ def image_tier(self, resolution: Optional[str]) -> Optional[ImageResolutionTier] "21:9": "1792x768", } +# WaveSpeed's aspect_ratio enum for Qwen Image 3.0 Pro (text-to-image and edit +# share it). WaveSpeed maps ratio + resolution tier to pixels itself, so both +# tiers use the same table and the values pass through unchanged. +_WAVESPEED_QWEN_IMAGE_ASPECT_RATIOS: dict[str, str] = { + ratio: ratio + for ratio in ( + "1:1", + "1:2", + "2:1", + "1:3", + "3:1", + "2:3", + "3:2", + "3:4", + "4:3", + "4:5", + "5:4", + "9:16", + "16:9", + "9:21", + "21:9", + ) +} + +# WaveSpeed's Krea 3 takes explicit ``width*height`` sizes (multiples of 16) +# and bills by them: $0.03 per started megapixel (1,048,576 px) of output, and +# the same again for each reference image. Every size in a table stays at or +# under its tier's billed area — 1 MP for 1K, 4 MP for 2K — so the flat tier +# price is never below what WaveSpeed charges. The ratio set is Krea's own. +_WAVESPEED_KREA_1K_SIZES: dict[str, str] = { + "1:1": "1024*1024", + "4:3": "1152*864", + "3:4": "864*1152", + "3:2": "1248*832", + "2:3": "832*1248", + "16:9": "1360*768", + "9:16": "768*1360", + "5:4": "1136*912", + "4:5": "912*1136", + "21:9": "1536*656", +} +_WAVESPEED_KREA_2K_SIZES: dict[str, str] = { + "1:1": "2048*2048", + "4:3": "2304*1728", + "3:4": "1728*2304", + "3:2": "2496*1664", + "2:3": "1664*2496", + "16:9": "2720*1536", + "9:16": "1536*2720", + "5:4": "2272*1824", + "4:5": "1824*2272", + "21:9": "3072*1312", +} + _GEMINI_IMAGE_ASPECT_RATIOS: dict[str, str] = { ratio: ratio for ratio in ( @@ -417,6 +492,20 @@ class SupportedModel(Enum): responses_api_for_tools=True, supports_temperature=False, ) + # GPT-6.1 Terra (released 2026-10-26) replaces gpt-5.6-terra: $1/$5 per MTok + # (https://openai.com/index/introducing-gpt-6-1-terra/), under half the old + # $2.5/$15 the registry bills for it. gpt-5.6-terra stays registered so + # existing callers keep resolving. Parameter restrictions are not listed in + # the launch docs; it is a GPT-6.1 reasoning model like Sol, so it takes the + # same supports_temperature=False and Responses-API-for-tools settings. + GPT_6_1_TERRA = ModelConfig( + provider="openai", + api_name="gpt-6.1-terra", + input_price_usd=Decimal("0.000001"), + output_price_usd=Decimal("0.000005"), + responses_api_for_tools=True, + supports_temperature=False, + ) # GPT-6.1 Sol (released 2026-09-29) replaces gpt-6-sol after a week: same # $2/$10 per MTok, better coding/agentic scores. gpt-6-sol stays registered # so existing callers keep resolving. Docs published so far do not list @@ -561,6 +650,20 @@ class SupportedModel(Enum): output_price_usd=Decimal("0.000015"), supports_temperature=False, ) + # Claude Haiku 5.5 — successor to Haiku 4.5, released 2026-10-07. Priced by + # prompt length: $0.10/$0.50 per MTok for prompts up to 100k tokens, + # $0.50/$2.50 above (platform.claude.com/docs/en/about-claude/pricing). + # The registry has one rate per model, so this carries the <=100k tier; + # prompts over 100k tokens are under-billed by 5x until a length-tiered + # rate exists. Adaptive thinking; a non-default temperature/top_p/top_k + # returns HTTP 400, so supports_temperature=False. + CLAUDE_HAIKU_5_5 = ModelConfig( + provider="anthropic", + api_name="claude-haiku-5-5", + input_price_usd=Decimal("0.0000001"), + output_price_usd=Decimal("0.0000005"), + supports_temperature=False, + ) CLAUDE_HAIKU_4_5 = ModelConfig( provider="anthropic", api_name="claude-haiku-4-5-20251001", @@ -719,6 +822,20 @@ class SupportedModel(Enum): image_output_price_usd=Decimal("0.00006"), image_aspect_ratios=_GEMINI_IMAGE_ASPECT_RATIOS, ) + # Native image generation ("nano banana 2.1"), successor to Nano Banana 2, + # released 2026-10-06. Google bills output at two rates: text/thinking at + # $7.50/MTok and images at $30/MTok ($0.0336/0.0504/0.113 per 1K/2K/4K + # image); input (text/image) is $1.50/MTok + # (ai.google.dev/gemini-api/docs/pricing). The API id has no `-image` suffix. + GEMINI_NANO_BANANA_2_1 = ModelConfig( + provider="google", + api_name="gemini-nano-banana-2.1", + input_price_usd=Decimal("0.0000015"), + output_price_usd=Decimal("0.0000075"), + image_output=True, + image_output_price_usd=Decimal("0.00003"), + image_aspect_ratios=_GEMINI_IMAGE_ASPECT_RATIOS, + ) GEMINI_3_5_FLASH = ModelConfig( provider="google", api_name="gemini-3.5-flash", @@ -750,6 +867,21 @@ class SupportedModel(Enum): input_price_usd=Decimal("0.00000075"), output_price_usd=Decimal("0.00000375"), ) + # Gemini 4 Argon — Google's flagship (generally available on the paid + # Gemini API 2026-10-13; text/image/video/audio/PDF in, text out, 1M + # context). Priced at Google's standard $4/$20 per MTok, NOT the + # introductory $2/$10 that runs to 2026-11-30 (never price at promotional + # rates; ai.google.dev/gemini-api/docs/pricing). Google's + # model page lists no sampling-parameter support, so per the add-gateway-model + # rule for undocumented new models it is registered with + # supports_temperature=False rather than risk a 400 on every request. + GEMINI_4_ARGON = ModelConfig( + provider="google", + api_name="gemini-4-argon", + input_price_usd=Decimal("0.000004"), + output_price_usd=Decimal("0.00002"), + supports_temperature=False, + ) # ── xAI Grok ──────────────────────────────────────────────────────── # Grok 4.7 — xAI's newest flagship (released 2026-09-21), a larger base @@ -907,12 +1039,6 @@ class SupportedModel(Enum): input_price_usd=Decimal("0.00000014"), output_price_usd=Decimal("0.00000028"), ) - DEEPSEEK_V4_PRO = ModelConfig( - provider="bytedance", - api_name="deepseek-v4-pro-260425", - input_price_usd=Decimal("0.00000174"), - output_price_usd=Decimal("0.00000348"), - ) # Seedream 4.0 image generation via ModelArk's OpenAI-compatible # /images/generations endpoint. Billed at a flat $0.03 per generated image; # token prices unused. @@ -1004,6 +1130,14 @@ class SupportedModel(Enum): ) # ── OpenRouter (OpenAI-compatible) ────────────────────────────────── + # DeepSeek V4 Pro 0813, routed through OpenRouter, pinned to Phala. + DEEPSEEK_V4_PRO = ModelConfig( + provider="openrouter", + api_name="deepseek/deepseek-v4-pro-0813", + input_price_usd=Decimal("0.00000145"), + output_price_usd=Decimal("0.00000436"), + openrouter_provider="phala", + ) # Tencent Hy4 Preview — 770B-parameter (49B active) open-weight MoE, # open-sourced 2026-08-28, aimed at coding agents and long-horizon # tool-use with a 1,048,576-token context window. Routed through @@ -1071,6 +1205,88 @@ class SupportedModel(Enum): image_aspect_ratio_param="size", ) + # ── WaveSpeed (async prediction API) ──────────────────────────────── + # Krea 3 — Krea's image editing model, served through WaveSpeed as two + # model endpoints (text-to-image, and edit with 1-8 references). Every + # request runs the quality variant ("berry", not the faster "jam") at + # reasoning_level "medium" ("high" is not offered on WaveSpeed); neither + # choice changes the price. WaveSpeed's published billing formula is + # $0.03 x output megapixels x (1 + reference count), with an omitted + # `size` billed as 1024*1024, so references are priced per tier at the + # tier's output price: 1K $0.03 + $0.03/ref, 2K $0.12 + $0.12/ref. 1K + # omits `size` when no ratio is asked for (text-to-image defaults to + # 1024*1024; edits keep the first reference's framing); 2K has no such + # automatic mode, so without a ratio it is square. + KREA_3 = ModelConfig( + provider="wavespeed", + api_name="krea/v3/text-to-image", + input_price_usd=Decimal("0"), + output_price_usd=Decimal("0"), + image_generation=True, + per_image_price_usd=Decimal("0.03"), + image_response_format=None, + image_send_n=False, + image_supports_reference=True, + image_edit_model="krea/v3/edit", + image_max_references=8, + per_reference_image_price_usd=Decimal("0.03"), + image_extra_params={"model": "berry", "reasoning_level": "medium"}, + image_aspect_ratios=_WAVESPEED_KREA_1K_SIZES, + image_aspect_ratio_param="size", + image_resolutions={ + "1K": ImageResolutionTier( + per_image_price_usd=Decimal("0.03"), + aspect_ratios=_WAVESPEED_KREA_1K_SIZES, + per_reference_image_price_usd=Decimal("0.03"), + ), + "2K": ImageResolutionTier( + per_image_price_usd=Decimal("0.12"), + aspect_ratios=_WAVESPEED_KREA_2K_SIZES, + extra_params={"size": "2048*2048"}, + per_reference_image_price_usd=Decimal("0.12"), + ), + }, + image_default_resolution="1K", + ) + # Qwen Image 3.0 Pro — Alibaba's flagship image model, served through + # WaveSpeed rather than Alibaba Cloud Model Studio. WaveSpeed exposes it as + # two model endpoints, text-to-image and edit (1-3 reference images), both + # run as async predictions the gateway submits and polls (see + # image_generation._generate_wavespeed). One image per prediction, so `n` + # is not sent. Priced from WaveSpeed's model pages: $0.04 per 1k image, + # $0.075 per 2k image, plus $0.003 per input image on edits. WaveSpeed's + # `resolution` field ("1k" / "2k") is what it bills by, so it is exposed as + # resolution tiers; omitted means 1k. + QWEN_IMAGE_3_0_PRO = ModelConfig( + provider="wavespeed", + api_name="alibaba/qwen-image-3.0-pro/text-to-image", + input_price_usd=Decimal("0"), + output_price_usd=Decimal("0"), + image_generation=True, + per_image_price_usd=Decimal("0.04"), + image_response_format=None, + image_send_n=False, + image_supports_reference=True, + image_edit_model="alibaba/qwen-image-3.0-pro/edit", + image_max_references=3, + per_reference_image_price_usd=Decimal("0.003"), + image_extra_params={"enable_prompt_expansion": False}, + image_aspect_ratios=_WAVESPEED_QWEN_IMAGE_ASPECT_RATIOS, + image_resolutions={ + "1K": ImageResolutionTier( + per_image_price_usd=Decimal("0.04"), + aspect_ratios=_WAVESPEED_QWEN_IMAGE_ASPECT_RATIOS, + extra_params={"resolution": "1k"}, + ), + "2K": ImageResolutionTier( + per_image_price_usd=Decimal("0.075"), + aspect_ratios=_WAVESPEED_QWEN_IMAGE_ASPECT_RATIOS, + extra_params={"resolution": "2k"}, + ), + }, + image_default_resolution="1K", + ) + # ── Legacy models (not in current SDK — retained for older SDK versions) ── # grok-3 was retired on 2026-05-15 alongside the slugs above and is # redirected to grok-4.3 (reasoning effort "none"), billed at grok-4.3's @@ -1116,6 +1332,8 @@ class SupportedModel(Enum): "gpt-5.6-terra": SupportedModel.GPT_5_6_TERRA, "gpt-5.6-luna": SupportedModel.GPT_5_6_LUNA, "gpt-6-astra": SupportedModel.GPT_6_ASTRA, + "gpt-6.1-terra": SupportedModel.GPT_6_1_TERRA, + "gpt-6-1-terra": SupportedModel.GPT_6_1_TERRA, "gpt-6.1-sol": SupportedModel.GPT_6_1_SOL, "gpt-6-1-sol": SupportedModel.GPT_6_1_SOL, "gpt-6-sol": SupportedModel.GPT_6_SOL, @@ -1132,6 +1350,8 @@ class SupportedModel(Enum): "claude-sonnet-5-5": SupportedModel.CLAUDE_SONNET_5_5, "claude-sonnet-5.5": SupportedModel.CLAUDE_SONNET_5_5, "claude-sonnet-5": SupportedModel.CLAUDE_SONNET_5, + "claude-haiku-5-5": SupportedModel.CLAUDE_HAIKU_5_5, + "claude-haiku-5.5": SupportedModel.CLAUDE_HAIKU_5_5, "claude-haiku-4-5": SupportedModel.CLAUDE_HAIKU_4_5, "claude-opus-4-5": SupportedModel.CLAUDE_OPUS_4_5, "claude-opus-4-6": SupportedModel.CLAUDE_OPUS_4_6, @@ -1151,11 +1371,15 @@ class SupportedModel(Enum): "gemini-3.1-pro-preview": SupportedModel.GEMINI_3_1_PRO_PREVIEW, "gemini-2.5-flash-image": SupportedModel.GEMINI_2_5_FLASH_IMAGE, "gemini-3.1-flash-image": SupportedModel.GEMINI_3_1_FLASH_IMAGE, + "gemini-nano-banana-2.1": SupportedModel.GEMINI_NANO_BANANA_2_1, + "gemini-nano-banana-2-1": SupportedModel.GEMINI_NANO_BANANA_2_1, "gemini-3.5-flash": SupportedModel.GEMINI_3_5_FLASH, "gemini-3.5-flash-lite": SupportedModel.GEMINI_3_5_FLASH_LITE, "gemini-3.6-flash": SupportedModel.GEMINI_3_6_FLASH, "gemini-3.7-flash": SupportedModel.GEMINI_3_7_FLASH, "gemini-3.8-flash": SupportedModel.GEMINI_3_8_FLASH, + "gemini-4-argon": SupportedModel.GEMINI_4_ARGON, + "gemini-4.0-argon": SupportedModel.GEMINI_4_ARGON, # xAI "grok-4.7": SupportedModel.GROK_4_7, "grok-4.6": SupportedModel.GROK_4_6, @@ -1188,8 +1412,6 @@ class SupportedModel(Enum): "dola-seed-2-0-mini": SupportedModel.DOLA_SEED_2_0_MINI, "deepseek-v4-flash-260425": SupportedModel.DEEPSEEK_V4_FLASH, "deepseek-v4-flash": SupportedModel.DEEPSEEK_V4_FLASH, - "deepseek-v4-pro-260425": SupportedModel.DEEPSEEK_V4_PRO, - "deepseek-v4-pro": SupportedModel.DEEPSEEK_V4_PRO, "seedream-4-0-250828": SupportedModel.SEEDREAM_4_0, "seedream-4.0": SupportedModel.SEEDREAM_4_0, "seedream-4-0": SupportedModel.SEEDREAM_4_0, @@ -1201,7 +1423,12 @@ class SupportedModel(Enum): "ep-20260803211347-hq9k8": SupportedModel.SEEDANCE_5_0, "seedance-5.0": SupportedModel.SEEDANCE_5_0, "seedance-5-0": SupportedModel.SEEDANCE_5_0, - # Nous Research models routed through OpenRouter + # Routed through OpenRouter + "deepseek-v4-pro": SupportedModel.DEEPSEEK_V4_PRO, + "deepseek-v4-pro-0813": SupportedModel.DEEPSEEK_V4_PRO, + "deepseek/deepseek-v4-pro-0813": SupportedModel.DEEPSEEK_V4_PRO, + # Former ModelArk id + "deepseek-v4-pro-260425": SupportedModel.DEEPSEEK_V4_PRO, "hermes-4-405b": SupportedModel.HERMES_4_405B, "nousresearch/hermes-4-405b": SupportedModel.HERMES_4_405B, "hy3": SupportedModel.HY3, @@ -1214,6 +1441,17 @@ class SupportedModel(Enum): "glm-5.2": SupportedModel.GLM_5_2, "ep-20260803211658-fwpzs": SupportedModel.GLM_5_2, "glm-image": SupportedModel.GLM_IMAGE, + # WaveSpeed + "krea-3": SupportedModel.KREA_3, + "krea-v3": SupportedModel.KREA_3, + "krea3": SupportedModel.KREA_3, + "krea/v3": SupportedModel.KREA_3, + "krea/v3/text-to-image": SupportedModel.KREA_3, + "qwen-image-3.0-pro": SupportedModel.QWEN_IMAGE_3_0_PRO, + "qwen-image-3-0-pro": SupportedModel.QWEN_IMAGE_3_0_PRO, + "qwen-image-3-pro": SupportedModel.QWEN_IMAGE_3_0_PRO, + "alibaba/qwen-image-3.0-pro": SupportedModel.QWEN_IMAGE_3_0_PRO, + "alibaba/qwen-image-3.0-pro/text-to-image": SupportedModel.QWEN_IMAGE_3_0_PRO, # Legacy — not in current SDK, retained for older SDK versions "grok-3-mini-beta": SupportedModel.GROK_3_MINI, # old beta alias "grok-3-mini": SupportedModel.GROK_3_MINI, diff --git a/tee_gateway/pricing.py b/tee_gateway/pricing.py index 527e0ae..bf1b91b 100644 --- a/tee_gateway/pricing.py +++ b/tee_gateway/pricing.py @@ -50,7 +50,11 @@ def _serialize_decimal(self, value: Decimal) -> str: def compute_session_cost( - model: str, usage: dict, image_count: int = 0, resolution: str | None = None + model: str, + usage: dict, + image_count: int = 0, + resolution: str | None = None, + reference_count: int = 0, ) -> SessionCost | None: """Compute the settled cost for a completed inference request. @@ -110,6 +114,19 @@ def compute_session_cost( if images and per_image is not None else Decimal(0) ) + # Providers that bill input images on edits (WaveSpeed) charge per + # reference forwarded, on top of the output images. Only billed when + # the request produced images, like the output price itself. + # A tier may set its own reference price (WaveSpeed's Krea 3 bills a + # reference at the output's size). + references = max(0, int(reference_count)) + per_reference = ( + tier.per_reference_image_price_usd + if tier is not None and tier.per_reference_image_price_usd is not None + else cfg.per_reference_image_price_usd + ) + if images and references and per_reference is not None: + image_usd += Decimal(references) * per_reference raw_usd += image_usd token_price_usd = get_price_feed().get_price() diff --git a/tee_gateway/test/test_image_generation.py b/tee_gateway/test/test_image_generation.py index 59a6fb0..f21aac0 100644 --- a/tee_gateway/test/test_image_generation.py +++ b/tee_gateway/test/test_image_generation.py @@ -1,5 +1,6 @@ """Tests for endpoint-based image generation (OpenAI gpt-image, xAI Grok, -ByteDance Seedream, ByteDance Seedance, Z.ai GLM-Image). +ByteDance Seedream, ByteDance Seedance, Z.ai GLM-Image, Qwen Image via +WaveSpeed's async prediction API). Unlike Gemini's inline-image chat models (see test_image_billing.py), these models are served via a dedicated OpenAI-compatible ``/images/generations`` @@ -32,6 +33,8 @@ SEEDANCE_5 = "seedance-5.0" GLM_IMAGE = "glm-image" GPT_IMAGE = "gpt-image-2" +QWEN_IMAGE = "qwen-image-3.0-pro" +KREA_3 = "krea-3" def _mock_response(data: list[dict]) -> MagicMock: @@ -443,6 +446,292 @@ def test_uninitialized_client_raises(self): generate_images(GROK_IMAGE, "p", n=1) +_WS_BASE = "https://api.wavespeed.ai/api/v3" +_WS_PREDICTION = "pred_abc123" +_WS_RESULT_PATH = f"/predictions/{_WS_PREDICTION}/result" +_WS_OUTPUT = "https://cdn.wavespeed.ai/outputs/out.png" + + +def _ws_response(data: dict, *, status: int = 200, code: int = 200) -> httpx.Response: + return httpx.Response( + status, + json={"code": code, "message": "success", "data": data}, + request=httpx.Request("POST", _WS_BASE), + ) + + +def _ws_prediction(status: str, **extra) -> httpx.Response: + return _ws_response({"id": _WS_PREDICTION, "status": status, **extra}) + + +class _WaveSpeedClient: + """Fake WaveSpeed API client: records posts, replays result polls.""" + + def __init__(self, polls: list, submit: httpx.Response | None = None): + self.posts: list[tuple[str, dict]] = [] + self.gets: list[str] = [] + self._polls = list(polls) + self._submit = submit or _ws_prediction("created") + + def post(self, path, json=None): + self.posts.append((path, json)) + if path == "/media/uploads": + n = sum(1 for p, _ in self.posts if p == "/media/uploads") + return _ws_response( + { + "download_url": f"https://cdn.wavespeed.ai/media/ref{n}.png", + "upload": { + "method": "PUT", + "url": f"https://storage.example/signed/{n}", + "headers": {"Content-Type": json["content_type"]}, + }, + } + ) + if path == "/predictions/delete": + return _ws_response({"deleted_count": 1, "deleted_ids": json["ids"]}) + return self._submit + + def get(self, path): + self.gets.append(path) + item = self._polls.pop(0) + if isinstance(item, Exception): + raise item + return item + + def submitted(self) -> tuple[str, dict]: + return next( + (path, body) + for path, body in self.posts + if path not in ("/media/uploads", "/predictions/delete") + ) + + +class TestWaveSpeedImageGeneration(unittest.TestCase): + """Qwen Image 3.0 Pro through WaveSpeed: submit, poll, fetch, delete.""" + + def setUp(self): + sleep = patch.object(image_generation.time, "sleep") + self.sleep = sleep.start() + self.addCleanup(sleep.stop) + fetch = patch.object( + image_generation, + "_fetch_url_as_data_uri", + return_value="data:image/png;base64,RkVUQ0hFRA==", + ) + self.fetch = fetch.start() + self.addCleanup(fetch.stop) + self.uploader = MagicMock() + self.uploader.request.return_value = httpx.Response(200) + uploader = patch.object( + image_generation, "_wavespeed_upload_client", self.uploader + ) + uploader.start() + self.addCleanup(uploader.stop) + + def _run(self, client: _WaveSpeedClient, **kwargs): + with patch.object(llm_backend, "wavespeed_http_client", client): + return generate_images(QWEN_IMAGE, "a lighthouse at dusk", **kwargs) + + def test_text_to_image_submits_polls_fetches_and_deletes(self): + client = _WaveSpeedClient( + [ + _ws_prediction("processing"), + _ws_prediction("completed", outputs=[_WS_OUTPUT]), + ] + ) + images, count = self._run(client, n=4) + + self.assertEqual((images, count), (["data:image/png;base64,RkVUQ0hFRA=="], 1)) + path, body = client.submitted() + self.assertEqual(path, "/alibaba/qwen-image-3.0-pro/text-to-image") + # One image per prediction: WaveSpeed has no `n`, and the default + # (1k) tier is always sent so the billed tier is never left implicit. + # Prompt expansion is off on every request. + self.assertEqual( + body, + { + "prompt": "a lighthouse at dusk", + "resolution": "1k", + "enable_prompt_expansion": False, + }, + ) + self.assertEqual(client.gets, [_WS_RESULT_PATH, _WS_RESULT_PATH]) + self.fetch.assert_called_once_with(_WS_OUTPUT) + self.assertIn(("/predictions/delete", {"ids": [_WS_PREDICTION]}), client.posts) + # WaveSpeed asks for no more than one poll every 2 seconds. + for call in self.sleep.call_args_list: + self.assertGreaterEqual(call.args[0], 2.0) + + def test_resolution_and_aspect_ratio_pass_through(self): + client = _WaveSpeedClient([_ws_prediction("completed", outputs=[_WS_OUTPUT])]) + self._run(client, aspect_ratio="9:21", resolution="2K") + _, body = client.submitted() + self.assertEqual(body["resolution"], "2k") + self.assertEqual(body["aspect_ratio"], "9:21") + + def test_unsupported_ratio_and_resolution_are_rejected_before_submit(self): + client = _WaveSpeedClient([]) + with self.assertRaisesRegex(ValueError, "Unsupported aspect_ratio"): + self._run(client, aspect_ratio="7:3") + with self.assertRaisesRegex(ValueError, "Unsupported resolution"): + self._run(client, resolution="4K") + self.assertEqual(client.posts, []) + + def test_krea_3_sends_berry_at_medium_reasoning(self): + client = _WaveSpeedClient([_ws_prediction("completed", outputs=[_WS_OUTPUT])]) + with patch.object(llm_backend, "wavespeed_http_client", client): + generate_images(KREA_3, "a lighthouse at dusk") + path, body = client.submitted() + self.assertEqual(path, "/krea/v3/text-to-image") + # The 1K default sends no size, which WaveSpeed bills as 1024*1024. + self.assertEqual( + body, + { + "prompt": "a lighthouse at dusk", + "model": "berry", + "reasoning_level": "medium", + }, + ) + + def test_krea_3_ratio_and_tier_select_the_size(self): + cases = [ + ({"aspect_ratio": "16:9"}, "1360*768"), + ({"resolution": "2K"}, "2048*2048"), + ({"aspect_ratio": "16:9", "resolution": "2K"}, "2720*1536"), + ] + for kwargs, size in cases: + with self.subTest(**kwargs): + client = _WaveSpeedClient( + [_ws_prediction("completed", outputs=[_WS_OUTPUT])] + ) + with patch.object(llm_backend, "wavespeed_http_client", client): + generate_images(KREA_3, "p", **kwargs) + _, body = client.submitted() + self.assertEqual(body["size"], size) + self.assertNotIn("aspect_ratio", body) + + def test_krea_3_edit_uses_the_edit_model_and_keeps_framing(self): + client = _WaveSpeedClient([_ws_prediction("completed", outputs=[_WS_OUTPUT])]) + with patch.object(llm_backend, "wavespeed_http_client", client): + generate_images( + KREA_3, "p", reference_images=["https://example.com/photo.jpg"] + ) + path, body = client.submitted() + self.assertEqual(path, "/krea/v3/edit") + self.assertEqual(body["images"], ["https://example.com/photo.jpg"]) + # No size at 1K, so the edit keeps the reference's framing. + self.assertNotIn("size", body) + + def test_edit_uploads_inline_references_and_uses_the_edit_model(self): + client = _WaveSpeedClient([_ws_prediction("completed", outputs=[_WS_OUTPUT])]) + refs = [ + "data:image/png;base64,iVBORw0KGgo=", + "https://example.com/photo.jpg", + "data:image/jpeg;base64,/9j/4AAQ", + "https://example.com/fourth.jpg", # over WaveSpeed's 3-image cap + ] + self._run(client, reference_images=refs) + + path, body = client.submitted() + self.assertEqual(path, "/alibaba/qwen-image-3.0-pro/edit") + # Inline references go up through the upload ticket; a plain URL is + # handed to WaveSpeed as-is; the fourth is dropped. + self.assertEqual( + body["images"], + [ + "https://cdn.wavespeed.ai/media/ref1.png", + "https://example.com/photo.jpg", + "https://cdn.wavespeed.ai/media/ref2.png", + ], + ) + # No ratio picked: the edit keeps the first reference's shape. + self.assertNotIn("aspect_ratio", body) + # Edits carry the same fixed params as text-to-image. + self.assertIs(body["enable_prompt_expansion"], False) + uploads = [b for p, b in client.posts if p == "/media/uploads"] + self.assertEqual( + uploads, + [ + {"filename": "image_0.png", "size": 8, "content_type": "image/png"}, + {"filename": "image_2.jpg", "size": 6, "content_type": "image/jpeg"}, + ], + ) + # The bytes go to the signed URL on the keyless client, never with + # the API key. + method, url = self.uploader.request.call_args_list[0].args + kwargs = self.uploader.request.call_args_list[0].kwargs + self.assertEqual((method, url), ("PUT", "https://storage.example/signed/1")) + self.assertEqual(kwargs["headers"], {"Content-Type": "image/png"}) + self.assertNotIn("Authorization", kwargs["headers"]) + + def test_failed_upload_does_not_leak_the_signed_url(self): + self.uploader.request.return_value = httpx.Response(403) + client = _WaveSpeedClient([]) + with self.assertRaises(image_generation.ProviderTaskError) as ctx: + self._run(client, reference_images=["data:image/png;base64,iVBORw0KGgo="]) + self.assertNotIn("storage.example", str(ctx.exception)) + + def test_failed_prediction_is_a_provider_error_and_still_deleted(self): + from tee_gateway.errors import describe_exception, http_status_for + + client = _WaveSpeedClient( + [_ws_prediction("failed", error="content policy violation")] + ) + with self.assertRaisesRegex( + image_generation.ProviderTaskError, "failed: content policy violation" + ) as ctx: + self._run(client) + self.assertEqual(describe_exception(ctx.exception)["source"], "provider") + self.assertEqual(http_status_for(ctx.exception), 502) + self.assertIn(("/predictions/delete", {"ids": [_WS_PREDICTION]}), client.posts) + + def test_prediction_past_the_deadline_is_a_provider_timeout(self): + from tee_gateway.errors import http_status_for + + client = _WaveSpeedClient([_ws_prediction("processing")] * 3) + with patch.object(image_generation, "_WAVESPEED_DEADLINE_SECONDS", 0.0): + with self.assertRaises(httpx.TimeoutException) as ctx: + self._run(client) + self.assertEqual(http_status_for(ctx.exception), 504) + self.fetch.assert_not_called() + + def test_transient_poll_errors_are_retried(self): + client = _WaveSpeedClient( + [ + httpx.ConnectError("reset"), + _ws_prediction("processing"), + httpx.ReadTimeout("slow"), + _ws_prediction("completed", outputs=[_WS_OUTPUT]), + ] + ) + _, count = self._run(client) + self.assertEqual(count, 1) + + def test_persistent_poll_errors_give_up(self): + client = _WaveSpeedClient([httpx.ConnectError("reset")] * 3) + with self.assertRaises(httpx.ConnectError): + self._run(client) + self.assertEqual(len(client.gets), 3) + + def test_error_code_in_body_is_raised(self): + client = _WaveSpeedClient([], submit=_ws_response({}, code=400, status=200)) + with self.assertRaisesRegex(image_generation.ProviderTaskError, "error 400"): + self._run(client) + + def test_undecodable_inline_reference_is_not_forwarded_or_billed(self): + cfg = get_model_config(QWEN_IMAGE) + refs = ["data:image/png;base64,***", "https://example.com/a.jpg"] + self.assertEqual( + image_generation._forwarded_references(cfg, refs), + ["https://example.com/a.jpg"], + ) + + def test_uninitialized_client_raises(self): + with patch.object(llm_backend, "wavespeed_http_client", None): + with self.assertRaises(RuntimeError): + generate_images(QWEN_IMAGE, "p") + + class TestProviderErrorDetail(unittest.TestCase): """Provider error bodies must surface in the raised error message. @@ -1046,7 +1335,16 @@ def _zero_usage() -> dict: return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} def test_single_image_charged_flat_price(self): - for model in (GROK_IMAGE, SEEDREAM, SEEDANCE, SEEDANCE_5, GLM_IMAGE, GPT_IMAGE): + for model in ( + GROK_IMAGE, + SEEDREAM, + SEEDANCE, + SEEDANCE_5, + GLM_IMAGE, + GPT_IMAGE, + QWEN_IMAGE, + KREA_3, + ): with self.subTest(model=model): cfg = get_model_config(model) cost = compute_session_cost(model, self._zero_usage(), image_count=1) @@ -1078,6 +1376,52 @@ def test_default_tier_mirrors_the_flat_price_on_every_tiered_model(self): default = cfg.image_resolutions[cfg.image_default_resolution] self.assertEqual(default.per_image_price_usd, cfg.per_image_price_usd) self.assertEqual(default.aspect_ratios, cfg.image_aspect_ratios) + if default.per_reference_image_price_usd is not None: + self.assertEqual( + default.per_reference_image_price_usd, + cfg.per_reference_image_price_usd, + ) + + def test_qwen_image_tiers_and_reference_images_are_billed(self): + # WaveSpeed bills $0.04 / $0.075 per 1k / 2k output image, plus + # $0.003 per input image on an edit. + usage = self._zero_usage() + cases = [ + ({}, Decimal("0.04")), + ({"resolution": "2K"}, Decimal("0.075")), + ({"reference_count": 3}, Decimal("0.049")), + ({"resolution": "2K", "reference_count": 1}, Decimal("0.078")), + ] + for kwargs, expected in cases: + with self.subTest(**kwargs): + cost = compute_session_cost(QWEN_IMAGE, usage, image_count=1, **kwargs) + self.assertEqual(cost.cost_usd, expected) + + def test_krea_3_references_are_billed_at_the_tier_output_price(self): + # WaveSpeed: $0.03 x output MP x (1 + references); 1K = 1 MP, 2K = 4 MP. + usage = self._zero_usage() + cases = [ + ({}, Decimal("0.03")), + ({"reference_count": 1}, Decimal("0.06")), + ({"reference_count": 8}, Decimal("0.27")), + ({"resolution": "2K"}, Decimal("0.12")), + ({"resolution": "2K", "reference_count": 1}, Decimal("0.24")), + ] + for kwargs, expected in cases: + with self.subTest(**kwargs): + cost = compute_session_cost(KREA_3, usage, image_count=1, **kwargs) + self.assertEqual(cost.cost_usd, expected) + + def test_reference_images_free_on_models_without_a_reference_price(self): + usage = self._zero_usage() + cost = compute_session_cost(SEEDREAM, usage, image_count=1, reference_count=3) + self.assertEqual(cost.cost_usd, get_model_config(SEEDREAM).per_image_price_usd) + + def test_reference_images_not_billed_when_no_image_was_produced(self): + cost = compute_session_cost( + QWEN_IMAGE, self._zero_usage(), image_count=0, reference_count=3 + ) + self.assertEqual(cost.cost_opg, 0) def test_cost_scales_with_image_count(self): cfg = get_model_config(GROK_IMAGE) diff --git a/tee_gateway/test/test_provider_usage_integration.py b/tee_gateway/test/test_provider_usage_integration.py index e9f057b..144e5d6 100644 --- a/tee_gateway/test/test_provider_usage_integration.py +++ b/tee_gateway/test/test_provider_usage_integration.py @@ -76,6 +76,7 @@ class _ProviderCase: _ProviderCase("xAI", "grok-imagine-image", "XAI_API_KEY"), _ProviderCase("ByteDance", "seedream-5.0-lite", "ARK_API_KEY"), _ProviderCase("Z.ai", "glm-image", "ZAI_API_KEY"), + _ProviderCase("WaveSpeed", "qwen-image-3.0-pro", "WAVESPEED_API_KEY"), ) CHAT_MODELS = PROVIDER_SMOKE_MODELS + NEW_CHAT_MODELS @@ -258,6 +259,7 @@ def setUpClass(cls) -> None: bytedance_api_key=os.environ["ARK_API_KEY"], openrouter_api_key=os.environ["OPENROUTER_API_KEY"], zai_api_key=os.environ["ZAI_API_KEY"], + wavespeed_api_key=os.environ["WAVESPEED_API_KEY"], ) ) set_price_feed(_FixedPriceFeed()) # type: ignore[arg-type] diff --git a/tee_gateway/test/test_tee_core.py b/tee_gateway/test/test_tee_core.py index c4b693c..7745116 100644 --- a/tee_gateway/test/test_tee_core.py +++ b/tee_gateway/test/test_tee_core.py @@ -1017,6 +1017,14 @@ def test_opus_omits_temperature(self): model = get_chat_model_cached("claude-opus-5", 0.0, 1024) self.assertIsNone(model.temperature) + def test_gemini_4_argon_omits_temperature(self): + """langchain-google-genai rejects temperature=None at construction, so + the Google path has to clear it afterwards — and it must then be absent + from the generation config actually sent.""" + model = get_chat_model_cached("gemini-4-argon", 0.0, 1024) + self.assertIsNone(model.temperature) + self.assertIsNone(model._prepare_params(stop=None).temperature) + def test_temperature_supporting_model_keeps_it(self): model = get_chat_model_cached("gpt-4.1", 0.0, 1024) self.assertEqual(model.temperature, 0.0) @@ -1035,3 +1043,32 @@ def test_every_supports_temperature_false_model_builds_without_it(self): with self.subTest(model=cfg.api_name): model = get_chat_model_cached(cfg.api_name, 0.0, 1024) self.assertIsNone(model.temperature) + + +class TestOpenRouterProviderRouting(unittest.TestCase): + """OpenRouter models route by price unless the registry pins a provider.""" + + _saved_config: Optional[ProviderConfig] = None + + @classmethod + def setUpClass(cls) -> None: + cls._saved_config = get_provider_config() + set_provider_config(ProviderConfig(openrouter_api_key="test-openrouter-key")) + + @classmethod + def tearDownClass(cls) -> None: + if cls._saved_config is not None: + set_provider_config(cls._saved_config) + get_chat_model_cached.cache_clear() + + def test_deepseek_v4_pro_pinned_to_phala(self): + model = get_chat_model_cached("deepseek-v4-pro", 0.0, 1024) + self.assertEqual(model.model_name, "deepseek/deepseek-v4-pro-0813") + self.assertEqual( + model.extra_body, + {"provider": {"order": ["phala"], "allow_fallbacks": False}}, + ) + + def test_unpinned_model_sorts_by_price(self): + model = get_chat_model_cached("hy3", 0.0, 1024) + self.assertEqual(model.extra_body, {"provider": {"sort": "price"}}) diff --git a/tests/test_pricing.py b/tests/test_pricing.py index d5600ed..03f62db 100644 --- a/tests/test_pricing.py +++ b/tests/test_pricing.py @@ -90,6 +90,24 @@ def test_claude_sonnet_4_6_resolves(self): self.assertEqual(cfg.input_price_usd, Decimal("0.000003")) self.assertEqual(cfg.output_price_usd, Decimal("0.000015")) + def test_claude_haiku_5_5_resolves(self): + cfg = get_model_config("claude-haiku-5-5") + self.assertEqual(cfg.provider, "anthropic") + self.assertEqual(cfg.api_name, "claude-haiku-5-5") + self.assertEqual(cfg.input_price_usd, Decimal("0.0000001")) + self.assertEqual(cfg.output_price_usd, Decimal("0.0000005")) + # Non-default temperature/top_p/top_k return HTTP 400 + self.assertFalse(cfg.supports_temperature) + + def test_gemini_nano_banana_2_1_resolves(self): + cfg = get_model_config("gemini-nano-banana-2.1") + self.assertEqual(cfg.provider, "google") + self.assertEqual(cfg.api_name, "gemini-nano-banana-2.1") + self.assertEqual(cfg.input_price_usd, Decimal("0.0000015")) + self.assertEqual(cfg.output_price_usd, Decimal("0.0000075")) + self.assertTrue(cfg.image_output) + self.assertEqual(cfg.image_output_price_usd, Decimal("0.00003")) + def test_claude_sonnet_5_5_resolves(self): cfg = get_model_config("claude-sonnet-5-5") self.assertEqual(cfg.provider, "anthropic") @@ -289,6 +307,16 @@ def test_gpt_6_astra_resolves(self): self.assertEqual(cfg.output_price_usd, Decimal("0.00005")) self.assertTrue(cfg.responses_api_for_tools) + def test_gpt_6_1_terra_resolves(self): + for name in ("gpt-6.1-terra", "gpt-6-1-terra"): + cfg = get_model_config(name) + self.assertEqual(cfg.provider, "openai") + self.assertEqual(cfg.api_name, "gpt-6.1-terra") + self.assertEqual(cfg.input_price_usd, Decimal("0.000001")) + self.assertEqual(cfg.output_price_usd, Decimal("0.000005")) + self.assertTrue(cfg.responses_api_for_tools) + self.assertFalse(cfg.supports_temperature) + def test_gpt_6_1_sol_resolves(self): for name in ("gpt-6.1-sol", "gpt-6-1-sol"): cfg = get_model_config(name) @@ -395,6 +423,14 @@ def test_gemini_3_8_flash_resolves(self): self.assertEqual(cfg.input_price_usd, Decimal("0.00000075")) self.assertEqual(cfg.output_price_usd, Decimal("0.00000375")) + def test_gemini_4_argon_resolves(self): + cfg = get_model_config("gemini-4-argon") + self.assertEqual(cfg.provider, "google") + self.assertEqual(cfg.api_name, "gemini-4-argon") + self.assertEqual(cfg.input_price_usd, Decimal("0.000004")) + self.assertEqual(cfg.output_price_usd, Decimal("0.00002")) + self.assertFalse(cfg.supports_temperature) + def test_gemini_3_1_flash_image_resolves(self): cfg = get_model_config("gemini-3.1-flash-image") self.assertEqual(cfg.provider, "google") @@ -579,16 +615,22 @@ def test_deepseek_v4_flash_dated_alias_resolves(self): def test_deepseek_v4_pro_resolves(self): cfg = get_model_config("deepseek-v4-pro") - self.assertEqual(cfg.provider, "bytedance") - self.assertEqual(cfg.api_name, "deepseek-v4-pro-260425") - self.assertEqual(cfg.input_price_usd, Decimal("0.00000174")) - self.assertEqual(cfg.output_price_usd, Decimal("0.00000348")) + self.assertEqual(cfg.provider, "openrouter") + self.assertEqual(cfg.api_name, "deepseek/deepseek-v4-pro-0813") + self.assertEqual(cfg.openrouter_provider, "phala") + self.assertEqual(cfg.input_price_usd, Decimal("0.00000145")) + self.assertEqual(cfg.output_price_usd, Decimal("0.00000436")) def test_deepseek_v4_pro_aliases_resolve(self): - self.assertEqual( - get_model_config("deepseek-v4-pro-260425"), - get_model_config("deepseek-v4-pro"), - ) + for alias in ( + "deepseek-v4-pro-0813", + "deepseek/deepseek-v4-pro-0813", + "deepseek-v4-pro-260425", + ): + with self.subTest(alias=alias): + self.assertEqual( + get_model_config(alias), get_model_config("deepseek-v4-pro") + ) def test_seedream_5_0_lite_resolves(self): cfg = get_model_config("seedream-5.0-lite") @@ -696,6 +738,87 @@ def test_glm_image_resolves(self): self.assertTrue(cfg.image_generation) self.assertEqual(cfg.per_image_price_usd, Decimal("0.015")) + # ── WaveSpeed ──────────────────────────────────────────────────────────── + + def test_qwen_image_3_0_pro_resolves(self): + cfg = get_model_config("qwen-image-3.0-pro") + self.assertEqual(cfg.provider, "wavespeed") + self.assertEqual(cfg.api_name, "alibaba/qwen-image-3.0-pro/text-to-image") + self.assertEqual(cfg.image_edit_model, "alibaba/qwen-image-3.0-pro/edit") + self.assertTrue(cfg.image_generation) + # WaveSpeed's model pages: $0.04 per 1k image, $0.075 per 2k image, + # $0.003 per input image on edits (1-3 of them). + self.assertEqual(cfg.per_image_price_usd, Decimal("0.04")) + self.assertEqual(cfg.image_default_resolution, "1K") + self.assertEqual( + {k: t.per_image_price_usd for k, t in cfg.image_resolutions.items()}, + {"1K": Decimal("0.04"), "2K": Decimal("0.075")}, + ) + self.assertEqual(cfg.per_reference_image_price_usd, Decimal("0.003")) + self.assertEqual(cfg.image_max_references, 3) + + def test_qwen_image_3_0_pro_aliases_resolve(self): + for alias in ( + "qwen-image-3-0-pro", + "qwen-image-3-pro", + "alibaba/qwen-image-3.0-pro", + "alibaba/qwen-image-3.0-pro/text-to-image", + "Qwen-Image-3.0-Pro", + ): + with self.subTest(alias=alias): + self.assertEqual( + get_model_config(alias), get_model_config("qwen-image-3.0-pro") + ) + + def test_krea_3_resolves(self): + cfg = get_model_config("krea-3") + self.assertEqual(cfg.provider, "wavespeed") + self.assertEqual(cfg.api_name, "krea/v3/text-to-image") + self.assertEqual(cfg.image_edit_model, "krea/v3/edit") + self.assertTrue(cfg.image_generation) + # Quality variant at medium reasoning on every request; neither is billed. + self.assertEqual( + dict(cfg.image_extra_params), + {"model": "berry", "reasoning_level": "medium"}, + ) + # WaveSpeed's formula: $0.03 x output MP x (1 + references). + self.assertEqual(cfg.per_image_price_usd, Decimal("0.03")) + self.assertEqual(cfg.per_reference_image_price_usd, Decimal("0.03")) + self.assertEqual(cfg.image_default_resolution, "1K") + self.assertEqual( + { + k: (t.per_image_price_usd, t.per_reference_image_price_usd) + for k, t in cfg.image_resolutions.items() + }, + { + "1K": (Decimal("0.03"), Decimal("0.03")), + "2K": (Decimal("0.12"), Decimal("0.12")), + }, + ) + self.assertEqual(cfg.image_max_references, 8) + + def test_krea_3_aliases_resolve(self): + for alias in ( + "krea-v3", + "krea3", + "krea/v3", + "krea/v3/text-to-image", + "Krea-3", + ): + with self.subTest(alias=alias): + self.assertEqual(get_model_config(alias), get_model_config("krea-3")) + + def test_krea_3_sizes_stay_within_the_billed_tier_area(self): + # WaveSpeed rounds each side up to 16 and bills per 1,048,576 px, so a + # size over the tier's area would be billed above the tier price. + cfg = get_model_config("krea-3") + for tier, megapixels in (("1K", 1), ("2K", 4)): + for ratio, size in cfg.image_resolutions[tier].aspect_ratios.items(): + with self.subTest(tier=tier, ratio=ratio): + w, h = (int(v) for v in size.split("*")) + self.assertEqual((w % 16, h % 16), (0, 0)) + self.assertLessEqual(w * h, megapixels * 1_048_576) + # ── Errors ─────────────────────────────────────────────────────────────── def test_unknown_model_raises(self): @@ -822,6 +945,13 @@ def test_gpt_6_astra_cost(self): # 1000*0.00001 + 500*0.00005 = 0.01 + 0.025 = 0.035 USD = 3.5e16 wei self.assertEqual(cost, 35_000_000_000_000_000) + def test_gpt_6_1_terra_cost(self): + cost = self._calc("gpt-6.1-terra", 1000, 500) + expected = _expected_cost_opg("gpt-6.1-terra", 1000, 500) + self.assertEqual(cost, expected) + # 1000*0.000001 + 500*0.000005 = 0.001 + 0.0025 = 0.0035 USD = 3.5e15 wei + self.assertEqual(cost, 3_500_000_000_000_000) + def test_gpt_6_1_sol_cost(self): cost = self._calc("gpt-6.1-sol", 1000, 500) expected = _expected_cost_opg("gpt-6.1-sol", 1000, 500) @@ -893,6 +1023,13 @@ def test_claude_opus_5_cost(self): # Same price tier as opus-4-5/4-6/4-7/4-8: 1000*0.000005 + 500*0.000025 = 0.0175 USD self.assertEqual(cost, 17_500_000_000_000_000) + def test_claude_haiku_5_5_cost(self): + cost = self._calc("claude-haiku-5-5", 1000, 500) + expected = _expected_cost_opg("claude-haiku-5-5", 1000, 500) + self.assertEqual(cost, expected) + # 1000*0.0000001 + 500*0.0000005 = 0.0001 + 0.00025 = 0.00035 USD = 3.5e14 wei + self.assertEqual(cost, 350_000_000_000_000) + def test_claude_sonnet_5_5_cost(self): cost = self._calc("claude-sonnet-5-5", 1000, 500) expected = _expected_cost_opg("claude-sonnet-5-5", 1000, 500) @@ -969,6 +1106,11 @@ def test_gemini_3_8_flash_cost(self): self.assertEqual(cost, _expected_cost_opg("gemini-3.8-flash", 1000, 500)) self.assertEqual(cost, 2_625_000_000_000_000) + def test_gemini_4_argon_cost(self): + cost = self._calc("gemini-4-argon", 1000, 500) + self.assertEqual(cost, _expected_cost_opg("gemini-4-argon", 1000, 500)) + self.assertEqual(cost, 14_000_000_000_000_000) + # ── xAI Grok ──────────────────────────────────────────────────────────── def test_grok_4_7_cost(self): @@ -1080,8 +1222,8 @@ def test_deepseek_v4_pro_cost(self): cost = self._calc("deepseek-v4-pro", 1000, 500) expected = _expected_cost_opg("deepseek-v4-pro", 1000, 500) self.assertEqual(cost, expected) - # 1000*0.00000174 + 500*0.00000348 = 0.00174 + 0.00174 = 0.00348 USD - self.assertEqual(cost, 3_480_000_000_000_000) + # 1000*0.00000145 + 500*0.00000436 = 0.00145 + 0.00218 = 0.00363 USD + self.assertEqual(cost, 3_630_000_000_000_000) # ── OpenRouter ─────────────────────────────────────────────────────────