From 55ef161fb17e5f23dc21e7cec331a695e7d6ff31 Mon Sep 17 00:00:00 2001 From: Anass Date: Mon, 21 Sep 2026 18:30:13 +0200 Subject: [PATCH 1/3] Add a Workers AI client and price partner models from the AI Gateway cost table Cloudflare's partner models (typesafe/jev) are not chat-shaped, and no client library can address any Workers AI model, so there is nothing for wrap() to patch. sdk.workers_ai() is a one-method client: @cf/ ids go through the gateway host (cache-hit skip, cf_log_id dimension), partner ids through the unified /ai/run path with cf-aig-gateway-id, the only route that consults the gateway's BYOK key. New adapter extract_workers_ai_native, backed by eight captured fixtures. Partner models are priced from GET .../ai-gateway/costs, fetched one row per id on the queue's background tick, served stale-while-revalidate; warm_pricing(workers_ai_models=[...]) fetches them up front so even the first call prices. Cloudflare log entries now expose `byok` in extras, because a BYOK row carries a list-price `cost` Cloudflare never charged. Verified live against a local Lago: typesafe/jev through BYOK, 446 in / 73 out, and 0.000018732 USD equals the gateway's own log cost (new money_golden.json row, identical in the JS repo). --- CHANGELOG.md | 11 + CONTRIBUTING.md | 19 ++ README.md | 8 + docs/cloudflare.md | 40 +++ src/lago_agent_sdk/__init__.py | 3 + src/lago_agent_sdk/adapters/__init__.py | 2 + .../adapters/workers_ai_native.py | 123 ++++++++ .../gateway/adapters/cloudflare_gateway.py | 6 + src/lago_agent_sdk/pricing.py | 146 ++++++++- src/lago_agent_sdk/sdk.py | 56 +++- src/lago_agent_sdk/workers_ai.py | 186 ++++++++++++ .../adapters/fixtures/capture_workers_ai.py | 158 ++++++++++ .../fixtures/workers_ai/01_chat_direct.json | 55 ++++ .../workers_ai/02_reasoning_direct.json | 24 ++ .../workers_ai/03_gpt_oss_direct.json | 54 ++++ .../workers_ai/04_mistral_small_direct.json | 53 ++++ .../workers_ai/05_jev_byok_gateway.json | 56 ++++ .../workers_ai/06_chat_gateway_miss.json | 56 ++++ .../workers_ai/07_chat_gateway_hit.json | 56 ++++ .../workers_ai/08_jev_typesafe_direct.json | 44 +++ tests/unit/adapters/test_workers_ai_native.py | 165 ++++++++++ tests/unit/fixtures/pricing/money_golden.json | 20 +- .../adapters/test_cloudflare_gateway.py | 12 +- tests/unit/gateway/test_ramp_router.py | 4 + tests/unit/test_auto_prime_pricing.py | 3 + tests/unit/test_drift.py | 54 ++++ tests/unit/test_pricing.py | 238 +++++++++++++++ tests/unit/test_workers_ai_client.py | 284 ++++++++++++++++++ 28 files changed, 1928 insertions(+), 8 deletions(-) create mode 100644 src/lago_agent_sdk/adapters/workers_ai_native.py create mode 100644 src/lago_agent_sdk/workers_ai.py create mode 100644 tests/unit/adapters/fixtures/capture_workers_ai.py create mode 100644 tests/unit/adapters/fixtures/workers_ai/01_chat_direct.json create mode 100644 tests/unit/adapters/fixtures/workers_ai/02_reasoning_direct.json create mode 100644 tests/unit/adapters/fixtures/workers_ai/03_gpt_oss_direct.json create mode 100644 tests/unit/adapters/fixtures/workers_ai/04_mistral_small_direct.json create mode 100644 tests/unit/adapters/fixtures/workers_ai/05_jev_byok_gateway.json create mode 100644 tests/unit/adapters/fixtures/workers_ai/06_chat_gateway_miss.json create mode 100644 tests/unit/adapters/fixtures/workers_ai/07_chat_gateway_hit.json create mode 100644 tests/unit/adapters/fixtures/workers_ai/08_jev_typesafe_direct.json create mode 100644 tests/unit/adapters/test_workers_ai_native.py create mode 100644 tests/unit/test_workers_ai_client.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 646a24a..ad31ed2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,17 @@ All notable changes to this project will be documented here. Format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow [SemVer](https://semver.org). +## [Unreleased] + +### Added + +- **Workers AI client: `sdk.workers_ai(account_id, api_token, gateway_id=..., gateway_auth=...)`.** The one provider with no client to wrap: the official `cloudflare` package percent-encodes the slash in every model name and "No route for that URI" is the answer for all of them, and partner models such as `typesafe/jev` reject the `messages` array the gateway's `/compat` endpoint requires. The SDK now ships a one-method client, `run(model, input, extra_lago=...)`, that reaches every Workers AI model and bills the response's own `usage`. Verified live on 2026-09-21: `typesafe/jev` answered through Cloudflare under BYOK and its 446 input / 73 output tokens landed in Lago. + - **Two routes, chosen per model id — measured, not chosen.** `@cf/...` ids go to the gateway host's path route, which answers with `cf-aig-cache-status` (a HIT is not billed) and `cf-aig-log-id` (emitted as the `cf_log_id` dimension for reconciliation against the Logs API); its model-in-body variant logs `model: "run"`, so it is not used for them. Partner ids (`typesafe/jev`) go to the unified `api.cloudflare.com/.../ai/run` path with the gateway named in `cf-aig-gateway-id` — the only route where the partner key stored under the gateway's BYOK is consulted; the gateway host answers 402 "Insufficient balance" for the same call even with the key stored. That path returns no `cf-aig-*` headers, so a cached replay is indistinguishable from a fresh call and the client sends `cf-aig-skip-cache: true` there rather than bill the same answer twice. + - **New adapter `extract_workers_ai_native`** for the two usage vocabularies the one endpoint returns: OpenAI-shaped `prompt_tokens`/`completion_tokens` plus Cloudflare's `neurons` (kept in `extras`, never a metric) for catalog models, and `input_tokens`/`output_tokens` one level deeper, under `result.result`, for partner models. Eight captured fixtures, all successful responses, including a real gateway cache MISS/HIT pair and Jev both through Cloudflare and straight from TypeSafe's API. + - **Partner models are priced from AI Gateway's own cost table.** `GET .../ai-gateway/costs` is the rate the gateway itself bills by — 2,839 rows across every provider it fronts, `token_pricing` in USD per million — and it lists the partner models the Workers AI catalog omits. A bare `vendor/model` id that misses the catalog is fetched from it one row at a time on the queue's next tick, so the first call to such a model in a process bills tokens and reports the miss, and every later call bills dollars; an id the gateway does not price is remembered as a miss for the TTL rather than re-queried every tick. `warm_pricing(workers_ai_models=[...])` fetches named partner rows up front so even that first call prices, and a fetched row keeps serving past its TTL while it refreshes. Verified on `typesafe/jev`: 446 input tokens at $0.042 per million is 0.000018732 USD, the `cost` Cloudflare stamped on the live log entry, now a `money_golden.json` row in both repos. The row's `cost_in`/`cost_out` per-token fields are ignored — they are 0 on rows whose `token_pricing` is not. + - **The requested id is the billing key for Workers AI, not the served name.** Cloudflare's price catalog is keyed by the id you request, and the served name drifts past the catalog's version-strip fallback: `@cf/mistralai/mistral-small-3.1-24b-instruct` answers as `...-24b-v2`, which misses the catalog while the requested id prices. The served name is kept in `extras["served_model"]`. + - **Cloudflare gateway log entries now surface `byok` in `extras`.** A row served under BYOK still carries Cloudflare's list-price `cost` (measured: 446 in × $0.042/M on `typesafe/jev`) although Cloudflare charged nothing — the partner bills the customer directly. A backfill that bills `cost` for such a row double-charges; this is the field that lets it not. + ## [0.3.1] - 2026-09-07 ### Added diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 549a787..f1991c3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -82,6 +82,25 @@ uv lock --upgrade-package X # bump a single package 5. Update `sdk.py::wrap()` to dispatch to the new wrapper. 6. Add unit tests against the captured fixtures. +## Adding a first-party client + +Some models have no client library to wrap — Cloudflare Workers AI is the case that forced this: +the official `cloudflare` package cannot address any Workers AI model, and partner models such as +`typesafe/jev` are not chat-shaped, so the OpenAI-compatible route cannot carry them either. For +these the SDK ships its own minimal client (`src/lago_agent_sdk/workers_ai.py`, built by +`LagoSDK.workers_ai()`). The rules are the wrappers' rules, applied to code we own: + +1. Capture real fixtures exactly as for a provider — successful responses only, one file per + shape, with a capture script alongside them. +2. Keep the adapter a pure function in `adapters/`; the client only does HTTP and calls `emit()`. +3. Raise the provider's own error before any instrumentation, so a failed call bills nothing and + reaches the caller unchanged. Wrap everything after the response in the same catch-all the + wrappers use — instrumentation never breaks the customer's call. +4. Write down, in the module docstring, every routing decision that was *measured* rather than + chosen (which route logs the model, which header the gateway honours, what a cache hit looks + like). A first-party client has no upstream to blame for its choices. +5. Mirror it in the JS repo in the same PR pair, tests and fixtures included. + ## Adding a gateway `gateway/` is a **second front door** into the same kernel, separate from the provider-native diff --git a/README.md b/README.md index df6b0f8..85236a4 100644 --- a/README.md +++ b/README.md @@ -148,6 +148,14 @@ client = sdk.wrap(Anthropic( )) ``` +Models with no client to wrap — every Workers AI model, including partner models like `typesafe/jev` that are not chat-shaped — go through the SDK's own one-method client. Store the partner's key under the gateway's **Provider Keys** (BYOK) and name the gateway; Cloudflare then bills nothing and the partner bills you: + +```python +ai = sdk.workers_ai(account_id, cf_api_token, gateway_id=gateway_id, gateway_auth=gateway_auth) +out = ai.run("typesafe/jev", {"state": ticket_text, "questions": {...}}, extra_lago={"subscription": "sub_acme"}) +out = ai.run("@cf/meta/llama-3.2-3b-instruct", {"messages": [{"role": "user", "content": "Hello"}]}) +``` + Full guide, including backfill from the gateway's Logs API: [docs/cloudflare.md](docs/cloudflare.md). ### Databricks AI Gateway diff --git a/docs/cloudflare.md b/docs/cloudflare.md index 21095b5..a812a4c 100644 --- a/docs/cloudflare.md +++ b/docs/cloudflare.md @@ -38,3 +38,43 @@ This page is the complete picture; runnable notebooks are kept out of the repo ( **Gateway-routed calls are billed at the gateway's metered cost.** Cloudflare reports its own `cost` per log entry and the backfill passes that straight through, so Lago reconciles against the dashboard you actually look at. One measured consequence to be aware of: that field excludes additive *reasoning* tokens, so a thinking-heavy Gemini call bills about 4% of what Google charges (verified live at 22.8x on one call, 39.6x on another — the ratio tracks each prompt's thinking-to-output ratio). Cloudflare is exact on input, output, cache-read and cache-write. **If you hand-roll a poller, don't use `urllib`.** `gateway.ai.cloudflare.com` returns `403` with body `error code: 1010` to `Python-urllib` — its bot-signature check. Any other User-Agent passes, and `requests` (which this SDK uses) is fine. The failure looks like an auth error because the body is otherwise empty. + +## Workers AI models with no client to wrap + +Cloudflare-hosted models can be reached through the OpenAI-compatible `/compat` endpoint above, but only when they are chat-shaped. Partner models are not: `typesafe/jev` takes `{state, questions}` and refuses a `messages` array, and the official `cloudflare` package cannot address any Workers AI model at all (it percent-encodes the slash in the model name). For these the SDK ships its own one-method client: + +```python +ai = sdk.workers_ai( + account_id, + cf_api_token, # a Cloudflare API token with Workers AI access + gateway_id=gateway_id, # optional — see what it buys below + gateway_auth=gateway_auth, # the gateway's cf-aig-authorization token, if authentication is on + subscription="sub_acme", # default for every call; extra_lago={"subscription": ...} per call +) + +# a partner model: the partner's key must be stored under the gateway's Provider Keys (BYOK) +out = ai.run( + "typesafe/jev", + { + "state": "I was charged twice and need the duplicate refunded before Friday.", + "questions": { + "department": {"type": "choice", "instructions": "Which team should handle this?", + "criteria": {"billing": "Payments, refunds", "technical": "Bugs, outages"}}, + }, + }, + extra_lago={"dimensions": {"ticket": "T-4821"}}, +) +out["result"]["result"]["answers"]["department"]["choice"] # "billing" + +# a catalog model: same client, same billing +ai.run("@cf/meta/llama-3.2-3b-instruct", {"messages": [{"role": "user", "content": "Hello"}], "max_tokens": 50}) +sdk.flush() +``` + +`run()` returns Cloudflare's full response envelope unchanged and raises `WorkersAIError` on an error status (a 402 for a partner model whose key is not stored, a 403 for a model not on your Workers plan), before anything is billed. + +**What the gateway buys.** With `gateway_id` set, `@cf/...` calls go through the gateway host: a cache hit (`cf-aig-cache-status: HIT`) is not billed, and every event carries a `cf_log_id` dimension that matches the entry's `id` in the Logs API. Partner models go through the unified `api.cloudflare.com/.../ai/run` path with the gateway named in a header, which is the only route where the gateway's stored partner key is consulted; that path returns no cache header, so the client asks the gateway to skip its cache for those calls rather than risk billing a cached replay. Pass `extra_headers={"cf-aig-skip-cache": "false"}` to opt back in. + +**Billing.** Token events from the response's own `usage`: `prompt_tokens`/`completion_tokens` for catalog models, `input_tokens`/`output_tokens` for partner models. Catalog models price from Cloudflare's published Workers AI rates in price mode as usual. Partner models are not in that catalog; their rates come from AI Gateway's own cost table (`GET .../ai-gateway/costs`), fetched per model in the background the first time one is seen. Left alone, the very first call to a partner model in a process bills tokens and reports the miss via `on_error`, and every call after it bills dollars; name the ids up front with `sdk.warm_pricing(["workers-ai"], workers_ai_models=["typesafe/jev"])` and even the first call prices. A fetched rate keeps serving past its TTL while it refreshes, so an expiry never bills a call as tokens. Jev lists input at $0.042 per million with free output, the same rate the gateway stamps as `cost` on its log entries. Under BYOK the gateway's Logs API still reports a `cost` for the partner call at Cloudflare's list price although Cloudflare charged nothing — the entry's `byok` field (surfaced in `extras["byok"]` by `extract_cloudflare_log`) is what tells a backfill not to bill it. The id billed is the one you requested, because that is what Cloudflare's price catalog is keyed by; the name the model reports (`jev-1.13.0`, `...-24b-v2`) is kept in `extras["served_model"]`. + +Streaming is not supported by this client; use the `/compat` endpoint through a wrapped OpenAI client for streamed chat. diff --git a/src/lago_agent_sdk/__init__.py b/src/lago_agent_sdk/__init__.py index 2205f3c..31e7d11 100644 --- a/src/lago_agent_sdk/__init__.py +++ b/src/lago_agent_sdk/__init__.py @@ -17,9 +17,12 @@ compute_cost, ) from .sdk import LagoSDK +from .workers_ai import WorkersAI, WorkersAIError __all__ = [ "LagoSDK", + "WorkersAI", + "WorkersAIError", "LagoConfig", "CanonicalUsage", "LagoApiError", diff --git a/src/lago_agent_sdk/adapters/__init__.py b/src/lago_agent_sdk/adapters/__init__.py index 88d3a71..266b71a 100644 --- a/src/lago_agent_sdk/adapters/__init__.py +++ b/src/lago_agent_sdk/adapters/__init__.py @@ -4,6 +4,7 @@ from .gemini_native import extract_gemini_native from .mistral_native import extract_mistral_native from .openai_native import extract_openai_native +from .workers_ai_native import extract_workers_ai_native __all__ = [ "extract_anthropic_native", @@ -13,4 +14,5 @@ "extract_gemini_native", "extract_mistral_native", "extract_openai_native", + "extract_workers_ai_native", ] diff --git a/src/lago_agent_sdk/adapters/workers_ai_native.py b/src/lago_agent_sdk/adapters/workers_ai_native.py new file mode 100644 index 0000000..49e4bcf --- /dev/null +++ b/src/lago_agent_sdk/adapters/workers_ai_native.py @@ -0,0 +1,123 @@ +"""Workers AI `/ai/run` adapter — maps a run response to CanonicalUsage. + +Verified against real captures (fixtures/workers_ai/, 2026-09-21) of the model-in-body +route `POST /accounts/{id}/ai/run {"model": ..., "input": {...}}`. That route matters +because it is the only one that reaches every Workers AI model: partner models such as +`typesafe/jev` have no `@cf/` prefix and the path-style `/ai/run/{model}` answers "No +route for that URI" for them, as does the gateway's `/compat` endpoint, which requires +a `messages` array the model rejects. + +Two usage vocabularies come back from the one endpoint: + + chat models result.usage.prompt_tokens / completion_tokens / total_tokens + result.usage.prompt_tokens_details.cached_tokens + result.usage.neurons (01-04, 06, 07) + typesafe/jev result.result.usage.input_tokens / output_tokens (05: one level + deeper — a partner model's answer is wrapped as + `result: {state, result: {model, answers, usage}, gatewayMetadata}`; + 08 is the same object straight from TypeSafe's API, unwrapped) + +The REQUESTED model id is the one carried, not the served name — the opposite of the +native adapters' rule, for a measured reason. Cloudflare's price catalog is keyed by the +id you request, and the served name drifts from it in ways the catalog's version-strip +fallback does not cover: `@cf/mistralai/mistral-small-3.1-24b-instruct` answers as +`...-24b-v2` (04) and that name MISSES the catalog while the requested id prices; +`typesafe/jev` answers as `jev-1.13.0` (05, 08). The served name is kept in +`extras["served_model"]` when it differs, so nothing is lost — only the billing key stays +the one Cloudflare itself bills by. + +`neurons` is Cloudflare's own billing unit, not a token count: it lands in `extras` +and is never a metric. `gatewayMetadata` (`keySource: "BYOK"` on 05) says whose key paid +for the call and lands in `extras` too — under BYOK Cloudflare charged nothing and the +partner bills the customer directly, which matters to anyone reconciling against the +Cloudflare dashboard. A reasoning model (02, deepseek-r1-distill) bundles its thinking +into `completion_tokens` with no separate field, so `reasoning` stays 0 — the same +shape Magistral has on Mistral. `cached_tokens` is a subset of `prompt_tokens`, which +is why "workers-ai" sits in `INPUT_INCLUDES_CACHE_READ`. + +A failure body (402 no credits / 403 not on plan / 400 no such model — seen live, not +kept as fixtures) carries `result: {}` and `errors: [...]`. The adapter yields an all-zero +usage for it rather than raising — it is a pure function and cannot know the HTTP status; +the client decides what a failure means. +""" + +from __future__ import annotations + +from typing import Any + +from ..canonical import CanonicalUsage + +# Every `usage` key this adapter maps or deliberately ignores. Anything else is drift +# and is swept into `extras["usage"]` — never silently dropped, never miscounted. +_KNOWN_USAGE_KEYS = frozenset( + { + # chat models + "prompt_tokens", + "completion_tokens", + "total_tokens", + "prompt_tokens_details", + # typesafe/jev + "input_tokens", + "output_tokens", + # Cloudflare's billing unit — kept in extras, not a metric + "neurons", + } +) +_KNOWN_DETAIL_KEYS = frozenset({"cached_tokens"}) + + +def _safe_dict(v: Any) -> dict[str, Any]: + return v if isinstance(v, dict) else {} + + +def _safe_int(v: Any) -> int: + try: + return max(0, int(v or 0)) + except (TypeError, ValueError): + return 0 + + +def extract_workers_ai_native(response: Any, model_id: str = "") -> CanonicalUsage: + """Translate a Workers AI `/ai/run` response body → CanonicalUsage. + + Accepts the full envelope (`{"result": {...}, "success": true, ...}`) or the bare + `result` object. `model_id` is the model the caller requested and is the id carried + (see the module docstring); the served name, when it differs, lands in + `extras["served_model"]`. + """ + payload = _safe_dict(response) + result = _safe_dict(payload.get("result")) or payload + gateway_meta = _safe_dict(result.get("gatewayMetadata")) + inner = _safe_dict(result.get("result")) + if inner and ("usage" in inner or "answers" in inner): + # Partner-model envelope (05): the model's own object sits one level down. + result = inner + # The envelope may also carry `usage` at the top level; check both so a shape change + # moves nothing to zero. + usage = _safe_dict(result.get("usage")) or _safe_dict(payload.get("usage")) + details = _safe_dict(usage.get("prompt_tokens_details")) + + extras: dict[str, Any] = {} + if gateway_meta: + extras["gateway_metadata"] = gateway_meta + served = result.get("model") + if isinstance(served, str) and served and served != model_id: + extras["served_model"] = served + if "neurons" in usage: + extras["neurons"] = usage["neurons"] + drift = {k: v for k, v in usage.items() if k not in _KNOWN_USAGE_KEYS} + detail_drift = {k: v for k, v in details.items() if k not in _KNOWN_DETAIL_KEYS} + if detail_drift: + drift["prompt_tokens_details"] = detail_drift + if drift: + extras["usage"] = drift + + return CanonicalUsage( + input=_safe_int(usage.get("prompt_tokens")) or _safe_int(usage.get("input_tokens")), + output=_safe_int(usage.get("completion_tokens")) or _safe_int(usage.get("output_tokens")), + cache_read=_safe_int(details.get("cached_tokens")), + model=model_id or (served if isinstance(served, str) else ""), + provider="workers-ai", + api="workers_ai_run", + extras=extras, + ) diff --git a/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py b/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py index 0a0236a..7b6c4b0 100644 --- a/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py +++ b/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py @@ -220,6 +220,12 @@ def extract_cloudflare_log(entry: dict[str, Any]) -> CanonicalUsage: "cached": entry.get("cached"), "step": entry.get("step"), "log_id": entry.get("id"), + # Which key paid: None (Cloudflare credits / the customer's own header key) or the + # BYOK alias (`"default"`). Measured 2026-09-21 on a `typesafe/jev` row served under + # BYOK: Cloudflare still fills `cost` with its list price (446 in × $0.042/M) although + # it charged nothing — the partner bills the customer directly. A backfill that bills + # `cost` for such a row double-charges; this is the field that lets it not. + "byok": entry.get("byok"), # Drift sweep — the same contract `adapters/openai_native.py` enforces, and # for the same reason: a counter this adapter does not map must not vanish # without an error or an on_error. `extras` used to be exactly the three diff --git a/src/lago_agent_sdk/pricing.py b/src/lago_agent_sdk/pricing.py index cf7ecc7..37e32ef 100644 --- a/src/lago_agent_sdk/pricing.py +++ b/src/lago_agent_sdk/pricing.py @@ -71,6 +71,15 @@ AWS_PRICING_HOST = "https://pricing.us-east-1.amazonaws.com" AWS_BEDROCK_REGION_INDEX = f"{AWS_PRICING_HOST}/offers/v1.0/aws/AmazonBedrock/current/region_index.json" CLOUDFLARE_MODELS_URL_TEMPLATE = "https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/models/search" +# AI Gateway's own price list — every provider the gateway fronts, including the partner +# models Workers AI serves under a bare `vendor/model` id (`typesafe/jev`), which the +# `/ai/models/search` catalog above does not list at all. Measured 2026-09-21: 2,839 rows, +# `per_page` capped at 100, `search=` filters by model id; the `typesafe/jev` row is +# `token_pricing: {input_tokens: 0.042, input_cached_tokens: 0, output_tokens: 0}` (USD per +# 1M) and 446 x 0.042e-6 is exactly the `cost` the gateway stamped on that call's log entry. +CLOUDFLARE_GATEWAY_COSTS_URL_TEMPLATE = ( + "https://api.cloudflare.com/client/v4/accounts/{account_id}/ai-gateway/costs" +) MISTRAL_MODELS_URL = "https://api.mistral.ai/v1/models" RAMP_ROUTER_MODELS_URL = "https://api.router.com/v1/models" @@ -188,6 +197,16 @@ # catalog of 64. _CF_PER_PAGE = 50 _CF_MAX_PAGES = 40 +# `token_pricing` keys on an `ai-gateway/costs` row → ModelPrice fields (USD per 1M tokens). +# The sibling `cost_in` / `cost_out` per-token fields are NOT used: on the same row they are +# frequently 0 where `token_pricing` is not (typesafe/jev, every Fireworks entry), and the +# gateway's own `cost` on a log entry reconciles against `token_pricing`, not against them. +_CF_COSTS_FIELD_MAP = { + "input_tokens": "input", + "output_tokens": "output", + "input_cached_tokens": "cache_read", + "input_cache_creation_tokens": "cache_write", +} # Bedrock cross-region inference prefix -> a representative AWS region. _BEDROCK_REGION_PREFIX = { @@ -846,6 +865,52 @@ def lookup_cloudflare_workers_ai(table: dict[str, ModelPrice], model: str) -> Mo return None +def parse_cloudflare_gateway_cost(rows: Any, model: str) -> ModelPrice | None: + """One `ai-gateway/costs?search=` response → the price of exactly `model`, or None. + + Only rows whose `model` equals the requested id and whose `cost_type` is `tokens` + count. The same id can appear under several providers at DIFFERENT rates + (`stealth/union-alpha` is listed by openrouter, unbiased and stealth), so when more + than one row matches, the row whose `provider` is the id's own namespace + (`typesafe/jev` → `typesafe`) wins; failing that, rows that all agree are one price + and rows that disagree are a refused lookup — the honest miss, same rule as Ramp + Router's foreign-backend aliases. A published 0 is kept as a real $0 rate (Jev's + output and cached input are free), not turned into "no rate": the gateway bills the + row literally, and so must we. + """ + if not isinstance(rows, list): + return None + matches: list[dict[str, Any]] = [] + for r in rows: + if not isinstance(r, dict) or r.get("model") != model or r.get("cost_type") != "tokens": + continue + if isinstance(r.get("token_pricing"), dict): + matches.append(r) + if not matches: + return None + namespace = model.split("/", 1)[0] if "/" in model else None + own = [r for r in matches if namespace and r.get("provider") == namespace] + candidates = own or matches + + def _fields(row: dict[str, Any]) -> dict[str, Decimal]: + out: dict[str, Decimal] = {} + for key, field in _CF_COSTS_FIELD_MAP.items(): + if key not in row["token_pricing"]: + continue + per_million = _parse_price(row["token_pricing"][key]) + if per_million is None: + continue + out[field] = (per_million / Decimal(1_000_000)).quantize(_Q, rounding=ROUND_DOWN) + return out + + first = _fields(candidates[0]) + if any(_fields(r) != first for r in candidates[1:]): + return None + if not first: + return None + return ModelPrice(source="cloudflare_gateway_costs", **first) + + # ---------------------------------------------------------------------- # Ramp Router parsing + matching # @@ -1168,6 +1233,7 @@ class PricingFetcher(Protocol): def fetch_openrouter(self) -> dict[str, Any]: ... def fetch_bedrock(self, region: str) -> dict[str, ModelPrice]: ... def fetch_cloudflare_workers_ai(self) -> dict[str, ModelPrice]: ... + def fetch_cloudflare_gateway_cost(self, model: str) -> ModelPrice | None: ... def fetch_mistral_aliases(self, api_key: str | None = None) -> dict[str, str]: ... def fetch_ramp_router(self, api_key: str | None = None) -> dict[str, ModelPrice]: ... @@ -1271,6 +1337,27 @@ def fetch_cloudflare_workers_ai(self) -> dict[str, ModelPrice]: page += 1 return parse_cloudflare_workers_ai(models) + def fetch_cloudflare_gateway_cost(self, model: str) -> ModelPrice | None: + """The gateway's own rate for one model id, or None when it lists none. + + One request per model, on demand — the full table is 2,839 rows across every + provider, and the only ids that reach this path are partner models the Workers AI + catalog omits, a handful per account. Same credentials as the catalog fetch. + """ + import requests + + if not self._cf_account_id or not self._cf_api_token: + return None + params: dict[str, str | int] = {"search": model, "per_page": 100} + resp = requests.get( + CLOUDFLARE_GATEWAY_COSTS_URL_TEMPLATE.format(account_id=self._cf_account_id), + headers={"Authorization": f"Bearer {self._cf_api_token}"}, + params=params, + timeout=self._timeout, + ) + resp.raise_for_status() + return parse_cloudflare_gateway_cost(resp.json().get("result"), model) + def fetch_mistral_aliases(self, api_key: str | None = None) -> dict[str, str]: import requests @@ -1340,6 +1427,14 @@ def __init__( self._cloudflare_workers_ai: dict[str, ModelPrice] | None = None self._cloudflare_fetched = 0.0 self._cloudflare_stale = False + # Partner models on Workers AI (`typesafe/jev`), priced from AI Gateway's own cost + # table one model at a time. Reactive like Bedrock: a miss in `lookup()` queues the + # id here, `maybe_refresh()` fetches it on the next tick, and every later call hits. + # A None value is a remembered "the gateway lists no rate" — kept for the TTL so an + # unpriced id does not cost one HTTP request per flush tick. + self._cf_gateway_costs: dict[str, ModelPrice | None] = {} + self._cf_gateway_costs_fetched: dict[str, float] = {} + self._cf_gateway_pending: set[str] = set() self._mistral_aliases: dict[str, str] | None = None self._mistral_fetched = 0.0 self._mistral_stale = False @@ -1375,7 +1470,7 @@ def _heal_fork(self) -> None: self._ramp_router_stale = self._ramp_router is not None or self._ramp_router_stale self._refreshing = set() - def prime(self, providers: Iterable[str] = ()) -> None: + def prime(self, providers: Iterable[str] = (), *, workers_ai_models: Iterable[str] = ()) -> None: """Flag OpenRouter for an eager background warm (used when price mode is the global default) to shrink the cold-start window. @@ -1413,6 +1508,17 @@ def prime(self, providers: Iterable[str] = ()) -> None: with self._lock: if self._is_cold(self._openrouter, self._openrouter_fetched): self._openrouter_stale = True + # Partner models on Workers AI are priced one row at a time from the gateway's + # cost table, and the SDK only learns an id when a call for it arrives — so the + # first call to each such model in a process is a cold miss. Naming the ids here + # (via `warm_pricing(workers_ai_models=[...])`) fetches their rows up front, the + # way `providers=["workers-ai"]` fetches the catalog, so even the first call + # prices. Same "only if cold" gate as everything else in this method. + for m in workers_ai_models: + if m and not m.startswith("@") and not m.startswith(WORKERS_AI_COMPAT_PREFIX): + fetched_at = self._cf_gateway_costs_fetched.get(m) + if fetched_at is None or (time.time() - fetched_at) >= self._ttl: + self._cf_gateway_pending.add(m) for p in providers: key = (p or "").lower() if key == "workers-ai": @@ -1512,7 +1618,20 @@ def lookup(self, provider: str, model: str, api: str) -> ModelPrice | None: fresh_cf = table_cf is not None and (time.time() - self._cloudflare_fetched) < self._ttl if not fresh_cf: self._cloudflare_stale = True - return lookup_cloudflare_workers_ai(table_cf, model) if table_cf is not None else None + hit = lookup_cloudflare_workers_ai(table_cf, model) if table_cf is not None else None + if hit is not None or model.startswith("@") or model.startswith(WORKERS_AI_COMPAT_PREFIX): + return hit + # A bare `vendor/model` id the catalog does not list: a partner model. Its + # rate lives in the gateway's cost table — fetched per id, in the background. + with self._lock: + fetched_at = self._cf_gateway_costs_fetched.get(model) + if fetched_at is None or (time.time() - fetched_at) >= self._ttl: + # Cold or past the TTL: queue a (re)fetch for the next tick. A row we + # already hold keeps serving meanwhile — stale-while-revalidate, the + # same as the catalog table — so a TTL expiry never bills a call as + # tokens. Only a genuinely cold id (never fetched) misses. + self._cf_gateway_pending.add(model) + return self._cf_gateway_costs.get(model) resolved_model = model is_mistral = (provider or "").lower() == "mistral" with self._lock: @@ -1549,6 +1668,7 @@ def maybe_refresh(self) -> None: and not self._cloudflare_stale and not self._mistral_stale and not self._ramp_router_stale + and not self._cf_gateway_pending ): return with self._lock: @@ -1589,6 +1709,13 @@ def _ready(source: str) -> bool: ] for r in regions: self._refreshing.add(f"bedrock:{r}") + partner_models = [ + m + for m in self._cf_gateway_pending + if f"cf_costs:{m}" not in self._refreshing and _ready(f"cf_costs:{m}") + ] + for m in partner_models: + self._refreshing.add(f"cf_costs:{m}") if do_openrouter: try: @@ -1620,6 +1747,21 @@ def _ready(source: str) -> bool: with self._lock: self._refreshing.discard("cloudflare_workers_ai") + for m in partner_models: + try: + price = self._fetcher.fetch_cloudflare_gateway_cost(m) + with self._lock: + self._cf_gateway_costs[m] = price + self._cf_gateway_costs_fetched[m] = time.time() + self._cf_gateway_pending.discard(m) + self._note_success(f"cf_costs:{m}") + except Exception as exc: # noqa: BLE001 + self._note_failure(f"cf_costs:{m}") + self._report(exc, "pricing.fetch_cloudflare_gateway_cost") + finally: + with self._lock: + self._refreshing.discard(f"cf_costs:{m}") + if do_mistral: try: with self._lock: diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index 0c395fc..de3115c 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -8,11 +8,14 @@ import uuid from collections.abc import Iterable from datetime import datetime, timezone -from typing import Any +from typing import TYPE_CHECKING, Any from .canonical import CanonicalUsage from .config import LagoConfig from .detector import detect_client_kind + +if TYPE_CHECKING: + from .workers_ai import WorkersAI from .exceptions import PricingUnavailableError, UnknownClientError from .gateway.adapters.snowflake_cortex import SNOWFLAKE_EVENT_ID_PREFIX from .lago_client import LagoClient @@ -326,6 +329,43 @@ def wrap( "Implemented: 'bedrock', 'mistral', 'anthropic', 'openai', 'gemini'." ) + def workers_ai( + self, + account_id: str, + api_token: str, + *, + gateway_id: str | None = None, + gateway_auth: str | None = None, + timeout: float = 60.0, + dimensions: dict[str, Any] | None = None, + subscription: str | None = None, + ) -> WorkersAI: + """Build an instrumented Workers AI client — the one provider with no client to wrap. + + Reaches every Workers AI model through the model-in-body `/ai/run` route, including + partner models (`typesafe/jev`) that neither the path-style route nor the gateway's + OpenAI-compatible `/compat` endpoint can address. Pass `gateway_id` + `gateway_auth` + to go through an AI Gateway: cache hits are then skipped and each call carries a + `cf_log_id` dimension for reconciliation against the Logs API. See `workers_ai.py`. + """ + from .workers_ai import WorkersAI + + if self.config.pricing_mode == "price": + # Same warm-up wrap() gives an OpenAI client pointed at the gateway: prime the + # Workers AI catalog now, in memory only, so the first call is not a cold miss. + self._pricing.prime(["workers-ai"]) + self._queue.wake() + return WorkersAI( + self, + account_id, + api_token, + gateway_id=gateway_id, + gateway_auth=gateway_auth, + timeout=timeout, + dimensions=dimensions, + subscription=subscription, + ) + # ------------------------------------------------------------------ # Emit # ------------------------------------------------------------------ @@ -675,7 +715,7 @@ def _report_error(self, exc: Exception, where: str) -> None: pass logger.warning("lago %s failed: %s", where, exc) - def warm_pricing(self, providers: Iterable[str] = ()) -> None: + def warm_pricing(self, providers: Iterable[str] = (), *, workers_ai_models: Iterable[str] = ()) -> None: """Block until the given price table(s) are fetched, instead of waiting for the queue's background thread to pick them up on its next tick (up to `flush_interval` seconds later, by default ~1s). @@ -705,8 +745,16 @@ def warm_pricing(self, providers: Iterable[str] = ()) -> None: say so and skip that one-time cost too: `providers=["mistral"]` and/or `["workers-ai"]`. A no-op for any source that isn't stale (e.g. the SDK isn't in price mode, was already warmed, or the - provider name wasn't recognized).""" - self._pricing.prime(providers) + provider name wasn't recognized). + + `workers_ai_models`: partner models on Workers AI (`typesafe/jev` — any id + without the `@cf/` prefix) are priced from the gateway's cost table one row per + id, and the SDK only learns an id when a call for it arrives, so the first call + to each in a process bills tokens. Name the ids you are about to call here to + fetch their rows now, so even that first call prices — the client's `run()` for + them then has a warm row from the start. + """ + self._pricing.prime(providers, workers_ai_models=workers_ai_models) self._pricing.maybe_refresh() def backfill_databricks( diff --git a/src/lago_agent_sdk/workers_ai.py b/src/lago_agent_sdk/workers_ai.py new file mode 100644 index 0000000..41cda93 --- /dev/null +++ b/src/lago_agent_sdk/workers_ai.py @@ -0,0 +1,186 @@ +"""Workers AI client — the SDK's own, because there is no third-party one to wrap. + +Every other provider is instrumented by patching its official client in place. Workers +AI has no such target: the official `cloudflare` package (5.7.0, measured 2026-09-21) +percent-encodes the slash in every model name, so `client.ai.run("@cf/meta/...")` hits +`/ai/run/@cf%2Fmeta%2F...` and Cloudflare answers "No route for that URI" — for every +model, not just partner ones. Patching its generic `client.post(...)` instead would +instrument every Cloudflare API call the customer makes. So the SDK ships this client: +one method, two routes, chosen per model id. + +Two kinds of model, two routes — measured, not chosen: + +* `@cf/...` models go to the gateway host, `gateway.ai.cloudflare.com/v1/{acct}/{gw}/ + workers-ai/{model}`. That route answers with `cf-aig-cache-status` (a HIT means the + model never ran and nothing is billed) and `cf-aig-log-id` (emitted as the `cf_log_id` + dimension, so a Lago row can be put beside its Logs API entry — see + `gateway/adapters/cloudflare_gateway.py`). The model-in-body variant of that host, + `.../workers-ai/run`, logs `model: "run"`, which is why it is not used for them. +* Partner models (`typesafe/jev` — no `@`) can only be reached on the *unified* path, + `api.cloudflare.com/.../ai/run` with `{"model", "input"}` in the body and the gateway + named in a `cf-aig-gateway-id` header. That is the one route where the partner key the + customer stored under the gateway's BYOK is consulted; the gateway host answers 402 + "Insufficient balance" for the same call even with the key stored (fixture 08). The + unified path returns NO `cf-aig-*` headers — a cached replay looks exactly like a fresh + call — so this client sends `cf-aig-skip-cache: true` there rather than risk billing + the same answer twice. Pass `extra_headers={"cf-aig-skip-cache": "false"}` to opt back + in, knowing a replay then bills again. + +Billing is by token count from the response's own `usage` block. In price mode a +`@cf/...` model prices from Cloudflare's Workers AI catalog as usual; a partner model is +not listed there and prices from AI Gateway's own cost table instead, fetched per id on +the queue's next tick — so the first call to a partner model in a process bills tokens +and reports the miss through `on_error`, and every call after it bills dollars (see +`pricing.parse_cloudflare_gateway_cost`). +""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +import requests + +from .adapters.workers_ai_native import extract_workers_ai_native + +logger = logging.getLogger("lago_agent_sdk.workers_ai") + +_DIRECT_BASE = "https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run" +_GATEWAY_BASE = "https://gateway.ai.cloudflare.com/v1/{account_id}/{gateway_id}/workers-ai" + + +def _is_catalog_model(model: str) -> bool: + """`@cf/...` (and `@hf/...`) ids are Cloudflare-hosted; anything else is a partner model.""" + return model.startswith("@") + + +class WorkersAIError(Exception): + """Cloudflare answered the run with an error status or `success: false`. + + Raised *before* any instrumentation, so a 402 (partner model with no BYOK key and no + gateway credits), a 403 (model not on the account's Workers plan) or a 400 (no such + model) reaches the caller exactly as the API reported it, and nothing is billed for it. + """ + + def __init__(self, status_code: int, errors: list[dict[str, Any]], body: dict[str, Any]) -> None: + self.status_code = status_code + self.errors = errors + self.body = body + detail = "; ".join(str(e.get("message", e)) for e in errors) if errors else "no error detail" + super().__init__(f"Workers AI HTTP {status_code}: {detail}") + + +class WorkersAI: + """Minimal Workers AI client with Lago instrumentation. Build one via `LagoSDK.workers_ai()`.""" + + def __init__( + self, + sdk: Any, + account_id: str, + api_token: str, + *, + gateway_id: str | None = None, + gateway_auth: str | None = None, + timeout: float = 60.0, + dimensions: dict[str, Any] | None = None, + subscription: str | None = None, + ) -> None: + self._sdk = sdk + self._timeout = timeout + self._base_dims = dict(dimensions or {}) + self._base_sub = subscription + self._gateway_id = gateway_id + self._direct = _DIRECT_BASE.format(account_id=account_id) + self._gateway = ( + _GATEWAY_BASE.format(account_id=account_id, gateway_id=gateway_id) if gateway_id else None + ) + self._auth = {"Authorization": f"Bearer {api_token}"} + self._gateway_auth = {"cf-aig-authorization": f"Bearer {gateway_auth}"} if gateway_auth else {} + self._session = requests.Session() + + def url_for(self, model: str) -> str: + """The endpoint a `run(model, ...)` posts to — see the module docstring for why two.""" + if _is_catalog_model(model): + return f"{self._gateway}/{model}" if self._gateway else f"{self._direct}/{model}" + return self._direct + + def _request_for(self, model: str, input: dict[str, Any]) -> tuple[dict[str, str], dict[str, Any]]: + """(headers, body) for the route `url_for(model)` picks.""" + if _is_catalog_model(model): + # Path route: the body IS the input. Gateway auth only on the gateway host. + return ({**self._auth, **(self._gateway_auth if self._gateway else {})}, input) + headers = dict(self._auth) + if self._gateway_id: + headers["cf-aig-gateway-id"] = self._gateway_id + headers["cf-aig-skip-cache"] = "true" # no cache header on this path — see module docstring + return (headers, {"model": model, "input": input}) + + def run( + self, + model: str, + input: dict[str, Any], + *, + extra_lago: dict[str, Any] | None = None, + extra_headers: dict[str, str] | None = None, + ) -> dict[str, Any]: + """Run `model` on `input` and bill the response's usage. + + `input` is whatever the model takes: `{"messages": [...]}` for a chat model, + `{"state": ..., "questions": {...}}` for `typesafe/jev`. Returns Cloudflare's full + response envelope (`result`, `success`, `errors`, `messages`) unchanged. + + `extra_lago` takes the same keys as the wrappers' per-call kwarg: `subscription`, + `dimensions`, `mode`, `markup`. `extra_headers` reaches the request as-is and wins + over the client's own — e.g. `{"cf-aig-cache-ttl": "300"}` to have the gateway + cache a `@cf/...` call. + + Streaming (`input["stream"] = True`) is refused up front: the body would be an SSE + stream this method does not parse, and mis-reading it as JSON would bill zero for + a call that ran. Use the OpenAI-compatible `/compat` endpoint through + `sdk.wrap(OpenAI(...))` for streamed chat. + """ + if input.get("stream"): + raise ValueError( + "WorkersAI.run() does not support stream=True; for streamed chat completions wrap an " + "OpenAI client against the gateway's /compat endpoint instead." + ) + lago_opts = extra_lago or {} + headers, body = self._request_for(model, input) + headers.update(extra_headers or {}) + sub = self._sdk._resolve_subscription(lago_opts.get("subscription") or self._base_sub) + if self._gateway_id and sub and "cf-aig-metadata" not in headers: + # Attribution travels with the call: the gateway stores this on the log entry, so + # the Logs API backfill resolves the same subscription this emit() uses. Honoured on + # both routes (measured on the unified path: the log row carried it). + headers["cf-aig-metadata"] = json.dumps({"lago_subscription": sub}) + + resp = self._session.post(self.url_for(model), headers=headers, json=body, timeout=self._timeout) + try: + payload: dict[str, Any] = resp.json() + except ValueError: + payload = {} + if resp.status_code >= 400 or payload.get("success") is False: + raise WorkersAIError(resp.status_code, list(payload.get("errors") or []), payload) + + try: + if resp.headers.get("cf-aig-cache-status") == "HIT": + # The gateway answered from its cache; the model never ran and Cloudflare + # billed nothing. Billing it would charge for a call that did not happen. + return payload + usage = extract_workers_ai_native(payload, model_id=model) + dims = {**self._base_dims, **(lago_opts.get("dimensions") or {})} + log_id = resp.headers.get("cf-aig-log-id") + if log_id: + dims["cf_log_id"] = log_id + self._sdk.emit( + usage, + subscription=sub, + dimensions=dims, + mode=lago_opts.get("mode"), + markup=lago_opts.get("markup"), + ) + except Exception as exc: # noqa: BLE001 — instrumentation never breaks the customer's call + logger.warning("lago: workers_ai.run instrumentation failed: %s", exc) + self._sdk._report_error(exc, "emit") + return payload diff --git a/tests/unit/adapters/fixtures/capture_workers_ai.py b/tests/unit/adapters/fixtures/capture_workers_ai.py new file mode 100644 index 0000000..7f7b820 --- /dev/null +++ b/tests/unit/adapters/fixtures/capture_workers_ai.py @@ -0,0 +1,158 @@ +"""Capture real Workers AI `/ai/run` responses for the Workers AI adapter. + +Saves to `workers_ai/.json` as `{_model_id, _status, _headers, _response}` — +successful responses only; the client's error path is unit-tested against inline bodies. + +Two routes are exercised, because they differ in what reaches a partner model: + + direct POST api.cloudflare.com/.../ai/run/{model} (`@cf/...` ids) + POST api.cloudflare.com/.../ai/run body {"model", "input"} (partner ids) + With `cf-aig-gateway-id: ` the call is routed through that gateway, + which is what lets a partner model (`typesafe/jev`) use the partner key stored + under the gateway's BYOK — without it the call bills gateway credits. + gateway POST gateway.ai.cloudflare.com/v1/.../workers-ai/{model} + Exposes `cf-aig-cache-status` / `cf-aig-log-id` response headers. + +Reads CF_ACCOUNT_ID, CF_API_TOKEN (or CF_LOGS_TOKEN), CF_GATEWAY_ID, CF_GATEWAY_AUTH and, +for scenario 8, TYPESAFE_API_KEY from env. Bodies carry no account identifiers; only the +gateway headers the client reads (`cf-aig-cache-status`) are kept. + +Pass scenario numbers to recapture a subset: `capture_workers_ai.py 3 4`. +""" + +from __future__ import annotations + +import json +import os +import pathlib +import sys +import time + +import requests + +OUT = pathlib.Path(__file__).parent / "workers_ai" +OUT.mkdir(parents=True, exist_ok=True) +KEPT_HEADERS = ("cf-aig-cache-status", "content-type") +LLAMA = "@cf/meta/llama-3.2-3b-instruct" + + +def save(name: str, model: str, resp: requests.Response, source: str | None = None) -> None: + payload: dict = { + "_model_id": model, + "_status": resp.status_code, + "_headers": {k: v for k, v in resp.headers.items() if k.lower() in KEPT_HEADERS}, + } + if source: + payload["_source"] = source + payload["_response"] = resp.json() + (OUT / f"{name}.json").write_text(json.dumps(payload, indent=2) + "\n") + print(f" ✓ saved {name}.json HTTP {resp.status_code}") + + +JEV_INPUT = { + "state": "Hi, I was charged twice for my subscription this month and I need the duplicate refunded " + "before my card statement closes on Friday. This is the second time this has happened.", + "questions": { + "is_urgent": { + "type": "noul", + "instructions": "Does this convey urgency?", + "criteria": {"true": "Explicitly time-sensitive", "false": "No urgency expressed"}, + }, + "department": { + "type": "choice", + "instructions": "Which team should handle this?", + "criteria": { + "billing": "Payments, invoicing, refunds", + "technical": "Bugs, outages, integrations", + "sales": "Pricing questions, upgrades", + }, + }, + "frustration": { + "type": "score", + "instructions": "How frustrated is the customer?", + "criteria": ["Calm", "Frustrated", "Very angry"], + }, + }, +} +CHAT_INPUT = { + "messages": [{"role": "user", "content": "Write one sentence about dolphins."}], + "max_tokens": 40, +} + + +def main() -> int: + only = {int(a) for a in sys.argv[1:] if a.isdigit()} + + def want(n: int) -> bool: + return not only or n in only + + acct = os.environ.get("CF_ACCOUNT_ID") + token = os.environ.get("CF_API_TOKEN") or os.environ.get("CF_LOGS_TOKEN") + gw, gw_auth = os.environ.get("CF_GATEWAY_ID"), os.environ.get("CF_GATEWAY_AUTH") + if not (acct and token): + print("error: set CF_ACCOUNT_ID and CF_API_TOKEN", file=sys.stderr) + return 2 + direct = f"https://api.cloudflare.com/client/v4/accounts/{acct}/ai/run" + hdr = {"Authorization": f"Bearer {token}"} + + def post(url: str, body: dict, extra: dict | None = None) -> requests.Response: + return requests.post(url, headers={**hdr, **(extra or {})}, json=body, timeout=90) + + if want(1): + print("\n[1] chat model, direct API, path route") + save("01_chat_direct", LLAMA, post(f"{direct}/{LLAMA}", CHAT_INPUT)) + + if want(2): + print("\n[2] reasoning model, direct API — does usage break reasoning out?") + m = "@cf/deepseek-ai/deepseek-r1-distill-qwen-32b" + save("02_reasoning_direct", m, post(f"{direct}/{m}", {**CHAT_INPUT, "max_tokens": 60})) + + if want(3): + print("\n[3] OpenAI-vendor model on Workers AI, direct API") + m = "@cf/openai/gpt-oss-120b" + save("03_gpt_oss_direct", m, post(f"{direct}/{m}", CHAT_INPUT)) + + if want(4): + print("\n[4] Mistral-vendor model on Workers AI, direct API") + m = "@cf/mistralai/mistral-small-3.1-24b-instruct" + save("04_mistral_small_direct", m, post(f"{direct}/{m}", CHAT_INPUT)) + + if want(5): + print("\n[5] partner model typesafe/jev, direct API routed through the BYOK gateway") + extra = {"cf-aig-gateway-id": gw} if gw else None + save( + "05_jev_byok_gateway", + "typesafe/jev", + post(direct, {"model": "typesafe/jev", "input": JEV_INPUT}, extra), + ) + + if gw and gw_auth and (want(6) or want(7)): + gwbase = f"https://gateway.ai.cloudflare.com/v1/{acct}/{gw}/workers-ai" + cache = {"cf-aig-authorization": f"Bearer {gw_auth}", "cf-aig-cache-ttl": "300"} + print("\n[6] chat model via gateway host (cache MISS), cache enabled for the pair") + save("06_chat_gateway_miss", LLAMA, post(f"{gwbase}/{LLAMA}", CHAT_INPUT, cache)) + time.sleep(2) + print("\n[7] identical call again — gateway cache HIT") + save("07_chat_gateway_hit", LLAMA, post(f"{gwbase}/{LLAMA}", CHAT_INPUT, cache)) + + if want(8) and os.environ.get("TYPESAFE_API_KEY"): + # The model's response body from its own API, for comparison with what Cloudflare + # wraps in `result.result` on scenario 5. + print("\n[8] typesafe/jev from TypeSafe's own API (not Cloudflare)") + r = requests.post( + "https://api.typesafe.ai/v1/systemone", + headers={"Authorization": f"Bearer {os.environ['TYPESAFE_API_KEY']}"}, + json={"model": "jev-latest", **JEV_INPUT}, + timeout=90, + ) + save( + "08_jev_typesafe_direct", + "typesafe/jev", + r, + source="POST https://api.typesafe.ai/v1/systemone (model=jev-latest)", + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unit/adapters/fixtures/workers_ai/01_chat_direct.json b/tests/unit/adapters/fixtures/workers_ai/01_chat_direct.json new file mode 100644 index 0000000..d7ffe43 --- /dev/null +++ b/tests/unit/adapters/fixtures/workers_ai/01_chat_direct.json @@ -0,0 +1,55 @@ +{ + "_model_id": "@cf/meta/llama-3.2-3b-instruct", + "_status": 200, + "_headers": { + "Content-Type": "application/json" + }, + "_response": { + "result": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": null, + "message": { + "annotations": null, + "audio": null, + "content": "Dolphins are highly intelligent, social marine mammals known for their playful and curious nature, often observed swimming and interacting with each other in complex behaviors.", + "function_call": null, + "reasoning": null, + "refusal": null, + "role": "assistant" + }, + "routed_experts": null, + "stop_reason": null, + "token_ids": null + } + ], + "created": 1789987074, + "ec_transfer_params": null, + "id": "chatcmpl-c205f782-079c-4867-a46e-a332a008834c", + "kv_transfer_params": null, + "metrics": null, + "model": "@cf/meta/llama-3.2-3b-instruct-v2", + "object": "chat.completion", + "prompt_logprobs": null, + "prompt_text": null, + "prompt_token_ids": null, + "response": "Dolphins are highly intelligent, social marine mammals known for their playful and curious nature, often observed swimming and interacting with each other in complex behaviors.", + "service_tier": null, + "tool_calls": [], + "usage": { + "prompt_tokens": 41, + "completion_tokens": 31, + "total_tokens": 72, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "neurons": 1.1343257427215576 + } + }, + "success": true, + "errors": [], + "messages": [] + } +} diff --git a/tests/unit/adapters/fixtures/workers_ai/02_reasoning_direct.json b/tests/unit/adapters/fixtures/workers_ai/02_reasoning_direct.json new file mode 100644 index 0000000..37b0ae7 --- /dev/null +++ b/tests/unit/adapters/fixtures/workers_ai/02_reasoning_direct.json @@ -0,0 +1,24 @@ +{ + "_model_id": "@cf/deepseek-ai/deepseek-r1-distill-qwen-32b", + "_status": 200, + "_headers": { + "Content-Type": "application/json" + }, + "_response": { + "result": { + "response": "\nOkay, so the user wants me to write one sentence about dolphins. Let me think about what's most interesting and accurate. Dolphins are intelligent, right? They have those big brains and they communicate. Maybe I should mention their intelligence and social behavior. Also, they're mammals, so they", + "usage": { + "prompt_tokens": 13, + "completion_tokens": 60, + "total_tokens": 73, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "neurons": 27.212594786540038 + } + }, + "success": true, + "errors": [], + "messages": [] + } +} diff --git a/tests/unit/adapters/fixtures/workers_ai/03_gpt_oss_direct.json b/tests/unit/adapters/fixtures/workers_ai/03_gpt_oss_direct.json new file mode 100644 index 0000000..9fa05a8 --- /dev/null +++ b/tests/unit/adapters/fixtures/workers_ai/03_gpt_oss_direct.json @@ -0,0 +1,54 @@ +{ + "_model_id": "@cf/openai/gpt-oss-120b", + "_status": 200, + "_headers": { + "Content-Type": "application/json" + }, + "_response": { + "result": { + "choices": [ + { + "finish_reason": "length", + "index": 0, + "logprobs": null, + "message": { + "annotations": null, + "audio": null, + "content": "Dolphins are highly intelligent, social marine mammals known for their", + "function_call": null, + "reasoning": "User asks: \"Write one sentence about dolphins.\" Straightforward. Provide a single sentence.", + "reasoning_content": "User asks: \"Write one sentence about dolphins.\" Straightforward. Provide a single sentence.", + "refusal": null, + "role": "assistant" + }, + "routed_experts": null, + "stop_reason": null, + "token_ids": null + } + ], + "created": 1789987522, + "ec_transfer_params": null, + "id": "chatcmpl-46141cfa-6441-4f51-b214-ad73ec255ec0", + "kv_transfer_params": null, + "metrics": null, + "model": "@cf/openai/gpt-oss-120b", + "object": "chat.completion", + "prompt_logprobs": null, + "prompt_text": null, + "prompt_token_ids": null, + "service_tier": null, + "usage": { + "prompt_tokens": 73, + "completion_tokens": 40, + "total_tokens": 113, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "neurons": 5.049985885620117 + } + }, + "success": true, + "errors": [], + "messages": [] + } +} diff --git a/tests/unit/adapters/fixtures/workers_ai/04_mistral_small_direct.json b/tests/unit/adapters/fixtures/workers_ai/04_mistral_small_direct.json new file mode 100644 index 0000000..88da437 --- /dev/null +++ b/tests/unit/adapters/fixtures/workers_ai/04_mistral_small_direct.json @@ -0,0 +1,53 @@ +{ + "_model_id": "@cf/mistralai/mistral-small-3.1-24b-instruct", + "_status": 200, + "_headers": { + "Content-Type": "application/json" + }, + "_response": { + "result": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": null, + "message": { + "annotations": null, + "audio": null, + "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", + "function_call": null, + "reasoning": null, + "refusal": null, + "role": "assistant", + "tool_calls": [] + }, + "stop_reason": null, + "token_ids": null + } + ], + "created": 1789987523, + "id": "chatcmpl-94af9498-eea6-4b55-bce8-1e0678ed6bbe", + "kv_transfer_params": null, + "model": "@cf/mistralai/mistral-small-3.1-24b-v2", + "object": "chat.completion", + "prompt_logprobs": null, + "prompt_token_ids": null, + "response": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", + "service_tier": null, + "system_fingerprint": null, + "tool_calls": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 19, + "total_tokens": 29, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "neurons": 1.2780319452285767 + } + }, + "success": true, + "errors": [], + "messages": [] + } +} diff --git a/tests/unit/adapters/fixtures/workers_ai/05_jev_byok_gateway.json b/tests/unit/adapters/fixtures/workers_ai/05_jev_byok_gateway.json new file mode 100644 index 0000000..1ea29dc --- /dev/null +++ b/tests/unit/adapters/fixtures/workers_ai/05_jev_byok_gateway.json @@ -0,0 +1,56 @@ +{ + "_model_id": "typesafe/jev", + "_status": 200, + "_headers": { + "Content-Type": "application/json" + }, + "_response": { + "result": { + "state": "Completed", + "result": { + "model": "jev-1.13.0", + "answers": { + "is_urgent": { + "type": "noul", + "noul": 0.97 + }, + "department": { + "type": "choice", + "choice": "billing", + "probabilities": { + "billing": 1, + "sales": 0, + "technical": 0 + }, + "confidence": 1 + }, + "frustration": { + "type": "score", + "score": 1, + "legend": { + "0": "Calm", + "1": "Frustrated", + "2": "Very angry" + }, + "probabilities": { + "0": 0.01, + "1": 0.99, + "2": 0 + }, + "confidence": 0.99 + } + }, + "usage": { + "input_tokens": 446, + "output_tokens": 73 + } + }, + "gatewayMetadata": { + "keySource": "BYOK" + } + }, + "success": true, + "errors": [], + "messages": [] + } +} diff --git a/tests/unit/adapters/fixtures/workers_ai/06_chat_gateway_miss.json b/tests/unit/adapters/fixtures/workers_ai/06_chat_gateway_miss.json new file mode 100644 index 0000000..d7b0e67 --- /dev/null +++ b/tests/unit/adapters/fixtures/workers_ai/06_chat_gateway_miss.json @@ -0,0 +1,56 @@ +{ + "_model_id": "@cf/meta/llama-3.2-3b-instruct", + "_status": 200, + "_headers": { + "Content-Type": "application/json", + "cf-aig-cache-status": "MISS" + }, + "_response": { + "result": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": null, + "message": { + "annotations": null, + "audio": null, + "content": "Dolphins are highly intelligent, social marine mammals known for their playful and curious nature, often observed swimming in large groups and communicating with each other using a variety of clicks and whistles.", + "function_call": null, + "reasoning": null, + "refusal": null, + "role": "assistant" + }, + "routed_experts": null, + "stop_reason": null, + "token_ids": null + } + ], + "created": 1789987080, + "ec_transfer_params": null, + "id": "chatcmpl-bd52fd8c-7173-427b-bcb9-79d7d42e4b3a", + "kv_transfer_params": null, + "metrics": null, + "model": "@cf/meta/llama-3.2-3b-instruct-v2", + "object": "chat.completion", + "prompt_logprobs": null, + "prompt_text": null, + "prompt_token_ids": null, + "response": "Dolphins are highly intelligent, social marine mammals known for their playful and curious nature, often observed swimming in large groups and communicating with each other using a variety of clicks and whistles.", + "service_tier": null, + "tool_calls": [], + "usage": { + "prompt_tokens": 41, + "completion_tokens": 39, + "total_tokens": 80, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "neurons": 1.3781238794326782 + } + }, + "success": true, + "errors": [], + "messages": [] + } +} diff --git a/tests/unit/adapters/fixtures/workers_ai/07_chat_gateway_hit.json b/tests/unit/adapters/fixtures/workers_ai/07_chat_gateway_hit.json new file mode 100644 index 0000000..ffe6c2a --- /dev/null +++ b/tests/unit/adapters/fixtures/workers_ai/07_chat_gateway_hit.json @@ -0,0 +1,56 @@ +{ + "_model_id": "@cf/meta/llama-3.2-3b-instruct", + "_status": 200, + "_headers": { + "Content-Type": "application/json", + "cf-aig-cache-status": "HIT" + }, + "_response": { + "result": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": null, + "message": { + "annotations": null, + "audio": null, + "content": "Dolphins are highly intelligent, social marine mammals known for their playful and curious nature, often observed swimming in large groups and communicating with each other using a variety of clicks and whistles.", + "function_call": null, + "reasoning": null, + "refusal": null, + "role": "assistant" + }, + "routed_experts": null, + "stop_reason": null, + "token_ids": null + } + ], + "created": 1789987080, + "ec_transfer_params": null, + "id": "chatcmpl-bd52fd8c-7173-427b-bcb9-79d7d42e4b3a", + "kv_transfer_params": null, + "metrics": null, + "model": "@cf/meta/llama-3.2-3b-instruct-v2", + "object": "chat.completion", + "prompt_logprobs": null, + "prompt_text": null, + "prompt_token_ids": null, + "response": "Dolphins are highly intelligent, social marine mammals known for their playful and curious nature, often observed swimming in large groups and communicating with each other using a variety of clicks and whistles.", + "service_tier": null, + "tool_calls": [], + "usage": { + "prompt_tokens": 41, + "completion_tokens": 39, + "total_tokens": 80, + "prompt_tokens_details": { + "cached_tokens": 0 + }, + "neurons": 1.3781238794326782 + } + }, + "success": true, + "errors": [], + "messages": [] + } +} diff --git a/tests/unit/adapters/fixtures/workers_ai/08_jev_typesafe_direct.json b/tests/unit/adapters/fixtures/workers_ai/08_jev_typesafe_direct.json new file mode 100644 index 0000000..4fa11d4 --- /dev/null +++ b/tests/unit/adapters/fixtures/workers_ai/08_jev_typesafe_direct.json @@ -0,0 +1,44 @@ +{ + "_model_id": "typesafe/jev", + "_status": 200, + "_headers": {}, + "_source": "POST https://api.typesafe.ai/v1/systemone (model=jev-latest)", + "_response": { + "model": "jev-1.13.0", + "answers": { + "is_urgent": { + "type": "noul", + "noul": 0.97 + }, + "department": { + "type": "choice", + "choice": "billing", + "confidence": 1.0, + "probabilities": { + "sales": 0.0, + "technical": 0.0, + "billing": 1.0 + } + }, + "frustration": { + "type": "score", + "score": 1.0, + "confidence": 0.99, + "legend": { + "0": "Calm", + "1": "Frustrated", + "2": "Very angry" + }, + "probabilities": { + "0": 0.01, + "1": 0.99, + "2": 0.0 + } + } + }, + "usage": { + "input_tokens": 446, + "output_tokens": 73 + } + } +} diff --git a/tests/unit/adapters/test_workers_ai_native.py b/tests/unit/adapters/test_workers_ai_native.py new file mode 100644 index 0000000..b1ea3e4 --- /dev/null +++ b/tests/unit/adapters/test_workers_ai_native.py @@ -0,0 +1,165 @@ +"""Workers AI `/ai/run` adapter — verified against real captured fixtures.""" + +from __future__ import annotations + +import json +import pathlib + +import pytest + +from lago_agent_sdk.adapters import extract_workers_ai_native + +FIX = pathlib.Path(__file__).parent / "fixtures" / "workers_ai" + + +def _all() -> list[pathlib.Path]: + return sorted(FIX.glob("*.json")) if FIX.exists() else [] + + +def _load(name: str) -> dict: + return json.loads((FIX / name).read_text()) + + +# -------------------------------------------------------------------------- +# Catalog models — OpenAI-shaped usage plus Cloudflare's `neurons` +# -------------------------------------------------------------------------- +def test_chat_direct_maps_openai_shaped_usage(): + d = _load("01_chat_direct.json") + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.input == 41 + assert u.output == 31 + assert u.cache_read == 0 + assert u.reasoning == 0 + assert u.provider == "workers-ai" + assert u.api == "workers_ai_run" + + +def test_requested_id_is_the_billing_key_and_served_name_is_kept(): + """`...-3b-instruct` answers as `...-3b-instruct-v2`. The requested id is what the price + catalog is keyed by, so it is the one carried; the served name is not lost.""" + d = _load("01_chat_direct.json") + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.model == "@cf/meta/llama-3.2-3b-instruct" + assert u.extras["served_model"] == "@cf/meta/llama-3.2-3b-instruct-v2" + + +def test_served_name_that_would_miss_the_catalog_does_not_become_the_model(): + """The case that decided the rule: Mistral small answers as `...-24b-v2` — `-instruct` + dropped, `-v2` added — and that name is NOT in Cloudflare's catalog (measured + 2026-09-21) while the requested id is.""" + d = _load("04_mistral_small_direct.json") + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.model == "@cf/mistralai/mistral-small-3.1-24b-instruct" + assert u.extras["served_model"] == "@cf/mistralai/mistral-small-3.1-24b-v2" + assert (u.input, u.output) == (10, 19) + + +def test_served_name_equal_to_requested_adds_nothing_to_extras(): + d = _load("03_gpt_oss_direct.json") + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.model == "@cf/openai/gpt-oss-120b" + assert "served_model" not in u.extras + assert (u.input, u.output) == (73, 40) + + +def test_neurons_land_in_extras_not_in_a_metric(): + d = _load("01_chat_direct.json") + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.extras["neurons"] == pytest.approx(1.1343257427215576) + assert "usage" not in u.extras # every other key was recognised — no drift reported + + +def test_reasoning_model_bundles_thinking_into_completion(): + """deepseek-r1-distill reasons but reports no separate field — do not invent one.""" + d = _load("02_reasoning_direct.json") + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.input == 13 + assert u.output == 60 + assert u.reasoning == 0 + + +def test_silent_response_keeps_the_requested_model(): + d = _load("02_reasoning_direct.json") + assert d["_response"]["result"].get("model") is None + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.model == "@cf/deepseek-ai/deepseek-r1-distill-qwen-32b" + assert "served_model" not in u.extras + + +def test_gateway_hit_and_miss_bodies_are_identical(): + """The gateway replays the cached body byte-for-byte, usage included. Nothing in the + body distinguishes a HIT — only the header does, which is why the client, not the + adapter, decides to skip billing.""" + miss = _load("06_chat_gateway_miss.json") + hit = _load("07_chat_gateway_hit.json") + assert miss["_headers"]["cf-aig-cache-status"] == "MISS" + assert hit["_headers"]["cf-aig-cache-status"] == "HIT" + um = extract_workers_ai_native(miss["_response"], model_id=miss["_model_id"]) + uh = extract_workers_ai_native(hit["_response"], model_id=hit["_model_id"]) + assert (um.input, um.output) == (uh.input, uh.output) == (41, 39) + + +# -------------------------------------------------------------------------- +# Partner model — typesafe/jev via Cloudflare (05, BYOK) and straight from TypeSafe (08) +# -------------------------------------------------------------------------- +def test_jev_via_cloudflare_is_wrapped_one_level_deeper(): + """`result: {state, result: {model, answers, usage}, gatewayMetadata}` — the partner's + own object sits under `result.result`. Captured through the unified `/ai/run` path with + `cf-aig-gateway-id` naming the gateway whose BYOK holds the TypeSafe key.""" + d = _load("05_jev_byok_gateway.json") + r = d["_response"]["result"] + assert r["state"] == "Completed" and "usage" in r["result"] + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.input == 446 + assert u.output == 73 + assert u.cache_read == 0 and u.reasoning == 0 + assert u.provider == "workers-ai" and u.api == "workers_ai_run" + + +def test_jev_keeps_its_catalog_id_and_reports_the_served_version(): + for name in ("05_jev_byok_gateway.json", "08_jev_typesafe_direct.json"): + d = _load(name) + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.model == "typesafe/jev" + assert u.extras["served_model"] == "jev-1.13.0" + + +def test_jev_byok_key_source_lands_in_extras(): + """Under BYOK Cloudflare charged nothing — TypeSafe bills the customer directly. Whoever + reconciles against the Cloudflare dashboard needs to know which.""" + d = _load("05_jev_byok_gateway.json") + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.extras["gateway_metadata"] == {"keySource": "BYOK"} + assert "usage" not in u.extras and "neurons" not in u.extras + + +def test_jev_from_typesafe_directly_is_the_same_object_unwrapped(): + d = _load("08_jev_typesafe_direct.json") + assert "typesafe.ai" in d["_source"] + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert (u.input, u.output) == (446, 73) + assert set(u.extras) == {"served_model"} + + +# -------------------------------------------------------------------------- +# Every capture, and the failure envelope +# -------------------------------------------------------------------------- +@pytest.mark.skipif( + not _all(), reason="Workers AI fixtures not captured (run fixtures/capture_workers_ai.py)" +) +@pytest.mark.parametrize("path", _all(), ids=lambda p: p.stem) +def test_every_capture_is_a_success_that_bills(path: pathlib.Path): + d = json.loads(path.read_text()) + assert d["_status"] == 200, f"{path.stem}: only successful responses are kept as fixtures" + u = extract_workers_ai_native(d["_response"], model_id=d["_model_id"]) + assert u.provider == "workers-ai" + assert u.input > 0 and u.output > 0, f"{path.stem}: adapter broken or capture stale" + + +def test_failure_envelope_yields_zero_usage_without_raising(): + """The shape Cloudflare returns on 402/403/400 (seen live: `result: {}` + `errors`). The + client raises before billing; the adapter must still be safe to call on it.""" + body = {"errors": [{"message": "Insufficient balance", "code": 2021}], "success": False, "result": {}} + u = extract_workers_ai_native(body, model_id="typesafe/jev") + assert not u.nonzero_numeric() + assert u.model == "typesafe/jev" diff --git a/tests/unit/fixtures/pricing/money_golden.json b/tests/unit/fixtures/pricing/money_golden.json index fc70ec4..4ab6847 100644 --- a/tests/unit/fixtures/pricing/money_golden.json +++ b/tests/unit/fixtures/pricing/money_golden.json @@ -243,7 +243,7 @@ "total_cents": "0.423" }, { - "name": "ramp_router: OpenAI-served cache write at the catalog write rate (gpt-5.6-luna, 2026-09-07; Router billed 1.1x this \u2014 documented mismatch, not corrected)", + "name": "ramp_router: OpenAI-served cache write at the catalog write rate (gpt-5.6-luna, 2026-09-07; Router billed 1.1x this — documented mismatch, not corrected)", "provider": "ramp_router", "api": "ramp_router", "prices": { @@ -323,6 +323,24 @@ "base": "0.0020523", "total": "0.0020523", "total_cents": "0.20523" + }, + { + "name": "workers-ai partner model typesafe/jev via ai-gateway/costs: $0.042/M input, free output — equals the gateway's own log `cost` 1.8732e-05 (2026-09-21)", + "provider": "workers-ai", + "api": "workers_ai_run", + "prices": { + "input": "0.000000042", + "output": "0", + "cache_read": "0" + }, + "counts": { + "input": 446, + "output": 73 + }, + "markup": "1", + "base": "0.000018732", + "total": "0.000018732", + "total_cents": "0.0018732" } ], "precomputed_cases": [ diff --git a/tests/unit/gateway/adapters/test_cloudflare_gateway.py b/tests/unit/gateway/adapters/test_cloudflare_gateway.py index f8ac9a4..d670d43 100644 --- a/tests/unit/gateway/adapters/test_cloudflare_gateway.py +++ b/tests/unit/gateway/adapters/test_cloudflare_gateway.py @@ -553,7 +553,7 @@ def test_drift_omits_the_key_entirely_when_there_is_none() -> None: }, } ) - assert u.extras == {"cached": False, "step": 0, "log_id": "log_3"} + assert u.extras == {"cached": False, "step": 0, "log_id": "log_3", "byok": None} assert "usage_metadata" not in u.extras @@ -612,3 +612,13 @@ def test_drift_no_captured_fixture_loses_a_counter() -> None: assert key in _MAPPED_USAGE_KEYS or key in swept, ( f"{path.name}: {key!r} is neither mapped nor swept into extras" ) + + +def test_byok_key_source_reaches_extras_on_every_entry(): + """A row served under BYOK still carries Cloudflare's list-price `cost` although it charged + nothing (measured 2026-09-21, `typesafe/jev`: 446 in x $0.042/M = 1.8732e-05 with + byok="default"). The poller needs the field to not bill that.""" + from lago_agent_sdk.gateway.adapters import extract_cloudflare_log + + assert extract_cloudflare_log({"id": "x", "byok": "default"}).extras["byok"] == "default" + assert extract_cloudflare_log({"id": "x"}).extras["byok"] is None diff --git a/tests/unit/gateway/test_ramp_router.py b/tests/unit/gateway/test_ramp_router.py index 2a2f3b4..b5d749f 100644 --- a/tests/unit/gateway/test_ramp_router.py +++ b/tests/unit/gateway/test_ramp_router.py @@ -16,6 +16,7 @@ from lago_agent_sdk.exceptions import PricingUnavailableError from lago_agent_sdk.pricing import ( TOKEN_BILLED_PROVIDERS, + ModelPrice, PricingProvider, lookup_ramp_router, parse_openrouter, @@ -476,6 +477,9 @@ def fetch_bedrock(self, region: str) -> dict[str, Any]: def fetch_cloudflare_workers_ai(self) -> dict[str, Any]: return {} + def fetch_cloudflare_gateway_cost(self, model: str) -> ModelPrice | None: + return None + def fetch_mistral_aliases(self, api_key: str | None = None) -> dict[str, str]: return {} diff --git a/tests/unit/test_auto_prime_pricing.py b/tests/unit/test_auto_prime_pricing.py index 2e2f7d8..6b8c98c 100644 --- a/tests/unit/test_auto_prime_pricing.py +++ b/tests/unit/test_auto_prime_pricing.py @@ -194,6 +194,9 @@ def fetch_mistral_aliases(self, api_key=None): class _RouterCallCountingFetcher(HttpPricingFetcher): + def fetch_cloudflare_gateway_cost(self, model: str) -> ModelPrice | None: + return None + def __init__(self): super().__init__() self.router_keys: list[str | None] = [] diff --git a/tests/unit/test_drift.py b/tests/unit/test_drift.py index 66f275a..ecfafe9 100644 --- a/tests/unit/test_drift.py +++ b/tests/unit/test_drift.py @@ -276,3 +276,57 @@ def test_databricks_gateway_token_details_drift_survives_the_json_string_path() u = extract_databricks_log(row) assert u.reasoning == 9 assert u.extras["token_details.output_audio_tokens"] == "42" + + +# ---------------------------------------------------------------------------------- +# Workers AI `/ai/run` — the model-in-body route, two usage vocabularies on one endpoint +# ---------------------------------------------------------------------------------- +def test_workers_ai_unknown_usage_key_reaches_extras() -> None: + from lago_agent_sdk.adapters import extract_workers_ai_native + + resp = { + "result": { + "model": "@cf/meta/llama-3.2-3b-instruct-v2", + "usage": { + "prompt_tokens": 41, + "completion_tokens": 34, + "total_tokens": 75, + "neurons": 1.22, + "future_counter": 9, + "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 3}, + }, + }, + "success": True, + } + u = extract_workers_ai_native(resp, model_id="@cf/meta/llama-3.2-3b-instruct") + assert u.input == 41 and u.output == 34 + assert u.extras["usage"] == {"future_counter": 9, "prompt_tokens_details": {"audio_tokens": 3}} + assert u.extras["neurons"] == 1.22 # Cloudflare's billing unit is kept, never counted as tokens + + +def test_workers_ai_mapped_keys_do_not_pollute_extras() -> None: + from lago_agent_sdk.adapters import extract_workers_ai_native + + resp = { + "result": { + "usage": { + "prompt_tokens": 1, + "completion_tokens": 2, + "total_tokens": 3, + "prompt_tokens_details": {"cached_tokens": 1}, + } + } + } + u = extract_workers_ai_native(resp, model_id="@cf/x/y") + assert "usage" not in u.extras + assert u.cache_read == 1 + + +def test_workers_ai_jev_vocabulary_is_known_not_drift() -> None: + from lago_agent_sdk.adapters import extract_workers_ai_native + + u = extract_workers_ai_native( + {"result": {"answers": {}, "usage": {"input_tokens": 5, "output_tokens": 7}}}, model_id="typesafe/jev" + ) + assert (u.input, u.output) == (5, 7) + assert u.extras == {} diff --git a/tests/unit/test_pricing.py b/tests/unit/test_pricing.py index c61c511..321e9e1 100644 --- a/tests/unit/test_pricing.py +++ b/tests/unit/test_pricing.py @@ -36,6 +36,7 @@ lookup_ramp_router, parse_bedrock_offer, parse_bedrock_region, + parse_cloudflare_gateway_cost, parse_cloudflare_workers_ai, parse_mistral_aliases, parse_openrouter, @@ -66,6 +67,8 @@ def __init__( self.openrouter_calls = 0 self.bedrock_calls: list[str] = [] self.cloudflare_workers_ai_calls = 0 + self.cloudflare_gateway_cost_calls: list[str] = [] + self._cloudflare_gateway_costs: dict[str, ModelPrice | None] = {} self.mistral_aliases_calls = 0 self.last_mistral_api_key: str | None = None self.ramp_router_calls = 0 @@ -83,6 +86,10 @@ def fetch_cloudflare_workers_ai(self) -> dict[str, ModelPrice]: self.cloudflare_workers_ai_calls += 1 return self._cloudflare_workers_ai + def fetch_cloudflare_gateway_cost(self, model: str) -> ModelPrice | None: + self.cloudflare_gateway_cost_calls.append(model) + return self._cloudflare_gateway_costs.get(model) + def fetch_mistral_aliases(self, api_key: str | None = None) -> dict[str, str]: self.mistral_aliases_calls += 1 self.last_mistral_api_key = api_key @@ -2697,3 +2704,234 @@ def test_ttl_split_is_inert_without_split_rates_openrouter_anthropic_unchanged() assert set(b.fields) == {"input", "cache_write"} assert b.fields["cache_write"]["tokens"] == "20113" assert b.base == "0.02515725" + + +# ---------------------------------------------------------------------- +# Partner models on Workers AI — priced from AI Gateway's own cost table +# ---------------------------------------------------------------------- +# The real `ai-gateway/costs?search=jev` row, captured 2026-09-21. `cost_in`/`cost_out` are 0 +# on it while `token_pricing` is not — the gateway's log `cost` (446 x 0.042e-6) proves which +# one it bills by. +_CF_COSTS_JEV_ROW = { + "id": "a5a44d54-d13a-405e-8dba-6911dc681300", + "provider": "typesafe", + "model": "typesafe/jev", + "model_rule": "equals", + "cost_type": "tokens", + "cost_in": 0, + "cost_out": 0, + "token_pricing": {"input_tokens": 0.042, "input_cached_tokens": 0, "output_tokens": 0}, +} + + +def test_parse_gateway_cost_reads_token_pricing_per_million_and_keeps_published_zeros() -> None: + price = parse_cloudflare_gateway_cost([_CF_COSTS_JEV_ROW], "typesafe/jev") + assert price is not None + assert price.source == "cloudflare_gateway_costs" + assert price.input == Decimal("0.000000042") + assert price.output == Decimal("0") # free output is a real $0 rate, not "no rate" + assert price.cache_read == Decimal("0") + assert price.cache_write is None + + +def test_parse_gateway_cost_ignores_other_models_and_non_token_rows() -> None: + rows = [ + {**_CF_COSTS_JEV_ROW, "model": "typesafe/jev-mini"}, + {**_CF_COSTS_JEV_ROW, "cost_type": "per_request"}, + {**_CF_COSTS_JEV_ROW, "token_pricing": None}, + ] + assert parse_cloudflare_gateway_cost(rows, "typesafe/jev") is None + assert parse_cloudflare_gateway_cost("not a list", "typesafe/jev") is None + assert parse_cloudflare_gateway_cost([], "typesafe/jev") is None + + +def test_parse_gateway_cost_prefers_the_ids_own_namespace_when_providers_disagree() -> None: + """`stealth/union-alpha` is listed by three providers at different rates; the id's own + vendor is the one Workers AI serves it through.""" + rows = [ + { + "provider": "openrouter", + "model": "stealth/union-alpha", + "cost_type": "tokens", + "token_pricing": {"input_tokens": 1, "output_tokens": 2}, + }, + { + "provider": "stealth", + "model": "stealth/union-alpha", + "cost_type": "tokens", + "token_pricing": {"input_tokens": 3, "output_tokens": 4}, + }, + ] + price = parse_cloudflare_gateway_cost(rows, "stealth/union-alpha") + assert price is not None and price.input == Decimal("0.000003") and price.output == Decimal("0.000004") + + +def test_parse_gateway_cost_refuses_disagreeing_rows_without_an_own_namespace_match() -> None: + rows = [ + { + "provider": "openrouter", + "model": "x/y", + "cost_type": "tokens", + "token_pricing": {"input_tokens": 1, "output_tokens": 2}, + }, + { + "provider": "groq", + "model": "x/y", + "cost_type": "tokens", + "token_pricing": {"input_tokens": 5, "output_tokens": 2}, + }, + ] + assert parse_cloudflare_gateway_cost(rows, "x/y") is None + agree = [rows[0], {**rows[0], "provider": "groq"}] + assert parse_cloudflare_gateway_cost(agree, "x/y") is not None + + +def test_partner_model_is_reactive_miss_then_hit_and_never_refetched_within_ttl() -> None: + fetcher = StubFetcher( + cloudflare_workers_ai={ + "@cf/meta/llama-3.2-3b-instruct": ModelPrice( + source="cloudflare_workers_ai", input=Decimal("0.00000005"), output=Decimal("0.0000003") + ) + } + ) + fetcher._cloudflare_gateway_costs["typesafe/jev"] = parse_cloudflare_gateway_cost( + [_CF_COSTS_JEV_ROW], "typesafe/jev" + ) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.prime(["workers-ai"]) + p.maybe_refresh() + assert fetcher.cloudflare_workers_ai_calls == 1 + # 1st lookup: catalog has no partner ids -> miss, id queued; nothing fetched on the hot path. + assert p.lookup("workers-ai", "typesafe/jev", "workers_ai_run") is None + assert fetcher.cloudflare_gateway_cost_calls == [] + p.maybe_refresh() + assert fetcher.cloudflare_gateway_cost_calls == ["typesafe/jev"] + hit = p.lookup("workers-ai", "typesafe/jev", "workers_ai_run") + assert hit is not None and hit.input == Decimal("0.000000042") and hit.output == Decimal("0") + # Warm: neither the catalog nor the cost row is fetched again. + p.maybe_refresh() + p.lookup("workers-ai", "typesafe/jev", "workers_ai_run") + p.maybe_refresh() + assert fetcher.cloudflare_gateway_cost_calls == ["typesafe/jev"] + assert fetcher.cloudflare_workers_ai_calls == 1 + # A catalog model never takes the partner path. + assert p.lookup("workers-ai", "@cf/meta/llama-3.2-3b-instruct", "workers_ai_run") is not None + assert p.lookup("workers-ai", "@cf/nobody/unlisted", "workers_ai_run") is None + p.maybe_refresh() + assert fetcher.cloudflare_gateway_cost_calls == ["typesafe/jev"] + + +def test_partner_model_the_gateway_does_not_price_is_remembered_as_a_miss() -> None: + """One HTTP request per unpriced id per TTL — not one per flush tick.""" + fetcher = StubFetcher(cloudflare_workers_ai={}) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.prime(["workers-ai"]) + p.maybe_refresh() + assert p.lookup("workers-ai", "acme/unknown", "workers_ai_run") is None + p.maybe_refresh() + assert fetcher.cloudflare_gateway_cost_calls == ["acme/unknown"] + for _ in range(3): + assert p.lookup("workers-ai", "acme/unknown", "workers_ai_run") is None + p.maybe_refresh() + assert fetcher.cloudflare_gateway_cost_calls == ["acme/unknown"] + + +def test_partner_model_fetch_failure_is_reported_and_backed_off() -> None: + errors: list[str] = [] + + class _Boom(StubFetcher): + def fetch_cloudflare_gateway_cost(self, model: str) -> ModelPrice | None: + self.cloudflare_gateway_cost_calls.append(model) + raise RuntimeError("HTTP 500") + + fetcher = _Boom(cloudflare_workers_ai={}) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600, on_error=lambda e, w: errors.append(w)) + p.prime(["workers-ai"]) + p.maybe_refresh() + p.lookup("workers-ai", "typesafe/jev", "workers_ai_run") + p.maybe_refresh() + p.maybe_refresh() # inside the 1s backoff — no second attempt + assert fetcher.cloudflare_gateway_cost_calls == ["typesafe/jev"] + assert errors == ["pricing.fetch_cloudflare_gateway_cost"] + + +def test_http_fetcher_queries_costs_by_model_and_parses_the_real_row() -> None: + import responses as _responses + + with _responses.RequestsMock() as rsps: + rsps.get( + "https://api.cloudflare.com/client/v4/accounts/acct/ai-gateway/costs", + json={ + "success": True, + "result": [_CF_COSTS_JEV_ROW], + "result_info": {"count": 1, "total_count": 1}, + }, + ) + f = HttpPricingFetcher(cloudflare_account_id="acct", cloudflare_api_token="tok") + price = f.fetch_cloudflare_gateway_cost("typesafe/jev") + assert price is not None and price.input == Decimal("0.000000042") + req = rsps.calls[0].request + assert "search=typesafe%2Fjev" in req.url and "per_page=100" in req.url + assert req.headers["Authorization"] == "Bearer tok" + assert ( + HttpPricingFetcher().fetch_cloudflare_gateway_cost("typesafe/jev") is None + ) # no credentials -> no request + + +def test_jev_end_to_end_cost_matches_the_gateway_log() -> None: + """446 in / 73 out at the gateway's rate = 0.000018732 USD — the `cost` Cloudflare itself + stamped on the live log entry (2026-09-21).""" + price = parse_cloudflare_gateway_cost([_CF_COSTS_JEV_ROW], "typesafe/jev") + assert price is not None + usage = CanonicalUsage( + model="typesafe/jev", provider="workers-ai", api="workers_ai_run", input=446, output=73 + ) + b = compute_cost(usage, price, Decimal("1")) + assert b.total == "0.000018732" + assert b.total_cents == "0.0018732" + + +def test_partner_model_row_keeps_serving_past_the_ttl_while_it_refetches() -> None: + """Stale-while-revalidate, same as the catalog table: a TTL expiry must never bill a call + as tokens. Only a never-fetched id misses.""" + fetcher = StubFetcher(cloudflare_workers_ai={}) + fetcher._cloudflare_gateway_costs["typesafe/jev"] = parse_cloudflare_gateway_cost( + [_CF_COSTS_JEV_ROW], "typesafe/jev" + ) + p = PricingProvider(fetcher=fetcher, ttl_seconds=0.05) + p.prime(["workers-ai"]) + p.maybe_refresh() + assert p.lookup("workers-ai", "typesafe/jev", "workers_ai_run") is None # cold: the one honest miss + p.maybe_refresh() + assert p.lookup("workers-ai", "typesafe/jev", "workers_ai_run") is not None + time.sleep(0.08) # past the TTL + stale_hit = p.lookup("workers-ai", "typesafe/jev", "workers_ai_run") + assert stale_hit is not None and stale_hit.input == Decimal("0.000000042") # still served + p.maybe_refresh() # ... and refetched in the background + assert fetcher.cloudflare_gateway_cost_calls == ["typesafe/jev", "typesafe/jev"] + + +def test_warm_pricing_by_partner_model_id_prices_the_very_first_call() -> None: + fetcher = StubFetcher(cloudflare_workers_ai={}) + fetcher._cloudflare_gateway_costs["typesafe/jev"] = parse_cloudflare_gateway_cost( + [_CF_COSTS_JEV_ROW], "typesafe/jev" + ) + sdk = LagoSDK( + api_key="k", + default_subscription_id="sub", + config=LagoConfig( + api_key="k", pricing_mode="price", pricing_provider=PricingProvider(fetcher=fetcher) + ), + ) + try: + sdk.warm_pricing(["workers-ai"], workers_ai_models=["typesafe/jev", "@cf/meta/llama-3.2-3b-instruct"]) + # The catalog id is ignored here (it is in the catalog); only the partner id was fetched. + assert fetcher.cloudflare_gateway_cost_calls == ["typesafe/jev"] + assert fetcher.cloudflare_workers_ai_calls == 1 + hit = sdk._pricing.lookup("workers-ai", "typesafe/jev", "workers_ai_run") + assert hit is not None and hit.input == Decimal("0.000000042") + # Warm: a second warm_pricing for the same id fetches nothing. + sdk.warm_pricing(["workers-ai"], workers_ai_models=["typesafe/jev"]) + assert fetcher.cloudflare_gateway_cost_calls == ["typesafe/jev"] + finally: + sdk.shutdown(timeout=1.0) diff --git a/tests/unit/test_workers_ai_client.py b/tests/unit/test_workers_ai_client.py new file mode 100644 index 0000000..9c5c22f --- /dev/null +++ b/tests/unit/test_workers_ai_client.py @@ -0,0 +1,284 @@ +"""WorkersAI client tests — mocked HTTP, no live API.""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest +import responses + +from lago_agent_sdk import LagoSDK, WorkersAIError + +ACCT, GW = "acct_test", "gw_test" +LLAMA = "@cf/meta/llama-3.2-3b-instruct" +# `@cf/...` ids take the gateway host's path route (cache + log headers, model logged); +# partner ids take the unified `/ai/run` path with `cf-aig-gateway-id`, the only route +# where the gateway's BYOK key is consulted. See the workers_ai module docstring. +DIRECT_RUN = f"https://api.cloudflare.com/client/v4/accounts/{ACCT}/ai/run" +DIRECT_LLAMA = f"{DIRECT_RUN}/{LLAMA}" +GATEWAY_BASE = f"https://gateway.ai.cloudflare.com/v1/{ACCT}/{GW}/workers-ai" +GATEWAY_LLAMA = f"{GATEWAY_BASE}/{LLAMA}" + +CHAT_BODY = { + "result": { + "response": "Hello there!", + "model": "@cf/meta/llama-3.2-3b-instruct-v2", + "usage": {"prompt_tokens": 41, "completion_tokens": 34, "total_tokens": 75, "neurons": 1.22}, + }, + "success": True, + "errors": [], + "messages": [], +} +JEV_INPUT = {"state": "charged twice", "questions": {"is_urgent": {"type": "noul", "instructions": "?"}}} +JEV_BODY = { # the partner-model envelope: the model's object one level down (fixture 03) + "result": { + "state": "Completed", + "result": { + "model": "jev-1.13.0", + "answers": {"is_urgent": {"type": "noul", "noul": 0.97}}, + "usage": {"input_tokens": 446, "output_tokens": 73}, + }, + "gatewayMetadata": {"keySource": "BYOK"}, + }, + "success": True, + "errors": [], + "messages": [], +} +JEV_402 = { + "errors": [{"message": "Insufficient balance; add money to your gateway or use BYOK", "code": 2021}], + "success": False, + "result": {}, + "messages": [], +} + + +def _new_sdk(default_sub: str | None = "sub_test") -> tuple[LagoSDK, list[dict], list[str]]: + received: list[dict] = [] + errors: list[str] = [] + + def sender(batch: list[dict]) -> None: + received.extend(batch) + + sdk = LagoSDK(api_key="dummy", default_subscription_id=default_sub) + sdk._queue._sender = sender # type: ignore[attr-defined] + sdk.config.on_error = lambda exc, where: errors.append(f"{where}: {exc}") + return sdk, received, errors + + +def _by_code(received: list[dict]) -> dict[str, int]: + return {e["code"]: int(e["properties"]["value"]) for e in received} + + +# -------------------------------------------------------------------------- +# Catalog models (`@cf/...`) — gateway host path route +# -------------------------------------------------------------------------- +@responses.activate +def test_run_bills_tokens_with_model_provider_api_and_log_id() -> None: + sdk, received, errors = _new_sdk() + responses.post( + GATEWAY_LLAMA, json=CHAT_BODY, headers={"cf-aig-cache-status": "MISS", "cf-aig-log-id": "01LOG"} + ) + ai = sdk.workers_ai(ACCT, "tok", gateway_id=GW, gateway_auth="gwtok") + out = ai.run(LLAMA, {"messages": [{"role": "user", "content": "hi"}]}) + assert out["result"]["response"] == "Hello there!" # envelope returned unchanged + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert errors == [] + assert _by_code(received) == {"llm_input_tokens": 41, "llm_output_tokens": 34} + props = received[0]["properties"] + assert props["model"] == LLAMA # requested id = catalog key; served "-v2" name stays in extras + assert props["provider"] == "workers-ai" + assert props["api"] == "workers_ai_run" + assert props["cf_log_id"] == "01LOG" + assert received[0]["external_subscription_id"] == "sub_test" + + +@responses.activate +def test_catalog_model_takes_the_gateway_host_with_gateway_auth() -> None: + sdk, _, _ = _new_sdk("sub_acme") + responses.post(GATEWAY_LLAMA, json=CHAT_BODY) + ai = sdk.workers_ai(ACCT, "tok", gateway_id=GW, gateway_auth="gwtok") + assert ai.url_for(LLAMA) == GATEWAY_LLAMA + ai.run(LLAMA, {"messages": []}, extra_headers={"cf-aig-cache-ttl": "300"}) + sdk.shutdown(timeout=1.0) + req = responses.calls[0].request + assert json.loads(req.body) == {"messages": []} # path route: the body IS the input + assert req.headers["Authorization"] == "Bearer tok" + assert req.headers["cf-aig-authorization"] == "Bearer gwtok" + assert req.headers["cf-aig-cache-ttl"] == "300" + assert "cf-aig-gateway-id" not in req.headers and "cf-aig-skip-cache" not in req.headers + # The resolved subscription rides along, so the Logs API backfill attributes the same way. + assert json.loads(req.headers["cf-aig-metadata"]) == {"lago_subscription": "sub_acme"} + + +@responses.activate +def test_direct_route_when_no_gateway_sets_no_gateway_headers() -> None: + sdk, received, _ = _new_sdk() + responses.post(DIRECT_LLAMA, json=CHAT_BODY) + ai = sdk.workers_ai(ACCT, "tok") + assert ai.url_for(LLAMA) == DIRECT_LLAMA + assert ai.url_for("typesafe/jev") == DIRECT_RUN + ai.run(LLAMA, {"messages": []}) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + req = responses.calls[0].request + assert json.loads(req.body) == {"messages": []} + for h in ("cf-aig-authorization", "cf-aig-gateway-id", "cf-aig-skip-cache", "cf-aig-metadata"): + assert h not in req.headers # nothing stores or reads these without a gateway + assert _by_code(received) == {"llm_input_tokens": 41, "llm_output_tokens": 34} + assert "cf_log_id" not in received[0]["properties"] + + +@responses.activate +def test_gateway_cache_hit_is_not_billed() -> None: + """The gateway replays the identical body, usage included (fixtures 06/07). Only the + header says the model never ran — so only the header can stop the bill.""" + sdk, received, errors = _new_sdk() + responses.post( + GATEWAY_LLAMA, json=CHAT_BODY, headers={"cf-aig-cache-status": "HIT", "cf-aig-log-id": "01HIT"} + ) + ai = sdk.workers_ai(ACCT, "tok", gateway_id=GW, gateway_auth="gwtok") + out = ai.run(LLAMA, {"messages": []}) + assert out["result"]["usage"]["prompt_tokens"] == 41 # the caller still gets the body + sdk.shutdown(timeout=1.0) + assert received == [] + assert errors == [] + + +# -------------------------------------------------------------------------- +# Partner models (`typesafe/jev`) — unified path through the BYOK gateway +# -------------------------------------------------------------------------- +@responses.activate +def test_partner_model_takes_the_unified_path_through_the_byok_gateway() -> None: + """Model in the body, `cf-aig-gateway-id` names the gateway whose BYOK holds the partner + key, cache skipped because this path returns no cache header, and the nested + `result.result.usage` bills.""" + sdk, received, errors = _new_sdk("sub_acme") + responses.post(DIRECT_RUN, json=JEV_BODY) + ai = sdk.workers_ai(ACCT, "tok", gateway_id=GW, gateway_auth="gwtok") + assert ai.url_for("typesafe/jev") == DIRECT_RUN + out = ai.run("typesafe/jev", JEV_INPUT) + assert out["result"]["result"]["answers"]["is_urgent"]["noul"] == 0.97 + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + req = responses.calls[0].request + assert json.loads(req.body) == {"model": "typesafe/jev", "input": JEV_INPUT} + assert req.headers["Authorization"] == "Bearer tok" + assert req.headers["cf-aig-gateway-id"] == GW + assert req.headers["cf-aig-skip-cache"] == "true" + assert "cf-aig-authorization" not in req.headers # gateway auth belongs to the gateway host only + assert json.loads(req.headers["cf-aig-metadata"]) == {"lago_subscription": "sub_acme"} + assert errors == [] + assert _by_code(received) == {"llm_input_tokens": 446, "llm_output_tokens": 73} + assert received[0]["properties"]["model"] == "typesafe/jev" + assert received[0]["properties"]["provider"] == "workers-ai" + assert "cf_log_id" not in received[0]["properties"] # the unified path returns no log id + + +@responses.activate +def test_extra_headers_win_over_the_client_defaults() -> None: + sdk, _, _ = _new_sdk() + responses.post(DIRECT_RUN, json=JEV_BODY) + ai = sdk.workers_ai(ACCT, "tok", gateway_id=GW) + ai.run("typesafe/jev", JEV_INPUT, extra_headers={"cf-aig-skip-cache": "false"}) + sdk.shutdown(timeout=1.0) + assert responses.calls[0].request.headers["cf-aig-skip-cache"] == "false" + + +@responses.activate +def test_402_raises_workers_ai_error_and_bills_nothing() -> None: + sdk, received, errors = _new_sdk() + responses.post(DIRECT_RUN, json=JEV_402, status=402) + ai = sdk.workers_ai(ACCT, "tok", gateway_id=GW, gateway_auth="gwtok") + with pytest.raises(WorkersAIError) as ei: + ai.run("typesafe/jev", JEV_INPUT) + assert ei.value.status_code == 402 + assert ei.value.errors[0]["code"] == 2021 + assert "Insufficient balance" in str(ei.value) + sdk.shutdown(timeout=1.0) + assert received == [] + assert errors == [] # the customer's error, not an instrumentation failure + + +@responses.activate +def test_success_false_with_200_is_still_an_error() -> None: + sdk, received, _ = _new_sdk() + responses.post(DIRECT_RUN, json={**JEV_402}, status=200) + ai = sdk.workers_ai(ACCT, "tok") + with pytest.raises(WorkersAIError): + ai.run("typesafe/jev", JEV_INPUT) + sdk.shutdown(timeout=1.0) + assert received == [] + + +# -------------------------------------------------------------------------- +# Options, attribution, failure isolation +# -------------------------------------------------------------------------- +@responses.activate +def test_per_call_extra_lago_overrides_subscription_and_adds_dimensions() -> None: + sdk, received, _ = _new_sdk("sub_default") + responses.post(GATEWAY_LLAMA, json=CHAT_BODY) + ai = sdk.workers_ai(ACCT, "tok", gateway_id=GW, gateway_auth="gwtok", dimensions={"team": "billing"}) + ai.run( + LLAMA, + {"messages": []}, + extra_lago={"subscription": "sub_override", "dimensions": {"ticket": "T-1"}}, + ) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert all(e["external_subscription_id"] == "sub_override" for e in received) + assert received[0]["properties"]["team"] == "billing" + assert received[0]["properties"]["ticket"] == "T-1" + req = responses.calls[0].request + assert json.loads(req.headers["cf-aig-metadata"]) == {"lago_subscription": "sub_override"} + + +@responses.activate +def test_no_resolvable_subscription_drops_with_on_error() -> None: + sdk, received, errors = _new_sdk(default_sub=None) + responses.post(DIRECT_LLAMA, json=CHAT_BODY) + ai = sdk.workers_ai(ACCT, "tok") + ai.run(LLAMA, {"messages": []}) + sdk.shutdown(timeout=1.0) + assert received == [] + assert any("no resolvable subscription" in e for e in errors) + + +def test_stream_is_refused_before_any_request() -> None: + sdk, _, _ = _new_sdk() + ai = sdk.workers_ai(ACCT, "tok") + with pytest.raises(ValueError, match="stream=True"): + ai.run(LLAMA, {"messages": [], "stream": True}) + sdk.shutdown(timeout=1.0) + + +@responses.activate +def test_instrumentation_failure_does_not_break_the_call(monkeypatch: pytest.MonkeyPatch) -> None: + import lago_agent_sdk.workers_ai as mod + + def boom(*a: Any, **k: Any) -> Any: + raise RuntimeError("adapter bug") + + monkeypatch.setattr(mod, "extract_workers_ai_native", boom) + sdk, received, errors = _new_sdk() + responses.post(DIRECT_LLAMA, json=CHAT_BODY) + ai = sdk.workers_ai(ACCT, "tok") + out = ai.run(LLAMA, {"messages": []}) + assert out["result"]["response"] == "Hello there!" + sdk.shutdown(timeout=1.0) + assert received == [] + assert errors and "adapter bug" in errors[0] + + +@responses.activate +def test_non_json_error_body_still_raises_cleanly() -> None: + sdk, received, _ = _new_sdk() + responses.post(DIRECT_LLAMA, body="bad gateway", status=502) + ai = sdk.workers_ai(ACCT, "tok") + with pytest.raises(WorkersAIError) as ei: + ai.run(LLAMA, {"messages": []}) + assert ei.value.status_code == 502 + assert ei.value.errors == [] + sdk.shutdown(timeout=1.0) + assert received == [] From 5958be45f48968e3b9f8edc06ed87e4ca01b0856 Mon Sep 17 00:00:00 2001 From: Anass Date: Mon, 21 Sep 2026 18:33:23 +0200 Subject: [PATCH 2/3] Scope the Cloudflare pagination fake to the catalog URL The fake swapped in for requests.get counted every call made while it was active, including a background price refresh from another test's still-running queue thread, so the bounded-pagination test saw 41 pages where the loop fetched 40 (CI, ubuntu/py3.10). Only the catalog's own URL counts now; anything else goes to the real function. --- tests/unit/test_pricing.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/unit/test_pricing.py b/tests/unit/test_pricing.py index 321e9e1..216796c 100644 --- a/tests/unit/test_pricing.py +++ b/tests/unit/test_pricing.py @@ -709,6 +709,13 @@ def json(self): return self._b def fake(url, **kw): + # `requests.get` is swapped module-wide, and another test's SDK may still have a + # queue thread refreshing prices in the background. Count only the catalog's own + # URL and hand anything else to the real function (which pytest-socket blocks), or + # one stray background call lands in `seen` and the page count is off by one — + # observed once in CI on ubuntu/py3.10 (41 pages where the loop fetched 40). + if "/ai/models/search" not in str(url): + return orig(url, **kw) page = int((kw.get("params") or {}).get("page", 1)) seen.append(page) return _Resp(pages[page - 1] if page - 1 < len(pages) else {"result": [], "result_info": {}}) From 692093cbe9f7a783fb051bb87fca12a659ec3821 Mon Sep 17 00:00:00 2001 From: Anass Date: Mon, 21 Sep 2026 19:54:53 +0200 Subject: [PATCH 3/3] Bill BYOK rows as tokens in the Logs API backfill example The example passed the entry's `cost` straight through. On a row served with the customer's own provider key that field is Cloudflare's list price for a call Cloudflare never charged, so following the example overbilled every BYOK-served partner-model call. The example now checks `extras["byok"]`, which this PR surfaces, and bills token counts for those rows. --- docs/cloudflare.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/cloudflare.md b/docs/cloudflare.md index a812a4c..5c88506 100644 --- a/docs/cloudflare.md +++ b/docs/cloudflare.md @@ -29,6 +29,12 @@ from lago_agent_sdk.gateway.adapters import extract_cloudflare_log, resolve_subs for entry in fetch_gateway_logs(): # GET .../ai-gateway/gateways/{id}/logs usage = extract_cloudflare_log(entry) sub = resolve_subscription(entry) or "sub_default" # from the call's cf-aig-metadata, if set + if usage.extras.get("byok"): + # Served with the customer's own provider key (BYOK): Cloudflare charged nothing and the + # partner bills them directly, yet `cost` still carries Cloudflare's list price. Bill the + # tokens, never that number. + sdk.emit(usage, subscription=sub, mode="tokens", event_id=f"cf_{entry['id']}") + continue sdk.emit(usage, subscription=sub, mode="price", usd_cost=entry.get("cost") or 0, event_id=f"cf_{entry['id']}") sdk.flush() ```