From 24eee43ebfd9e27cd5c53707c8f198e24e8d2c98 Mon Sep 17 00:00:00 2001 From: Anass Date: Tue, 25 Aug 2026 19:51:20 +0200 Subject: [PATCH 1/2] Account for additive cache counts in the total_tokens guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An OpenAI-compatible wire says nothing about the token CONVENTION behind it, and the consistency guard assumed it did. The guard folds any positive `total_tokens - (input + output + reasoning)` remainder into `output`, on the premise that cache counts always sit INSIDE `prompt_tokens` and so can never appear in that remainder. That held for all three surfaces this adapter served when it was written: native OpenAI (zero deltas across every capture), Databricks (112 of 112 rows with total == input + output), and Cloudflare (cache outside `input`, but outside `total_tokens` too, so it never inflated the delta). Snowflake Cortex breaks it. Cortex answers on an OpenAI wire — /api/v2/cortex/v1/chat/completions, byte-for-byte the OpenAI payload shape — with Anthropic's ADDITIVE convention. Measured live 2026-08-25: prompt_tokens 7 prompt_tokens_details.cached_tokens 4805 completion_tokens 6 total_tokens 4818 = 7 + 4805 + 6 The cached block sits outside `prompt_tokens` and inside `total_tokens`, so 4,805 tokens looked unaccounted and were added to `output`: 4,811 reported for a call that generated 6, while the same 4,805 also shipped as cache_read. On the first capture that was 17,503 tokens billed for 8,758 consumed (2.0x), with the output line inflated ~800x. Not Snowflake-specific — any OpenAI-compatible proxy with additive caching hits it. The reconciliation now subtracts cache_read and the raw prompt_tokens_details.cache_write_tokens as well. cache_write_tokens is read straight off the payload rather than through a canonical field because it is deliberately NOT mapped to CanonicalUsage.cache_write — for OpenAI it sits inside prompt_tokens and billing it separately over-charges 2.24x — so it has no other route into the arithmetic, and an additive cache WRITE would inflate `output` exactly the way the read did. Subtracting cannot suppress a genuine fold: on a subtractive surface those counts are already inside `input`, so removing them again only drives the delta further negative, where the `> 0` guard already no-ops. Re-verified across every captured OpenAI, Databricks and Cloudflare fixture — all still 0 — and a payload carrying both an additive cache block and hidden thinking tokens still folds the thinking remainder alone. Deliberately NOT gated on the details objects being empty, which was the other candidate fix: that suppresses genuine folds on any proxy that reports a details block alongside unreported tokens, reintroducing the silent under-bill the guard exists for. This hazard was raised in review on #14 (2026-08-17) with the exact payload shape, and answered with a three-surface census returning unaccounted = 0 everywhere. That census was correct; it expired when a surface with additive caching arrived. It is now pinned by live fixtures rather than by an argument: 11_snowflake_cortex_plain_chat.json and 12_snowflake_cortex_cache_chat.json, captured by capture_snowflake_cortex.py and byte-identical to the JS port's copies. Reverting the subtraction fails three tests in each repo. --- CHANGELOG.md | 4 + src/lago_agent_sdk/adapters/openai_native.py | 28 +++- .../fixtures/capture_snowflake_cortex.py | 124 ++++++++++++++++++ .../11_snowflake_cortex_plain_chat.json | 50 +++++++ .../12_snowflake_cortex_cache_chat.json | 50 +++++++ tests/unit/adapters/test_openai_native.py | 101 ++++++++++++++ 6 files changed, 356 insertions(+), 1 deletion(-) create mode 100644 tests/unit/adapters/fixtures/capture_snowflake_cortex.py create mode 100644 tests/unit/adapters/fixtures/openai_native/11_snowflake_cortex_plain_chat.json create mode 100644 tests/unit/adapters/fixtures/openai_native/12_snowflake_cortex_cache_chat.json diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b8d529..bc20feb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **A cached token on an OpenAI-compatible endpoint with additive cache semantics was billed twice, the second time at the output rate.** The `total_tokens` consistency guard folds any positive `total_tokens - (input + output + reasoning)` remainder into `output`, on the premise — verified at the time across all three surfaces this adapter served — that cache counts always sit INSIDE `prompt_tokens`, so they can never appear in that remainder. Snowflake Cortex breaks the premise: it answers on an OpenAI **wire** (`/api/v2/cortex/v1/chat/completions`) with Anthropic's **additive** convention. Measured live against the real endpoint on 2026-08-25: `prompt_tokens: 7`, `prompt_tokens_details.cached_tokens: 4805`, `completion_tokens: 6`, `total_tokens: 4818` — the cached block sits outside `prompt_tokens` and inside `total_tokens`. The adapter therefore read 4,805 tokens as unaccounted and folded them into `output`, reporting **4,811 output tokens for a call that generated 6**, while the same 4,805 also shipped as `cache_read`: 17,503 tokens billed for 8,758 consumed on the original capture (**2.0x**), with the output line inflated ~800x. The reconciliation now subtracts `cache_read` and the raw `prompt_tokens_details.cache_write_tokens` as well. + - **`cache_write_tokens` is accounted for without being mapped.** It is deliberately absent from `CanonicalUsage.cache_write` — for OpenAI it sits inside `prompt_tokens` and billing it separately over-charges 2.24x — so it has no canonical route into the arithmetic and is read straight off the payload instead. Without that, an additive cache *write* inflates `output` exactly the way the read did. + - **The correction cannot suppress a genuine fold.** On a subtractive surface the cache counts are already inside `input`, so removing them again only drives the delta further negative, where the existing `> 0` guard already no-ops. Re-verified across every captured OpenAI, Databricks and Cloudflare fixture — all still 0 — and a payload carrying both an additive cache block and hidden thinking tokens still folds the thinking remainder alone. + - **Reported in review and answered "unreachable" — correctly, at the time.** PR #14 (2026-08-17) raised that `reasoning` was not a priori the only field able to inflate `total_tokens`, and named the exact payload shape. The answer was a three-surface census returning `unaccounted = 0` everywhere, which held until an OpenAI-compatible surface with additive caching existed. The regression fixtures are captured from that surface (`11_snowflake_cortex_plain_chat.json`, `12_snowflake_cortex_cache_chat.json`, via `capture_snowflake_cortex.py`) so the census is now pinned by a live payload rather than by an argument. - **An explicitly-passed falsy `api_url` silently resolved to PRODUCTION Lago.** Preferring the config value over `""` is right — `requests` raises `MissingSchema`, which is not a `LagoApiError`, so the queue classified it transient, re-prepended the batch and retried at the 60s ceiling forever, stopping all billing with only a growing buffer as the symptom. But `LagoConfig`'s default is the production URL, so `api_url=os.environ.get("LAGO_API_URL", "")` with the var unset resolved to production with **no `on_error` and no log** — verified live: 0 reports, 0 log lines, and a client posting to `api.getlago.com`. For a CI job or a developer holding a real production key that writes live billing data, and ingested events cannot be un-ingested. The fallback is unchanged, so the original config-clobber bug stays fixed; it is now reported under `config.api_url` through the same log-plus-callback floor as every other drop path. An *unpassed* `api_url` stays silent — `None` means the caller never mentioned it, and reporting the common case would train customers to ignore the channel this fix depends on. - **`usage_metadata` from the Cloudflare gateway got no drift sweep, and had already lost two counters.** `extras` was a fixed three-key dict, so any counter the adapter does not map vanished with no error and no `on_error` — the one place violating the drift contract `test_drift.py` enforces for the native adapters. Not hypothetical: replaying the **14 captured fixtures** through the adapter drops **`neurons`** (Cloudflare's Workers AI billing unit) in 4 entries and **`input_text_tokens`** in 1, and a live Logs API pull also returns **`units`**, a cost quantity that appears in no fixture at all — the hand-maintained key enumeration in the module docstring had already drifted past reality, which is exactly the failure mode a snapshot invites. Unmapped keys are now swept into `extras["usage_metadata"]` against an explicit `_MAPPED_USAGE_KEYS` set, so a ninth spelling surfaces on its own instead of needing another 14-fixture audit. Deliberately **nested** rather than merged flat into `extras`: the poller reads `extras["cached"]` to decide whether to skip billing a request Cloudflare served for free, so a future `usage_metadata` key called `cached` or `step` must not be able to shadow it. The regression test iterates the fixture directory rather than a fixed key list, so a recapture that introduces a new counter fails it with no test edit. Closes #16. - **`prime()` re-downloaded the ~400-model OpenRouter catalogue regardless of the TTL.** It set `_openrouter_stale` unconditionally, and it is reached from `_auto_prime_pricing_for` on a matching `wrap()` and from `warm_pricing()` — both of which a server can run per request — so `pricing_ttl_seconds` never applied on that path at all, and the catalogue was refetched on essentially every flush tick, on the thread the queue drains events from. Measured with the shipped 1-hour TTL: **4 `prime()`+`maybe_refresh()` cycles produced 4 full downloads where 1 was correct; now 1.** Gated on the same "no table, or past the TTL" test `lookup()` already uses, so priming and looking up cannot disagree about what needs fetching — and a table that genuinely ages out is still re-primed, so prices do not freeze at the first fetch. diff --git a/src/lago_agent_sdk/adapters/openai_native.py b/src/lago_agent_sdk/adapters/openai_native.py index 547de8e..2534fb6 100644 --- a/src/lago_agent_sdk/adapters/openai_native.py +++ b/src/lago_agent_sdk/adapters/openai_native.py @@ -254,10 +254,36 @@ def extract_openai_native(response: Any, model_id: str = "", provider_hint: str # breakdown at all (measured: prompt 57, completion 47, total 1253) — still # recovers its 1,149 tokens, because reasoning is 0 there. # + # The cache counts are subtracted for the SAME reason as reasoning, and this was + # the half that was missing. The guard assumed every OpenAI-shaped surface reports + # cache_read INSIDE prompt_tokens, which held for all three surfaces that existed + # when it was written (native OpenAI: zero deltas; Databricks: 112/112 rows with + # total == input + output; Cloudflare: cache outside `input` but outside `total` + # too, so it never inflated the delta). Snowflake Cortex is the surface that broke + # it — an OpenAI-WIRE endpoint with Anthropic's ADDITIVE convention: measured + # 2026-08-25, prompt_tokens=7, cached_tokens=4805, completion_tokens=6, + # total_tokens=4818, i.e. the cached block sits OUTSIDE prompt_tokens and INSIDE + # total_tokens. Accounting for only input+output+reasoning made those 4,805 cached + # tokens look unaccounted, so they were folded into `output` — 4,811 reported for a + # call that generated 6, while the same tokens also shipped as cache_read. 2.0x on + # the call, 800x on the output line. See 12_snowflake_cortex_cache_chat.json. + # + # `cache_write_tokens` is read straight from the payload rather than from a + # canonical field because it is deliberately NOT mapped to CanonicalUsage.cache_write + # (for OpenAI it sits inside prompt_tokens and billing it separately over-charges + # 2.24x — see _MAPPED_DETAIL_FIELDS). It still has to be accounted for here, or an + # additive cache WRITE would inflate `output` exactly the way the read did. + # + # Subtracting them cannot suppress a genuine fold on a subtractive surface: there + # the cache counts are already inside `input`, so removing them again only drives + # the delta further negative, where the `> 0` guard already no-ops. Verified against + # every captured OpenAI, Databricks and Cloudflare fixture — all still 0. + # # A no-op for real OpenAI either way: total always equals prompt + completion. declared_total = _safe_int(usage.get("total_tokens")) if declared_total: - unaccounted = declared_total - (input_tokens + output_tokens + reasoning) + cache_write = _safe_int(_safe_dict(usage.get("prompt_tokens_details")).get("cache_write_tokens")) + unaccounted = declared_total - (input_tokens + output_tokens + reasoning + cache_read + cache_write) if unaccounted > 0: output_tokens += unaccounted extras["unaccounted_output_tokens"] = unaccounted diff --git a/tests/unit/adapters/fixtures/capture_snowflake_cortex.py b/tests/unit/adapters/fixtures/capture_snowflake_cortex.py new file mode 100644 index 0000000..701dad0 --- /dev/null +++ b/tests/unit/adapters/fixtures/capture_snowflake_cortex.py @@ -0,0 +1,124 @@ +"""Capture Snowflake Cortex responses off the OpenAI-compatible endpoint. + +Saves to: + tests/unit/adapters/fixtures/openai_native/11_snowflake_cortex_plain_chat.json + tests/unit/adapters/fixtures/openai_native/12_snowflake_cortex_cache_chat.json + +These live under `openai_native/` on purpose: Cortex answers on an OpenAI-wire +endpoint, so `extract_openai_native` / `extractOpenAINative` is the adapter that +serves them. They are the surface that proves the `total_tokens` reconciliation +cannot assume OpenAI's subtractive cache convention — on Cortex, `cached_tokens` +sits OUTSIDE `prompt_tokens` and INSIDE `total_tokens`. + +Two things about Cortex that this script encodes, both measured 2026-08-25: + + * Caching only happens with an explicit Anthropic-style `cache_control` part. + The same 4,800-token prompt sent twice WITHOUT it reports `cached_tokens: 0` + both times, so the "call1 then call2" pattern the OpenAI cache fixtures use + captures nothing here. + * `max_tokens` is rejected outright ("deprecated in favor of + max_completion_tokens"), unlike OpenAI which still accepts it. + +Requires a Snowflake account with the Cortex REST endpoint entitled — it returns +403 `003001` otherwise, which is an account-level grant no config can work around. + + SNOWFLAKE_HOST=-.snowflakecomputing.com \ + SNOWFLAKE_PAT= \ + python3 capture_snowflake_cortex.py + +Idempotent: skips files that already exist. Re-run after deleting one to refresh it. +""" + +from __future__ import annotations + +import json +import os +import pathlib +import sys + +import requests + +MODEL = "claude-sonnet-4-5" +OUT = pathlib.Path(__file__).parent / "openai_native" + +# Long enough to clear Anthropic's minimum cacheable prefix. Deliberately dull, +# fixed text: the fixture must not carry anything account- or person-identifying. +CACHEABLE_PREFIX = "Reference notes on widget calibration tolerances, revision seven. " * 400 + + +def call(host: str, pat: str, body: dict) -> dict: + r = requests.post( + f"https://{host}/api/v2/cortex/v1/chat/completions", + headers={ + "Authorization": f"Bearer {pat}", + "X-Snowflake-Authorization-Token-Type": "PROGRAMMATIC_ACCESS_TOKEN", + "Content-Type": "application/json", + "Accept": "application/json", + }, + json=body, + timeout=120, + ) + if r.status_code != 200: + sys.exit(f"Cortex returned {r.status_code}: {r.text[:300]}") + return r.json() + + +def save(name: str, response: dict) -> None: + path = OUT / name + if path.exists(): + print(f"skip {name} (exists)") + return + path.write_text(json.dumps({"_model_id": MODEL, "_response": response}, indent=2) + "\n") + print(f"wrote {name}") + + +def main() -> None: + host = os.environ.get("SNOWFLAKE_HOST") + pat = os.environ.get("SNOWFLAKE_PAT") + if not host or not pat: + sys.exit("set SNOWFLAKE_HOST and SNOWFLAKE_PAT") + + save( + "11_snowflake_cortex_plain_chat.json", + call( + host, + pat, + { + "model": MODEL, + "messages": [{"role": "user", "content": "What is 2 + 2? Answer in one word."}], + "max_completion_tokens": 32, + }, + ), + ) + + # The regression fixture. `cache_control` is what makes Cortex report a cached + # block at all, and the resulting payload is the one that used to inflate + # `output` by the whole cached count. + save( + "12_snowflake_cortex_cache_chat.json", + call( + host, + pat, + { + "model": MODEL, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": CACHEABLE_PREFIX, + "cache_control": {"type": "ephemeral"}, + }, + {"type": "text", "text": "Reply with one word."}, + ], + } + ], + "max_completion_tokens": 32, + }, + ), + ) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/adapters/fixtures/openai_native/11_snowflake_cortex_plain_chat.json b/tests/unit/adapters/fixtures/openai_native/11_snowflake_cortex_plain_chat.json new file mode 100644 index 0000000..2dc742a --- /dev/null +++ b/tests/unit/adapters/fixtures/openai_native/11_snowflake_cortex_plain_chat.json @@ -0,0 +1,50 @@ +{ + "_model_id": "claude-sonnet-4-5", + "_response": { + "choices": [ + { + "finish_reason": "", + "index": 0, + "message": { + "annotations": null, + "audio": { + "data": "", + "expires_at": 0, + "id": "", + "transcript": "" + }, + "content": "Four", + "function_call": { + "arguments": "", + "name": "" + }, + "refusal": "", + "role": "assistant", + "tool_calls": null + } + } + ], + "created": 1787679484, + "id": "", + "model": "claude-sonnet-4-5", + "object": "chat.completion", + "service_tier": "", + "system_fingerprint": "", + "usage": { + "completion_tokens": 4, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0 + }, + "prompt_tokens": 21, + "prompt_tokens_details": { + "audio_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 0 + }, + "total_tokens": 25 + } + } +} diff --git a/tests/unit/adapters/fixtures/openai_native/12_snowflake_cortex_cache_chat.json b/tests/unit/adapters/fixtures/openai_native/12_snowflake_cortex_cache_chat.json new file mode 100644 index 0000000..4a1f8a3 --- /dev/null +++ b/tests/unit/adapters/fixtures/openai_native/12_snowflake_cortex_cache_chat.json @@ -0,0 +1,50 @@ +{ + "_model_id": "claude-sonnet-4-5", + "_response": { + "choices": [ + { + "finish_reason": "", + "index": 0, + "message": { + "annotations": null, + "audio": { + "data": "", + "expires_at": 0, + "id": "", + "transcript": "" + }, + "content": "Acknowledged.", + "function_call": { + "arguments": "", + "name": "" + }, + "refusal": "", + "role": "assistant", + "tool_calls": null + } + } + ], + "created": 1787679488, + "id": "", + "model": "claude-sonnet-4-5", + "object": "chat.completion", + "service_tier": "", + "system_fingerprint": "", + "usage": { + "completion_tokens": 6, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0 + }, + "prompt_tokens": 7, + "prompt_tokens_details": { + "audio_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 4805 + }, + "total_tokens": 4818 + } + } +} diff --git a/tests/unit/adapters/test_openai_native.py b/tests/unit/adapters/test_openai_native.py index fc776e5..4c3da25 100644 --- a/tests/unit/adapters/test_openai_native.py +++ b/tests/unit/adapters/test_openai_native.py @@ -350,6 +350,47 @@ def test_total_tokens_guard_recovers_unaccounted_output() -> None: assert u.extras["unaccounted_output_tokens"] == 1149 +def test_total_tokens_guard_does_not_fold_an_additive_cache_write() -> None: + """The payload shape raised in review on PY #14: a proxy reporting cache-creation + tokens outside `prompt_tokens` but inside `total_tokens`. It was answered + "unreachable on the three surfaces we have", which was true at the time — + Snowflake Cortex then shipped the same class of payload with `cached_tokens`. + Accounted for now whether or not a live surface reports it this way, because + `cache_write_tokens` is deliberately never mapped to CanonicalUsage.cache_write + and so has no other route into the accounting.""" + u = extract_openai_native( + { + "usage": { + "prompt_tokens": 13, + "completion_tokens": 4, + "total_tokens": 1829, + "prompt_tokens_details": {"cache_write_tokens": 1812}, + } + } + ) + assert u.output == 4, "was 1816" + assert "unaccounted_output_tokens" not in u.extras + assert u.extras["prompt_tokens_details.cache_write_tokens"] == 1812 + + +def test_total_tokens_guard_still_recovers_a_remainder_beside_a_cache_count() -> None: + """The two corrections must not cancel each other: an additive cache block AND + hidden thinking tokens in the same payload. 20 + 5 + 100 = 125 accounted, + total 200, so 75 are real unreported output and must still fold.""" + u = extract_openai_native( + { + "usage": { + "prompt_tokens": 20, + "completion_tokens": 5, + "total_tokens": 200, + "prompt_tokens_details": {"cached_tokens": 100}, + } + } + ) + assert u.output == 80 + assert u.extras["unaccounted_output_tokens"] == 75 + + def test_total_tokens_guard_is_a_noop_for_genuine_openai() -> None: """For real OpenAI total_tokens == prompt + completion always holds, because reasoning is a SUBSET of completion rather than additive. Verified across @@ -381,3 +422,63 @@ def test_total_tokens_guard_ignores_a_negative_delta() -> None: u = extract_openai_native({"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 10}}) assert u.output == 50 assert "unaccounted_output_tokens" not in u.extras + + +# -------------------------------------------------------------------------- +# Snowflake Cortex — an OpenAI-wire endpoint with ADDITIVE cache +# +# Cortex answers on `/api/v2/cortex/v1/chat/completions` with OpenAI's exact +# payload shape, so this adapter serves it — but it does NOT follow OpenAI's +# token convention. Captured live 2026-08-25 by capture_snowflake_cortex.py; +# never hand-edit these numbers, recapture instead. +# -------------------------------------------------------------------------- +def test_snowflake_cortex_plain_call() -> None: + """No cache: total reconciles to prompt + completion, guard never fires.""" + model_id, resp = _load("11_snowflake_cortex_plain_chat.json") + u = extract_openai_native(resp, model_id=model_id, provider_hint="snowflake") + assert u.input == 21 + assert u.output == 4 + assert u.cache_read == 0 + assert u.provider == "snowflake" + assert u.api == "chat_completions" + assert "unaccounted_output_tokens" not in u.extras + + +def test_snowflake_cortex_cached_tokens_are_additive() -> None: + """THE regression. 7 + 4805 + 6 = 4818, so under the old accounting (input + + output + reasoning only) the 4,805 cached tokens looked unaccounted and were + folded into `output`: 4,811 reported for a call that generated 6, while the + same tokens also shipped as cache_read — 2.0x on the call, 800x on the output + line. Revert the cache subtraction in openai_native.py and this fails on + `output`. + + This exact hazard was raised in review on PY #14 (2026-08-17) and answered + "measured 0 on all three surfaces we have" — true then. Cortex is the surface + that did not exist yet.""" + model_id, resp = _load("12_snowflake_cortex_cache_chat.json") + usage = resp["usage"] + assert usage["prompt_tokens"] + 4805 + usage["completion_tokens"] == usage["total_tokens"] + + u = extract_openai_native(resp, model_id=model_id, provider_hint="snowflake") + assert u.input == 7 + assert u.output == 6, "NOT 4811" + assert u.cache_read == 4805 + assert u.cache_write == 0 + assert u.reasoning == 0 + assert "unaccounted_output_tokens" not in u.extras + + +def test_snowflake_cortex_keeps_the_customers_model_spelling() -> None: + """A Cortex fine-tune answers as `database.schema.model`. CanonicalUsage.model + keeps it verbatim — normalising here would report a model the customer cannot + find in their own Snowflake account.""" + u = extract_openai_native( + { + "model": "mydb.myschema.my_tuned_model", + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + }, + model_id="mydb.myschema.my_tuned_model", + provider_hint="snowflake", + ) + assert u.model == "mydb.myschema.my_tuned_model" + assert u.provider == "snowflake" From 67c6766b42bf25347f7fdaca8dfd5f41f8800816 Mon Sep 17 00:00:00 2001 From: Anass Date: Thu, 27 Aug 2026 10:58:15 +0200 Subject: [PATCH 2/2] Key the total_tokens guard off the provider's token convention MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review on this branch showed the unconditional subtraction fixes Cortex by disarming the guard everywhere else, in both directions: * A SUBSET-convention surface with a cached block AND a genuine remainder under-folds by exactly the cached count. Gemini through Google's own OpenAI-compat layer — the surface the guard was written for — reports cached_tokens inside prompt_tokens: at prompt 1200 / cached 1000 / completion 47 / total 2396 the fold shrank from 1,149 to 149, so 1,000 generated tokens went unbilled, silently, with no on_error. Live-verified 2026-08-27: a real gemini-2.5-flash call through that layer carried a 155-token thinking remainder in total_tokens. * OpenAI's own measured cache_write shape (prompt_tokens=3025 WITH cache_write_tokens=3022 inside it) behind an under-reporting proxy lost the guard entirely: the 1,171-token remainder stopped folding and unaccounted_output_tokens vanished from extras. The payload cannot answer which convention is in play — `accounted <= total` admits a small subtractive cache, `cache_read > input` rejects a small additive one; both were run against real shapes and both leak. The only thing that knows the convention is the table pricing already bills from, so the guard now reads it: each subset field joins the accounted sum exactly when the provider reports it OUTSIDE its parent count. The sets move to token_semantics.py so the adapter does not import pricing's HTTP machinery — the guard, compute_cost and deoverlapped_token_total are three readings of one convention and now cannot disagree. One new set, INPUT_INCLUDES_CACHE_WRITE, records the measured OpenAI write-inside-prompt fact; a KNOWN_PROVIDERS roster plus test_token_semantics.py make absence from a set a recorded decision instead of a default. Also from the review, all verified before changing anything: * The Responses branch read cache_write from the chat branch's container, so the two API shapes disagreed about one provider's convention — it now reads per-branch, next to cache_read. * A test pins what UNHINTED Cortex traffic does (stamped openai, folds): the runtime fix is PR #26's /api/v2/cortex/ hint entry, and merge order matters — this branch alone does not fix live Cortex traffic, and never did: even with the unconditional subtraction, deoverlapped_token_total re-zeroed the cache_read of a payload stamped openai one layer down (unit=13 for a 4,818-token call). * The fixture test reads the cached count off the fixture instead of a literal, asserts the raw cache_write key in extras instead of the can-never-fail canonical field, and the fine-tune spelling test moved to the model-attribution block it belongs to. * capture_snowflake_cortex.py checks for an existing fixture BEFORE firing the request — save(call(...)) evaluated the call first, so a re-run on a complete checkout made two live calls and then printed "skip". Also documented: one cold call is enough, Cortex reports a cache CREATION under cached_tokens (measured on a matched pair against the account-usage view), so there is no warm-up race on recapture. * The "verified against every captured OpenAI, Databricks and Cloudflare fixture" claim narrowed to the 12 fixtures that actually traverse this function — the Databricks and Cloudflare sets go through the gateway adapters and never reach this guard. Measured for the snowflake entries, on top of the 2026-08-25 captures: llama on Cortex accepts cache_control and ignores it (cached_tokens 0 on a matched pair, total = prompt + completion), and reasoning_tokens is always 0 on this wire — thinking exists only on Cortex's Anthropic wire. The OpenAI family needs cross-region inference the capture account cannot enable; the set comments say to re-measure the day it becomes reachable, since Cortex documents caching per model family. Live-verified end to end against every adapter (openai cache pair, gemini compat + native, anthropic cache pair, mistral, bedrock, workers-ai via the CF gateway, Cortex hinted and unhinted) with the events accepted by a real Lago — values exact for all 30. --- CHANGELOG.md | 8 +- src/lago_agent_sdk/adapters/openai_native.py | 105 ++++++---- src/lago_agent_sdk/pricing.py | 104 ++------- src/lago_agent_sdk/token_semantics.py | 198 ++++++++++++++++++ .../fixtures/capture_snowflake_cortex.py | 74 ++++--- tests/unit/adapters/test_openai_native.py | 153 +++++++++++--- tests/unit/test_token_semantics.py | 100 +++++++++ 7 files changed, 555 insertions(+), 187 deletions(-) create mode 100644 src/lago_agent_sdk/token_semantics.py create mode 100644 tests/unit/test_token_semantics.py diff --git a/CHANGELOG.md b/CHANGELOG.md index bc20feb..53704fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,9 +6,11 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed -- **A cached token on an OpenAI-compatible endpoint with additive cache semantics was billed twice, the second time at the output rate.** The `total_tokens` consistency guard folds any positive `total_tokens - (input + output + reasoning)` remainder into `output`, on the premise — verified at the time across all three surfaces this adapter served — that cache counts always sit INSIDE `prompt_tokens`, so they can never appear in that remainder. Snowflake Cortex breaks the premise: it answers on an OpenAI **wire** (`/api/v2/cortex/v1/chat/completions`) with Anthropic's **additive** convention. Measured live against the real endpoint on 2026-08-25: `prompt_tokens: 7`, `prompt_tokens_details.cached_tokens: 4805`, `completion_tokens: 6`, `total_tokens: 4818` — the cached block sits outside `prompt_tokens` and inside `total_tokens`. The adapter therefore read 4,805 tokens as unaccounted and folded them into `output`, reporting **4,811 output tokens for a call that generated 6**, while the same 4,805 also shipped as `cache_read`: 17,503 tokens billed for 8,758 consumed on the original capture (**2.0x**), with the output line inflated ~800x. The reconciliation now subtracts `cache_read` and the raw `prompt_tokens_details.cache_write_tokens` as well. - - **`cache_write_tokens` is accounted for without being mapped.** It is deliberately absent from `CanonicalUsage.cache_write` — for OpenAI it sits inside `prompt_tokens` and billing it separately over-charges 2.24x — so it has no canonical route into the arithmetic and is read straight off the payload instead. Without that, an additive cache *write* inflates `output` exactly the way the read did. - - **The correction cannot suppress a genuine fold.** On a subtractive surface the cache counts are already inside `input`, so removing them again only drives the delta further negative, where the existing `> 0` guard already no-ops. Re-verified across every captured OpenAI, Databricks and Cloudflare fixture — all still 0 — and a payload carrying both an additive cache block and hidden thinking tokens still folds the thinking remainder alone. +- **A cached token on an OpenAI-compatible endpoint with additive cache semantics was billed twice, the second time at the output rate.** The `total_tokens` consistency guard folds any positive `total_tokens - (input + output + reasoning)` remainder into `output`, on the premise — verified at the time across all three surfaces this adapter served — that cache counts always sit INSIDE `prompt_tokens`, so they can never appear in that remainder. Snowflake Cortex breaks the premise: it answers on an OpenAI **wire** (`/api/v2/cortex/v1/chat/completions`) with Anthropic's **additive** convention. Measured live against the real endpoint on 2026-08-25: `prompt_tokens: 7`, `prompt_tokens_details.cached_tokens: 4805`, `completion_tokens: 6`, `total_tokens: 4818` — the cached block sits outside `prompt_tokens` and inside `total_tokens`. The adapter therefore read 4,805 tokens as unaccounted and folded them into `output`, reporting **4,811 output tokens for a call that generated 6**, while the same 4,805 also shipped as `cache_read`: 17,503 tokens billed for 8,758 consumed on the original capture (**2.0x**), with the output line inflated ~800x. The reconciliation now builds its accounted sum from the provider's own convention — each of `reasoning`, `cache_read` and the raw `cache_write_tokens` joins it exactly when that provider reports the count OUTSIDE its parent — read from the same table `compute_cost` and `deoverlapped_token_total` bill from, so the guard and the money paths cannot answer the convention question differently. + - **The convention tables moved to their own module, `token_semantics.py`**, because the adapter layer now needs them and importing `pricing`'s HTTP machinery from a pure extraction function was the wrong direction. `pricing` re-reads them from there; the sets themselves are unchanged except for the new `INPUT_INCLUDES_CACHE_WRITE` (OpenAI's cache-write is measured INSIDE `prompt_tokens`, so it must never join the accounted sum for OpenAI — an unconditional subtraction would have disarmed the guard on the one cache-write shape this codebase has actually measured). A roster (`KNOWN_PROVIDERS`) plus `test_token_semantics.py` pin that every provider the SDK can stamp has a RECORDED decision, so absence from a set is a choice, not a default. + - **Keyed on the provider, not inferred from the payload, deliberately.** Both payload-only gates fail on real shapes: `accounted <= total` still under-folds a subtractive surface with a small cached block, and `cache_read > input` folds tokens an additive surface never generated. On a subset surface the cache is already inside `input`, so subtracting it *unconditionally* (the first version of this fix) eats a genuine remainder — a Gemini-through-OpenAI-compat payload with a cached block and hidden thinking would have under-folded by exactly the cached count, silently. Both directions are now pinned by tests; re-verified across all 12 fixtures this adapter actually serves (the 10 OpenAI captures plus the two Cortex ones — the Databricks and Cloudflare fixture sets go through the gateway adapters and never reach this guard), and a payload carrying both an additive cache block and hidden thinking tokens still folds the thinking remainder alone. + - **`cache_write_tokens` is accounted for without being mapped.** It is deliberately absent from `CanonicalUsage.cache_write` — for OpenAI it sits inside `prompt_tokens` and billing it separately over-charges 2.24x — so it has no canonical route into the arithmetic and is read straight off the payload instead, from whichever details container the API shape uses (`prompt_tokens_details` on chat, `input_tokens_details` on Responses). Without that, an additive cache *write* inflates `output` exactly the way the read did. + - **An unhinted additive proxy still folds** — the payload carries no convention, so a surface the wrapper cannot identify from its `base_url` is stamped `openai` and gets OpenAI's subset semantics; the runtime fix for Cortex traffic is the `/api/v2/cortex/` hint-table entry shipping separately (PR #26), and a test pins what unhinted traffic does so the dependency is visible. - **Reported in review and answered "unreachable" — correctly, at the time.** PR #14 (2026-08-17) raised that `reasoning` was not a priori the only field able to inflate `total_tokens`, and named the exact payload shape. The answer was a three-surface census returning `unaccounted = 0` everywhere, which held until an OpenAI-compatible surface with additive caching existed. The regression fixtures are captured from that surface (`11_snowflake_cortex_plain_chat.json`, `12_snowflake_cortex_cache_chat.json`, via `capture_snowflake_cortex.py`) so the census is now pinned by a live payload rather than by an argument. - **An explicitly-passed falsy `api_url` silently resolved to PRODUCTION Lago.** Preferring the config value over `""` is right — `requests` raises `MissingSchema`, which is not a `LagoApiError`, so the queue classified it transient, re-prepended the batch and retried at the 60s ceiling forever, stopping all billing with only a growing buffer as the symptom. But `LagoConfig`'s default is the production URL, so `api_url=os.environ.get("LAGO_API_URL", "")` with the var unset resolved to production with **no `on_error` and no log** — verified live: 0 reports, 0 log lines, and a client posting to `api.getlago.com`. For a CI job or a developer holding a real production key that writes live billing data, and ingested events cannot be un-ingested. The fallback is unchanged, so the original config-clobber bug stays fixed; it is now reported under `config.api_url` through the same log-plus-callback floor as every other drop path. An *unpassed* `api_url` stays silent — `None` means the caller never mentioned it, and reporting the common case would train customers to ignore the channel this fix depends on. - **`usage_metadata` from the Cloudflare gateway got no drift sweep, and had already lost two counters.** `extras` was a fixed three-key dict, so any counter the adapter does not map vanished with no error and no `on_error` — the one place violating the drift contract `test_drift.py` enforces for the native adapters. Not hypothetical: replaying the **14 captured fixtures** through the adapter drops **`neurons`** (Cloudflare's Workers AI billing unit) in 4 entries and **`input_text_tokens`** in 1, and a live Logs API pull also returns **`units`**, a cost quantity that appears in no fixture at all — the hand-maintained key enumeration in the module docstring had already drifted past reality, which is exactly the failure mode a snapshot invites. Unmapped keys are now swept into `extras["usage_metadata"]` against an explicit `_MAPPED_USAGE_KEYS` set, so a ninth spelling surfaces on its own instead of needing another 14-fixture audit. Deliberately **nested** rather than merged flat into `extras`: the poller reads `extras["cached"]` to decide whether to skip billing a request Cloudflare served for free, so a future `usage_metadata` key called `cached` or `step` must not be able to shadow it. The regression test iterates the fixture directory rather than a fixed key list, so a recapture that introduces a new counter fails it with no test edit. Closes #16. diff --git a/src/lago_agent_sdk/adapters/openai_native.py b/src/lago_agent_sdk/adapters/openai_native.py index 2534fb6..fe255dc 100644 --- a/src/lago_agent_sdk/adapters/openai_native.py +++ b/src/lago_agent_sdk/adapters/openai_native.py @@ -40,6 +40,7 @@ from typing import Any, cast from ..canonical import WORKERS_AI_COMPAT_PREFIX, CanonicalUsage +from ..token_semantics import token_semantics from ._common import resolve_model # Cloudflare Workers AI names every model "@cf//". Reaching one @@ -189,12 +190,19 @@ def extract_openai_native(response: Any, model_id: str = "", provider_hint: str # Responses API uses input_tokens. They never both appear. is_responses_api = "input_tokens" in usage and "prompt_tokens" not in usage + # `cache_write` here is the RAW reported count, kept out of CanonicalUsage + # (see _MAPPED_DETAIL_FIELDS) and read per-branch from that branch's own + # details container — the Responses API spells it input_tokens_details, so a + # single prompt_tokens_details lookup would leave the Responses branch + # answering the total_tokens reconciliation below differently from the chat + # branch for the same convention. if is_responses_api: input_tokens = _safe_int(usage.get("input_tokens")) output_tokens = _safe_int(usage.get("output_tokens")) input_details = _safe_dict(usage.get("input_tokens_details")) output_details = _safe_dict(usage.get("output_tokens_details")) cache_read = _safe_int(input_details.get("cached_tokens")) + cache_write = _safe_int(input_details.get("cache_write_tokens")) reasoning = _safe_int(output_details.get("reasoning_tokens")) audio_input = _safe_int(input_details.get("audio_tokens")) audio_output = 0 # not exposed by Responses API today @@ -206,6 +214,7 @@ def extract_openai_native(response: Any, model_id: str = "", provider_hint: str prompt_details = _safe_dict(usage.get("prompt_tokens_details")) completion_details = _safe_dict(usage.get("completion_tokens_details")) cache_read = _safe_int(prompt_details.get("cached_tokens")) + cache_write = _safe_int(prompt_details.get("cache_write_tokens")) reasoning = _safe_int(completion_details.get("reasoning_tokens")) audio_input = _safe_int(prompt_details.get("audio_tokens")) audio_output = _safe_int(completion_details.get("audio_tokens")) @@ -225,11 +234,14 @@ def extract_openai_native(response: Any, model_id: str = "", provider_hint: str if k not in mapped: extras[f"{container}.{k}"] = v + resolved_model = resolve_model(resp.get("model"), model_id) + provider = provider_hint or _infer_provider(resolved_model) + # Consistency guard: for genuine OpenAI, total_tokens always equals # prompt + completion (reasoning is a SUBSET of completion, never additive). - # Verified across every captured real OpenAI-shaped response — zero deltas. - # So a POSITIVE delta means tokens exist that neither named bucket accounts - # for, which only happens behind an OpenAI-COMPATIBLE proxy that under-reports. + # Verified across every fixture under openai_native/ — zero deltas. So a + # POSITIVE delta means tokens exist that neither named bucket accounts for, + # which only happens behind an OpenAI-COMPATIBLE proxy that under-reports. # # Measured on Gemini through Google's own OpenAI-compat layer: # prompt_tokens=57, completion_tokens=47, total_tokens=1253 — 1149 real @@ -239,56 +251,67 @@ def extract_openai_native(response: Any, model_id: str = "", provider_hint: str # provider's own total proves those tokens were generated. # # Deliberately NOT assigned to `reasoning`: compute_cost zeroes reasoning for - # providers in _OUTPUT_INCLUDES_REASONING, so for real OpenAI that would set the + # providers in OUTPUT_INCLUDES_REASONING, so for real OpenAI that would set the # field and immediately discard it, recovering nothing. # - # `reasoning` is subtracted from the accounted total, and that subtraction is - # load-bearing rather than cosmetic. This adapter no longer only ever emits - # provider="openai" — it also emits "workers-ai" (Cloudflare `/compat`) and - # "databricks" (via provider_hint), and for those compute_cost bills reasoning - # ADDITIVELY. A payload reporting both `reasoning_tokens` and an inflated - # `total_tokens` would then be charged for them twice: once inside the grown - # `output` and again as a separate reasoning line. Subtracting first means a - # provider that already broke reasoning out gets no second bill, while the case - # this guard exists for — a thinking model behind a proxy that reports NO - # breakdown at all (measured: prompt 57, completion 47, total 1253) — still - # recovers its 1,149 tokens, because reasoning is 0 there. + # WHAT COUNTS AS ACCOUNTED is a per-provider fact, not a payload fact. The + # wire is one shape, but the convention behind it splits: OpenAI puts the + # cache and reasoning counts INSIDE prompt/completion, while Snowflake + # Cortex answers on the same wire with Anthropic's ADDITIVE convention — + # measured 2026-08-25, prompt_tokens=7, cached_tokens=4805, + # completion_tokens=6, total_tokens=4818: the cached block sits OUTSIDE + # prompt_tokens and INSIDE total_tokens. Accounting for input+output only + # made those 4,805 cached tokens look unaccounted, so they were folded into + # `output` — 4,811 reported for a call that generated 6, while the same + # tokens also shipped as cache_read. 2.0x on the call, 800x on the output + # line. See 12_snowflake_cortex_cache_chat.json. # - # The cache counts are subtracted for the SAME reason as reasoning, and this was - # the half that was missing. The guard assumed every OpenAI-shaped surface reports - # cache_read INSIDE prompt_tokens, which held for all three surfaces that existed - # when it was written (native OpenAI: zero deltas; Databricks: 112/112 rows with - # total == input + output; Cloudflare: cache outside `input` but outside `total` - # too, so it never inflated the delta). Snowflake Cortex is the surface that broke - # it — an OpenAI-WIRE endpoint with Anthropic's ADDITIVE convention: measured - # 2026-08-25, prompt_tokens=7, cached_tokens=4805, completion_tokens=6, - # total_tokens=4818, i.e. the cached block sits OUTSIDE prompt_tokens and INSIDE - # total_tokens. Accounting for only input+output+reasoning made those 4,805 cached - # tokens look unaccounted, so they were folded into `output` — 4,811 reported for a - # call that generated 6, while the same tokens also shipped as cache_read. 2.0x on - # the call, 800x on the output line. See 12_snowflake_cortex_cache_chat.json. + # So the accounted sum adds each subset field exactly when the provider + # reports it OUTSIDE its parent count, read from the same token_semantics + # table compute_cost and deoverlapped_token_total bill from — the guard and + # the money paths cannot answer the convention question differently. It + # cannot be decided from the payload instead: `accounted <= total` admits a + # small subtractive cache (folds too little), `cache_read > input` rejects a + # small additive one (folds tokens never generated) — both were tried + # against real shapes and both leak. And subtracting unconditionally + # disarms the guard where it is load-bearing: on a SUBSET surface the cache + # is already inside `input`, so also adding it to the accounted sum eats a + # genuine remainder — Gemini-compat's own cached+thinking payload would + # under-fold by exactly the cached count, silently, with no on_error. # - # `cache_write_tokens` is read straight from the payload rather than from a - # canonical field because it is deliberately NOT mapped to CanonicalUsage.cache_write - # (for OpenAI it sits inside prompt_tokens and billing it separately over-charges - # 2.24x — see _MAPPED_DETAIL_FIELDS). It still has to be accounted for here, or an - # additive cache WRITE would inflate `output` exactly the way the read did. + # `cache_write` is the raw prompt/input_tokens_details count because it is + # deliberately NOT mapped to CanonicalUsage.cache_write (for OpenAI it sits + # inside prompt_tokens and billing it separately over-charges 2.24x — see + # _MAPPED_DETAIL_FIELDS), yet an additive cache WRITE would inflate `output` + # exactly the way the read did. `cache_write_tokens` is the only spelling + # accounted for — an OpenAI-compat proxy re-reporting Anthropic's + # `cache_creation_input_tokens` (or `cache_creation.*`) inside a details + # block would still fold. Known limit: those spellings land in `extras` via + # the drift sweep, which is the signal to add them HERE, deliberately — + # deriving the accounting from the sweep itself would assume every unmapped + # count is additive, the same payload-only guess ruled out above. # - # Subtracting them cannot suppress a genuine fold on a subtractive surface: there - # the cache counts are already inside `input`, so removing them again only drives - # the delta further negative, where the `> 0` guard already no-ops. Verified against - # every captured OpenAI, Databricks and Cloudflare fixture — all still 0. + # The residual: an UNRECOGNIZED additive proxy arrives as provider="openai" + # and folds its cached block, exactly as Cortex did before its base_url rule + # existed. The payload carries no convention, so identification (a + # _provider_hint_for entry) is the fix — not loosening this arithmetic. # # A no-op for real OpenAI either way: total always equals prompt + completion. declared_total = _safe_int(usage.get("total_tokens")) if declared_total: - cache_write = _safe_int(_safe_dict(usage.get("prompt_tokens_details")).get("cache_write_tokens")) - unaccounted = declared_total - (input_tokens + output_tokens + reasoning + cache_read + cache_write) + inc_cache_read, inc_cache_write, inc_reasoning = token_semantics(provider, api) + accounted = input_tokens + output_tokens + if not inc_reasoning: + accounted += reasoning + if not inc_cache_read: + accounted += cache_read + if not inc_cache_write: + accounted += cache_write + unaccounted = declared_total - accounted if unaccounted > 0: output_tokens += unaccounted extras["unaccounted_output_tokens"] = unaccounted - resolved_model = resolve_model(resp.get("model"), model_id) return CanonicalUsage( input=input_tokens, output=output_tokens, @@ -298,7 +321,7 @@ def extract_openai_native(response: Any, model_id: str = "", provider_hint: str audio_output=audio_output, tool_calls=tool_calls, model=resolved_model, - provider=provider_hint or _infer_provider(resolved_model), + provider=provider, api=api, extras=extras, ) diff --git a/src/lago_agent_sdk/pricing.py b/src/lago_agent_sdk/pricing.py index 6849f31..fe102a1 100644 --- a/src/lago_agent_sdk/pricing.py +++ b/src/lago_agent_sdk/pricing.py @@ -52,6 +52,7 @@ from typing import Any, Protocol from .canonical import WORKERS_AI_COMPAT_PREFIX, CanonicalUsage +from .token_semantics import token_semantics logger = logging.getLogger("lago_agent_sdk.pricing") @@ -69,72 +70,12 @@ # Canonical usage fields we know how to price. PRICED_FIELDS = ("input", "output", "cache_read", "cache_write", "reasoning") -# Providers whose reported `input` token count ALREADY includes the cached -# (`cache_read`) tokens — i.e. cache_read is a subset of input, not additive. -# For these, the cached portion must be billed at the cache-read rate, not the -# full prompt rate, so compute_cost moves it out of `input`. Anthropic reports -# input EXCLUSIVE of cache (cache_read/cache_write are additive), so it's absent. -# -# "workers-ai" belongs here because it is only ever reached through Cloudflare's -# OpenAI-COMPATIBLE endpoint (`.../compat`), so its usage payload is the OpenAI -# shape: `prompt_tokens` includes `prompt_tokens_details.cached_tokens`. It is a -# distinct provider only because it prices against Cloudflare's own catalog -# (see _infer_provider in adapters/openai_native.py) — the token semantics are -# still OpenAI's. Omitting it billed the cached tokens twice: once at the full -# input rate because they were never subtracted, and again at the cache-read -# rate, which Cloudflare's catalog does publish for some models. -# -# "mistral" belongs here for the same reason: the API is OpenAI-shaped and reports -# `prompt_tokens_details.cached_tokens` as a SUBSET of `prompt_tokens`. Mistral's own -# documented example is unambiguous — prompt_tokens=1013, cached_tokens=1008, and -# total_tokens=1043 = prompt + completion, which only reconciles if the cached tokens -# sit inside the prompt count. Omitting it double-billed the cached portion by 6.15x -# on that payload. 13 of 18 Mistral models on OpenRouter publish a cache-read rate, -# so the wrong path was reachable for most of them, including Mistral traffic routed -# through a Cloudflare gateway (the gateway adapter leaves provider="mistral" as-is). -_INPUT_INCLUDES_CACHE_READ = frozenset({"openai", "gemini", "workers-ai", "mistral"}) - -# Providers whose reported `output` token count ALREADY includes the reasoning -# tokens (reasoning is a subset of output). For these, reasoning is billed as -# part of output and must NOT be billed again separately. (Gemini's `thoughts` -# are additive to output, so it's absent here.) -# "workers-ai" belongs here for the same reason it is in _INPUT_INCLUDES_CACHE_READ -# above: it is only ever reached through Cloudflare's OpenAI-COMPATIBLE endpoint, so -# its usage payload is the OpenAI shape — and in that shape -# `completion_tokens_details.reasoning_tokens` is a SUBSET of `completion_tokens`, -# exactly as it is for real OpenAI. `extract_openai_native` fills `reasoning` from that -# key with no provider gate, so omitting it counted the subset twice: measured, a -# 100/1000/reasoning-800 call reported unit=1900 against 1100 consumed. `compute_cost` -# would double-BILL the same tokens and does not today only because -# _CLOUDFLARE_UNIT_FIELD_MAP happens to carry no reasoning unit — an accident, not a -# guard, and Cloudflare hosts reasoning models (deepseek-r1, qwen, glm). -_OUTPUT_INCLUDES_REASONING = frozenset({"openai", "workers-ai"}) - -# Gateway SURFACES that re-report every vendor's usage in the OpenAI shape: `input` -# already contains cache_read AND cache_write, and `output` already contains -# reasoning, no matter which vendor actually served the call. -# -# This keys on `CanonicalUsage.api` rather than the provider because on a gateway -# it is the SURFACE that decides the shape, and a surface row reuses the live -# vendor names. A `provider="anthropic"` row read from Databricks' system table -# needs the correction; a `provider="anthropic"` response from Anthropic's own API -# must NOT get it. The vendor name cannot tell those two apart, so it is the wrong -# key — unlike "workers-ai" above, which names a vendor reachable through exactly -# one surface and so works as a provider entry. -# -# Measured on `system.ai_gateway.usage`, 246 rows across 6 vendors: `total_tokens -# == input + output` for EVERY vendor group, with cache_read and cache_write inside -# input and reasoning inside output. Anthropic's own API reports the exact opposite -# (cache_read=3962 against input=9, additive), which is why keying on the vendor -# over-billed a real backfill 1.570x — 48,798 tokens reported against 31,091 -# consumed, the excess being exactly cache_read + cache_write. -# -# Cloudflare AI Gateway is deliberately ABSENT: its logs preserve each vendor's -# native shape instead of normalising them. A real Anthropic entry there reads -# input=10, output=4, total=14 with input_cached_tokens=3429 sitting OUTSIDE that -# total — additive, exactly like the native API — so the provider-keyed sets are -# already right for it and adding it here would UNDER-bill the cached portion. -_OPENAI_SHAPED_APIS = frozenset({"databricks_gateway"}) +# The subset-vs-additive convention sets (_INPUT_INCLUDES_CACHE_READ and friends) +# used to live here. They moved to token_semantics.py the day the total_tokens +# guard in adapters/openai_native.py needed the same answers: the guard, the cost +# split and the token total are three readings of one convention, and keeping the +# sets in this module would have forced the adapter layer to import pricing's +# HTTP machinery to reach them. # Providers this SDK bills as TOKEN COUNTS by design, even in price mode — because # no per-token rate for them exists anywhere the SDK could read it. @@ -376,24 +317,14 @@ class CostBreakdown: def _token_semantics(usage: Any) -> tuple[bool, bool, bool]: - """Which of a record's subsets are ALREADY inside their parent count. - - Returns ``(input_includes_cache_read, input_includes_cache_write, - output_includes_reasoning)``, the three overlaps the billing paths have to - remove. The SURFACE wins over the vendor: a gateway that re-reports usage in - its own shape has already decided the convention, so `api` is checked first - and the provider-keyed sets only answer for a native call. + """`token_semantics` read off a CanonicalUsage — see token_semantics.py. - `cache_write` is surface-only by design and has no provider set to consult: - Anthropic is the one vendor whose native API bills cache writes at all, and it - reports them additively, so no native response needs the correction. + Kept as the module-internal spelling so the billing paths keep reading the + convention from the record they are billing, not from loose strings. """ - provider = (getattr(usage, "provider", "") or "").lower() - shaped = (getattr(usage, "api", "") or "").lower() in _OPENAI_SHAPED_APIS - return ( - shaped or provider in _INPUT_INCLUDES_CACHE_READ, - shaped, - shaped or provider in _OUTPUT_INCLUDES_REASONING, + return token_semantics( + getattr(usage, "provider", "") or "", + getattr(usage, "api", "") or "", ) @@ -468,10 +399,11 @@ def deoverlapped_token_total(usage: Any) -> int: applied, because a subset counted twice inflates the reported quantity exactly as it would inflate a price: - * reasoning ⊆ output — providers in _OUTPUT_INCLUDES_REASONING, or any - row from a surface in _OPENAI_SHAPED_APIS - * cache_read ⊆ input — providers in _INPUT_INCLUDES_CACHE_READ, likewise - * cache_write ⊆ input — surfaces in _OPENAI_SHAPED_APIS only + * reasoning ⊆ output — providers in OUTPUT_INCLUDES_REASONING, or any + row from a surface in OPENAI_SHAPED_APIS + * cache_read ⊆ input — providers in INPUT_INCLUDES_CACHE_READ, likewise + * cache_write ⊆ input — providers in INPUT_INCLUDES_CACHE_WRITE, likewise + (all in token_semantics.py) Deliberately NOT gated on a unit price existing, unlike `compute_cost`'s subtraction — this is a token count, so whether a rate happens to be published diff --git a/src/lago_agent_sdk/token_semantics.py b/src/lago_agent_sdk/token_semantics.py new file mode 100644 index 0000000..0b0f04a --- /dev/null +++ b/src/lago_agent_sdk/token_semantics.py @@ -0,0 +1,198 @@ +"""Token-count conventions per provider — the single source of truth. + +Providers disagree about whether a subset count (cached tokens, reasoning +tokens) is reported INSIDE its parent count or beside it. OpenAI reports +`cached_tokens` inside `prompt_tokens`; Anthropic reports `cache_read_input_tokens` +outside `input_tokens`. The wire shape says nothing about which convention is in +play — Snowflake Cortex answers on a byte-for-byte OpenAI wire with Anthropic's +additive convention (measured live 2026-08-25: prompt 7, cached 4805, completion 6, +total 4818) — so the convention can only be KEYED, never inferred from a payload. +Two payload-only inference gates were tried and both were shown unsound on real +shapes (see PR #23 review); do not reintroduce one. + +Three places have to answer the same question, and this module exists so they +cannot answer it differently: + + - `adapters/openai_native.extract_openai_native` — deciding which reported + counts sit inside `total_tokens` when reconciling it, + - `pricing.compute_cost` — deciding which subsets to move out of their parent + before pricing each field, + - `pricing.deoverlapped_token_total` — deciding which subsets to drop from the + single-event token total. + +The two failure directions are not symmetric and both have happened here: +treat a subtractive surface as additive and real tokens are silently never +billed; treat an additive surface as subtractive and the same tokens bill twice +(measured at 1.570x, 2.0x, and 6.15x on three different providers). Entries are +added on MEASUREMENT of a real payload, never on the wire shape or vendor +documentation alone. +""" + +from __future__ import annotations + +# Providers whose reported `input` token count ALREADY includes the cached +# (`cache_read`) tokens — i.e. cache_read is a subset of input, not additive. +# For these, the cached portion must be billed at the cache-read rate, not the +# full prompt rate, so compute_cost moves it out of `input`. Anthropic reports +# input EXCLUSIVE of cache (cache_read/cache_write are additive), so it's absent. +# +# "workers-ai" belongs here because it is only ever reached through Cloudflare's +# OpenAI-COMPATIBLE endpoint (`.../compat`), so its usage payload is the OpenAI +# shape: `prompt_tokens` includes `prompt_tokens_details.cached_tokens`. It is a +# distinct provider only because it prices against Cloudflare's own catalog +# (see _infer_provider in adapters/openai_native.py) — the token semantics are +# still OpenAI's. Omitting it billed the cached tokens twice: once at the full +# input rate because they were never subtracted, and again at the cache-read +# rate, which Cloudflare's catalog does publish for some models. +# +# "mistral" belongs here for the same reason: the API is OpenAI-shaped and reports +# `prompt_tokens_details.cached_tokens` as a SUBSET of `prompt_tokens`. Mistral's own +# documented example is unambiguous — prompt_tokens=1013, cached_tokens=1008, and +# total_tokens=1043 = prompt + completion, which only reconciles if the cached tokens +# sit inside the prompt count. Omitting it double-billed the cached portion by 6.15x +# on that payload. 13 of 18 Mistral models on OpenRouter publish a cache-read rate, +# so the wrong path was reachable for most of them, including Mistral traffic routed +# through a Cloudflare gateway (the gateway adapter leaves provider="mistral" as-is). +# +# "snowflake" is deliberately ABSENT even though Cortex answers on an OpenAI-WIRE +# endpoint: measured live 2026-08-25, `prompt_tokens: 7`, `cached_tokens: 4805`, +# `completion_tokens: 6`, `total_tokens: 4818` — the cached block sits OUTSIDE +# `prompt_tokens` and INSIDE the total, Anthropic's additive convention on OpenAI's +# wire. Measured for the Claude family, which is the only family on that surface +# that caches at all today: llama accepts `cache_control` and ignores it +# (cached_tokens 0 on a matched pair, total = prompt + completion, measured +# 2026-08-27), and the OpenAI family needs cross-region inference the capture +# account cannot enable — re-verify the day an OpenAI-family model becomes +# reachable, since Cortex documents its caching behaviour per model family. +INPUT_INCLUDES_CACHE_READ = frozenset({"openai", "gemini", "workers-ai", "mistral"}) + +# Providers whose reported `input` token count ALREADY includes the cache-WRITE +# tokens. OpenAI is measured: a live gpt-5.6-sol response carries +# prompt_tokens=3025 with prompt_tokens_details.cache_write_tokens=3022, and +# Databricks' metered spend for that call matched billing all 3025 at the plain +# input rate — the write sits inside the prompt count (OpenAI's own docs now say +# the same: weighted input = ordinary + cached + cache-write portions). That is +# also why the field is deliberately NOT mapped to CanonicalUsage.cache_write — +# billing it separately over-charged 2.24x (see adapters/openai_native.py). +# +# The other three subset-cache providers are carried here on the SHAPE argument +# that earned "workers-ai" its cache_read entry: their surfaces re-report usage +# in the OpenAI shape, where every prompt_tokens_details member is a subset of +# prompt_tokens. None of the three emits the key today (Gemini and Mistral have +# no cache_write concept on these wires), so for them membership decides only +# what the total_tokens guard does if the key ever appears — and for a +# subset-convention surface the guard must stay live (a genuine remainder still +# folds), which membership preserves. Snowflake is absent: cache_write_tokens +# exists on its wire (a key OpenAI never sends) and was 0 in every capture +# including cache-creation calls — creation reports under cached_tokens there. +INPUT_INCLUDES_CACHE_WRITE = frozenset({"openai", "gemini", "workers-ai", "mistral"}) + +# Providers whose reported `output` token count ALREADY includes the reasoning +# tokens (reasoning is a subset of output). For these, reasoning is billed as +# part of output and must NOT be billed again separately. (Gemini's `thoughts` +# are additive to output, so it's absent here.) +# +# "workers-ai" belongs here for the same reason it is in INPUT_INCLUDES_CACHE_READ +# above: it is only ever reached through Cloudflare's OpenAI-COMPATIBLE endpoint, so +# its usage payload is the OpenAI shape — and in that shape +# `completion_tokens_details.reasoning_tokens` is a SUBSET of `completion_tokens`, +# exactly as it is for real OpenAI. `extract_openai_native` fills `reasoning` from that +# key with no provider gate, so omitting it counted the subset twice: measured, a +# 100/1000/reasoning-800 call reported unit=1900 against 1100 consumed. `compute_cost` +# would double-BILL the same tokens and does not today only because +# _CLOUDFLARE_UNIT_FIELD_MAP happens to carry no reasoning unit — an accident, not a +# guard, and Cloudflare hosts reasoning models (deepseek-r1, qwen, glm). +# +# "snowflake" absence is a measured no-op rather than an open question: +# `reasoning_tokens` is always 0 on Cortex's OpenAI-compat wire (reasoning_effort +# is accepted and ignored; extended thinking exists only on Cortex's Anthropic +# wire, which this adapter never serves). Re-measure if that wire ever starts +# reporting it. +OUTPUT_INCLUDES_REASONING = frozenset({"openai", "workers-ai"}) + +# Gateway SURFACES that re-report every vendor's usage in the OpenAI shape: `input` +# already contains cache_read AND cache_write, and `output` already contains +# reasoning, no matter which vendor actually served the call. +# +# This keys on `CanonicalUsage.api` rather than the provider because on a gateway +# it is the SURFACE that decides the shape, and a surface row reuses the live +# vendor names. A `provider="anthropic"` row read from Databricks' system table +# needs the correction; a `provider="anthropic"` response from Anthropic's own API +# must NOT get it. The vendor name cannot tell those two apart, so it is the wrong +# key — unlike "workers-ai" above, which names a vendor reachable through exactly +# one surface and so works as a provider entry. +# +# Measured on `system.ai_gateway.usage`, 246 rows across 6 vendors: `total_tokens +# == input + output` for EVERY vendor group, with cache_read and cache_write inside +# input and reasoning inside output. Anthropic's own API reports the exact opposite +# (cache_read=3962 against input=9, additive), which is why keying on the vendor +# over-billed a real backfill 1.570x — 48,798 tokens reported against 31,091 +# consumed, the excess being exactly cache_read + cache_write. +# +# Cloudflare AI Gateway is deliberately ABSENT: its logs preserve each vendor's +# native shape instead of normalising them. A real Anthropic entry there reads +# input=10, output=4, total=14 with input_cached_tokens=3429 sitting OUTSIDE that +# total — additive, exactly like the native API — so the provider-keyed sets are +# already right for it and adding it here would UNDER-bill the cached portion. +OPENAI_SHAPED_APIS = frozenset({"databricks_gateway"}) + +# Every provider name the SDK's own code can stamp on a CanonicalUsage, so that +# absence from the sets above is always a recorded DECISION and never a default +# nobody made. The convention for a name not in any set is "everything additive" — +# correct for Anthropic-style reporters and for gateway vendors whose logs +# preserve the native shape, and the conservative direction for an unknown (it +# can over-count a subset into the total but never silently drop generated +# tokens). test_token_semantics.py pins this list against the stamps in the +# adapters and wrappers; when adding a provider, add it here IN THE SAME CHANGE +# as its (measured) set entries, per the recipe in CONTRIBUTING.md. +# +# The bedrock_* names are the vendor spellings `_provider_from_model` can emit +# for Bedrock model ids. Bedrock reports cache counts ADDITIVELY for every +# vendor it hosts (its `inputTokens` excludes `cacheRead/WriteInputTokens`), and +# today only its Anthropic and Nova families cache at all — both stamped with +# names absent from the subset sets, so the additive default is the measured +# answer. If Bedrock ever enables caching for a vendor whose name IS in a subset +# set ("openai" via gpt-oss, "mistral"), the bedrock adapters must start +# stamping a surface-distinct api the way databricks_gateway does — the vendor +# name alone would answer wrongly there. +KNOWN_PROVIDERS = frozenset( + { + # adapters/, by inference or wrapper hint + "openai", + "workers-ai", + "anthropic", + "gemini", + "mistral", + "databricks", + "snowflake", + # adapters/bedrock_*, from _provider_from_model + "amazon", + "meta", + "cohere", + "qwen", + "google", + "minimax", + "nvidia", + "zai", + "bedrock", + } +) + + +def token_semantics(provider: str, api: str) -> tuple[bool, bool, bool]: + """Which of a record's subsets are ALREADY inside their parent count. + + Returns ``(input_includes_cache_read, input_includes_cache_write, + output_includes_reasoning)``, the three overlaps the billing paths have to + remove and the total_tokens guard has to leave alone. The SURFACE wins over + the vendor: a gateway that re-reports usage in its own shape has already + decided the convention, so `api` is checked first and the provider-keyed + sets only answer for a native call. + """ + p = (provider or "").lower() + shaped = (api or "").lower() in OPENAI_SHAPED_APIS + return ( + shaped or p in INPUT_INCLUDES_CACHE_READ, + shaped or p in INPUT_INCLUDES_CACHE_WRITE, + shaped or p in OUTPUT_INCLUDES_REASONING, + ) diff --git a/tests/unit/adapters/fixtures/capture_snowflake_cortex.py b/tests/unit/adapters/fixtures/capture_snowflake_cortex.py index 701dad0..efffe9b 100644 --- a/tests/unit/adapters/fixtures/capture_snowflake_cortex.py +++ b/tests/unit/adapters/fixtures/capture_snowflake_cortex.py @@ -10,12 +10,18 @@ cannot assume OpenAI's subtractive cache convention — on Cortex, `cached_tokens` sits OUTSIDE `prompt_tokens` and INSIDE `total_tokens`. -Two things about Cortex that this script encodes, both measured 2026-08-25: +Three things about Cortex that this script encodes, all measured 2026-08-25: * Caching only happens with an explicit Anthropic-style `cache_control` part. The same 4,800-token prompt sent twice WITHOUT it reports `cached_tokens: 0` both times, so the "call1 then call2" pattern the OpenAI cache fixtures use captures nothing here. + * One cold call is enough for fixture 12: unlike Anthropic's own wire, Cortex + reports a cache CREATION under `cached_tokens` too (`cache_write_tokens` + stays 0), so the first `cache_control` call already carries the cached + block. Measured on a matched pair — both calls returned identical usage + while the account-usage view logged one `cache_write_input` row and one + `cache_read_input` row. No warm-up call, no 5-minute-TTL race on recapture. * `max_tokens` is rejected outright ("deprecated in favor of max_completion_tokens"), unlike OpenAI which still accepts it. @@ -63,11 +69,16 @@ def call(host: str, pat: str, body: dict) -> dict: return r.json() -def save(name: str, response: dict) -> None: +def save(name: str, body: dict, host: str, pat: str) -> None: + # The existence check runs BEFORE the request: `save(..., call(...))` would + # evaluate the call first and fire two live Cortex requests — the 4,800-token + # cacheable one included — on a checkout where both fixtures already exist, + # then print "skip". Idempotent means no request, not just no write. path = OUT / name if path.exists(): print(f"skip {name} (exists)") return + response = call(host, pat, body) path.write_text(json.dumps({"_model_id": MODEL, "_response": response}, indent=2) + "\n") print(f"wrote {name}") @@ -80,43 +91,40 @@ def main() -> None: save( "11_snowflake_cortex_plain_chat.json", - call( - host, - pat, - { - "model": MODEL, - "messages": [{"role": "user", "content": "What is 2 + 2? Answer in one word."}], - "max_completion_tokens": 32, - }, - ), + { + "model": MODEL, + "messages": [{"role": "user", "content": "What is 2 + 2? Answer in one word."}], + "max_completion_tokens": 32, + }, + host, + pat, ) # The regression fixture. `cache_control` is what makes Cortex report a cached # block at all, and the resulting payload is the one that used to inflate - # `output` by the whole cached count. + # `output` by the whole cached count. One call suffices — see the docstring's + # creation-reports-as-cached_tokens note. save( "12_snowflake_cortex_cache_chat.json", - call( - host, - pat, - { - "model": MODEL, - "messages": [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": CACHEABLE_PREFIX, - "cache_control": {"type": "ephemeral"}, - }, - {"type": "text", "text": "Reply with one word."}, - ], - } - ], - "max_completion_tokens": 32, - }, - ), + { + "model": MODEL, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": CACHEABLE_PREFIX, + "cache_control": {"type": "ephemeral"}, + }, + {"type": "text", "text": "Reply with one word."}, + ], + } + ], + "max_completion_tokens": 32, + }, + host, + pat, ) diff --git a/tests/unit/adapters/test_openai_native.py b/tests/unit/adapters/test_openai_native.py index 4c3da25..0657559 100644 --- a/tests/unit/adapters/test_openai_native.py +++ b/tests/unit/adapters/test_openai_native.py @@ -168,6 +168,23 @@ def test_model_falls_back_to_request_when_response_is_silent() -> None: assert u.model == "gpt-4o-mini" +def test_snowflake_cortex_keeps_the_customers_model_spelling() -> None: + """A Cortex fine-tune answers as `database.schema.model`. CanonicalUsage.model + keeps it verbatim — normalising here would report a model the customer cannot + find in their own Snowflake account. (The hint is what a wrapped client whose + base_url matches the Cortex path supplies — see _provider_hint_for.)""" + u = extract_openai_native( + { + "model": "mydb.myschema.my_tuned_model", + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + }, + model_id="mydb.myschema.my_tuned_model", + provider_hint="snowflake", + ) + assert u.model == "mydb.myschema.my_tuned_model" + assert u.provider == "snowflake" + + # -------------------------------------------------------------------------- # Robustness # -------------------------------------------------------------------------- @@ -355,9 +372,11 @@ def test_total_tokens_guard_does_not_fold_an_additive_cache_write() -> None: tokens outside `prompt_tokens` but inside `total_tokens`. It was answered "unreachable on the three surfaces we have", which was true at the time — Snowflake Cortex then shipped the same class of payload with `cached_tokens`. - Accounted for now whether or not a live surface reports it this way, because - `cache_write_tokens` is deliberately never mapped to CanonicalUsage.cache_write - and so has no other route into the accounting.""" + The write is accounted for from the raw payload because `cache_write_tokens` + is deliberately never mapped to CanonicalUsage.cache_write and so has no + other route into the accounting — but only under a provider whose convention + IS additive: for OpenAI itself the write sits inside prompt_tokens (see + test_total_tokens_guard_survives_openai_cache_write_beside_a_remainder).""" u = extract_openai_native( { "usage": { @@ -366,13 +385,36 @@ def test_total_tokens_guard_does_not_fold_an_additive_cache_write() -> None: "total_tokens": 1829, "prompt_tokens_details": {"cache_write_tokens": 1812}, } - } + }, + provider_hint="snowflake", ) assert u.output == 4, "was 1816" assert "unaccounted_output_tokens" not in u.extras assert u.extras["prompt_tokens_details.cache_write_tokens"] == 1812 +def test_total_tokens_guard_handles_an_additive_cache_write_on_the_responses_shape() -> None: + """Same convention, other API branch: the Responses shape spells the container + `input_tokens_details`, so a chat-only `prompt_tokens_details` lookup would + leave this exact payload folding 1,812 cached-write tokens into `output` — + the two API shapes must not disagree about one provider's convention.""" + u = extract_openai_native( + { + "usage": { + "input_tokens": 13, + "output_tokens": 4, + "total_tokens": 1829, + "input_tokens_details": {"cache_write_tokens": 1812}, + } + }, + provider_hint="snowflake", + ) + assert u.api == "responses" + assert u.output == 4, "was 1816" + assert "unaccounted_output_tokens" not in u.extras + assert u.extras["input_tokens_details.cache_write_tokens"] == 1812 + + def test_total_tokens_guard_still_recovers_a_remainder_beside_a_cache_count() -> None: """The two corrections must not cancel each other: an additive cache block AND hidden thinking tokens in the same payload. 20 + 5 + 100 = 125 accounted, @@ -385,12 +427,61 @@ def test_total_tokens_guard_still_recovers_a_remainder_beside_a_cache_count() -> "total_tokens": 200, "prompt_tokens_details": {"cached_tokens": 100}, } - } + }, + provider_hint="snowflake", ) assert u.output == 80 assert u.extras["unaccounted_output_tokens"] == 75 +def test_total_tokens_guard_keeps_a_subtractive_fold_beside_a_cache_count() -> None: + """The case that rules out subtracting the cache unconditionally, raised in + review on #23: a SUBSET-convention surface reporting a cached block AND a + genuine remainder. Gemini through Google's own OpenAI-compat layer reports + `cached_tokens` inside `prompt_tokens` (that is why "gemini"/"openai" are in + INPUT_INCLUDES_CACHE_READ) while thinking tokens appear only in the total — + so the 1,000 cached tokens are ALREADY accounted for by prompt_tokens, and + also adding them to the accounted sum would shrink the fold to 149: 1,000 + generated tokens unbilled, silently, with no on_error. The full 1,149 must + fold.""" + u = extract_openai_native( + { + "model": "gemini-2.5-flash", + "usage": { + "prompt_tokens": 1200, + "completion_tokens": 47, + "total_tokens": 2396, + "prompt_tokens_details": {"cached_tokens": 1000}, + }, + } + ) + assert u.provider == "openai", "no hint: OpenAI-compat traffic is stamped openai" + assert u.cache_read == 1000 + assert u.output == 1196, "47 reported + 1149 unaccounted — NOT 196" + assert u.extras["unaccounted_output_tokens"] == 1149 + + +def test_total_tokens_guard_survives_openai_cache_write_beside_a_remainder() -> None: + """The documented-real OpenAI shape (see the NOTE on _MAPPED_DETAIL_FIELDS: + prompt_tokens=3025 measured WITH cache_write_tokens=3022 inside it) behind a + proxy that under-reports 1,171 tokens. OpenAI's write sits inside + prompt_tokens, so it must NOT join the accounted sum — subtracting it + unconditionally would swallow the delta and disarm the guard on the one + payload shape this file documents as measured.""" + u = extract_openai_native( + { + "usage": { + "prompt_tokens": 3025, + "completion_tokens": 4, + "total_tokens": 4200, + "prompt_tokens_details": {"cache_write_tokens": 3022}, + } + } + ) + assert u.output == 1175, "4 reported + 1171 unaccounted — NOT 4" + assert u.extras["unaccounted_output_tokens"] == 1171 + + def test_total_tokens_guard_is_a_noop_for_genuine_openai() -> None: """For real OpenAI total_tokens == prompt + completion always holds, because reasoning is a SUBSET of completion rather than additive. Verified across @@ -457,28 +548,42 @@ def test_snowflake_cortex_cached_tokens_are_additive() -> None: that did not exist yet.""" model_id, resp = _load("12_snowflake_cortex_cache_chat.json") usage = resp["usage"] - assert usage["prompt_tokens"] + 4805 + usage["completion_tokens"] == usage["total_tokens"] + # Read the cached count off the fixture rather than pinning a literal: the + # assertion is the additive IDENTITY, so a recapture with a different cached + # count must keep passing instead of nudging someone toward the hand-edit + # the header above forbids. + cached = usage["prompt_tokens_details"]["cached_tokens"] + assert cached > 0, "recapture produced no cached block — see capture_snowflake_cortex.py" + assert usage["prompt_tokens"] + cached + usage["completion_tokens"] == usage["total_tokens"] u = extract_openai_native(resp, model_id=model_id, provider_hint="snowflake") - assert u.input == 7 - assert u.output == 6, "NOT 4811" - assert u.cache_read == 4805 - assert u.cache_write == 0 + assert u.input == usage["prompt_tokens"] + assert u.output == usage["completion_tokens"], "NOT completion + cached" + assert u.cache_read == cached + # cache_write stays unmapped BY DESIGN (u.cache_write is 0 on every path, so + # asserting it proves nothing) — the load-bearing check is that the raw key + # is still visible in extras rather than silently consumed by the guard. + assert u.extras["prompt_tokens_details.cache_write_tokens"] == 0 assert u.reasoning == 0 assert "unaccounted_output_tokens" not in u.extras -def test_snowflake_cortex_keeps_the_customers_model_spelling() -> None: - """A Cortex fine-tune answers as `database.schema.model`. CanonicalUsage.model - keeps it verbatim — normalising here would report a model the customer cannot - find in their own Snowflake account.""" - u = extract_openai_native( - { - "model": "mydb.myschema.my_tuned_model", - "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, - }, - model_id="mydb.myschema.my_tuned_model", - provider_hint="snowflake", - ) - assert u.model == "mydb.myschema.my_tuned_model" - assert u.provider == "snowflake" +def test_snowflake_cortex_without_a_hint_is_stamped_openai_and_folds() -> None: + """Pins what REAL traffic does until the wrapper carries a Cortex base_url + rule (PR #26 adds `/api/v2/cortex/` → "snowflake" to _provider_hint_for): + the response body has no marker of its own, so an unhinted Cortex payload is + stamped "openai", whose SUBSET convention folds the additive cached block + into `output` again. This is deliberate — the payload cannot carry the + convention, so identification is the fix, not looser arithmetic. If this + test starts failing because the fold stopped, the guard has been loosened + for every genuine OpenAI-compat proxy; if it fails on `provider`, the hint + now reaches this adapter by default and the test should assert the fixed + behaviour instead.""" + model_id, resp = _load("12_snowflake_cortex_cache_chat.json") + usage = resp["usage"] + cached = usage["prompt_tokens_details"]["cached_tokens"] + + u = extract_openai_native(resp, model_id=model_id) + assert u.provider == "openai" + assert u.output == usage["completion_tokens"] + cached + assert u.extras["unaccounted_output_tokens"] == cached diff --git a/tests/unit/test_token_semantics.py b/tests/unit/test_token_semantics.py new file mode 100644 index 0000000..346c44f --- /dev/null +++ b/tests/unit/test_token_semantics.py @@ -0,0 +1,100 @@ +"""token_semantics — the one table three billing paths answer from. + +These tests pin the DECISIONS, not the mechanism: each provider's entry (or +deliberate absence) traces to a measurement recorded in token_semantics.py, and +a provider the SDK can stamp without a recorded decision is a red test here — +absence must always be a choice somebody made. +""" + +from __future__ import annotations + +from lago_agent_sdk.token_semantics import ( + INPUT_INCLUDES_CACHE_READ, + INPUT_INCLUDES_CACHE_WRITE, + KNOWN_PROVIDERS, + OPENAI_SHAPED_APIS, + OUTPUT_INCLUDES_REASONING, + token_semantics, +) + +# Every provider string the SDK's own code can stamp on a CanonicalUsage today. +# Kept explicit rather than scraped from the source: when an adapter or wrapper +# grows a new stamp, add it BOTH here and (with its measured decision) to +# KNOWN_PROVIDERS — this list failing to cover a stamp is exactly the silent +# default the roster exists to prevent. Gateway backfills additionally pass +# vendor names through from their logs verbatim; those arrive with a surface +# `api` and are decided by OPENAI_SHAPED_APIS (or the vendor's own entry), not +# by this list. +_STAMPABLE = { + # adapters/openai_native.py: _infer_provider + "openai", + "workers-ai", + # wrappers/openai.py: _provider_hint_for (base_url table) + "databricks", + "snowflake", + # adapters/anthropic_native.py, gemini_native.py, mistral_native.py + "anthropic", + "gemini", + "mistral", + # adapters/bedrock_converse.py / bedrock_invoke.py: _provider_from_model + "amazon", + "meta", + "cohere", + "qwen", + "google", + "minimax", + "nvidia", + "zai", + "bedrock", +} + + +def test_every_stampable_provider_has_a_recorded_decision() -> None: + missing = _STAMPABLE - KNOWN_PROVIDERS + assert not missing, f"providers stamped by the SDK with no semantics decision: {sorted(missing)}" + + +def test_the_subset_sets_only_name_known_providers() -> None: + """A set entry for a name nothing can stamp is dead weight at best and a + typo silently reverting a measured decision at worst.""" + for s in (INPUT_INCLUDES_CACHE_READ, INPUT_INCLUDES_CACHE_WRITE, OUTPUT_INCLUDES_REASONING): + assert s <= KNOWN_PROVIDERS + + +def test_openai_convention_is_subset_on_all_three_dimensions() -> None: + assert token_semantics("openai", "chat_completions") == (True, True, True) + + +def test_snowflake_convention_is_additive_on_all_three_dimensions() -> None: + """Measured live 2026-08-25: prompt 7 / cached 4805 / completion 6 / + total 4818 — Anthropic's additive convention on OpenAI's wire. This single + row is what the total_tokens guard, compute_cost and + deoverlapped_token_total all read; if it ever flips, all three flip + together or 4,805 tokens bill twice.""" + assert token_semantics("snowflake", "chat_completions") == (False, False, False) + + +def test_anthropic_convention_is_additive() -> None: + assert token_semantics("anthropic", "messages") == (False, False, False) + + +def test_gemini_cache_is_subset_but_reasoning_is_additive() -> None: + """cachedContentTokenCount ⊆ promptTokenCount, thoughtsTokenCount additive — + Google documents totalTokenCount = prompt + thoughts + candidates.""" + assert token_semantics("gemini", "generate_content") == (True, True, False) + + +def test_an_openai_shaped_surface_overrides_the_vendor() -> None: + """A provider="anthropic" ROW from Databricks' system table is in the + gateway's re-reported shape — everything subset — while the same vendor + name from its own API is additive. The surface wins.""" + assert token_semantics("anthropic", "databricks_gateway") == (True, True, True) + assert "databricks_gateway" in OPENAI_SHAPED_APIS + + +def test_an_unknown_provider_defaults_to_additive() -> None: + """No overlap is removed for a name nobody measured: the conservative + direction — it can over-count a subset into a token total, but it can never + silently drop generated tokens or zero a real cache line.""" + assert token_semantics("some-new-gateway", "chat_completions") == (False, False, False) + assert token_semantics("", "") == (False, False, False)