Workers AI client + partner-model pricing (typesafe/jev) - #40
anassg-lago wants to merge 3 commits into
Conversation
…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).
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.
sarkissianraffi
left a comment
sarkissianraffi
left a comment
There was a problem hiding this comment.
Reviewed against the harness: billing correctness, security sweep, gates, and the mirrored JS PR (getlago/lago-agent-sdk-js#60).
Blocker
- See inline comment on
docs/cloudflare.md. The newbyokfield is exposed and tested at the extraction layer, but the pre-existing generic Logs-API backfill example on the same page still passesentry.get("cost")straight through with nobyokcheck, so a customer following the documented pattern will overbill every BYOK-served partner-model row. The fix is local to this same file.
What I checked
uv sync --all-extras --all-groups && uv run pytest tests/unit -q: 979 passed, matches the PR description.uv run ruff check .: clean.uv run mypy --strict src: clean, matches the claimed "ruff + mypy strict".- Verified the math: 446 input tokens x $0.042/1M = $0.000018732, matches the new
money_golden.jsonrow and the PR's "verified live" claim. - Diffed every new fixture (
workers_ai/01..08,cloudflare_gateway/03_real_workers_ai_failed.json) and the newmoney_golden.jsonrow against the JS twin: byte-identical, as claimed. - Security sweep on
git log -p main..pr40, the capture script, and all fixtures: no real account ids, tokens, or hostnames beyond the documented Cloudflare/TypeSafe endpoints.capture_workers_ai.pyreads credentials from env vars only. - Route selection (
@cf/...-> gateway/direct host, partner id -> unified/ai/runwithcf-aig-gateway-id), the cache-HIT skip, and thecf-aig-skip-cache: trueguard on the unified path (which has no cache header to check) are all correctly reasoned and match the module docstring's measured claims. - Test coverage is good: cache hit not billed, 402/403 bill nothing,
success: falsewith HTTP 200 treated as an error, stream refused up front, instrumentation failure never breaks the caller's call.
Non-blocking
- The
_request_for/requestFortwo-route split is a clean way to keep onerun()method honest about two very different wire shapes. Good use of measured, dated comments instead of guessing at Cloudflare's behavior.
Merge order
- Only other open PR in this repo is dependabot #4 (CI bump), no overlap.
- Mirrors getlago/lago-agent-sdk-js#60 fixture-for-fixture and row-for-row. The same blocker applies there; fix both in the same pass so they stay identical.
REQUEST_CHANGES for the byok/backfill-example gap. Everything else here is solid.
Generated by Claude Code
|
|
||
| **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"]`. |
There was a problem hiding this comment.
Blocker: this paragraph states the invariant (byok tells a backfill not to bill the row) but the docs page's own generic backfill example above ("Backfill from the Logs API", the for entry in fetch_gateway_logs(): ... usd_cost=entry.get("cost") or 0" snippet) does not check it. That snippet is the money path most integrators will copy verbatim. As written, running it against a window that includes a BYOK-served typesafe/jevcall bills Cloudflare's phantom list-pricecost` even though Cloudflare charged nothing and the partner already bills the customer directly, an overbill.
Small fix: have that example skip or zero the cost when usage.extras.get("byok") is truthy (mirroring what this PR already tests in test_byok_key_source_reaches_extras_on_every_entry), or at minimum add a one-line warning right at the example, not only in this paragraph 40 lines below it.
Generated by Claude Code
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.
|
Fixed in 692093c: the backfill example now checks |
sarkissianraffi
left a comment
sarkissianraffi
left a comment
There was a problem hiding this comment.
Re-review after the byok/backfill fix. Only docs/cloudflare.md changed since my last review, same file I flagged.
The fix adds an extras.get("byok") check next to the loop and routes those rows through mode="tokens" instead of usd_cost=entry.get("cost"). Mirrors the JS fix exactly (extras.byok / mode: "tokens"), and the underlying extraction/token-population logic hasn't changed since I already verified it last round.
Re-ran the gates
uv sync --all-extras --all-groups && uv run pytest tests/unit -q: 979 passed, same as before. (The twoopenrouter.aifetch failures in the output are the sandbox proxy blocking that outbound call inside an instrumentation-isolation test, not a code failure, no test reported failed.)uv run ruff check .: clean.uv run mypy --strict src: clean.- CI on 692093c: green across py3.10/3.11/3.12 x ubuntu/macos.
- Diffed the two new commits (python 692093c, js 4208b25): each touches only
docs/cloudflare.md, no drift beyond the byok guard.
Blocker resolved. Nothing else outstanding.
Merge order (unchanged from last review)
- Only other open PR in this repo is dependabot #4 (CI bump), no overlap.
APPROVE.
Generated by Claude Code
2 participants
Why
Cloudflare's partner models on Workers AI —
typesafe/jevtoday — can't be wrapped: they aren't chat-shaped, and no client library can call any Workers AI model (the officialcloudflarepackage percent-encodes the slash in every model id). Until now the only way to bill them was the Logs API backfill.What
sdk.workers_ai(account_id, token, gateway_id=…, gateway_auth=…)— a one-method client,run(model, input, extra_lago=…).@cf/…models use the gateway host (cache hits skipped,cf_log_iddimension for reconciliation); partner models use the unified/ai/runpath withcf-aig-gateway-id, the only route where the gateway's stored BYOK key applies.extract_workers_ai_native, backed by 8 captured fixtures: chat models, a reasoning model, a real gateway cache MISS/HIT pair, and Jev both via Cloudflare and straight from TypeSafe. The requested id is the billing key (served names like…-24b-v2miss the price catalog); the served name is kept inextras.ai-gateway/coststable, one row per id, fetched in the background;warm_pricing(workers_ai_models=[…])prices the very first call too. Rows keep serving past the TTL while they refresh.extract_cloudflare_logexposesbyok— a BYOK row carries a list-pricecostCloudflare never charged, and a backfill must not bill it.Verified live (local Lago)
$0.000018732= the gateway's own logcost→ newmoney_golden.jsonrow.JS counterpart: getlago/lago-agent-sdk-js#60 (same fixtures and golden row, byte-identical).