diff --git a/.gitignore b/.gitignore index c836f0e..6500a62 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ .env .env.* +!.env.example .envrc .ca-bundle.pem logs.log diff --git a/CHANGELOG.md b/CHANGELOG.md index ceafed4..fd388fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,104 @@ All notable changes to this project will be documented here. Format follows [Kee ## [Unreleased] +### Added +- **Price mode can now price `workers-ai` calls live, via Cloudflare's own model catalog** (`/accounts/{id}/ai/models/search`) — a third pricing source alongside OpenRouter and AWS Bedrock. This is the real rate the gateway bills at, not a third party's price for hosting the same open-weight model elsewhere: verified live that a real call's actual charged cost matched this catalog's rate exactly, while the closest OpenRouter listing for the same underlying model (`meta-llama/llama-3.3-70b-instruct`) came out ~3.5x lower — a genuinely different price, not a naming mismatch, and the reason OpenRouter can never be the right source for Workers AI regardless of how well its model names are matched. New `LagoConfig.cloudflare_account_id`/`cloudflare_api_token` (needed because, unlike OpenRouter/AWS, this catalog isn't public/no-auth — without both set, this source is simply empty and behaves exactly like any other pricing miss). Same non-blocking design as the existing sources: the fetch runs on the queue's background thread on the existing TTL cycle, never on the customer's call path. +- **Fixed a real bug this surfaced**: `extract_openai_native` hardcoded `provider="openai"` unconditionally — correct for a real OpenAI response, but also stamped on any call made through Cloudflare's OpenAI-compatible endpoint (`.../compat`) to a non-OpenAI backend, since the response shape looks identical either way. A Workers AI call routed through it was permanently unpriceable, silently, at the extraction layer — OpenAI's price table has no Workers AI entries, so it always missed, before the new catalog source could ever be reached. Now infers `provider="workers-ai"` from the resolved model string (Cloudflare's `@cf/...` naming is unambiguous) instead of assuming the SDK shape implies the provider. +- **`LagoConfig.verify_ssl`** (default `True`) — threads through to the internal `LagoClient`'s `requests.post(..., verify=...)`. A local dev Lago instance behind a self-signed certificate (Traefik's default) is a real, common setup; without this the only option was routing every request through a public tunnel (ngrok) purely to get a browser-trusted cert — which turned out to be unreliable enough on its own (repeated `SSLEOFError`s, for both me and the person actually using the example) to cause real, confusing failures unrelated to any of the SDK's own code. Suppresses `requests`'s `InsecureRequestWarning` when explicitly set to `False` (the customer already accepted the risk by setting it; the warning on every single request is noise, not new information) — never touches it otherwise. `examples/cloudflare_gateway_demo.ipynb` now reads this from `LAGO_VERIFY_SSL` and can hit a local instance directly with zero tunnel dependency. +- **Mistral alias resolution for price mode**, via Mistral's own `/v1/models`. Mistral has no per-token price table of its own (confirmed: their pricing page lists one FAQ example, not a structured/JSON list) — but a customer request commonly uses a moving alias (`mistral-small-latest`) that Mistral's response never resolves (unlike Anthropic/OpenAI, which report the dated snapshot that actually answered), so the existing OpenRouter lookup missed even though OpenRouter *does* list the resolved id with real pricing (verified live: `mistral-small-latest` resolves via `/v1/models`'s `aliases` array to `mistral-small-2603`, which OpenRouter lists as `mistralai/mistral-small-2603`). New `LagoConfig.mistral_api_key` (needed because, unlike OpenRouter, this endpoint isn't public/no-auth — without it, alias resolution is simply skipped and lookups fall back to the pre-existing behavior: a safe miss for an alias, a hit for anything already an exact id). Same non-blocking background-refresh design as the other sources. + - **Found and fixed a real bug in this same feature before it ever shipped correctly**: the first implementation mapped "each alias in this entry's `aliases` array -> this entry's `id`", which is wrong for Mistral's actual response shape — every name in an alias family (`mistral-small-2603`, `mistral-small-latest`, `magistral-small-latest`, `mistral-vibe-cli-fast`) appears as its OWN top-level `id` too, each listing the other three as `aliases`. A directional last-write-wins map is order-dependent and resolved `mistral-small-latest` to `magistral-small-latest` (whichever entry got parsed last) instead of the real dated snapshot — confirmed live against a real notebook run, where this exact miss showed up as `lago pricing failed: no price for provider='mistral' model='mistral-small-latest'`. Replaced with union-find: every name in a mutually-aliasing family is grouped regardless of who mentions whom, then one canonical name per group is picked deterministically (prefer a dated id like `-2603` over any `-latest` moniker). Re-verified live end-to-end against real Mistral + OpenRouter data after the fix: resolves correctly and finds real pricing. + - **`PricingProvider.prime()` no longer eagerly force-fetches Cloudflare Workers AI or Mistral alias resolution** — only OpenRouter. Both are credential-gated and provider-specific; most price-mode customers never call `workers-ai` or `mistral` at all in a given session, and the original design (added for the Cloudflare catalog above, then copied for Mistral) eagerly hit both APIs at SDK-construction time regardless of whether that provider was ever actually used — real, wasted network calls on every construction, every TTL cycle. Both now stay purely reactive: the session's first real call to that specific provider is what flags it stale (this already existed in `lookup()`); `maybe_refresh()` fetches it on the queue's very next tick; every call after that, even a moment later, hits the cache with zero further network calls until the TTL expires. Only that first per-provider call can race a cold cache — a provider a session never calls now costs nothing at all, instead of one unconditional fetch per SDK instance regardless of use. `warm_pricing()`'s docstring updated to describe the narrower (OpenRouter-only) guarantee; it also now accepts an optional `providers=["mistral", "workers-ai"]` to eagerly warm one or both when you already know you'll call them, closing even that first-call race if you want to. + - **`wrap()` now automatically (and non-blockingly) warms Cloudflare Workers AI/Mistral pricing the moment it sees a client that needs them** — no `warm_pricing(providers=[...])` call required at all. Wrapping a `mistralai` client learns that client's own `api_key` (`LagoSDK._extract_mistral_api_key`, reading `client.sdk_configuration.security.api_key` — verified against a real client instance) and feeds it straight to alias resolution via a new `PricingProvider.learn_mistral_api_key()`, so **no separate `LagoConfig.mistral_api_key` is needed at all** for the common case — the credential the customer already has to provide to make the real call is reused for pricing it too. Wrapping an OpenAI-shaped client checks its `base_url` for `gateway.ai.cloudflare.com` to distinguish "real OpenAI" from "Workers AI via Cloudflare's `.../compat` endpoint" (the client kind alone can't tell them apart) and warms the Cloudflare catalog the same way. Because `wrap()` normally happens some real time before the actual completion call (building the prompt, setting up messages), this closes the one-time cold-start race from the previous entry for the common case too — verified live: a fresh session's very first Mistral call, with no `warm_pricing()` call anywhere in the code, correctly billed as `llm_cost` instead of falling back to token events. An explicitly configured `mistral_api_key` still wins over a learned one if both are present. + +### Fixed + +- **A recovery path silently reversed FIFO order.** `_send_individually` re-queued each transiently-failing event as it went, and `_replay_failed` PREPENDS — so a 413 batch of `a,b,c,d,e` whose `b,c,d` failed while isolated came back as `d,c,b`. FIFO is the queue's contract: it is what makes the oldest-dropped-first overflow policy and Lago's own event ordering mean anything. Survivors are now collected and re-queued once. Present in **both** repos, not just JS as first reported. +- **A negative token count was dropped without a word.** `nonzero_numeric` correctly filters it (Lago would otherwise sum a negative billable quantity), but this was the last drop path in the SDK that never reached `on_error` — the same gap already closed for queue overflow and for an unresolvable subscription. It is reachable, not theoretical: `CanonicalUsage` is exported and `emit()` takes one directly, which is the documented way to backfill usage the SDK did not intercept, so a caller computing a delta wrongly really can hand us one. Now reported under `negative_tokens`, before the empty-check, so an event whose only fields were negative still reports instead of returning silently. +- **One log line per dropped event, not two.** `_report_error` already invokes `on_error` AND logs; an extra `logger.error` beside it emitted the same drop twice at two levels, so a customer grepping logs counted one lost call as two. The JS port logged **nothing** for the same event, so the two repos reported 2 lines vs 0 — `reportError` there now logs as well, since `onError` is opt-in and the log is the floor. +- **`verify_ssl=False` could crash `LagoSDK()` construction.** The InsecureRequestWarning suppression reached through `requests.packages`, a legacy compatibility alias with no guarantee of existing, in an unguarded attribute chain inside `__init__`. Now `import urllib3` directly, wrapped: suppressing a warning must never fail construction. This sits on an advertised path — `verify_ssl` is a first-class constructor argument the docstring recommends for local dev — so the crash would have hit exactly the setup the flag was added to serve. +- **`WORKERS_AI_COMPAT_PREFIX` had drifted into two definitions.** `adapters/openai_native` decides the *provider* from it and `pricing` strips it before a catalog lookup; those two must never import each other, so it now lives in `canonical` (which imports nothing from the package — no cycle either way, and no pulling `pricing`'s ~50KB into a lightweight adapter). A drift between the copies would have been a silently unpriced call rather than a crash, so a test now asserts there is exactly one definition in the tree. +- **Two comments in the Cloudflare gateway adapter were wrong, and the module docstring described behaviour Cloudflare does not have.** Both are corrected, because they were the justification for code and would have misled the next reader: + - The docstring claimed `usage_metadata`'s key casing "is NOT normalized by Cloudflare — it passes through whatever convention the underlying provider's own usage object used". It does not. Across all 14 captured fixtures (Anthropic, Workers AI, Mistral and Gemini, via every ingress method) the only keys that ever appear are Cloudflare's own: `input_tokens`, `output_tokens`, `total_tokens`, `input_cached_tokens`, `input_cache_creation_tokens`, `neurons`, `input_text_tokens`, `reasoningTokens`. **Not one provider-native key shows up.** The cited proof — camelCase `reasoningTokens` in the real Gemini entry — is Cloudflare's own inconsistency, not a leaked provider key: Gemini's native spelling for that quantity is `thoughtsTokenCount`, which appears nowhere. + - `_first_int`'s comment said a missed cache key is "an over-bill, not an omission". That is true only for a **subtractive** provider (`gemini`/`openai`/`workers-ai`, where `compute_cost` subtracts `cache_read` out of `input`). For **additive** Anthropic the same miss means those tokens are never billed at all — an under-bill, the direction this SDK treats as worse. The comment asserted one direction for a function used by both. +- **Added Anthropic's `cache_read_input_tokens` and Gemini's `thoughtsTokenCount` to the gateway adapter's spelling fallthrough**, alongside the `cachedContentTokenCount` / `cache_creation_input_tokens` entries that were already there. Labelled honestly in the code as **unobserved insurance**: neither has ever appeared in a captured fixture, so this guards against Cloudflare one day forwarding a provider's usage object instead of rewriting it, and is not handling for a case we have seen. Kept because `_first_int` fallthrough is free and a missed cache key mis-bills in one direction or the other for every provider. Pinned by synthetic tests marked as such. +- **Gemini streaming attributed the requested alias instead of the resolved model when a chunk carried usage without `model_version` — a real divergence between the two ports.** Python read `model_version` off whichever chunk carried usage; JS remembered it across chunks. On a stream that announces the version early and sends usage last, Python emitted `gemini-flash-latest` where JS emitted `gemini-2.5-flash-002`, so the two repos priced the same call differently. Python now persists it across chunks (both the sync and async stream wrappers), and accepts `modelVersion` as well as `model_version` since `model_dump()` yields snake_case while a raw REST dict is camelCase. **Not reachable with Gemini as it behaves today** — verified live that every streaming chunk carries both `model_version` and `usage_metadata`, which is why the existing fixture (faithful to that) could not catch it. The captured fixture was left faithful rather than edited to expose the bug; a separate, explicitly synthetic case pins the property. The alias hot-swap itself is real and live-verified: `gemini-flash-latest` resolved to `gemini-3.7-flash`. +- **A moving `~` alias could overwrite a real listing's price, decided purely by catalog order.** Stripping OpenRouter's `~` marker indexes an alias under its real vendor, which is what makes a plain `-latest` id priceable at all — but it wrote the alias-derived keys with plain assignment. Collision-freedom was verified against the live catalog and still holds, so nothing is mispriced today; it was a property of that day's response rather than of the code. If OpenRouter ever lists both `google/gemini-flash-latest` and `~google/gemini-flash-latest`, whichever arrived later won: measured on a synthetic pair, the same lookup returned `0.009` or `0.001` depending only on position in the response. Alias-derived keys are now written only when absent, so a real listing always wins regardless of order; the `~`-spelled id still resolves to its own entry, and non-alias entries keep plain assignment so genuine duplicates behave exactly as before. +- **The `-\d{3}` version-strip arm is now scoped to OpenRouter, where it was the only thing that ever needed it.** Gemini's `model_version` can report a `-002` revision that OpenRouter omits from its ids, so the shared `_strip_version` grew a 3-digit arm — but that helper also builds the AWS/Bedrock price keys, and there a shortened key does not merely miss: `bedrock_model_key` feeds `table.setdefault(key, {})[direction] = price`, so two distinct models collapsing onto one key silently overwrite each other's rate. All four live catalogs are clean (OpenRouter 415 ids, Cloudflare 64, AWS offer 77, captured Bedrock 39 — zero model parts ending in exactly three digits), so this was latent, not active. Splitting it into an OpenRouter-only strip makes the risk structurally zero rather than empirically zero, and `amazon.titan-text-001` / `-002` now stay distinct keys. +- **`apply_markup`'s two bad-input fallbacks were not equivalent, and the ports disagreed on one.** An unparseable `usd` means the cost is unusable, so 0 is right; an unparseable `markup` means only the multiplier is unusable, and returning 0 there discards a perfectly good cost. Python returned `"0"` for both, JS fell back to 1.0 for a bad markup, so identical input would have produced different bills. Python now matches JS. This is defence in depth rather than a live fix: every `emit()` path already runs the customer's markup through `coerce_markup` (which falls back to 1.0 and reports under `on_error`), and `CostBreakdown.markup` / `fields[*]["cost"]` are `_fmt_money` output, so neither argument can actually arrive unparseable — the divergence was latent. What *was* untested is the end-to-end consequence of that guard, now pinned: a customer sending `markup="1,5"` gets the cost billed at 1.0 rather than zeroed, and the lost markup reaches `on_error`. +- **A dropped event now always reaches `on_error`, not just a log line.** Two paths lost billable events while telling only the module logger, so a customer watching the error hook — the documented channel for every other billing gap — saw nothing: the queue's buffer **overflow**, and `emit()` dropping a call because no subscription resolved. The JS port already reported both, so the two repos disagreed on whether a lost event is even visible. + - **The overflow report is raised with the queue lock released, and is re-entrancy guarded per thread.** `_lock` is a plain `Lock`, not an `RLock`, so an `on_error` hook that itself emits — a plausible hook, since reporting a billing gap by emitting a metric is an obvious thing to do — deadlocked the *customer's* thread permanently if the report was made while holding it. Reporting after release fixes that but re-enters unboundedly, because the buffer is still full when the hook runs; a `threading.local` guard bounds it to one report per overflow per thread. Both failure modes are covered by a subprocess test, since an in-process deadlock wedges interpreter shutdown and hangs the whole suite rather than failing it. +- **A negative token count could be emitted as a billable quantity.** `nonzero_numeric` filtered on truthiness, so `input=-100` survived and was pushed as `value="-100"`. Nothing upstream should produce one — every adapter clamps at extraction — but this is the last gate before an event is built, and the JS port already filtered on `> 0`. Now positive-only in both. +- **`apply_markup` was the one money helper that could raise.** A non-numeric input hit a bare `Decimal("abc")` and threw `InvalidOperation` from inside `_push_cost_event`, under `emit()`'s catch-all — so the cost event was dropped and reported as an unknown `"emit"` error instead of taking the documented no-price path, past every caller relying on this module's `None`-on-bad-input convention. Now parsed through `_parse_price` and floored to `0`, matching the JS port and `compute_precomputed_cost`'s existing floor-to-zero behaviour. `money_golden.json` is unchanged and still byte-identical across repos. + +- **The Cloudflare catalog fetch had no test coverage at all, and two ways to silently under-price.** + - **`result_info.total_count` was trusted as a terminator, and defaulted to `len(models)` when absent** — which made a missing count break after page one, keeping 50 of the 64 models the endpoint actually serves, with no diagnostic. Measured live, `total_count` is worse than merely absent-able: it reports **291 while the endpoint serves 64** (50, then 14, then 0), so a `len(models) >= total` test can never fire at all. A short page is the only reliable end-of-catalog signal, and is now the only one used. The current behaviour was correct by luck — the short page terminated the loop — but any change to that count would have cost 14 models silently. + - **The `while True` had no page bound**, on a loop that runs on the queue's flush tick *ahead of the drain* — so an endpoint returning a full page indefinitely would stall event delivery, not merely waste bandwidth. Now capped at 40 pages (~2000 models against a real catalog of 64) and the truncation is logged, since a short catalog otherwise reads as "these models are unpriced". + - **Python only: one malformed entry unpriced every Workers AI model.** `m.get("properties", [])` only defaults when the key is *absent*, so an explicit JSON `null` returned `None` and `for p in None` raised `TypeError` straight out of `parse_cloudflare_workers_ai` into `maybe_refresh`'s handler — which reports and leaves `_cloudflare_workers_ai` at `None`. The whole table stayed empty, so every Workers AI call fell back to token events. The JS port already isinstance-guarded here and dropped only the bad entry. + - Five new tests per repo cover the fetch loop, which previously had none: the real 50-then-14 shape, a missing `total_count`, a wrong `total_count`, the page bound, and a null-`properties` sibling surviving. + +- **Every `-latest` model alias on OpenRouter was unpriceable — 11 models across 6 vendors, including `claude-sonnet-latest`, `claude-opus-latest` and `gpt-latest`.** OpenRouter marks a *moving* alias with a leading `~` on the vendor (`~anthropic/claude-sonnet-latest`, `~openai/gpt-latest`, `~google/gemini-flash-latest`). `parse_openrouter` split the id on `/` and took the left half as the vendor, so these indexed under `~anthropic` / `~openai` / `~google` — none of which appear in `_VENDOR_MAP`. A customer in price mode requesting a plain `-latest` alias therefore missed the table and fell back to token events: **nothing at all is billed in an `llm_cost`-only setup.** Verified live against the 415-model catalog: all 11 carry real token pricing, all 11 missed, all 11 now resolve. The `~` id stays indexed as well as the bare one, so nothing that already worked changes, and the strip is collision-free — no un-prefixed id duplicates a `~`-prefixed one. + - Found while checking the reviewer's separate question about Gemini's `-002` suffix, which turned out to be a *second, independent* reason that one model could not price. Fixing only the suffix would have left it missing anyway. +- **A 3-digit revision suffix is now stripped when matching OpenRouter.** Gemini's `model_version` can report `-002` where OpenRouter lists only the bare name, so preferring the resolved id (which is correct, and what the response reports) turned a hit into a miss. Latent rather than live: every captured real Gemini response reports a bare `gemini-2.5-flash` with no revision suffix, so the trigger does not currently occur — but it costs one regex arm to close, and it is a config change away. Verified safe against the live catalog: **zero** of the 415 ids have a model part ending in exactly three digits, so the new arm cannot shorten a real listing. + +- **A cached Gemini call through the gateway was over-billed 3.96x, because `cache_read` was read under one spelling only.** `cache_read` checked just `input_cached_tokens` and `cache_write` just `input_cache_creation_tokens`, while `reasoning` on the next line already checked two casings — and that asymmetry was the tell. Surveying every `usage_metadata` key across all 14 captured fixtures shows the gateway's OWN counters are consistently snake_case (`input_tokens`, `output_tokens`, `total_tokens`, `input_cached_tokens`, `input_cache_creation_tokens`), but a provider's native key can pass through untouched: the real Gemini entry carries camelCase `reasoningTokens`, plus an `input_text_tokens` this adapter maps nowhere. So the spelling of a cache key on a provider we have no *cached* capture for was genuinely unknown. + - **The consequence is an over-bill, not a lost field.** `_normalize_provider` now correctly maps these entries to `gemini`, and `gemini` is in `_INPUT_INCLUDES_CACHE_READ` — so `compute_cost` relies on `cache_read` being populated in order to SUBTRACT the cached portion out of `input`. A silent `0` leaves all 10,000 prompt tokens billed at the full input rate instead of 1,000 at input plus 9,000 at the cache rate. Measured against the live OpenRouter table on a `gemini-2.5-flash` call with 9,000 cached: **$0.00325 against a true $0.00082, a 3.96x over-bill**, and the error grows with cache hit rate — worst on the long-cached-system-prompt workload that caching exists for. + - Both fields now check the gateway's snake_case name, its camelCase form, and the provider's own native name (`cachedContentTokenCount` for Gemini, `cache_creation_input_tokens` for Anthropic). Checking extra spellings costs a dict lookup; missing one costs a 4x over-bill. + - **The JS port additionally used `??` where Python used `or`**, so it only fell through on null/undefined. A provider sending both its own key and the gateway's with one of them zeroed resolved to the **zero** and lost the real count in JS while Python resolved it correctly — the two repos disagreeing on live money again. Both now fall through on any falsy value. + +- **An explicitly supplied `usd_cost` was discarded with no log and no `on_error`.** `emit()` returns early whenever the effective mode isn't `"price"`, and never consulted `usd_cost` on the way out — so a caller who had gone to the trouble of passing a gateway's real metered price got token counts instead, silently. A hand-rolled backfill written from the module docstring (the pattern `examples/cloudflare_gateway_demo.ipynb` demonstrates) would drop every real cost it read and still look like it succeeded. The configured mode is still respected — that part is deliberate, since honouring a per-call `usd_cost` in token mode would emit an `llm_cost` event mapping to none of a token-mode customer's charges — but the discard is now reported through `on_error`, the same hook every other billing gap uses. Reported **per occurrence rather than deduped**: the count of discarded costs is precisely what a caller reconciling on `on_error` needs, and the documented backfill pattern passes an explicit `mode="price"`, so hitting this at volume means a real misconfiguration rather than normal operation. The common case — no `usd_cost` supplied at all — stays silent. + +- **The single-event cost path reported a different token quantity than the split path, for the same call.** `unit` was `str(usage.input + usage.output)`, which dropped `reasoning` and `cache_write` entirely and counted a cache-inclusive provider's cached tokens at full weight — while the per-`token_type` branch right below it reports the de-overlapped `parts["tokens"]`. Two branches of one method, two bases. On the captured `16_real_gemini_via_dedicated_endpoint.json` row (`tokens_in=9`, `tokens_out=21`, `reasoningTokens=852`) it published `unit="30"` for a call that consumed **882**. New `deoverlapped_token_total()` sums the same `PRICED_FIELDS` the split path emits one event each for, so the two now agree by construction. The charged amount was never affected — that comes from `precise_total_amount_cents` — but `unit` is the *reported quantity*, which is what a customer points a `sum` aggregation at. + - **Both `_INCLUDES_` sets are applied, and deliberately NOT gated on a price existing** (unlike `compute_cost`'s subtraction): this is a token count, so whether a rate happens to be published cannot change how many tokens were consumed. The two still agree — when a cache-inclusive provider has no `cache_read` price, `compute_cost` leaves the cached tokens inside `input` and emits no `cache_read` event, and this skips `cache_read` for the same reason. + - **Limited to the five `PRICED_FIELDS` on purpose.** `tool_calls` is a count of calls rather than tokens, and `cache_write_5m` / `cache_write_1h` are a breakdown OF `cache_write`, so including any of them would not be a token total. Mirrors price mode's documented five-field scope. + - Worth recording an under-report the review didn't name: for an **additive** provider the old basis was far wider than the reasoning case. An Anthropic call with `input=1000`, `cache_read=900`, `output=100` reported `unit=1100` against 2000 actually consumed — cache tokens are real extra consumption there, not a subset. Verified end-to-end against a live Lago instance: the `llm_cost` charge's `units` moved by exactly 2000, with the dynamic charge carrying the metered `$0.05`. + +- **A config-only `api_url` sent every event to PRODUCTION Lago.** `LagoSDK.__init__` defaulted `api_url` to the production URL, so the `if api_url:` guard meant to let explicit args win *always* fired — silently overwriting whatever the caller had set on `LagoConfig`. `LagoSDK(api_key=k, config=LagoConfig(api_url="http://localhost:3000/api/v1"))` therefore shipped a local-dev customer's usage data to `api.getlago.com`, with no error and nothing in the logs. The default is now `None`, and every explicit arg is guarded on "was it actually passed?" rather than on truthiness, so a config value survives when the arg was omitted while an explicit arg still wins when given. The production default is unchanged when nothing is passed at all. + - **`api_url` is the one field guarded on non-emptiness rather than on presence.** `api_url=os.environ.get("LAGO_API_URL", "")` with the variable unset — the ordinary way to write this — would otherwise store `""`, and an empty base URL is unrecoverable downstream: `requests` raises `MissingSchema`, which is not a `LagoApiError`, so the queue classifies it as transient and retries at the 60s ceiling forever. Every event is buffered and none is ever delivered, with a growing buffer as the only symptom. An empty string now falls through to config, then to the production default, exactly as an omitted argument does. + - **This was the shortest path to the bug, not an exotic one.** A custom `api_url` and `verify_ssl=False` go together in exactly one setup — a local dev Lago behind a self-signed cert (Traefik's default) — which is the setup `verify_ssl` was added for. Since `verify_ssl` was reachable *only* through a `LagoConfig`, the feature pushed callers straight into the clobber. `LagoSDK(..., verify_ssl=False)` is now accepted directly, so that setup needs no config object at all. + - **The JS port had the same bug inverted, so the two repos disagreed.** `...(opts.config || {})` was spread LAST, so a `config.apiUrl` overrode an *explicitly passed* `apiUrl` — the opposite of what Python documents and does. The identical call therefore billed a different Lago instance depending on which SDK you used. Both now apply explicit options over config, and `verifySsl` is accepted at the top level to match. + +- **A caller dimension was honoured on token events and silently discarded on cost events.** The cost path spread `dimensions` into `base_properties` — i.e. *before* `unit` / `value` / `base_cost` / `unit_price` — so those four SDK-computed keys overwrote a caller dimension of the same name, while the token emitter (which spreads `dimensions` last) honoured it. One customer config, two different outcomes depending on the mode, with no error either way. `dimensions` is now spread last in both emitters, so one rule holds everywhere: **a caller dimension always wins over an SDK-computed property of the same name.** Names like `unit`, `value` and `model` are ordinary English words a customer may well use for their own breakdown — a per-seat biller naturally writes `dimensions={"unit": "seat"}` — so this is a reachable config, not a contrived one. The accepted consequence, now pinned by a test rather than left implicit: a dimension named `value` overrides the *reported* quantity on token events. It cannot affect the amount actually charged on a cost event, because `precise_total_amount_cents` is a sibling of `properties`, not a member of it. + +- **Mistral alias resolution priced a whole model family at its OLDEST snapshot's rate.** `_pick_mistral_canonical` sorted by `(len(n), n)` — but every dated id in one family is the same length, so the length term always tied and the choice fell through to the alphabetical one, which for `-2402` / `-2407` / `-2411` *is* the date, ascending. A family of `mistral-large-2402` / `-2407` / `-2411` / `-latest` therefore collapsed onto `mistral-large-2402`, and every member was priced at a two-year-old rate. Now sorts by a normalized date descending, so the newest snapshot wins; 4-digit `YYMM` is widened to `YYYYMM00` first, because mixed suffix widths do not compare correctly as raw strings (`"20250929"` sorts *below* `"2411"`). Two further parts to the same bug: + - **An explicit dated snapshot is no longer remapped at all.** The mapping loop rewrote *every* non-canonical group member, including the real dated ids — so asking for `mistral-large-2411` by name was silently redirected to the group's canonical and priced at that snapshot's rate. Before alias resolution existed, `mistral-large-2411` matched OpenRouter directly, which makes this a regression rather than a gap. An exact snapshot request now passes through untouched. + - **The JS port used `localeCompare`, so the two repos disagreed on live money.** `localeCompare` is ICU/locale-dependent, i.e. not reproducible across environments — directly contradicting the docstring's own determinism claim — and for a group differing only by case/separator the two repos picked *different* canonicals (`Mistral-Small-2603` in Python, `mistral_small_2603` in JS). `_norm` lowercases and maps `.`→`-` but leaves `_` alone, so the JS pick normalized onto a name OpenRouter does not list and the whole group fell back to token events there while pricing correctly in Python. Both now order by Unicode code point and agree. + +- **A backfill re-run after the price table warmed up silently discarded the priced events.** `_emit_token_events` and the per-field branch of `_push_cost_event` both built `f"{event_id}_{field_name}"` over the same field vocabulary (`input`, `output`, `cache_read`, ...), and both are reachable for the SAME `event_id`: `emit()` falls back to token events when a price lookup misses, then takes the cost path once the table is warm. So run 1 (cold table) sent `backfill_X_input` under `llm_input_tokens`, and run 2 (warm) re-sent the identical id under `llm_cost` — Lago rejected it as a duplicate `transaction_id`, and because `/events/batch` is all-or-nothing that rejection failed **every other event in the batch** too. Net effect: the dollar amounts for that window were never billed, only the raw token counts, and nothing surfaced it — which defeats the very idempotency promise `event_id` exists to provide. The two multi-event paths now use disjoint namespaces (`_tok_` / `_cost_`); the single precomputed-cost event stays unsuffixed, since one event has nothing to disambiguate. No migration needed: `event_id` is absent from 0.2.0 entirely, so no released path has ever emitted the old format (they all use a random UUID). Pinned by a regression test that asserts a cold run's ids and a warm run's ids over the same `event_id` do not intersect — verified to fail without the fix, where all four ids collided. + +- **Rate-limited events were dropped for good, and the retry made the rate limit worse.** `_is_permanent_failure` classified the whole `400-499` range as unretryable, but two 4xx statuses mean "try again, later": **429** (rate limited) and **408** (request timeout). So a throttled 100-event batch took the isolate-and-drop path — 100 further requests aimed at the server that had just asked us to slow down, each of them also throttled, each then logged and discarded. 100 billable events lost, and the isolation actively deepened the throttle. This was a regression: before the permanent/transient split existed, the same 429 went through the 1s→60s backoff and eventually landed. "Permanent" is now an explicit set (`400, 401, 402, 403, 404, 409, 413, 415, 422`) rather than a range, so an unrecognized 4xx errs toward retrying — waiting on an event that would have been dropped costs latency, dropping one that would have been accepted costs revenue. The isolate-one-by-one recovery is unchanged for the validation statuses it was written for, so a single duplicate `transaction_id` still doesn't take its batch down with it. + - **Three of those statuses are permanent for the opposite reason to 429, and that is why they belong in the set.** `413` (payload too large), `402` (payment required) and `415` (unsupported media type) reject the batch *as a batch* — re-sending the identical bytes can never succeed, so retrying it re-prepends it at the head of the FIFO and blocks every event behind it until the 60s ceiling, indefinitely. Classifying them permanent routes them to the isolate-one-by-one path, which splits the batch and delivers what is individually deliverable — the case that path was built for, and the one it could not previously reach. `405` and `410` stay transient: both are plausibly a misconfigured or mid-deploy endpoint that will start working again. + +- **Every documented way to call Workers AI went unpriced.** `_infer_provider` matched only a bare `@cf/...`, but Cloudflare's OpenAI-compatible `/compat` endpoint requires the routing-prefixed form — `workers-ai/@cf/meta/llama-3.3-70b-instruct-fp8-fast` — which is exactly what `README` and `examples/cloudflare_gateway_demo.ipynb` prescribe, and the only form a **streaming** call can report (the synthetic usage payload carries no model, so `resolve_model` returns the requested string verbatim). Those calls were stamped `provider="openai"`, looked up against OpenRouter as `openai/workers-ai/@cf/...`, missed, and silently degraded to token events — with nothing to degrade *to* in an `llm_cost`-only billing setup. The `workers-ai` entry in `_INPUT_INCLUDES_CACHE_READ` never applied either, so on the paths that did price, cache reads were billed twice again. Fixing the provider check alone was not enough: Cloudflare's catalog keys models as bare `@cf/...`, so `lookup_cloudflare_workers_ai` now strips the routing prefix before matching. `CanonicalUsage.model` deliberately keeps the spelling the customer used, so reporting stays faithful to the request while pricing resolves. + +- **The model-attribution fix never reached the streaming path, in any of the three wrappers.** Each stream wrapper rebuilds a synthetic usage payload from the chunks, and all three discarded the model the response reported — so `resolve_model` fell back to the requested alias, which is precisely the bug the non-streaming path had just been fixed for. A streamed `gpt-5-chat-latest` stayed `gpt-5-chat-latest` instead of resolving to the dated snapshot OpenRouter lists, so price mode missed and degraded to token events, while the identical non-streaming call priced correctly. Each provider hides the resolved name somewhere different: OpenAI puts `model` on every chunk (and on `response` for the Responses API), Anthropic reports it **only** on `message_start` under `message.model` — so the wrapper now carries it across the whole accumulate-and-merge stream — and Gemini reports `model_version`, which it hot-swaps server-side for `-latest` aliases. It matters most on a gateway, where the resolved name is what decides which price table the call is looked up in at all. Pinned by a new attribution test per wrapper; the streaming fakes previously omitted the model field entirely, and no wrapper test asserted `properties["model"]`. + +- **Gateway-backfilled Gemini calls could never be priced, and their cached tokens were billed twice.** `extract_cloudflare_log` passed Cloudflare's own provider vocabulary through verbatim, but that is not the vocabulary the pricing and token-semantics tables key off — and not even Cloudflare's own URL slug (the logs say `workers-ai` where the endpoint path says `workersai`). A real captured entry reports `provider: "google-ai-studio"`, which matched no vendor in `_VENDOR_MAP`, so `lookup_openrouter` searched a vendor that does not exist and missed every time (verified against the live 400-model OpenRouter table: miss as `google-ai-studio`, hit as `gemini`). The same miss also kept it out of `_INPUT_INCLUDES_CACHE_READ`, so Gemini's `cache_read` — a **subset** of its input count — was billed on top of the full input rather than subtracted from it. Cloudflare's names are now mapped onto the SDK's (`google-ai-studio`/`google-vertex-ai`/`vertex` → `gemini`, `azure-openai`/`azureopenai` → `openai`, `workersai` → `workers-ai`); anything unrecognized passes through untouched, since a clean miss falling back to token events beats an invented mapping. AWS Bedrock is deliberately **not** mapped — its prices key off `api.startswith("bedrock")` and this connector always sets `api="cloudflare_gateway"`, so a mapping would route it to OpenRouter under a vendor that cannot match. +- **A model already carrying its vendor prefix never matched a price.** A real gateway log for a REST-path call reports `model: "anthropic/claude-opus-4.8"` alongside `provider: "anthropic"`, which `lookup_openrouter` turned into `"anthropic/anthropic/claude-opus-4.8"` — a guaranteed miss. The prefix is now stripped, but **only** when it agrees with the vendor resolved from `provider`, so the lookup stays vendor-gated as documented: a model naming a different vendor than the call claims is still a miss, not a cross-vendor mispricing. With both fixes, all 10 distinct (provider, model) pairs across the real captured fixtures now resolve to a live price; three of them previously missed. +- **Workers AI cached tokens were billed twice.** `provider="workers-ai"` was missing from `_INPUT_INCLUDES_CACHE_READ`, but Workers AI is only ever reached through Cloudflare's OpenAI-**compatible** endpoint (`.../compat`), so its usage payload is the OpenAI shape — `prompt_tokens` already **includes** `prompt_tokens_details.cached_tokens`. It is a distinct provider only because it prices against Cloudflare's own catalog, not because its token semantics differ. With the cached portion never subtracted from `input`, those tokens were charged once at the full input rate *and* again at the cache-read rate, which Cloudflare's catalog does publish (verified live: `@cf/moonshotai/kimi-k2.6`, `@cf/moonshotai/kimi-k2.7-code` and `@cf/zai-org/glm-5.2` all list a "per M cached input tokens" price). Measured against a real cached call (prompt 23233 / cached 23168) at live catalog rates, this overbilled by **+583%**; the error scales with cache hit rate, so a long cached system prompt — the standard agent workload — is the worst case. Pinned by two new golden cases (`workers-ai` subtracts, `anthropic` stays additive) carrying those real counts. + - **The same omission existed on the output side and is fixed with it.** `workers-ai` was missing from `_OUTPUT_INCLUDES_REASONING` too, and for the identical reason: the `/compat` endpoint returns the OpenAI shape, where `completion_tokens_details.reasoning_tokens` is a **subset** of `completion_tokens`, so counting it additively inflates the token basis — a 100-input/1000-output call with 800 reasoning tokens reported 1900 instead of 1100. **Unlike the cache-read side, this one is not reachable on a live call today, and the 1900 above is the arithmetic, not a measured bill.** Verified against real Workers AI on 2026-08-20: three reasoning-capable models (`@cf/openai/gpt-oss-120b`, `@cf/deepseek-ai/deepseek-r1-distill-qwen-32b`, `@cf/qwen/qwen3-30b-a3b-fp8`) returned no `completion_tokens_details` block at all, and the gateway Logs API reports only `input_tokens`/`output_tokens`/`total_tokens`/`input_cached_tokens`/`neurons`/`units` on `workers-ai` rows — so `reasoning` is never populated for this provider by either the native `/compat` adapter or the logs backfill. It is fixed regardless, because the two sets must agree about the same provider for the same reason: the cache-read side *is* live-reachable (`input_cached_tokens` is reported and non-null), and one set knowing that `/compat` implies OpenAI token semantics while the other does not is precisely the asymmetry that produced the original double-billing. +- **`_parse_price` raised instead of returning `None` on absurdly large values.** `.quantize()` sat outside the `try`, so any value ≥ 1e16 (16 integer + 12 fractional digits exceeds `Decimal`'s default 28-digit context precision) threw `InvalidOperation` straight out of a function documented as returning `None` on bad input — past every caller relying on that, and out of `compute_precomputed_cost` into `emit()`'s catch-all, where the event was dropped as an unknown error rather than taking the normal "no price" path. Now returns `None`, which also matches what the JS port returns for the same inputs. +- **Cross-repo golden fixture gained a `precomputed_cases` section**, asserted by both repos, carrying verbatim `cost` values from real Cloudflare AI Gateway log entries — including ones below 1e-6, where JS renders the number in exponential notation. Python has always parsed those correctly; the JS port silently billed them as $0 (fixed in that repo's matching release). `cases` also gained an optional `provider`, so per-provider token semantics are now pinned by the shared fixture rather than by each repo's own tests. + +- **Two real reliability bugs in `EventQueue`, both found live while backfilling a real, already-partially-backfilled Cloudflare window:** + 1. Any send failure — a permanent Lago 4xx (e.g. a duplicate `transaction_id` from replaying the same window twice) or a transient one — got identical treatment: re-queue the whole batch and retry with backoff, forever. Since Lago's `/events/batch` is all-or-nothing, one permanently-doomed duplicate at the front of the FIFO buffer blocked every event queued behind it — including brand new, perfectly valid ones — indefinitely. Now a `LagoApiError` with a 4xx status falls back to sending the batch one-by-one: individually-permanent failures are logged and dropped for good, individually-transient ones re-queue normally, and neither blocks the other. + 2. `shutdown()`'s final drain silently swallowed any failure at all (`except Exception: pass`) and only ever attempted a single batch, so a buffer holding more than `max_batch_size` events at shutdown time left the rest never even attempted — with no error, no log, nothing. Now drains every remaining batch (time-bounded so a persistently-down network can't spin the exiting thread forever), applies the same permanent/transient handling as the main loop, and reports every failure via `on_error`/a warning log instead of vanishing silently. +- **`LagoSDK.warm_pricing()`** — blocks until price mode's tables are fetched, instead of waiting for the queue's background thread's next tick (~1s later by default). Found live, the hard way: a real notebook constructing the SDK and immediately making a price-mode call raced a cold cache; with a single `llm_cost`-only billing setup (no token-metric charge left to fall back to — see the `token_type` change below), the event was silently lost rather than merely mis-priced. A long-running server's first real call naturally lands well after that first tick and never hits this; a script, notebook, or one-shot job making a call right after construction does. Call it once, right after constructing the SDK with `pricing_mode="price"`. +- **`LagoSDK.emit()` now accepts `usd_cost` and `event_id`** — the connector's one-call entrypoint for billing a gateway-reported cost directly, instead of hand-building `precise_total_amount_cents` events with a raw REST call. `usd_cost` skips the SDK's own OpenRouter/Bedrock price lookup entirely and bills that exact amount (for a gateway that already reports its own real, metered cost per call — e.g. Cloudflare AI Gateway's `cost` field). `event_id` sets Lago's idempotency key (`transaction_id`) instead of a random UUID, so replaying/backfilling the same log window twice never double-bills; in token mode (which can push several events from one call) each field's event gets a `f"{event_id}_{field_name}"` suffix so they don't collide with each other. New `compute_precomputed_cost()` in `pricing.py` mirrors `compute_cost()`'s money conventions (`Decimal`, floored to 12 dp) but skips the per-field breakdown since a gateway gives one lump sum, not a per-token table. +- **Price mode now bills one `llm_cost` event per `token_type` (input/output/cache_read/cache_write/reasoning) when a real per-field breakdown exists, instead of one event summing the whole call.** Lets a single `llm_cost` billable metric be `grouped_by: ["model", "token_type"]` in Lago — broken down by both dimensions from one metric, live wrap() calls and Cloudflare backfill alike. Markup is applied per field (previously only to the summed total — `compute_cost`'s per-field `cost` values are pre-markup, so this needed its own fix: `pricing.apply_markup()`). The `usd_cost`/precomputed path (Cloudflare's own lump cost per call) has no real per-field split to work with — it still emits a single event, grouped by `model` only; `token_type` is absent rather than a fabricated proportional guess. Verified empirically that this doesn't create a real pricing mismatch between the two paths for at least one model: Cloudflare's actual charged rate for `claude-sonnet-4-5` ($3/M input, $15/M output, solved from real invoiced amounts across several real calls) matches OpenRouter's listed price for the same model exactly. +- Live-verified end to end: backfilled 99 real historical Cloudflare AI Gateway log entries into Lago as `llm_cost` events priced from Cloudflare's own `cost` field (not our pricing tables), through a new `llm_cost` dynamic-charge-model billable metric — total backfilled cost $0.0175, matching Cloudflare's own numbers exactly. Confirmed idempotency for real: re-running the backfill against the same window has Lago reject every duplicate `transaction_id` (`"value_already_exist"`) — worth noting for connector design that `/events/batch` rejects the **whole batch** atomically on any single collision, not just the colliding entry, so a real poller needs cursor-based dedup rather than relying on idempotency alone to make replay safe. + +### Fixed +- **OpenAI/Anthropic adapters mis-tagged usage with the requested model instead of the model that actually answered.** `extract_openai_native`/`extract_anthropic_native` preferred the request's `model` kwarg over the response's own `model` field. Harmless calling a provider directly with a fully-qualified model id, but wrong the moment a provider resolves a short alias to a dated snapshot — confirmed live with no gateway involved at all: requesting `claude-sonnet-4-5` answered as `claude-sonnet-4-5-20250929`. Nearly every captured OpenAI fixture in this suite shows the same pattern (`gpt-4o-mini` → `gpt-4o-mini-2024-07-18`). Both adapters now prefer the response's own `model`, falling back to the request only when the response is silent about it (e.g. a synthetic streaming usage blob). Pricing and per-model attribution now key off what actually served the request. +- **`extract_gemini_native` had the same bug, but backwards from how it looked in OpenAI/Anthropic.** It preferred the requested `model_id` over the response's own `model_version`, even though `model_version` was already present in every response — it was just never used unless `model_id` was empty. Gemini resolves "-latest" aliases (`gemini-flash-latest`) to a dated snapshot server-side the same way OpenAI/Anthropic do (confirmed in [Google's docs](https://ai.google.dev/gemini-api/docs/models): "this alias will get hot-swapped with every new release"); every captured fixture happened to request an already-dated model, so `model_version` came back identical and this never showed. Flipped to prefer the response's `model_version`, matching the OpenAI/Anthropic adapters — no new fetch or credential needed, the resolved id was already being discarded. + +### Added +- **Gateway cache-hit detection for OpenAI/Anthropic wrappers.** Non-streaming `.create(...)` calls now go through `.with_raw_response.create(...)` so the wrapper can see response headers before parsing the body. If a gateway in front of the provider (e.g. Cloudflare AI Gateway) marks the response `cf-aig-cache-status: HIT`, the provider served it from cache at zero cost to the customer, and the wrapper skips billing it. `.parse()` on the raw response returns the identical object `.create()` would have, so this is invisible to the customer and a no-op with no gateway in the path. Streaming calls are not covered yet — gateways typically recommend `.with_streaming_response` for that, which behaves differently and hasn't been verified end-to-end; streaming keeps using the plain `.create()` path. Falls back to the pre-existing behavior if `.with_raw_response` isn't available on the client (older SDK versions). +- **`lago_agent_sdk.gateway.adapters.cloudflare_gateway`** — `extract_cloudflare_log()` maps a Cloudflare AI Gateway Logs API entry (`tokens_in`/`tokens_out`/`usage_metadata`/`model`/`provider`) to `CanonicalUsage`, and `resolve_subscription()` reads Lago attribution from the customer's `cf-aig-metadata` header. Lives in a new `lago_agent_sdk.gateway` namespace, separate from the provider-native `adapters/` used by `wrap()` — this is the extraction half of a standalone log-polling connector (not part of `wrap()`), verified against a real captured log entry whose token counts were independently confirmed to roll up correctly in a real Lago instance. The poller itself (scheduler, cursor store, credential store) is not part of this SDK and isn't built yet. +- **Verified `extract_cloudflare_log()` against all three of Cloudflare's ingress methods, live**, not just the provider-native `/{provider}` routes covered above: the REST API (`POST /accounts/{account}/ai/run`), the Unified/OpenAI-compat endpoint (`.../compat/chat/completions`, called with the real `openai` SDK), and the Native/binding method (`env.AI.run(model, input, {gateway: {id, metadata}})`, only reachable from inside a deployed Cloudflare Worker). Same extraction function, zero code changes, correct results and correct attribution (`resolve_subscription()`) across all three — confirms the log schema is normalized regardless of how the call reached the gateway. Also swept 26 real Workers AI models through the REST API in one pass (22 succeeded, 4 failed for real account/licensing reasons — Workers Paid plan required, or a model needing explicit license acceptance — none a compatibility gap); extraction had zero failures across the full spread, including an unusual moderation-model shape (`llama-guard-3-8b`: 203 input / 3 output tokens). New fixtures 06–11 in `tests/unit/gateway/adapters/fixtures/cloudflare_gateway/` capture this. +- **`wrap_gemini_client`/`wrap_mistral_client` verified through Cloudflare's dedicated per-provider passthrough endpoints, live, with real customer API keys** (`.../google-ai-studio` and `.../mistral`) — real calls, real Lago billing, same pattern already proven for Anthropic: the customer's own key is forwarded directly, no Cloudflare-side BYOK/wholesale credits needed. Both required an explicit `cf-aig-authorization` header the SDK doesn't add on its own (`http_options.headers` for `google-genai`, `http_headers=` per-call for `mistralai`). +- **Fixed a real gap this surfaced: `extract_cloudflare_log()` never mapped reasoning tokens.** The real Gemini call's log entry has `usage_metadata.reasoningTokens: 852` (camelCase) sitting right next to `tokens_out: 21` (the visible completion only) — Cloudflare doesn't normalize `usage_metadata`'s key casing across providers; it passes through whatever convention each provider's own usage object used (Anthropic: snake_case, Gemini: camelCase). Now checks both cases for the reasoning field. +- **Replaced two hand-built synthetic cache fixtures with real captured ones — and corrected a wrong assumption in the process.** The old synthetic gateway-cache-hit fixture assumed a `cached: true` entry still reports the token counts the call "would have" cost, leaving billing policy to decide whether to skip it. A real captured cache hit (same request sent twice with `cf-aig-cache-ttl` set; the second came back in 8ms vs 296ms) proves that's wrong: Cloudflare's own log already reports `tokens_in`/`tokens_out` as 0 on a real hit — no caller-side branching on `cached` is needed. Separately, real back-to-back Anthropic calls through the gateway with a >1024-token `cache_control: {"type": "ephemeral"}` block confirm `usage_metadata.input_cache_creation_tokens`/`input_cached_tokens` exactly match Anthropic's own `cache_creation_input_tokens`/`cache_read_input_tokens` (3429 tokens, both directions) — this mapping was previously untested against real data. + ## [0.2.0] - 2026-06-15 ### Added diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 368a059..62acb56 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -35,14 +35,12 @@ make test # Unit tests with coverage report uv run pytest tests/unit --cov=lago_agent_sdk --cov-report=term-missing - -# Integration tests (require credentials — see env vars in each test) -AWS_BEARER_TOKEN_BEDROCK="..." \ -MISTRAL_API_KEY="..." \ -LAGO_API_URL="..." LAGO_API_KEY="..." LAGO_EXTERNAL_SUBSCRIPTION_ID="..." \ -uv run pytest tests/integration -q ``` +There is no committed live-provider test tier. Adapter behaviour is pinned by +captured real responses under `tests/unit/adapters/fixtures/`, which is what the +unit tests assert against; re-capture a fixture rather than hand-editing one. + ## Linting and type checks ```bash @@ -72,7 +70,6 @@ uv lock --upgrade-package X # bump a single package - `src/lago_agent_sdk/lago_client.py` — thin HTTP client to `/events/batch` - `tests/unit/` — unit tests, organized to mirror `src/` - `tests/unit/adapters/fixtures/` — captured real provider responses, used by adapter tests -- `tests/integration/` — live tests, gated on credential env vars ## Adding a provider @@ -82,7 +79,6 @@ uv lock --upgrade-package X # bump a single package 4. Update `detector.py` to recognize the client class. 5. Update `sdk.py::wrap()` to dispatch to the new wrapper. 6. Add unit tests against the captured fixtures. -7. Add a live integration test gated on the provider's API key env var. ## Pull request checklist diff --git a/README.md b/README.md index 70562a2..104167c 100644 --- a/README.md +++ b/README.md @@ -129,9 +129,42 @@ sdk.flush() Wraps the modern `google-genai` SDK (`from google import genai`). Covers `client.models.generate_content` + `generate_content_stream`, sync + async (via `client.aio.models`). -**Reasoning tokens** populate automatically on Gemini 2.5 — the model reasons internally by default and surfaces `thoughts_token_count`. Note the semantic difference vs OpenAI: -- **OpenAI:** `reasoning_tokens` is a *subset* of `completion_tokens` (already counted in output) -- **Gemini:** `thoughts_token_count` is *additive* to `candidates_token_count` (total Google bill = output + reasoning) +**Reasoning tokens** populate automatically on Gemini 2.5 — the model reasons internally by default and surfaces `thoughts_token_count` (see the note on reasoning semantics below). + +## Cloudflare AI Gateway + +Point any of the clients above at your gateway instead of the provider directly — `wrap()` detects it and bills correctly, with two behaviors on top of the plain provider case: + +```python +from anthropic import Anthropic +from lago_agent_sdk import LagoSDK + +sdk = LagoSDK(api_key="...", default_subscription_id="sub_acme") +client = sdk.wrap(Anthropic( + api_key="...", + base_url=f"https://gateway.ai.cloudflare.com/v1/{account_id}/{gateway_id}/anthropic", + default_headers={"cf-aig-authorization": f"Bearer {gateway_auth}"}, +)) +client.messages.create(model="claude-sonnet-4-6", max_tokens=200, messages=[{"role": "user", "content": "Hello"}]) +sdk.flush() +``` + +- **Gateway cache hits aren't billed.** If the gateway serves a response from its own cache (`cf-aig-cache-status: HIT`), the provider was never called, so the SDK skips emitting for that response. +- **Workers AI gets priced automatically.** Wrap an OpenAI-shaped client against the gateway's `/compat` endpoint (`model="workers-ai/@cf/..."`) with `pricing_mode="price"`, and the SDK fetches Cloudflare's own published Workers AI rates in the background — no separate price table to maintain. + +For usage that already happened, backfill straight from the gateway's own Logs API instead of replaying calls — `lago_agent_sdk.gateway.adapters` extracts a log entry into `CanonicalUsage` and bills Cloudflare's own metered `cost` for it, so there's no separate price lookup and re-running over the same window never double-bills: + +```python +from lago_agent_sdk.gateway.adapters import extract_cloudflare_log, resolve_subscription + +for entry in fetch_gateway_logs(): # GET .../ai-gateway/gateways/{id}/logs + usage = extract_cloudflare_log(entry) + sub = resolve_subscription(entry) or "sub_default" # from the call's cf-aig-metadata, if set + sdk.emit(usage, subscription=sub, mode="price", usd_cost=entry.get("cost") or 0, event_id=f"cf_{entry['id']}") +sdk.flush() +``` + +See [`examples/cloudflare_gateway_demo.ipynb`](examples/cloudflare_gateway_demo.ipynb) for a runnable end-to-end version of both. ## Multi-tenant — pick a subscription per call @@ -161,7 +194,6 @@ Backed by `contextvars` for safe propagation across `asyncio` tasks. | Mistral | native SDK (`chat.complete` + `chat.stream`) | ✓ | | OpenAI | native SDK (`chat.completions.create` + `responses.create`, sync + async + stream) | ✓ | | Google Gemini | native SDK (`google-genai`: `models.generate_content` + `generate_content_stream`, sync + async) | ✓ | -| LiteLLM | callback bridge | Phase 4 | ## Token dimensions captured @@ -178,21 +210,15 @@ Backed by `contextvars` for safe propagation across `asyncio` tasks. | tool_calls | `llm_tool_calls` | ✓ | ✓ | ✓ | ✓ | ✓ | | audio_input | `llm_audio_input_tokens` | ✗ | ✗ | ✗ | ✓ (GPT-4o-audio) | ✓ (multimodal AUDIO) | | audio_output | `llm_audio_output_tokens` | ✗ | ✗ | ✗ | ✓ (GPT-4o-audio) | ✓ (multimodal AUDIO) | -| image_input | `llm_image_input_tokens` | ✗ | ✗ | ✗ | ✗ (Phase 3) | ✓ (multimodal IMAGE) | +| image_input | `llm_image_input_tokens` | ✗ | ✗ | ✗ | ✗ | ✓ (multimodal IMAGE) | -**Semantic note on `reasoning`:** -- **OpenAI's `reasoning_tokens` is a SUBSET of `output`** — already counted in `completion_tokens`. -- **Gemini's `thoughts_token_count` is ADDITIVE to `output`** — `candidates + thoughts = total billable output`. +**Reasoning:** OpenAI's `reasoning_tokens` is a *subset* of `output` (already counted in `completion_tokens`). Gemini's `thoughts_token_count` is *additive* to `output` (`candidates + thoughts = total billable output`). -**Semantic note on input breakdowns (avoid double-counting):** -For both OpenAI and Gemini, `cache_read`, `audio_input`, and `image_input` are **subsets of `input`**, not additive to it — they are a breakdown of tokens already counted in `llm_input_tokens`. For example, OpenAI reports `cached_tokens` under `prompt_tokens_details` *within* `prompt_tokens`, and Gemini's docs state `prompt_token_count` "includes the number of tokens in the cached content". A billable metric that sums `llm_input_tokens + llm_cached_input_tokens` (or `+ llm_audio_input_tokens`, `+ llm_image_input_tokens`) will **double-count**. Bill on `llm_input_tokens` as the total; use the breakdown fields only for cost attribution or discounted-rate tiers (e.g. cached input billed at a lower rate), subtracting them from `input` rather than adding. - -OpenAI's Predicted Outputs tokens (`accepted_prediction_tokens`, `rejected_prediction_tokens`) are not surfaced — see the OpenAI adapter docstring for details on this intentional gap. +**Cache/audio/image on OpenAI and Gemini are subsets of `input`, not additive.** Both providers count cached/audio/image tokens *within* their input total, so summing `llm_input_tokens + llm_cached_input_tokens` (or `+ audio/image`) double-counts. Bill on `llm_input_tokens` alone; use the breakdown fields only for cost attribution (e.g. a discounted cache rate). ## Pricing mode — send dollar cost instead of tokens -By default the SDK emits **token counts** (`pricing_mode="tokens"`). You can instead have it -compute and emit the **dollar cost** of each call: `Σ(unit_price_per_token × tokens) × markup`. +By default the SDK emits **token counts** (`pricing_mode="tokens"`). Set `pricing_mode="price"` to instead emit the **dollar cost** of each call: `Σ(unit_price_per_token × tokens) × markup`. ```python from lago_agent_sdk import LagoSDK, LagoConfig @@ -207,46 +233,18 @@ client = sdk.wrap(anthropic_client) # ... use the client normally ... ``` -In **price mode** the SDK emits **one event per call** with code `llm_cost`. The event carries a -top-level `precise_total_amount_cents` (the total cost in cents, after markup) for Lago's -**dynamic charge model**, plus a breakdown in `properties`: `unit` (total tokens), `value` (USD -total), `base_cost` (pre-markup), `markup`, `price_source`, and per-field `*_tokens` / -`*_unit_price` / `*_cost`. Set up in Lago a `sum`-aggregation billable metric `llm_cost` on -`field_name: "unit"` and a **dynamic** charge on it — Lago sums each event's -`precise_total_amount_cents` into a single fee (`unit` is the displayed usage quantity). See -`testing/lago_setup_pricing_plan.py` for a script that creates this. +Price mode emits one `llm_cost` event per priced field (input, output, cache, ...), each carrying `precise_total_amount_cents` for Lago's **dynamic charge model** plus a `token_type` property so a single billable metric can be grouped by both `model` and `token_type`. Prices come from public sources (OpenRouter for native providers, the AWS Bedrock price list for Bedrock), fetched and cached in the background — your LLM call is never blocked on pricing. If a price isn't available yet, the SDK falls back to token-count events and reports via `on_error` rather than under-billing. -Per-call override via `extra_lago` (mode and markup, in addition to subscription/dimensions): +Per-call override via `extra_lago`: ```python client.messages.create(model="claude-...", messages=[...], extra_lago={"mode": "price", "markup": 1.5}) ``` -**Live, public pricing sources (no API keys):** -- **OpenRouter** (`/api/v1/models`) for native `anthropic` / `openai` / `mistral` / `gemini` - clients — USD per token. -- **AWS Bedrock Price List Bulk API** (public) for Bedrock — parsed per region. - -Prices are fetched and cached in the background (TTL `pricing_ttl_seconds`, default 1h); the -refresh runs on the SDK's background thread, so **your LLM call is never blocked on pricing**. - -**Fallback (never under-bill):** if a price is unavailable (table not warm on the first call, -or the model isn't found in the source), the SDK **falls back to emitting token-count events** -and calls `on_error` so it's visible — it never silently drops the usage. - -**Bedrock note:** AWS's public bulk data lists many models (Titan, Llama, Mistral, Cohere, and -older Claude) but, at time of writing, **not the current Claude 3.5/3.7/4 models**. Bedrock -calls for models absent from AWS's data fall back to token events. Native Anthropic clients are -priced via OpenRouter and unaffected. - ## Error policy -The SDK never breaks your LLM call. If anything in instrumentation fails (adapter bug, Lago down, network error), the SDK swallows it, logs a warning, and your call returns normally. - -## Subscription resolution returns nothing → drop with `ERROR` log - -Configurable via `LagoConfig.on_error` callback to integrate with Sentry, Datadog, etc.: +The SDK never breaks your LLM call. If anything in instrumentation fails (adapter bug, Lago down, network error, no subscription resolved), it's swallowed, logged, and your call returns normally. Wire your own observability via `LagoConfig.on_error`: ```python from lago_agent_sdk import LagoConfig, LagoSDK @@ -274,17 +272,6 @@ pip install -e '.[dev]' pytest ``` -Run live integration tests (requires real credentials): - -```bash -AWS_BEARER_TOKEN_BEDROCK="..." \ -MISTRAL_API_KEY="..." \ -LAGO_API_URL="https://api.getlago.com/api/v1/" \ -LAGO_API_KEY="..." \ -LAGO_EXTERNAL_SUBSCRIPTION_ID="sub_..." \ -pytest tests/integration -``` - ## Security Found a vulnerability? See [SECURITY.md](SECURITY.md). diff --git a/examples/.env.example b/examples/.env.example new file mode 100644 index 0000000..247a32b --- /dev/null +++ b/examples/.env.example @@ -0,0 +1,23 @@ +# Copy this file to examples/.env and fill in real values. +# examples/.env is gitignored — never commit real credentials. + +CF_ACCOUNT_ID= +CF_GATEWAY_ID= +CF_LOGS_TOKEN= +# Also doubles as the Cloudflare API token for live workers-ai pricing (via +# Cloudflare's own model catalog) — no separate credential needed for that. +CF_GATEWAY_AUTH= + +LAGO_API_KEY= +LAGO_API_URL=https://api.getlago.com/api/v1 +LAGO_SUBSCRIPTION_ID=cloudflare_gateway_demo_sub +# Set to false ONLY for a local dev Lago instance behind a self-signed cert +# (e.g. LAGO_API_URL=https://api.lago.dev/api/v1). Never for a real Lago URL. +LAGO_VERIFY_SSL=true + +# Only needed for the provider you pick in Part 2 — workers-ai needs neither. +ANTHROPIC_API_KEY= +# wrap()-ing the Mistral client auto-detects this key for pricing's alias +# resolution too (see LagoConfig.mistral_api_key) — no separate credential +# needed for that. +MISTRAL_API_KEY= diff --git a/examples/cloudflare_gateway_demo.ipynb b/examples/cloudflare_gateway_demo.ipynb new file mode 100644 index 0000000..1b5bcd3 --- /dev/null +++ b/examples/cloudflare_gateway_demo.ipynb @@ -0,0 +1,581 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "3ae3c48f", + "metadata": {}, + "source": [ + "# Cloudflare AI Gateway ↔ Lago\n", + "\n", + "Two things, both landing on the same `llm_cost` metric — breakable by model in\n", + "Lago because every event already carries `model` in its properties, and the\n", + "plan's `llm_cost` charge has `grouped_by: [\"model\"]` set:\n", + "\n", + "1. **Backfill** — read every log entry from your Cloudflare AI Gateway and\n", + " bill each one's *real*, already-metered cost straight from Cloudflare's\n", + " own `cost` field. Idempotent: safe to re-run over the same window.\n", + "2. **Live call** — wrap a real provider SDK client, make one call through the\n", + " gateway, and let `sdk.wrap()` bill it automatically.\n", + "\n", + "### Setup\n", + "\n", + "Set these as environment variables before starting the kernel (never hardcode\n", + "real credentials into the notebook itself) — **or** put them in a `.env` file\n", + "next to this notebook (`examples/.env`); the next cell loads one automatically\n", + "if present. A `.env` file survives kernel restarts, unlike shell exports made\n", + "after Jupyter is already running — if you restart the kernel and still see a\n", + "missing-variable error, that's usually why.\n", + "\n", + "| Variable | What it is |\n", + "|---|---|\n", + "| `CF_ACCOUNT_ID` | Cloudflare account id |\n", + "| `CF_GATEWAY_ID` | the AI Gateway's id |\n", + "| `CF_LOGS_TOKEN` | Cloudflare API token scoped for AI Gateway logs read |\n", + "| `CF_GATEWAY_AUTH` | the gateway's own auth token (`cf-aig-authorization`) |\n", + "| `LAGO_API_KEY` | your Lago API key |\n", + "| `LAGO_API_URL` | defaults to `https://api.getlago.com/api/v1` |\n", + "| `LAGO_SUBSCRIPTION_ID` | defaults to `cloudflare_gateway_demo_sub` |\n", + "| `LAGO_VERIFY_SSL` | defaults to `true`. Set to `false` **only** for a local dev Lago instance behind a self-signed certificate — never for a real Lago URL. Lets you hit a local instance directly instead of needing a public tunnel just to get a browser-trusted cert. |\n", + "| `ANTHROPIC_API_KEY` / `MISTRAL_API_KEY` | only needed for the provider you pick in Part 2 — `workers-ai` needs none at all, it's billed directly by Cloudflare. `MISTRAL_API_KEY` doubles as pricing's alias-resolution credential (see below) — without it, a `-latest` Mistral call still bills, just as an unpriced token-event fallback. |\n", + "\n", + "Mistral has no per-token price table of its own, so `mistral-small-latest`\n", + "(what you actually call) can't be priced directly. Setting `MISTRAL_API_KEY`\n", + "lets price mode resolve it via Mistral's own `/v1/models` — which reports\n", + "`mistral-small-latest`'s real dated id (`mistral-small-2603`) — and look\n", + "*that* up against OpenRouter, which does list it with real pricing." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bfc3cff9", + "metadata": {}, + "outputs": [], + "source": "%load_ext autoreload\n%autoreload 2\n# Reloads lago_agent_sdk automatically whenever its source changes, so fixes\n# take effect on the next cell run — no kernel restart needed. Only helps for\n# edits made AFTER this cell has run once in the current kernel; the first\n# time you pull in a change to this cell itself (or add a brand new\n# top-level name the rest of the notebook needs), you still need one restart.\n\nimport os\nimport sys\n\nimport requests\n\nsys.path.insert(0, \"../src\") # run this notebook from examples/, or adjust to your install\n\n\ndef _load_dotenv(path: str) -> None:\n \"\"\"No extra dependency — just KEY=VALUE lines, same as python-dotenv's basics.\"\"\"\n if not os.path.exists(path):\n return\n for line in open(path):\n line = line.strip()\n if line and not line.startswith(\"#\") and \"=\" in line:\n key, _, value = line.partition(\"=\")\n os.environ.setdefault(key.strip(), value.strip().strip('\"').strip(\"'\"))\n\n\n_load_dotenv(os.path.join(os.getcwd(), \".env\"))\n\nfrom lago_agent_sdk import LagoSDK # noqa: E402\nfrom lago_agent_sdk.config import LagoConfig # noqa: E402\nfrom lago_agent_sdk.gateway.adapters import extract_cloudflare_log, resolve_subscription # noqa: E402\n\n_REQUIRED = [\"CF_ACCOUNT_ID\", \"CF_GATEWAY_ID\", \"CF_LOGS_TOKEN\", \"LAGO_API_KEY\"]\n_missing = [name for name in _REQUIRED if not os.environ.get(name)]\nif _missing:\n raise SystemExit(\n f\"Missing required environment variable(s): {', '.join(_missing)}.\\n\"\n \"Set them before starting the kernel, or put them in examples/.env — see the Setup cell above.\"\n )\n\nCF_ACCOUNT_ID = os.environ[\"CF_ACCOUNT_ID\"]\nCF_GATEWAY_ID = os.environ[\"CF_GATEWAY_ID\"]\nCF_LOGS_TOKEN = os.environ[\"CF_LOGS_TOKEN\"]\nCF_GATEWAY_AUTH = os.environ.get(\"CF_GATEWAY_AUTH\", \"\")\nLAGO_API_KEY = os.environ[\"LAGO_API_KEY\"]\nLAGO_API_URL = os.environ.get(\"LAGO_API_URL\", \"https://api.getlago.com/api/v1\")\nLAGO_SUBSCRIPTION_ID = os.environ.get(\"LAGO_SUBSCRIPTION_ID\", \"cloudflare_gateway_demo_sub\")\nLAGO_VERIFY_SSL = os.environ.get(\"LAGO_VERIFY_SSL\", \"true\").lower() != \"false\"\nMISTRAL_API_KEY = os.environ.get(\"MISTRAL_API_KEY\", \"\")\n\nsdk = LagoSDK(\n api_key=LAGO_API_KEY,\n api_url=LAGO_API_URL,\n default_subscription_id=LAGO_SUBSCRIPTION_ID,\n config=LagoConfig(\n api_key=LAGO_API_KEY, api_url=LAGO_API_URL, pricing_mode=\"price\", verify_ssl=LAGO_VERIFY_SSL,\n # Prices \"workers-ai\" calls from Cloudflare's own model catalog — the\n # real rate the gateway bills at, not a third party's guess. This is\n # just a credential declaration, not an eager fetch: Cloudflare's\n # catalog is only ever actually fetched lazily, on this session's\n # first real workers-ai call (see warm_pricing() below). Optional:\n # without it, Workers AI calls just fall back to token events.\n cloudflare_account_id=CF_ACCOUNT_ID, cloudflare_api_token=CF_GATEWAY_AUTH,\n # Mistral has no price table of its own — this resolves \"-latest\"\n # aliases (e.g. \"mistral-small-latest\") via Mistral's own /v1/models\n # to the dated id OpenRouter actually lists. Same as Cloudflare\n # above: declaring the key here doesn't fetch anything by itself —\n # it's fetched lazily on this session's first real Mistral call.\n # Optional: without it, an aliased Mistral call falls back to token\n # events instead.\n mistral_api_key=MISTRAL_API_KEY,\n ),\n)\n# Blocks until OpenRouter's table is fetched — closes the cold-start race for\n# the very first call, for whichever native provider (anthropic/openai/\n# mistral/gemini) that first call happens to use. Deliberately does NOT also\n# force-fetch Cloudflare/Mistral above: both are credential-gated and\n# provider-specific, and this demo (like most price-mode setups) may only\n# ever call one of the three PROVIDER options below in a given run — eagerly\n# hitting all their APIs regardless of which one gets used would be wasted\n# work. So: whichever provider your first live call below actually uses,\n# THAT one's table gets fetched lazily right then (and is cached for every\n# call after) — only that very first call for a given provider can race a\n# cold cache, and only if you picked workers-ai or mistral (never for\n# anthropic/openai/gemini, which OpenRouter already warmed here).\nsdk.warm_pricing()\nprint(\"SDK ready — billing to\", LAGO_SUBSCRIPTION_ID)" + }, + { + "cell_type": "markdown", + "id": "91c45afd", + "metadata": {}, + "source": [ + "## Part 1 — Backfill historic usage from Cloudflare\n", + "\n", + "Fetch every log entry the gateway has recorded, and bill each one's real\n", + "Cloudflare-reported cost. No price lookup on our side — Cloudflare already\n", + "metered it.\n", + "\n", + "`UNIFIED_BILLING = True` (the default here) bills every entry to\n", + "`LAGO_SUBSCRIPTION_ID`, ignoring any `cf-aig-metadata` attribution a call\n", + "might carry — the right choice when this gateway's traffic should all land on\n", + "one subscription. Set it `False` instead to respect real per-call\n", + "attribution and route each entry to whichever subscription it names,\n", + "falling back to `LAGO_SUBSCRIPTION_ID` only for entries with none — the right\n", + "choice when one gateway serves multiple customers/subscriptions." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "e1786a29", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "fetched 123 real log entries from Cloudflare\n" + ] + } + ], + "source": [ + "def fetch_all_logs():\n", + " entries, page = [], 1\n", + " while True:\n", + " body = requests.get(\n", + " f\"https://api.cloudflare.com/client/v4/accounts/{CF_ACCOUNT_ID}\"\n", + " f\"/ai-gateway/gateways/{CF_GATEWAY_ID}/logs\",\n", + " headers={\"Authorization\": f\"Bearer {CF_LOGS_TOKEN}\"},\n", + " params={\"per_page\": 50, \"page\": page},\n", + " timeout=30,\n", + " ).json()\n", + " entries.extend(body[\"result\"])\n", + " if len(body[\"result\"]) < 50 or len(entries) >= body[\"result_info\"][\"total_count\"]:\n", + " return entries\n", + " page += 1\n", + "\n", + "\n", + "logs = fetch_all_logs()\n", + "print(f\"fetched {len(logs)} real log entries from Cloudflare\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "ba489082", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "backfilled 123 entries\n" + ] + } + ], + "source": [ + "UNIFIED_BILLING = True\n", + "\n", + "for entry in logs:\n", + " usage = extract_cloudflare_log(entry)\n", + " sub = LAGO_SUBSCRIPTION_ID if UNIFIED_BILLING else (resolve_subscription(entry) or LAGO_SUBSCRIPTION_ID)\n", + " # transaction_id is unique across the whole ORG, not just this subscription —\n", + " # always scope it by subscription. Without this, switching LAGO_SUBSCRIPTION_ID\n", + " # to a second, different subscription later (unified or not) would have every\n", + " # entry collide with the ids already used for the first one and silently never\n", + " # land anywhere at all — this bit a real run: the first unified subscription's\n", + " # ids blocked every entry from ever reaching a second one.\n", + " event_id = f\"unified_{sub}_{entry['id']}\" if UNIFIED_BILLING else f\"backfill_{sub}_{entry['id']}\"\n", + " sdk.emit(\n", + " usage,\n", + " subscription=sub,\n", + " mode=\"price\",\n", + " usd_cost=entry.get(\"cost\") or 0, # Cloudflare's own metered price\n", + " event_id=event_id,\n", + " )\n", + "\n", + "assert sdk.flush(timeout=30.0), \"queue did not flush in time\"\n", + "print(f\"backfilled {len(logs)} entries\")" + ] + }, + { + "cell_type": "markdown", + "id": "753795a6", + "metadata": {}, + "source": [ + "## Part 2 — Live call through the gateway\n", + "\n", + "Pick a provider. `workers-ai` needs no external key at all — Cloudflare bills\n", + "it directly. `anthropic`/`mistral` need their own key set as an env var." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "28027ea8", + "metadata": {}, + "outputs": [], + "source": [ + "PROMPT = \"Tell me about getLago, the billing company - give as many details as you can find\"" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "a2f7c058", + "metadata": {}, + "outputs": [], + "source": [ + "from anthropic import Anthropic\n", + "client = sdk.wrap(Anthropic(\n", + " api_key=os.environ[\"ANTHROPIC_API_KEY\"],\n", + " base_url=f\"https://gateway.ai.cloudflare.com/v1/{CF_ACCOUNT_ID}/{CF_GATEWAY_ID}/anthropic\",\n", + " default_headers={\"cf-aig-authorization\": f\"Bearer {CF_GATEWAY_AUTH}\"},\n", + " ))\n", + "resp = client.messages.create(model=\"claude-sonnet-4-5\", max_tokens=20000,\n", + " messages=[{\"role\": \"user\", \"content\": PROMPT}])\n", + "text = resp.content[0].text" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "b360432e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "# GetLago - Open-Source Billing Platform\n", + "\n", + "## Overview\n", + "GetLago is an open-source billing and metering platform designed for product-led SaaS companies. It provides an alternative to proprietary billing solutions like Stripe Billing, Chargebee, and Recurly.\n", + "\n", + "## Key Information\n", + "\n", + "### Company Background\n", + "- **Founded**: 2021\n", + "- **Founders**: Anh-Tho Chuong and Raffi Sarkissian\n", + "- **Headquarters**: Paris, France (with remote-first culture)\n", + "- **Funding**: Raised significant seed funding from investors including Y Combinator (YC Winter 2022 batch), SignalFire, and others\n", + "\n", + "### Core Product Features\n", + "\n", + "**1. Usage-Based Billing**\n", + "- Real-time event ingestion and metering\n", + "- Supports complex pricing models (pay-as-you-go, tiered, graduated, package pricing)\n", + "- Aggregation capabilities for billing metrics\n", + "\n", + "**2. Subscription Management**\n", + "- Handles recurring subscriptions\n", + "- Supports hybrid models (combining subscriptions + usage)\n", + "- Plan versioning and management\n", + "\n", + "**3. Pricing Flexibility**\n", + "- Multiple charge models: standard, graduated, package, percentage, volume\n", + "- Support for in-arrears and in-advance billing\n", + "- Minimum commitments and spending caps\n", + "- Proration handling\n", + "\n", + "**4. Coupons & Credits**\n", + "- Discount management\n", + "- Prepaid credits/wallet system\n", + "- Credit notes\n", + "\n", + "**5. Invoicing**\n", + "- Automated invoice generation\n", + "- PDF invoicing\n", + "- Tax management\n", + "- Multiple currencies\n", + "\n", + "**6. Integrations**\n", + "- Payment processors: Stripe, GoCardless, Adyen\n", + "- Data warehouses and analytics tools\n", + "- Accounting software\n", + "- CRM systems\n", + "- REST API for custom integrations\n", + "\n", + "### Technical Architecture\n", + "\n", + "**Open Source**\n", + "- Available on GitHub under AGPL-3.0 license\n", + "- Community edition freely available\n", + "- Self-hostable option\n", + "\n", + "**Technology Stack**\n", + "- Backend: Ruby on Rails\n", + "- Frontend: React\n", + "- PostgreSQL database\n", + "- Event-driven architecture for metering\n", + "\n", + "**Deployment Options**\n", + "- Self-hosted (open-source version)\n", + "- Cloud-hosted (managed service)\n", + "\n", + "### Differentiators\n", + "\n", + "1. **Open Source**: Full transparency and customizability\n", + "2. **Developer-First**: API-first design, extensive documentation\n", + "3. **No Vendor Lock-in**: Can be self-hosted\n", + "4. **Real-time Metering**: Built for usage-based pricing from the ground up\n", + "5. **Pricing**: More cost-effective than traditional billing platforms, especially at scale\n", + "\n", + "### Use Cases\n", + "- B2B SaaS companies with usage-based pricing\n", + "- API-first companies (like Algolia, Segment model)\n", + "- Companies needing complex billing logic\n", + "- Businesses wanting to avoid vendor lock-in\n", + "- Startups to enterprises requiring scalable billing\n", + "\n", + "### Target Market\n", + "- Product-led growth companies\n", + "- Engineering teams that want control over billing infrastructure\n", + "- Companies with complex or hybrid pricing models\n", + "- Businesses scaling usage-based revenue\n", + "\n", + "### Competitive Position\n", + "Competes with:\n", + "- **Proprietary solutions**: Stripe Billing, Chargebee, Recurly, Zuora\n", + "- **Open-source alternatives**: Kill Bill (though Kill Bill is more Java-based and enterprise-focused)\n", + "\n", + "### Community & Growth\n", + "- Active GitHub community\n", + "- Regular updates and feature releases\n", + "- Growing adoption among YC companies and tech startups\n", + "- Developer-focused documentation and resources\n", + "\n", + "## Recent Developments\n", + "The company has been actively developing features like:\n", + "- Enhanced analytics and reporting\n", + "- More payment gateway integrations\n", + "- Improved tax handling\n", + "- Advanced dunning management\n", + "- Better webhook systems\n", + "\n", + "GetLago represents the trend toward open-source infrastructure for critical business functions, giving companies more control and flexibility over their billing operations while reducing costs compared to traditional SaaS billing platforms.\n" + ] + } + ], + "source": [ + "print(text)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "4ef82052", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "lago pricing failed: no price for provider='mistral' model='mistral-small-latest' api='native'\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "**GetLago** is an open-source **usage-based billing** and **metering** platform designed to help SaaS companies implement **pay-as-you-go** pricing models efficiently. It provides developers with the tools to track customer usage, apply custom pricing rules, and generate invoices—all while integrating seamlessly with existing billing systems.\n", + "\n", + "Here’s a detailed breakdown of **GetLago**, including its features, architecture, pricing, integrations, and more:\n", + "\n", + "---\n", + "\n", + "## **1. Overview & Key Features**\n", + "GetLago is built to solve the challenges of **usage-based billing**, which is becoming increasingly popular among SaaS companies (e.g., AWS, Stripe, Datadog). Key features include:\n", + "\n", + "### **🔹 Core Features**\n", + "✅ **Metering & Usage Tracking**\n", + "- Tracks API calls, feature usage, storage, compute time, etc.\n", + "- Supports **real-time** and **batch** event ingestion.\n", + "- Customizable **metric definitions** (e.g., \"API requests,\" \"database queries\").\n", + "\n", + "✅ **Pricing & Billing Engine**\n", + "- Supports **tiered pricing**, **volume discounts**, **overage charges**, and **custom formulas**.\n", + "- Handles **prepaid credits**, **subscription renewals**, and **one-time charges**.\n", + "- **Granular pricing rules** (e.g., \"$0.01 per API call after 10,000\").\n", + "\n", + "✅ **Subscription Management**\n", + "- Manages **customer subscriptions**, **plan changes**, and **downgrades**.\n", + "- Supports **free trials**, **usage-based add-ons**, and **commitment discounts**.\n", + "\n", + "✅ **Invoicing & Payments**\n", + "- Generates **invoices** based on usage.\n", + "- Integrates with **Stripe**, **Paddle**, **Lemon Squeezy**, and other payment processors.\n", + "- Supports **one-off invoices** and **recurring billing**.\n", + "\n", + "✅ **Multi-Tenancy & Security**\n", + "- **Role-based access control (RBAC)** for teams.\n", + "- **Customer isolation** (data is segmented per customer).\n", + "- **Audit logs** for compliance.\n", + "\n", + "✅ **Open-Source & Self-Hosted**\n", + "- **MIT-licensed** (free to use, modify, and self-host).\n", + "- **Cloud-hosted** option available (GetLago Cloud).\n", + "- **Docker & Kubernetes** support for easy deployment.\n", + "\n", + "✅ **API-First & Developer-Friendly**\n", + "- **REST API** for integration with apps.\n", + "- **Webhooks** for real-time event notifications.\n", + "- **SDKs** (JavaScript, Python, Ruby, etc.).\n", + "\n", + "---\n", + "\n", + "## **2. Architecture & Tech Stack**\n", + "GetLago is built with modern technologies:\n", + "\n", + "| **Component** | **Technology** |\n", + "|---------------------|---------------|\n", + "| **Backend** | Ruby on Rails (API) |\n", + "| **Database** | PostgreSQL (primary), Redis (caching) |\n", + "| **Event Processing** | Kafka (for high-throughput usage events) |\n", + "| **Frontend** | React (for the admin dashboard) |\n", + "| **Deployment** | Docker, Kubernetes, Helm |\n", + "| **Monitoring** | Prometheus, Grafana |\n", + "| **CI/CD** | GitHub Actions |\n", + "\n", + "### **🔹 How It Works**\n", + "1. **Events Ingestion** → Customers send usage events (e.g., API calls) via the API.\n", + "2. **Metering** → GetLago processes events and updates usage counters.\n", + "3. **Pricing Calculation** → Applies pricing rules to compute charges.\n", + "4. **Invoicing** → Generates an invoice (or sends to a payment processor).\n", + "5. **Payment Processing** → Charges the customer (via Stripe, etc.).\n", + "\n", + "---\n", + "\n", + "## **3. Pricing (GetLago Cloud)**\n", + "GetLago offers **two pricing models**:\n", + "\n", + "### **💰 Self-Hosted (Free)**\n", + "- **Open-source (MIT license)** – Free to use, modify, and self-host.\n", + "- **No usage limits** (but you pay for infrastructure costs).\n", + "\n", + "### **☁️ GetLago Cloud (Paid)**\n", + "- **Pay-as-you-go** pricing based on **usage volume**.\n", + "- **No upfront costs**, but charges apply per:\n", + " - **Events processed** (e.g., API calls, feature usage).\n", + " - **Customers managed**.\n", + " - **Invoices generated**.\n", + "- **Free tier** available (limited usage).\n", + "\n", + "*(Exact pricing details are not publicly listed—contact GetLago for a quote.)*\n", + "\n", + "---\n", + "\n", + "## **4. Integrations**\n", + "GetLago integrates with popular tools:\n", + "\n", + "| **Category** | **Integrations** |\n", + "|--------------------|------------------|\n", + "| **Payment Processors** | Stripe, Paddle, Lemon Squeezy, Adyen |\n", + "| **CRM & Analytics** | HubSpot, Segment, Mixpanel |\n", + "| **Dev Tools** | GitHub, Slack (for alerts) |\n", + "| **Databases** | PostgreSQL, MySQL (for custom metrics) |\n", + "| **Auth** | Auth0, Firebase Auth, Supabase |\n", + "\n", + "### **🔹 Example Use Cases**\n", + "- **API-based SaaS** (e.g., AI models, cloud services).\n", + "- **Feature-based billing** (e.g., \"Pay per API call\").\n", + "- **Multi-tenant apps** (e.g., per-user pricing).\n", + "- **Prepaid credits** (e.g., \"$100 credit = 10,000 API calls\").\n", + "\n", + "---\n", + "\n", + "## **5. Competitors Comparison**\n", + "| **Tool** | **Type** | **Open-Source** | **Usage-Based Billing** | **Self-Hosted** | **Pricing** |\n", + "|----------|---------|----------------|------------------------|----------------|------------|\n", + "| **GetLago** | Billing & Metering | ✅ Yes | ✅ Yes | ✅ Yes | Free (self-hosted) / Paid (cloud) |\n", + "| **Stripe Billing** | Billing Platform | ❌ No | ✅ Yes | ❌ No | Pay-per-use |\n", + "| **Chargebee** | Subscription Billing | ❌ No | ✅ Yes | ❌ No | Subscription-based |\n", + "| **Recurly** | Subscription Billing | ❌ No | ✅ Yes | ❌ No | Subscription-based |\n", + "| **Copper** | Usage-Based Billing | ✅ Yes | ✅ Yes | ✅ Yes | Free (self-hosted) |\n", + "| **Orb** | Usage-Based Billing | ❌ No | ✅ Yes | ❌ No | Pay-per-use |\n", + "\n", + "**Key Differentiators of GetLago:**\n", + "✔ **Open-source** (unlike Stripe/Chargebee).\n", + "✔ **Self-hostable** (unlike most competitors).\n", + "✔ **Developer-first** (API & SDKs).\n", + "✔ **Flexible pricing rules** (tiered, volume-based, etc.).\n", + "\n", + "---\n", + "\n", + "## **6. Getting Started with GetLago**\n", + "### **🚀 Self-Hosted Setup**\n", + "1. **Deploy with Docker**:\n", + " ```bash\n", + " docker-compose up -d\n", + " ```\n", + "2. **Configure via Admin Dashboard** (or API).\n", + "3. **Send Usage Events** (via API or SDK).\n", + "4. **Generate Invoices** (manually or automated).\n", + "\n", + "### **📖 Documentation & Resources**\n", + "- **[Official Website](https://getlago.com/)**\n", + "- **[GitHub Repository](https://github.com/getlago/lago)**\n", + "- **[Documentation](https://doc.lago.dev/)**\n", + "- **[Discord Community](https://discord.gg/9ae6K3dX7T)**\n", + "\n", + "### **🎓 Tutorials & Examples**\n", + "- [Building a Usage-Based SaaS with GetLago](https://getlago.com/blog/usage-based-billing-guide)\n", + "- [Integrating with Stripe](https://doc.lago.dev/docs/payment-processors/stripe)\n", + "- [Metering API Calls](https://doc.lago.dev/docs/metering)\n", + "\n", + "---\n", + "\n", + "## **7. Pros & Cons**\n", + "### **✅ Pros**\n", + "✔ **Open-source & self-hostable** (no vendor lock-in).\n", + "✔ **Flexible pricing models** (tiered, volume, overage).\n", + "✔ **Developer-friendly** (API-first, SDKs).\n", + "✔ **Real-time usage tracking**.\n", + "✔ **Multi-tenant support**.\n", + "\n", + "### **❌ Cons**\n", + "❌ **Self-hosting requires DevOps effort** (PostgreSQL, Redis, Kafka).\n", + "❌ **Cloud pricing is not transparent** (must contact sales).\n", + "❌ **Young project** (fewer integrations than Stripe/Chargebee).\n", + "❌ **Limited enterprise features** (e.g., dunning management).\n", + "\n", + "---\n", + "\n", + "## **8. Who Should Use GetLago?**\n", + "✅ **Startups & SMBs** needing **usage-based billing** without Stripe’s complexity.\n", + "✅ **Developers** who want **full control** over billing logic.\n", + "✅ **Open-source advocates** who prefer **self-hosted solutions**.\n", + "✅ **SaaS companies** with **custom pricing models** (tiered, volume-based).\n", + "\n", + "❌ **Not ideal for:**\n", + "- Companies needing **advanced dunning** (failed payment retries).\n", + "- Enterprises requiring **SOC 2 / ISO 27001 compliance** (self-hosted may need extra setup).\n", + "- Businesses that **prefer managed SaaS** (like Stripe Billing).\n", + "\n", + "---\n", + "\n", + "## **9. Recent Updates & Roadmap**\n", + "- **2024:** Major updates in **pricing engine**, **multi-currency support**, and **improved API performance**.\n", + "- **Upcoming:** **Webhook improvements**, **more payment processor integrations**, and **enhanced analytics**.\n", + "- **Community-driven** – GetLago is actively maintained with **regular GitHub contributions**.\n", + "\n", + "---\n", + "\n", + "## **10. Alternatives to Consider**\n", + "If GetLago doesn’t fit your needs, check out:\n", + "- **[Copper](https://github.com/coopnorge/copper)** (Open-source, usage-based billing).\n", + "- **[Orb](https://orb.com/)** (SaaS-focused, usage-based billing).\n", + "- **[Stripe Billing](https://stripe.com/billing)** (Managed, but expensive).\n", + "- **[Chargebee](https://www.chargebee.com/)** (Subscription-focused).\n", + "\n", + "---\n", + "\n", + "## **Final Verdict**\n", + "GetLago is a **powerful, open-source alternative** to Stripe Billing and Chargebee, ideal for **developers and startups** who need **flexible, usage-based billing** without vendor lock-in. While it requires some **DevOps effort** for self-hosting, its **API-first approach** and **custom pricing rules** make it a strong choice for modern SaaS businesses.\n", + "\n", + "🔗 **Try it out:**\n", + "- [GetLago Cloud (Free Tier)](https://app.getlago.com/)\n", + "- [GitHub Repository](https://github.com/getlago/lago)\n", + "\n", + "Would you like a deeper dive into any specific aspect (e.g., API integration, pricing engine, or deployment)?\n" + ] + } + ], + "source": [ + "from mistralai.client import Mistral\n", + "client = sdk.wrap(Mistral(\n", + " api_key=os.environ[\"MISTRAL_API_KEY\"],\n", + " server_url=f\"https://gateway.ai.cloudflare.com/v1/{CF_ACCOUNT_ID}/{CF_GATEWAY_ID}/mistral\",\n", + " ))\n", + "resp = client.chat.complete(model=\"mistral-small-latest\",\n", + " messages=[{\"role\": \"user\", \"content\": PROMPT}],\n", + " http_headers={\"cf-aig-authorization\": f\"Bearer {CF_GATEWAY_AUTH}\"})\n", + "text = resp.choices[0].message.content\n", + "print(text)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d74ec257", + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "client = sdk.wrap(OpenAI(\n", + " api_key=CF_GATEWAY_AUTH,\n", + " base_url=f\"https://gateway.ai.cloudflare.com/v1/{CF_ACCOUNT_ID}/{CF_GATEWAY_ID}/compat\",\n", + " ))\n", + "resp = client.chat.completions.create(\n", + " model=\"workers-ai/@cf/meta/llama-3.3-70b-instruct-fp8-fast\",\n", + " messages=[{\"role\": \"user\", \"content\": PROMPT}],\n", + " )\n", + "text = resp.choices[0].message.content\n", + "\n", + "print(text)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 3daae9d..b170c14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ dev = [ "ruff>=0.6", "mypy>=1.10", "types-requests>=2.31", - # every provider SDK (so unit + live integration tests can import them) + # every provider SDK (so the unit + wrapper tests can import them) "boto3>=1.34", "mistralai>=2.0", "anthropic>=0.30", diff --git a/src/lago_agent_sdk/adapters/_common.py b/src/lago_agent_sdk/adapters/_common.py new file mode 100644 index 0000000..1a57ea8 --- /dev/null +++ b/src/lago_agent_sdk/adapters/_common.py @@ -0,0 +1,24 @@ +"""Shared helpers used by more than one native provider adapter.""" + +from __future__ import annotations + +from typing import Any + + +def resolve_model(response_model: Any, requested_model: str) -> str: + """Prefer the model a response reports over the one requested. + + Every native provider can resolve a short alias/moniker to a more + specific snapshot id server-side, under different names — Anthropic and + OpenAI turn a short alias into a dated snapshot (e.g. + "claude-sonnet-4-5" -> "claude-sonnet-4-5-20250929"), Gemini hot-swaps + "-latest" aliases the same way (see + https://ai.google.dev/gemini-api/docs/models). Pricing/attribution must + key off what actually answered: OpenRouter lists the resolved snapshot, + never the alias. Falls back to the requested model only when the + response is silent about its own model (e.g. a synthetic streaming + usage blob). + """ + if isinstance(response_model, str) and response_model: + return response_model + return requested_model or "" diff --git a/src/lago_agent_sdk/adapters/anthropic_native.py b/src/lago_agent_sdk/adapters/anthropic_native.py index 5943676..b0b0e57 100644 --- a/src/lago_agent_sdk/adapters/anthropic_native.py +++ b/src/lago_agent_sdk/adapters/anthropic_native.py @@ -53,6 +53,20 @@ def _to_dict(obj: Any) -> dict[str, Any]: return {} +def _resolve_model(response_model: Any, requested_model: str) -> str: + """Prefer the model the response reports over the one requested. + + Anthropic can resolve a short alias to a more specific name — e.g. + "claude-sonnet-4-5" → "claude-sonnet-4-5-20250929" — with no gateway or + fallback involved at all. Pricing and attribution must key off what actually + answered. Falls back to the requested model only when the response is silent + about its own model (e.g. a synthetic streaming usage blob). + """ + if isinstance(response_model, str) and response_model: + return response_model + return requested_model or "" + + def extract_anthropic_native(response: Any, model_id: str = "") -> CanonicalUsage: """Translate an Anthropic native response (Message or dict) → CanonicalUsage. @@ -84,7 +98,7 @@ def extract_anthropic_native(response: Any, model_id: str = "") -> CanonicalUsag cache_write_5m=_safe_int(cache_creation.get("ephemeral_5m_input_tokens")), cache_write_1h=_safe_int(cache_creation.get("ephemeral_1h_input_tokens")), tool_calls=tool_calls, - model=model_id or (resp.get("model") if isinstance(resp.get("model"), str) else "") or "", + model=_resolve_model(resp.get("model"), model_id), provider="anthropic", api="native", extras=extras, diff --git a/src/lago_agent_sdk/adapters/gemini_native.py b/src/lago_agent_sdk/adapters/gemini_native.py index f3bdc96..d9770a6 100644 --- a/src/lago_agent_sdk/adapters/gemini_native.py +++ b/src/lago_agent_sdk/adapters/gemini_native.py @@ -35,6 +35,7 @@ from typing import Any, cast from ..canonical import CanonicalUsage +from ._common import resolve_model _KNOWN_USAGE_FIELDS = { "prompt_token_count", @@ -126,9 +127,7 @@ def extract_gemini_native(response: Any, model_id: str = "") -> CanonicalUsage: audio_output=_modality_token_count(candidates_details, "AUDIO"), image_input=_modality_token_count(prompt_details, "IMAGE"), tool_calls=_count_tool_calls(resp), - model=model_id - or (resp.get("model_version") if isinstance(resp.get("model_version"), str) else "") - or "", + model=resolve_model(resp.get("model_version"), model_id), provider="gemini", api="native", extras=extras, diff --git a/src/lago_agent_sdk/adapters/openai_native.py b/src/lago_agent_sdk/adapters/openai_native.py index 55bd09d..b10fee6 100644 --- a/src/lago_agent_sdk/adapters/openai_native.py +++ b/src/lago_agent_sdk/adapters/openai_native.py @@ -39,7 +39,16 @@ from typing import Any, cast -from ..canonical import CanonicalUsage +from ..canonical import WORKERS_AI_COMPAT_PREFIX, CanonicalUsage +from ._common import resolve_model + +# Cloudflare Workers AI names every model "@cf//". Reaching one +# through the gateway's OpenAI-compatible `/compat` endpoint additionally requires +# the "workers-ai/" routing prefix, so the same model arrives under two spellings +# depending on which surface the customer used. `pricing.lookup_cloudflare_workers_ai` +# strips the routing prefix before matching, because Cloudflare's own catalog lists +# only the bare form. +_WORKERS_AI_MODEL_PREFIX = "@cf/" # Top-level usage fields we recognize across BOTH chat completions and responses APIs. _KNOWN_USAGE_FIELDS = { @@ -101,6 +110,32 @@ def _count_responses_tool_calls(resp: dict[str, Any]) -> int: return sum(1 for item in output if isinstance(item, dict) and item.get("type") == "function_call") +def _infer_provider(resolved_model: str) -> str: + """The SDK shape only ever tells you "this looks like an OpenAI response" — + it can't tell you who actually served it. Going through a gateway's + OpenAI-compatible endpoint (e.g. Cloudflare's `.../compat`), the resolved + model string is the only real signal: "@cf/..." is Cloudflare Workers AI's + own naming convention, never a real OpenAI model. This isn't cosmetic — + `provider` is what price-mode keys pricing off of, and Workers AI has a + genuinely different price table (Cloudflare's own catalog) than real + OpenAI models (OpenRouter); stamping "openai" on a Workers AI call would + have made it permanently unpriceable, quietly, at the extraction layer. + + BOTH spellings have to match. Cloudflare's `/compat` endpoint takes the + provider-prefixed form — `workers-ai/@cf/meta/llama-3.3-70b-instruct-fp8-fast` + — which is what the README and the demo notebook prescribe, and what a + streaming call always reports (the synthetic usage payload carries no model, + so `resolve_model` falls back to the requested string verbatim). Matching + only the bare `@cf/` left every documented Workers AI call stamped "openai", + priced against OpenRouter, missed, and silently degraded to token events. + """ + if resolved_model.startswith(_WORKERS_AI_MODEL_PREFIX) or resolved_model.startswith( + f"{WORKERS_AI_COMPAT_PREFIX}{_WORKERS_AI_MODEL_PREFIX}" + ): + return "workers-ai" + return "openai" + + def extract_openai_native(response: Any, model_id: str = "") -> CanonicalUsage: """Translate an OpenAI response (chat completion or responses API) → CanonicalUsage. @@ -142,6 +177,7 @@ def extract_openai_native(response: Any, model_id: str = "") -> CanonicalUsage: if k not in _KNOWN_USAGE_FIELDS: extras[k] = v + resolved_model = resolve_model(resp.get("model"), model_id) return CanonicalUsage( input=input_tokens, output=output_tokens, @@ -150,8 +186,8 @@ def extract_openai_native(response: Any, model_id: str = "") -> CanonicalUsage: audio_input=audio_input, audio_output=audio_output, tool_calls=tool_calls, - model=model_id or (resp.get("model") if isinstance(resp.get("model"), str) else "") or "", - provider="openai", + model=resolved_model, + provider=_infer_provider(resolved_model), api=api, extras=extras, ) diff --git a/src/lago_agent_sdk/canonical.py b/src/lago_agent_sdk/canonical.py index 715a595..054aa22 100644 --- a/src/lago_agent_sdk/canonical.py +++ b/src/lago_agent_sdk/canonical.py @@ -5,6 +5,21 @@ from dataclasses import asdict, dataclass, field from typing import Any +# The routing prefix Cloudflare's OpenAI-compatible `/compat` endpoint requires: +# "workers-ai/@cf/meta/llama-3.3-70b-instruct-fp8-fast". The same model therefore +# arrives under two spellings depending on which surface the customer used, and two +# unrelated layers need to agree on this string — `adapters/openai_native` decides +# the PROVIDER from it, and `pricing.lookup_cloudflare_workers_ai` strips it before +# matching, because Cloudflare's own catalog lists only the bare "@cf/..." form. +# +# It lives here rather than in either of them because they must never import each +# other (an adapter is a pure function of a provider response; pricing is SDK state), +# and because a drift between two copies is a silent unpriced call, not a crash. This +# module is the natural shared floor: it imports nothing from the package, so there is +# no cycle in either direction, and depending on it does not pull `pricing`'s ~50KB +# into a lightweight adapter. +WORKERS_AI_COMPAT_PREFIX = "workers-ai/" + @dataclass class CanonicalUsage: @@ -39,7 +54,29 @@ class CanonicalUsage: ) def nonzero_numeric(self) -> dict[str, int]: - return {k: getattr(self, k) for k in self.NUMERIC_FIELDS if getattr(self, k)} + """Fields with a POSITIVE count, i.e. the ones worth billing. + + `> 0`, not just truthy: a negative slipped through and was emitted verbatim + as `value="-100"`, which Lago would sum into a negative billable quantity. + Nothing upstream should produce one — every adapter clamps at extraction — + but this is the last gate before an event is built, and the JS port already + filtered on `> 0`, so the two disagreed on the same input. + """ + return {k: v for k in self.NUMERIC_FIELDS if (v := getattr(self, k)) and v > 0} + + def negative_numeric(self) -> dict[str, int]: + """Fields `nonzero_numeric` DROPPED for being negative, so the caller can + report them. + + Reachable, unlike most defensive paths here: `CanonicalUsage` is exported and + `emit()` takes one directly, which is the documented way to backfill usage the + SDK did not intercept. A caller computing a delta wrongly can hand us a + negative, and silently dropping it is the one drop path that never reached + `on_error` — the same gap that was closed for queue overflow and for an + unresolvable subscription. Kept as a separate pure query so `CanonicalUsage` + stays a dumb dataclass with no notification channel of its own. + """ + return {k: v for k in self.NUMERIC_FIELDS if (v := getattr(self, k)) and v < 0} def to_dict(self) -> dict[str, Any]: return asdict(self) diff --git a/src/lago_agent_sdk/config.py b/src/lago_agent_sdk/config.py index ea81f64..26a66c1 100644 --- a/src/lago_agent_sdk/config.py +++ b/src/lago_agent_sdk/config.py @@ -49,6 +49,14 @@ class LagoConfig: request_timeout_seconds: float = 10.0 max_retry_seconds: float = 60.0 on_error: Callable[[Exception, str], None] | None = None + # TLS certificate verification for requests to `api_url`. Defaults to True + # (always verify — never disable this against a real Lago instance). The + # one legitimate reason to set False: a local dev Lago instance behind a + # self-signed certificate (e.g. Traefik's default local cert), where the + # alternative is routing through a public tunnel (ngrok, etc.) purely to + # get a browser-trusted cert — adding a flaky, unnecessary network hop for + # a problem this flag solves directly. + verify_ssl: bool = True # --- pricing (price mode) --- # Global default mode. "tokens" preserves the existing behavior exactly. @@ -61,6 +69,17 @@ class LagoConfig: pricing_ttl_seconds: float = 3600.0 # Region used for Bedrock pricing when the model id carries no region prefix. bedrock_default_region: str = "us-east-1" + # Cloudflare account id + API token for pricing "workers-ai" calls in price + # mode, via Cloudflare's own model catalog (not a public/no-auth source the + # way OpenRouter/AWS are — without both set, Workers AI pricing is simply + # unavailable and falls back to token events, same as any other miss). + cloudflare_account_id: str | None = None + cloudflare_api_token: str | None = field(default=None, repr=False) + # Usually NOT needed — wrap()-ing a mistralai client auto-detects this + # (see LagoSDK._auto_prime_pricing_for). Set it explicitly only when + # pricing Mistral usage without ever calling wrap() (e.g. a log-backfill + # path); an explicit value here always wins over an auto-detected one. + mistral_api_key: str | None = field(default=None, repr=False) # Optional injected PricingProvider (or a stub) — primarily for tests/overrides. # Typed Any to avoid a config→pricing import cycle. pricing_provider: Any | None = field(default=None, repr=False) @@ -79,5 +98,7 @@ def __repr__(self) -> str: f"markup={self.markup}, " f"cost_metric_code={self.cost_metric_code!r}, " f"pricing_ttl_seconds={self.pricing_ttl_seconds}, " - f"bedrock_default_region={self.bedrock_default_region!r})" + f"bedrock_default_region={self.bedrock_default_region!r}, " + f"cloudflare_account_id={self.cloudflare_account_id!r}, " + f"verify_ssl={self.verify_ssl})" ) diff --git a/src/lago_agent_sdk/gateway/__init__.py b/src/lago_agent_sdk/gateway/__init__.py new file mode 100644 index 0000000..3cca14d --- /dev/null +++ b/src/lago_agent_sdk/gateway/__init__.py @@ -0,0 +1,13 @@ +"""Gateway connector code — a second front door into the same billing kernel. + +Everything under `lago_agent_sdk.gateway` maps a third-party AI gateway's own +usage-reporting surface (Cloudflare's Logs API, Vercel's Reporting API, ...) +into the SDK's existing `CanonicalUsage` shape. It is consumed by a standalone +poller service, not by `wrap()` — there is no client to monkey-patch here. + +This is intentionally a separate namespace from `lago_agent_sdk.adapters` +(which extracts usage from a provider-native response inside a wrapped call). +The two never import from each other; both target `CanonicalUsage`. +""" + +from __future__ import annotations diff --git a/src/lago_agent_sdk/gateway/adapters/__init__.py b/src/lago_agent_sdk/gateway/adapters/__init__.py new file mode 100644 index 0000000..1e0177b --- /dev/null +++ b/src/lago_agent_sdk/gateway/adapters/__init__.py @@ -0,0 +1,6 @@ +from .cloudflare_gateway import extract_cloudflare_log, resolve_subscription + +__all__ = [ + "extract_cloudflare_log", + "resolve_subscription", +] diff --git a/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py b/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py new file mode 100644 index 0000000..8b4d21f --- /dev/null +++ b/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py @@ -0,0 +1,189 @@ +"""Cloudflare AI Gateway log adapter — maps a Logs API entry to CanonicalUsage. + +Verified against a real captured log entry (live account, real Anthropic call +routed through a real gateway, real Lago rollup confirmed exact). + +Field mapping (`GET .../ai-gateway/gateways/{id}/logs` and the single-entry +`GET .../logs/{log_id}`): + tokens_in → input + tokens_out → output + usage_metadata.input_cached_tokens → cache_read + usage_metadata.input_cache_creation_tokens → cache_write + usage_metadata.reasoningTokens → reasoning + model, provider → passed straight through + +Cloudflare reports its OWN counter vocabulary here, not the provider's. Across all +14 captured fixtures — Anthropic, Workers AI, Mistral and Gemini, via every ingress +method — the only keys that ever appear are `input_tokens`, `output_tokens`, +`total_tokens`, `input_cached_tokens`, `input_cache_creation_tokens`, `neurons`, +`input_text_tokens` and `reasoningTokens`. Not one provider-native key shows up: +no Anthropic `cache_read_input_tokens`, no Gemini `thoughtsTokenCount` or +`cachedContentTokenCount`. + +That vocabulary is *mostly* snake_case, with `reasoningTokens` as a camelCase +outlier — Cloudflare's own inconsistency, not a provider key leaking through +(Gemini's native spelling for the same quantity is `thoughtsTokenCount`, which +appears nowhere). The extra spellings checked below are therefore unobserved +insurance against a convention we have not seen, not handling for a known case. + +Unlike the provider-native adapters (`adapters/openai_native.py`, +`adapters/anthropic_native.py`), there is no request-side model kwarg to prefer +or fall back on here — a Cloudflare log entry always reports the model that +actually served the request. This adapter is immune, by construction, to the +alias-vs-resolved-model bug fixed in those two. + +Billing *policy* is deliberately not decided here — this module only extracts. +`cached`, `step`, and the log's own `id` land in `extras` because the caller +(the poller) needs them: `cached` to decide whether to skip billing a request +Cloudflare served for free, `id` as the idempotency key against replays. +`resolve_subscription()` is separate from extraction because attribution can be +absent, and dropping vs. warning on that is also a caller policy decision. +""" + +from __future__ import annotations + +from typing import Any + +from ...canonical import CanonicalUsage + + +def _safe_dict(v: Any) -> dict[str, Any]: + return v if isinstance(v, dict) else {} + + +def _safe_int(v: Any) -> int: + try: + return max(0, int(v or 0)) + except (TypeError, ValueError): + return 0 + + +def _safe_str(v: Any) -> str: + return v if isinstance(v, str) else "" + + +def _first_int(meta: dict[str, Any], *names: str) -> int: + """First of `names` present in `meta` with a usable value, as an int. + + Cloudflare's counter names are its own and mostly snake_case, but not + reliably so — `reasoningTokens` is camelCase in the real Gemini entry, right + next to snake_case `input_tokens` in the same object. Since the vocabulary is + internally inconsistent, the spelling it will use for a provider we have no + capture for is genuinely unknown. + + Checking every plausible spelling costs nothing and the downside is lopsided + — though it is lopsided in OPPOSITE DIRECTIONS depending on the provider, so + neither "over-bill" nor "under-bill" describes it alone: + + - For a SUBTRACTIVE provider (`gemini`, `openai`, `workers-ai` — in + `_INPUT_INCLUDES_CACHE_READ`), `compute_cost` subtracts `cache_read` out of + `input`. A missed cache key leaves those tokens billed at the full prompt + rate instead of the cache rate: an OVER-bill. + - For an ADDITIVE provider (`anthropic`), `cache_read` is billed as its own + line on top of `input`. A missed key means those tokens are not billed at + all: an UNDER-bill, which is the direction this SDK treats as worse. + + Uses `or`-style fallthrough (not "first key present"), so a provider that sends + both its own name and the gateway's with one of them zeroed still resolves to + the real count. + """ + for name in names: + v = _safe_int(meta.get(name)) + if v: + return v + return 0 + + +# Cloudflare AI Gateway logs its OWN provider vocabulary, which is not the name +# the pricing tables and token-semantics tables key off — and not always its own +# URL slug either (the logs say "workers-ai" where the endpoint path says +# "workersai"). Passed through verbatim, "google-ai-studio" matched no vendor in +# pricing's _VENDOR_MAP, so every Gemini call backfilled through the gateway +# missed on price; worse, it also missed _INPUT_INCLUDES_CACHE_READ, so Gemini's +# cache_read — a SUBSET of its input count, not additive — was billed twice. +# +# Only providers this SDK can actually price need an entry. Anything else passes +# through unchanged: an unrecognized provider is one we have no table for, and a +# clean miss falls back to token events, which is strictly better than inventing +# a mapping. AWS Bedrock is deliberately absent for that reason — Bedrock prices +# are keyed off `api.startswith("bedrock")`, and this connector always sets +# api="cloudflare_gateway", so mapping its provider name would route it to +# OpenRouter under a vendor that cannot match. A miss there is honest. +_PROVIDER_ALIASES = { + "google-ai-studio": "gemini", + "google-vertex-ai": "gemini", + "vertex": "gemini", + "azure-openai": "openai", + "azureopenai": "openai", + "workersai": "workers-ai", +} + + +def _normalize_provider(v: Any) -> str: + """Map Cloudflare's provider name onto the SDK's own provider vocabulary.""" + p = _safe_str(v).lower() + return _PROVIDER_ALIASES.get(p, p) + + +def extract_cloudflare_log(entry: dict[str, Any]) -> CanonicalUsage: + """Translate one Cloudflare AI Gateway log entry → CanonicalUsage. + + Accepts a single log entry dict as returned by the Logs API (either the + list endpoint or the single-entry endpoint — same shape). Missing/malformed + fields degrade to zero/empty rather than raising, matching the defensive + style of the other adapters — a poller processing a batch of log entries + must not have one malformed entry take down the whole run. + """ + usage_meta = _safe_dict(entry.get("usage_metadata")) + + return CanonicalUsage( + input=_safe_int(entry.get("tokens_in")), + output=_safe_int(entry.get("tokens_out")), + # Cloudflare's own key first — that is the only spelling ever observed + # (`input_cached_tokens` in 8 of the 14 captured fixtures). Everything after + # it is unobserved insurance: its camelCase form, then the two big providers' + # native names, in case Cloudflare ever forwards a provider's usage object + # rather than rewriting it into its own vocabulary. Kept because + # `_first_int` fallthrough is free and a missed cache key mis-bills in one + # direction or the other for EVERY provider (see `_first_int`) — but this is + # belt-and-braces, not handling for a case we have seen. + cache_read=_first_int( + usage_meta, + "input_cached_tokens", + "inputCachedTokens", + "cachedContentTokenCount", # Gemini native + "cache_read_input_tokens", # Anthropic native + ), + cache_write=_first_int( + usage_meta, + "input_cache_creation_tokens", + "inputCacheCreationTokens", + "cache_creation_input_tokens", # Anthropic native + ), + reasoning=_first_int( + usage_meta, + "reasoningTokens", # Cloudflare's own camelCase outlier — the observed one + "reasoning_tokens", + "thoughtsTokenCount", # Gemini native + ), + model=_safe_str(entry.get("model")), + provider=_normalize_provider(entry.get("provider")), + api="cloudflare_gateway", + extras={ + "cached": entry.get("cached"), + "step": entry.get("step"), + "log_id": entry.get("id"), + }, + ) + + +def resolve_subscription(entry: dict[str, Any]) -> str | None: + """Pull the Lago subscription id from the customer's `cf-aig-metadata` header. + + Returns None if the customer never set `lago_subscription` — the caller + decides what to do with an unattributed entry (drop it, log a warning, ...); + this function only reports whether attribution is present. + """ + metadata = _safe_dict(entry.get("metadata")) + value = metadata.get("lago_subscription") + return value if isinstance(value, str) and value else None diff --git a/src/lago_agent_sdk/lago_client.py b/src/lago_agent_sdk/lago_client.py index cf32d38..01f03f4 100644 --- a/src/lago_agent_sdk/lago_client.py +++ b/src/lago_agent_sdk/lago_client.py @@ -11,10 +11,29 @@ class LagoClient: - def __init__(self, api_key: str, api_url: str, timeout: float = 10.0) -> None: + def __init__(self, api_key: str, api_url: str, timeout: float = 10.0, verify_ssl: bool = True) -> None: self.api_key = api_key self.api_url = api_url.rstrip("/") self.timeout = timeout + self.verify_ssl = verify_ssl + if not verify_ssl: + # The customer explicitly opted out via config — they've already + # accepted the risk; requests/urllib3's warning on every single + # request would just be noise at that point, not new information. + # Import urllib3 directly rather than through `requests.packages`, + # which is a legacy compatibility shim that is not guaranteed to exist. + # Wrapped because this is an optional convenience: suppressing a warning + # must never be able to fail construction of the SDK itself. That is not + # hypothetical — `verify_ssl=False` is now a first-class constructor + # argument that the docstring recommends for local dev, so this line sits + # on an advertised path, and an ImportError/AttributeError here would + # take down `LagoSDK()` for the exact setup the flag was added to serve. + try: + import urllib3 + + urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + except Exception: # noqa: BLE001 + pass def __repr__(self) -> str: if not self.api_key: @@ -23,7 +42,10 @@ def __repr__(self) -> str: masked = "***" else: masked = f"***{self.api_key[-4:]}" - return f"LagoClient(api_key={masked!r}, api_url={self.api_url!r}, timeout={self.timeout})" + return ( + f"LagoClient(api_key={masked!r}, api_url={self.api_url!r}, " + f"timeout={self.timeout}, verify_ssl={self.verify_ssl})" + ) def send_batch(self, events: list[dict[str, Any]]) -> None: if not events: @@ -34,6 +56,8 @@ def send_batch(self, events: list[dict[str, Any]]) -> None: "Content-Type": "application/json", } payload = {"events": events} - resp = requests.post(url, headers=headers, data=json.dumps(payload), timeout=self.timeout) + resp = requests.post( + url, headers=headers, data=json.dumps(payload), timeout=self.timeout, verify=self.verify_ssl + ) if not (200 <= resp.status_code < 300): raise LagoApiError(resp.status_code, resp.text) diff --git a/src/lago_agent_sdk/pricing.py b/src/lago_agent_sdk/pricing.py index 3b3dfe7..7475d8f 100644 --- a/src/lago_agent_sdk/pricing.py +++ b/src/lago_agent_sdk/pricing.py @@ -7,6 +7,24 @@ - OpenRouter (``https://openrouter.ai/api/v1/models``) for native providers (anthropic / openai / mistral / gemini). Prices are USD per token. - AWS Bedrock Price List **Bulk** API (public, no credentials) for Bedrock. + - Cloudflare's own model catalog (``/accounts/{id}/ai/models/search``) for + ``workers-ai`` — the actual rate the gateway bills at, not a third party's + price for hosting the same open-weight model elsewhere (verified live: + Cloudflare's real charged cost for one call matched this catalog's rate + exactly; OpenRouter's listing for the same underlying model came out ~3.5x + lower — a genuinely different price, not just a naming mismatch). Needs + an account id + API token (Cloudflare's catalog isn't public/no-auth the + way OpenRouter/AWS are); without both set, this source is simply empty. + - Mistral's own ``/v1/models`` for *alias resolution*, not pricing directly. + Mistral has no per-token price table of its own (confirmed: their pricing + page lists one FAQ example, not a structured/JSON price list) — it genuinely + has no analogue to Cloudflare's catalog. But a customer request commonly + uses a moving alias (``mistral-small-latest``) and Mistral's response never + resolves it (unlike Anthropic/OpenAI, which report the dated snapshot that + answered) — so the OpenRouter lookup below misses even though OpenRouter + *does* list the resolved id (e.g. ``mistralai/mistral-small-2603``) with + real pricing. ``/v1/models`` exposes the resolution directly via each + model's ``aliases`` array; needs the customer's own Mistral API key. Design constraints (mirror the queue's non-blocking guarantee): - ``lookup()`` is pure in-memory and O(1); it NEVER does network I/O, so the @@ -28,18 +46,20 @@ import re import threading import time -from collections.abc import Callable +from collections.abc import Callable, Iterable from dataclasses import dataclass from decimal import ROUND_DOWN, Decimal, InvalidOperation from typing import Any, Protocol -from .canonical import CanonicalUsage +from .canonical import WORKERS_AI_COMPAT_PREFIX, CanonicalUsage logger = logging.getLogger("lago_agent_sdk.pricing") OPENROUTER_URL = "https://openrouter.ai/api/v1/models" AWS_PRICING_HOST = "https://pricing.us-east-1.amazonaws.com" AWS_BEDROCK_REGION_INDEX = f"{AWS_PRICING_HOST}/offers/v1.0/aws/AmazonBedrock/current/region_index.json" +CLOUDFLARE_MODELS_URL_TEMPLATE = "https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/models/search" +MISTRAL_MODELS_URL = "https://api.mistral.ai/v1/models" # Canonical usage fields we know how to price. PRICED_FIELDS = ("input", "output", "cache_read", "cache_write", "reasoning") @@ -49,13 +69,32 @@ # 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. -_INPUT_INCLUDES_CACHE_READ = frozenset({"openai", "gemini"}) +# +# "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. +_INPUT_INCLUDES_CACHE_READ = frozenset({"openai", "gemini", "workers-ai"}) # 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.) -_OUTPUT_INCLUDES_REASONING = frozenset({"openai"}) +# "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"}) # Canonical field -> OpenRouter pricing key. _OPENROUTER_FIELD_MAP = { @@ -75,6 +114,26 @@ "google": "google", } +# Cloudflare's catalog price unit -> canonical field. Real, surveyed units also +# include "per 1k characters", "per step", "per 512 by 512 tile", "per audio +# minute (websocket)", "per audio minute", "per inference request" — none of +# those are token-based, so they're deliberately absent: a model priced only in +# those units yields a ModelPrice with no input/output/cache_read at all, which +# `compute_cost` already treats as "unpriced field, skip it" — the same safe +# behavior as any other model with no usable price. +_CLOUDFLARE_UNIT_FIELD_MAP = { + "per M input tokens": "input", + "per M output tokens": "output", + "per M cached input tokens": "cache_read", +} + +# Cloudflare's catalog page size, and a hard bound on the paging loop. The loop runs +# on the queue's flush tick ahead of the drain, so it must terminate even if the +# endpoint keeps returning full pages. 40 pages covers ~2000 models against a real +# catalog of 64. +_CF_PER_PAGE = 50 +_CF_MAX_PAGES = 40 + # Bedrock cross-region inference prefix -> a representative AWS region. _BEDROCK_REGION_PREFIX = { "us": "us-east-1", @@ -98,8 +157,23 @@ _SCALE = 12 _Q = Decimal(1).scaleb(-_SCALE) # Decimal("1E-12") +# A trailing version/revision marker OpenRouter usually omits from its own ids. +# Shapes seen live: Anthropic's compact date ("-20250929") and an explicit "-v2". _VERSION_DATE_SUFFIX = re.compile(r"-(?:\d{8}|v\d+)$") +# Gemini's 3-digit revision ("-002", which `model_version` can report where +# OpenRouter lists only the bare name) is stripped for OpenRouter matching ONLY. +# It is deliberately NOT in the shared `_strip_version`: that helper also builds +# the AWS/Bedrock price keys, where a shortened key does not merely miss but +# silently MIS-prices — `bedrock_model_key` feeds +# `table.setdefault(key, {})[direction] = price`, so two distinct models +# collapsing to one key overwrite each other's rate. All four live catalogs are +# currently clean (OpenRouter 415 ids, Cloudflare 64, AWS offer 77, captured +# Bedrock 39: zero model parts end in exactly three digits), but the arm was only +# ever motivated by OpenRouter, and scoping it makes that risk structurally zero +# instead of empirically zero. +_OPENROUTER_VERSION_SUFFIX = re.compile(r"-(?:\d{8}|\d{3}|v\d+)$") + # ---------------------------------------------------------------------- # Money helpers (kept in lock-step with the JS implementation) @@ -112,7 +186,63 @@ def _parse_price(value: Any) -> Decimal | None: return None if d.is_nan() or d.is_infinite() or d < 0: return None - return d.quantize(_Q, rounding=ROUND_DOWN) + try: + return d.quantize(_Q, rounding=ROUND_DOWN) + except InvalidOperation: + # quantize raises once the result would exceed the default 28-digit + # context precision — i.e. at 1e16 and above (16 integer digits + the + # 12 fractional ones this always produces). Absurd as a price, but it + # must not ESCAPE: this function is documented as returning None on bad + # input, and callers rely on that. Uncaught, it propagated out of + # compute_precomputed_cost into emit()'s catch-all, so the event was + # dropped as an unknown error instead of taking the normal "no price" + # path. Returning None also keeps JS byte-identical, where parseScaled + # returns null for exactly these inputs. + return None + + +def money_str_to_cents(usd: str) -> str: + """A money string (already floored to 12dp) → the same amount in cents, + same floor-and-format conventions as everywhere else.""" + return _fmt_money(Decimal(usd) * 100) + + +def apply_markup(usd: str, markup: str) -> str: + """`compute_cost`'s per-field `cost` values are PRE-markup — only the + summed `total` has markup applied. Splitting a breakdown into one event + per field (per token_type) needs markup applied to each field individually, + with the same floor-to-12dp convention as everywhere else, or a markup + != 1.0 would silently vanish from every per-field/token_type event. + + Parsed through `_parse_price` rather than `Decimal()` directly. A bare + `Decimal("abc")` raises `InvalidOperation` — and this is called from inside + `_push_cost_event`, under `emit()`'s catch-all, so the whole cost event was + dropped and reported as an unknown "emit" error instead of taking the + documented no-price path. It was also the one money helper in this module that + could raise at all, past every caller relying on the `None`-on-bad-input + convention. + + Both fallbacks are DEFENCE IN DEPTH, not live behaviour: every `emit()` path + runs the customer's markup through `coerce_markup` first (which falls back to + 1.0 and reports under "pricing"), and `CostBreakdown.markup` / + `fields[*]["cost"]` are `_fmt_money` output, so neither argument can actually + arrive unparseable here. They are still not interchangeable, and the two ports + disagreed on them: + + - An unparseable `usd` means the cost itself is unusable — nothing to bill: 0. + - An unparseable `markup` means only the MULTIPLIER is unusable. Returning 0 + there would discard a good cost, an under-bill to nothing; 1.0 bills the real + cost with no markup, the smallest defensible error. JS already did this; + Python returned "0", so identical input produced different bills if anything + ever did reach it. Aligned rather than left as a latent divergence. + """ + base = _parse_price(usd) + mult = _parse_price(markup) + if base is None: + return _fmt_money(Decimal(0)) + if mult is None: + mult = Decimal(1) + return _fmt_money((base * mult).quantize(_Q, rounding=ROUND_DOWN)) def _fmt_money(d: Decimal) -> str: @@ -139,6 +269,12 @@ def _strip_version(model: str) -> str: return _VERSION_DATE_SUFFIX.sub("", model) +def _strip_version_openrouter(model: str) -> str: + """`_strip_version`, plus Gemini's 3-digit revision. OpenRouter matching only — + see `_OPENROUTER_VERSION_SUFFIX` for why this is not the shared helper.""" + return _OPENROUTER_VERSION_SUFFIX.sub("", model) + + # ---------------------------------------------------------------------- # Price tables # ---------------------------------------------------------------------- @@ -205,19 +341,74 @@ def compute_cost(usage: CanonicalUsage, price: ModelPrice, markup: Decimal) -> C "unit_price": _fmt_money(unit), "cost": _fmt_money(cost), } - # Floor the USD total to 12 dp FIRST, then derive cents from it, so cents == - # billed-USD × 100 exactly (matches the JS integer-division implementation). + return _finalize_breakdown(base, markup, price.source, fields) + + +def _finalize_breakdown( + base: Decimal, markup: Decimal, source: str, fields: dict[str, dict[str, str]] +) -> CostBreakdown: + """Shared tail for `compute_cost`/`compute_precomputed_cost`: floor the + USD total to 12 dp FIRST, then derive cents from it, so cents == + billed-USD × 100 exactly (matches the JS integer-division implementation).""" total = (base * markup).quantize(_Q, rounding=ROUND_DOWN) return CostBreakdown( total=_fmt_money(total), total_cents=_fmt_money(total * 100), base=_fmt_money(base), markup=_fmt_money(markup), - source=price.source, + source=source, fields=fields, ) +def deoverlapped_token_total(usage: Any) -> int: + """Total tokens a call actually consumed, with per-provider overlaps removed. + + Sums the same PRICED_FIELDS the split cost path emits one event each for, so + the single-event `unit` equals the sum of the split path's `unit`s instead of + reporting a different basis. Both `_INCLUDES_` sets are applied, because a + subset counted twice inflates the reported quantity exactly as it would inflate + a price: + + * reasoning ⊆ output for providers in _OUTPUT_INCLUDES_REASONING + * cache_read ⊆ input for providers in _INPUT_INCLUDES_CACHE_READ + + 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 + cannot change how many tokens were consumed. The two still agree: when a + cache-inclusive provider has no cache_read price, `compute_cost` leaves the + cached tokens inside `input` and emits no cache_read event, and this skips + cache_read for the same reason. + + Deliberately limited to PRICED_FIELDS — the five text fields. `tool_calls` is a + count of calls rather than tokens, and `cache_write_5m` / `cache_write_1h` are + a breakdown OF `cache_write`, so including any of them would not be a token + total. This mirrors price mode's documented five-field scope. + """ + provider = (getattr(usage, "provider", "") or "").lower() + counts = {f: (getattr(usage, f, 0) or 0) for f in PRICED_FIELDS} + if provider in _OUTPUT_INCLUDES_REASONING: + counts["reasoning"] = 0 + if provider in _INPUT_INCLUDES_CACHE_READ: + counts["cache_read"] = 0 + return sum(int(v or 0) for v in counts.values()) + + +def compute_precomputed_cost(usd_cost: Any, markup: Decimal) -> CostBreakdown: + """Build a CostBreakdown from a cost the CALLER already knows. + + For a gateway that reports its own real, metered price per call (e.g. + Cloudflare AI Gateway's `cost` field), computing our own per-token estimate + via the OpenRouter/Bedrock tables would be redundant AND less accurate than + the number the gateway already gives us. This skips `compute_cost` entirely + — there's one lump sum, not a per-field breakdown, so `fields` is empty and + the invalid/negative case floors to 0 the same way `_parse_price` always has, + rather than raising or silently mis-billing. + """ + base = _parse_price(usd_cost) or Decimal(0) + return _finalize_breakdown(base, markup, "precomputed", {}) + + def coerce_markup(markup: Any) -> tuple[Decimal, bool]: """Return (markup_decimal, ok). Falls back to 1.0 when invalid/non-positive.""" d = _parse_price(markup) @@ -251,16 +442,176 @@ def parse_openrouter(data: Any) -> dict[str, Any]: cache_write=_parse_price(pricing.get(_OPENROUTER_FIELD_MAP["cache_write"])), reasoning=_parse_price(pricing.get(_OPENROUTER_FIELD_MAP["reasoning"])), ) + # OpenRouter marks a MOVING alias with a leading "~" on the vendor — + # "~anthropic/claude-sonnet-latest", "~openai/gpt-latest", + # "~google/gemini-flash-latest". Measured live: 11 such ids across 6 + # vendors, every one a "-latest" moniker, every one carrying real token + # pricing. Indexed verbatim they were ALL unpriceable, because the vendor + # parsed as "~anthropic"/"~openai"/"~google" — none of which appear in + # _VENDOR_MAP — so a customer in price mode asking for a plain "-latest" + # alias missed and fell back to token events, billing nothing at all in an + # llm_cost-only setup. Stripping the marker indexes them under their real + # vendor. Verified collision-free against the live catalog: no un-prefixed + # id duplicates a "~"-prefixed one, so nothing is overwritten. + # + # `setdefault` for the alias-derived keys rather than assignment: the + # collision-freedom above is a property of TODAY's catalog, and with plain + # assignment the winner depended purely on iteration order — if OpenRouter + # ever ships both "google/gemini-flash-latest" and + # "~google/gemini-flash-latest", the moving alias could overwrite the real + # listing's rate (measured on a synthetic pair: 0.009 vs 0.001 for the same + # lookup, decided by nothing but position in the response). A real listing + # now always wins, whatever the order. Non-alias entries keep plain + # assignment so genuine duplicates behave exactly as before. + bare = mid[1:] if mid.startswith("~") else mid + is_alias = bare != mid exact[mid] = mp - if "/" in mid: - vendor, _, suffix = mid.partition("/") - norm[(vendor.lower(), _norm(suffix))] = mp + if is_alias: + exact.setdefault(bare, mp) + if "/" in bare: + vendor, _, suffix = bare.partition("/") + norm_key = (vendor.lower(), _norm(suffix)) + if is_alias: + norm.setdefault(norm_key, mp) + else: + norm[norm_key] = mp return {"exact": exact, "norm": norm} +# A real dated Mistral snapshot ends in a short numeric tag (e.g. "-2603", +# "-2411", "-2508") — never a "-latest"-style moniker. Used to pick the one +# genuine canonical name out of a family that mutually lists each other (see +# parse_mistral_aliases). +_MISTRAL_DATED_ID = re.compile(r"-\d{4,8}$") + + +def _mistral_date_key(name: str) -> int: + """Normalize a dated Mistral suffix to a comparable integer; newest = largest. + + Mistral's own convention is a 4-digit YYMM ("-2411", "-2603"), but the regex + admits 4-8 digits and mixed widths do NOT compare correctly as raw strings: + "20241101" sorts *below* "2411" lexicographically. Widening YYMM to YYYYMM00 + puts both shapes on one scale. + """ + m = _MISTRAL_DATED_ID.search(name) + if m is None: + return -1 + digits = m.group(0)[1:] # drop the leading "-" + if len(digits) == 4: # YYMM -> 20YY-MM, day unknown + return int(f"20{digits}00") + return int(digits) # YYYYMMDD, or an unexpected width taken at face value + + +def _pick_mistral_canonical(names: list[str]) -> str: + """Prefer the NEWEST dated snapshot id (what OpenRouter actually lists + models under) over a "-latest"-style moniker. + + Newest, not shortest. Every dated id in one family is the same length, so a + shortest-then-alphabetical tie-break silently resolved on the DATE — and + ascending: `mistral-large-2402` / `-2407` / `-2411` / `-latest` all collapsed + onto `mistral-large-2402`, the OLDEST, so the whole family got priced at a + two-year-old rate. `-2411` had matched OpenRouter directly before alias + resolution existed, which makes that a regression rather than a gap. + + Falls back to shortest-then-alphabetical only when the group has no dated + candidate at all, so the choice stays deterministic either way. Ordering is + by Unicode code point — the JS port must NOT use `localeCompare`, which is + ICU/locale-dependent and made the two repos pick different canonicals for + the same input. + """ + dated = [n for n in names if _MISTRAL_DATED_ID.search(n)] + if dated: + return sorted(dated, key=lambda n: (-_mistral_date_key(n), n))[0] + return sorted(names, key=lambda n: (len(n), n))[0] + + +def parse_mistral_aliases(data: Any) -> dict[str, str]: + """Parse Mistral's `/v1/models` response into {alias: canonical_id}. + + Naively mapping "each name in this entry's `aliases` -> this entry's + `id`" is wrong: Mistral's real response lists EVERY name in a family as + its own top-level entry, each one's `aliases` pointing at the others — + e.g. `id="mistral-small-2603"`, `id="mistral-small-latest"`, AND + `id="magistral-small-latest"` each appear separately, each listing the + other two as `aliases`. A directional last-write-wins map is then + order-dependent and can resolve an alias to ANOTHER alias instead of the + real dated snapshot (confirmed live: this resolved + "mistral-small-latest" -> "magistral-small-latest", which OpenRouter + doesn't list, instead of -> "mistral-small-2603", which it does). + + Union-find instead: treat a model's id + its aliases as one connected + group regardless of which entry mentions which, then pick a single + canonical name per group (see `_pick_mistral_canonical`) and map every + other member of the group to it. + """ + models = data.get("data") if isinstance(data, dict) else None + if not isinstance(models, list): + return {} + + parent: dict[str, str] = {} + + def find(x: str) -> str: + root = x + while parent.get(root, root) != root: + root = parent[root] + return root + + def union(a: str, b: str) -> None: + ra, rb = find(a), find(b) + if ra != rb: + parent[ra] = rb + + names: set[str] = set() + for m in models: + if not isinstance(m, dict): + continue + mid = m.get("id") + if not isinstance(mid, str) or not mid: + continue + parent.setdefault(mid, mid) + names.add(mid) + for alias in m.get("aliases") or []: + if isinstance(alias, str) and alias: + parent.setdefault(alias, alias) + names.add(alias) + union(mid, alias) + + groups: dict[str, list[str]] = {} + for name in names: + groups.setdefault(find(name), []).append(name) + + result: dict[str, str] = {} + for members in groups.values(): + if len(members) < 2: + continue # no aliasing at all — nothing to resolve + canonical = _pick_mistral_canonical(members) + for name in members: + if name == canonical: + continue + # An explicit dated snapshot is already the real id OpenRouter lists, + # so it must pass through untouched — never rewritten onto a sibling. + # Without this, requesting `mistral-large-2411` was remapped to the + # group's canonical and priced at THAT snapshot's rate instead of its + # own, which is a mispricing rather than a miss. + if _MISTRAL_DATED_ID.search(name): + continue + result[name] = canonical + return result + + def lookup_openrouter(table: dict[str, Any], provider: str, model: str) -> ModelPrice | None: """Match (provider, model) to an OpenRouter price. Conservative: vendor-gated.""" vendor = _VENDOR_MAP.get((provider or "").lower(), (provider or "").lower()) + # Some sources report the model ALREADY carrying its vendor prefix — a real + # Cloudflare AI Gateway log for a REST-path call says + # model="anthropic/claude-opus-4.8" with provider="anthropic" — which would + # otherwise build "anthropic/anthropic/claude-opus-4.8" and never match. + # Strip it only when the prefix agrees with the vendor we just resolved, so + # this stays vendor-gated as documented: a model naming a DIFFERENT vendor + # than the call claims is still a miss, not a cross-vendor mispricing. + head, sep, tail = model.partition("/") + if sep and head.lower() in (vendor, (provider or "").lower()): + model = tail exact: dict[str, ModelPrice] = table.get("exact", {}) norm: dict[tuple[str, str], ModelPrice] = table.get("norm", {}) # 1. exact id @@ -272,12 +623,96 @@ def lookup_openrouter(table: dict[str, Any], provider: str, model: str) -> Model if hit is not None: return hit # 3. date/version-stripped, normalized - hit = norm.get((vendor, _norm(_strip_version(model)))) + hit = norm.get((vendor, _norm(_strip_version_openrouter(model)))) if hit is not None: return hit return None +# ---------------------------------------------------------------------- +# Cloudflare Workers AI parsing + matching +# +# Unlike OpenRouter/Bedrock, this is the ACTUAL rate the gateway bills at — not +# a third party's price for hosting the same open-weight model elsewhere, which +# can (and does) differ meaningfully. Model strings (e.g. +# "@cf/meta/llama-3.3-70b-instruct-fp8-fast") are already exact and +# self-contained; no vendor-prefix mapping is needed the way OpenRouter needs +# one to disambiguate "anthropic" -> "anthropic" vs "mistral" -> "mistralai". +# ---------------------------------------------------------------------- +def parse_cloudflare_workers_ai(models: Any) -> dict[str, ModelPrice]: + """Parse `/ai/models/search` results into {model_name: ModelPrice}. + + A model with no `price` property at all, or whose price entries are all + non-token units (per-image, per-audio-minute, ...), is simply absent from + the table — `lookup` then returns None, same as any other priced-nowhere + model, and the caller safely falls back to token events. + """ + table: dict[str, ModelPrice] = {} + if not isinstance(models, list): + return table + for m in models: + if not isinstance(m, dict): + continue + name = m.get("name") + if not isinstance(name, str) or not name: + continue + # `.get("properties", [])` only defaults when the key is ABSENT — an + # explicit JSON null returns None, and `for p in None` raises TypeError + # straight out of this function into `maybe_refresh`'s handler, which leaves + # `_cloudflare_workers_ai` at None. One malformed entry would therefore + # unprice EVERY Workers AI model, not just its own. The JS port already + # isinstance-guarded here and dropped only the bad entry. + props = m.get("properties") + price_prop = next( + ( + p + for p in (props if isinstance(props, list) else []) + if isinstance(p, dict) and p.get("property_id") == "price" + ), + None, + ) + if not isinstance(price_prop, dict): + continue + entries = price_prop.get("value") + if not isinstance(entries, list): + continue + fields: dict[str, Decimal] = {} + for entry in entries: + if not isinstance(entry, dict) or entry.get("currency") != "USD": + continue + field = _CLOUDFLARE_UNIT_FIELD_MAP.get(str(entry.get("unit", ""))) + if field is None: + continue + per_million = _parse_price(entry.get("price")) + if per_million is None: + continue + fields[field] = (per_million / Decimal(1_000_000)).quantize(_Q, rounding=ROUND_DOWN) + if fields: + table[name] = ModelPrice(source="cloudflare_workers_ai", **fields) + return table + + +def lookup_cloudflare_workers_ai(table: dict[str, ModelPrice], model: str) -> ModelPrice | None: + """Exact match first; a version-suffix fallback covers the same drift we've + seen in practice — e.g. a live response naming a model + "...instruct-v2" when the catalog itself only lists "...instruct". + + The "workers-ai/" routing prefix comes off first. Cloudflare's catalog keys + models as bare "@cf/...", but calling one through the gateway's `/compat` + endpoint requires "workers-ai/@cf/..." — the form the README prescribes and + the only form a streaming call can report. Without the strip, recognising the + prefixed spelling as Workers AI upstream just moves the miss here. + """ + for candidate in (model, model.removeprefix(WORKERS_AI_COMPAT_PREFIX)): + hit = table.get(candidate) + if hit is not None: + return hit + hit = table.get(_strip_version(candidate)) + if hit is not None: + return hit + return None + + # ---------------------------------------------------------------------- # Bedrock parsing + matching # @@ -416,13 +851,37 @@ def lookup_bedrock(region_table: dict[str, ModelPrice], model: str) -> ModelPric class PricingFetcher(Protocol): def fetch_openrouter(self) -> dict[str, Any]: ... def fetch_bedrock(self, region: str) -> dict[str, ModelPrice]: ... + def fetch_cloudflare_workers_ai(self) -> dict[str, ModelPrice]: ... + def fetch_mistral_aliases(self, api_key: str | None = None) -> dict[str, str]: ... class HttpPricingFetcher: - """Default fetcher using ``requests`` (already a core dependency).""" + """Default fetcher using ``requests`` (already a core dependency). + + ``cloudflare_account_id``/``cloudflare_api_token``: unlike OpenRouter/AWS, + Cloudflare's model catalog is account-scoped and needs auth — there's no + public, no-credentials equivalent. Without both set, + ``fetch_cloudflare_workers_ai`` returns an empty table rather than raising, + so Workers AI pricing is simply unavailable (safe token-event fallback) + instead of breaking price mode for every other provider. + + ``mistral_api_key``: same story — Mistral's ``/v1/models`` needs the + customer's own key. Without it, ``fetch_mistral_aliases`` returns an + empty map, so alias resolution is simply skipped and lookups fall back to + whatever the request already spelled out (safe miss, not a break). + """ - def __init__(self, timeout: float = 10.0) -> None: + def __init__( + self, + timeout: float = 10.0, + cloudflare_account_id: str | None = None, + cloudflare_api_token: str | None = None, + mistral_api_key: str | None = None, + ) -> None: self._timeout = timeout + self._cf_account_id = cloudflare_account_id + self._cf_api_token = cloudflare_api_token + self._mistral_api_key = mistral_api_key def fetch_openrouter(self) -> dict[str, Any]: import requests @@ -444,6 +903,63 @@ def fetch_bedrock(self, region: str) -> dict[str, ModelPrice]: offer.raise_for_status() return parse_bedrock_offer(offer.json(), region) + def fetch_cloudflare_workers_ai(self) -> dict[str, ModelPrice]: + import requests + + if not self._cf_account_id or not self._cf_api_token: + return {} + url = CLOUDFLARE_MODELS_URL_TEMPLATE.format(account_id=self._cf_account_id) + headers = {"Authorization": f"Bearer {self._cf_api_token}"} + models: list[Any] = [] + page = 1 + while True: + resp = requests.get( + url, + headers=headers, + params={"per_page": _CF_PER_PAGE, "page": page}, + timeout=self._timeout, + ) + resp.raise_for_status() + body = resp.json() + batch = body.get("result") or [] + models.extend(batch) + # A SHORT page is the only reliable end-of-catalog signal here. + # `result_info.total_count` is not: measured live it reports 291 while + # the endpoint serves 64 (50 then 14 then 0), so a `len(models) >= total` + # test never fires. It must also never be defaulted to `len(models)` — + # that made an ABSENT total_count break after page one, silently keeping + # 50 of the 64 available. + if len(batch) < _CF_PER_PAGE: + break + if page >= _CF_MAX_PAGES: + # Bounded because this runs on the queue's flush tick, ahead of the + # drain — an endpoint that always returns a full page must not stall + # event delivery indefinitely. Truncation is reported rather than + # silent, since a short catalog reads as "these models are unpriced". + logger.warning( + "lago: cloudflare model catalog truncated at %d pages (%d models); " + "prices for later models are unavailable", + _CF_MAX_PAGES, + len(models), + ) + break + page += 1 + return parse_cloudflare_workers_ai(models) + + def fetch_mistral_aliases(self, api_key: str | None = None) -> dict[str, str]: + import requests + + # An explicitly configured key (LagoConfig.mistral_api_key) always + # wins over one learned from a wrapped client — a deliberate config + # value shouldn't be silently shadowed by an auto-detected one. + key = self._mistral_api_key or api_key + if not key: + return {} + headers = {"Authorization": f"Bearer {key}"} + resp = requests.get(MISTRAL_MODELS_URL, headers=headers, timeout=self._timeout) + resp.raise_for_status() + return parse_mistral_aliases(resp.json()) + # ---------------------------------------------------------------------- # PricingProvider — cache + background refresh + non-blocking lookup @@ -455,8 +971,15 @@ def __init__( ttl_seconds: float = 3600.0, default_region: str = "us-east-1", on_error: Callable[[Exception, str], None] | None = None, + cloudflare_account_id: str | None = None, + cloudflare_api_token: str | None = None, + mistral_api_key: str | None = None, ) -> None: - self._fetcher: PricingFetcher = fetcher or HttpPricingFetcher() + self._fetcher: PricingFetcher = fetcher or HttpPricingFetcher( + cloudflare_account_id=cloudflare_account_id, + cloudflare_api_token=cloudflare_api_token, + mistral_api_key=mistral_api_key, + ) self._ttl = ttl_seconds self._default_region = default_region self._on_error = on_error @@ -470,6 +993,17 @@ def __init__( self._bedrock: dict[str, dict[str, ModelPrice]] = {} self._bedrock_fetched: dict[str, float] = {} self._bedrock_stale: set[str] = set() + self._cloudflare_workers_ai: dict[str, ModelPrice] | None = None + self._cloudflare_fetched = 0.0 + self._cloudflare_stale = False + self._mistral_aliases: dict[str, str] | None = None + self._mistral_fetched = 0.0 + self._mistral_stale = False + # Learned from a wrapped Mistral client (see LagoSDK._auto_prime_pricing_for), + # not configured — the customer's own client already carries this key + # for making real calls, so alias resolution can reuse it without + # ever requiring a separate LagoConfig.mistral_api_key. + self._mistral_api_key_override: str | None = None self._refreshing: set[str] = set() def _heal_fork(self) -> None: @@ -483,13 +1017,57 @@ def _heal_fork(self) -> None: self._pid = os.getpid() self._openrouter_stale = self._openrouter is not None or self._openrouter_stale self._bedrock_stale = set(self._bedrock.keys()) + self._cloudflare_stale = self._cloudflare_workers_ai is not None or self._cloudflare_stale + self._mistral_stale = self._mistral_aliases is not None or self._mistral_stale self._refreshing = set() - def prime(self) -> None: - """Flag the OpenRouter table for an eager background warm (used when - price mode is the global default) to shrink the cold-start window.""" + def prime(self, providers: Iterable[str] = ()) -> None: + """Flag OpenRouter for an eager background warm (used when price mode + is the global default) to shrink the cold-start window. + + Deliberately does NOT also eagerly warm Cloudflare Workers AI or + Mistral alias resolution by default — both are credential-gated and + provider-specific; most price-mode customers never touch Workers AI + or Mistral at all, and eagerly hitting either's API at construction + time regardless of actual usage is real, unnecessary work (an extra + network round-trip per SDK instance, every TTL cycle, for a provider + that may never be called). Instead they stay purely reactive: the + first real `lookup()` for that provider flags it stale (see below), + `maybe_refresh()` fetches it on the queue's very next tick, and every + call after that — even the one a second later — hits the cache, with + zero further network calls until the TTL expires. Only that first + call for a given provider can race a cold cache; every provider that + session never calls costs nothing. + + Pass `providers=["mistral"]` and/or `["workers-ai"]` when you already + know, in advance, which of these two you're about to call this + session — this eagerly warms exactly that source too, so even ITS + first call prices correctly instead of paying the one-time lazy + cold-start cost. Unknown provider names are silently ignored (no + source is warmed) rather than raising, since this is a hint, not a + contract.""" with self._lock: self._openrouter_stale = True + for p in providers: + key = (p or "").lower() + if key == "workers-ai": + self._cloudflare_stale = True + elif key == "mistral": + self._mistral_stale = True + + def learn_mistral_api_key(self, api_key: str) -> None: + """Adopt a Mistral API key discovered from a wrapped client, so + alias resolution can run without ever requiring the customer to + also declare it in `LagoConfig` — their Mistral client already + carries the exact credential needed. Pure in-memory, no I/O. A key + explicitly set via `LagoConfig.mistral_api_key` always wins over one + learned this way (see `HttpPricingFetcher.fetch_mistral_aliases`); + this only fills the gap when no explicit key was configured.""" + if not api_key: + return + with self._lock: + if not self._mistral_api_key_override: + self._mistral_api_key_override = api_key # ---- non-blocking lookup (customer thread) ---- def lookup(self, provider: str, model: str, api: str) -> ModelPrice | None: @@ -506,12 +1084,32 @@ def lookup(self, provider: str, model: str, api: str) -> ModelPrice | None: if not fresh: self._bedrock_stale.add(region) return lookup_bedrock(table, model) if table is not None else None + if (provider or "").lower() == "workers-ai": + with self._lock: + table_cf = self._cloudflare_workers_ai + fresh_cf = table_cf is not None and (time.time() - self._cloudflare_fetched) < self._ttl + if not fresh_cf: + self._cloudflare_stale = True + return lookup_cloudflare_workers_ai(table_cf, model) if table_cf is not None else None + resolved_model = model + is_mistral = (provider or "").lower() == "mistral" with self._lock: + if is_mistral: + aliases = self._mistral_aliases + fresh_m = aliases is not None and (time.time() - self._mistral_fetched) < self._ttl + if not fresh_m: + self._mistral_stale = True + # Cold/miss: resolved_model stays the alias as-requested, + # and the OpenRouter lookup below misses safely, same as + # before this resolution step existed — never worse than + # the old behavior, only better once the table is warm. + if aliases: + resolved_model = aliases.get(model, model) table_or = self._openrouter fresh = table_or is not None and (time.time() - self._openrouter_fetched) < self._ttl if not fresh: self._openrouter_stale = True - return lookup_openrouter(table_or, provider, model) if table_or is not None else None + return lookup_openrouter(table_or, provider, resolved_model) if table_or is not None else None except Exception: # noqa: BLE001 — lookup must never raise return None @@ -523,12 +1121,23 @@ def maybe_refresh(self) -> None: # keeps the queue's background tick essentially free and avoids extra # cross-thread lock churn. The reads are racy but harmless: a missed flag # just defers a refresh by one tick. - if not self._openrouter_stale and not self._bedrock_stale: + if ( + not self._openrouter_stale + and not self._bedrock_stale + and not self._cloudflare_stale + and not self._mistral_stale + ): return with self._lock: do_openrouter = self._openrouter_stale and "openrouter" not in self._refreshing if do_openrouter: self._refreshing.add("openrouter") + do_cloudflare = self._cloudflare_stale and "cloudflare_workers_ai" not in self._refreshing + if do_cloudflare: + self._refreshing.add("cloudflare_workers_ai") + do_mistral = self._mistral_stale and "mistral_aliases" not in self._refreshing + if do_mistral: + self._refreshing.add("mistral_aliases") regions = [r for r in self._bedrock_stale if f"bedrock:{r}" not in self._refreshing] for r in regions: self._refreshing.add(f"bedrock:{r}") @@ -546,6 +1155,34 @@ def maybe_refresh(self) -> None: with self._lock: self._refreshing.discard("openrouter") + if do_cloudflare: + try: + table_cf = self._fetcher.fetch_cloudflare_workers_ai() + with self._lock: + self._cloudflare_workers_ai = table_cf + self._cloudflare_fetched = time.time() + self._cloudflare_stale = False + except Exception as exc: # noqa: BLE001 + self._report(exc, "pricing.fetch_cloudflare_workers_ai") + finally: + with self._lock: + self._refreshing.discard("cloudflare_workers_ai") + + if do_mistral: + try: + with self._lock: + learned_key = self._mistral_api_key_override + aliases = self._fetcher.fetch_mistral_aliases(learned_key) + with self._lock: + self._mistral_aliases = aliases + self._mistral_fetched = time.time() + self._mistral_stale = False + except Exception as exc: # noqa: BLE001 + self._report(exc, "pricing.fetch_mistral_aliases") + finally: + with self._lock: + self._refreshing.discard("mistral_aliases") + for r in regions: try: table = self._fetcher.fetch_bedrock(r) diff --git a/src/lago_agent_sdk/queue.py b/src/lago_agent_sdk/queue.py index ac62274..18c59d0 100644 --- a/src/lago_agent_sdk/queue.py +++ b/src/lago_agent_sdk/queue.py @@ -1,9 +1,21 @@ """Async batched event queue. Thread-safe, in-memory. Background thread flushes every `flush_interval` -seconds or immediately when buffer reaches `max_batch_size`. On send -failure, re-prepends the batch and applies exponential backoff -(1s, 2s, 4s, 8s, capped at 60s). Resets on next success. +seconds or immediately when buffer reaches `max_batch_size`. On a TRANSIENT +send failure (network error, 5xx), re-prepends the batch and applies +exponential backoff (1s, 2s, 4s, 8s, capped at 60s). Resets on next success. + +A PERMANENT failure (a Lago *validation* 4xx — e.g. a duplicate +`transaction_id` from replaying/backfilling the same window twice) is +different: retrying it will never succeed, so it is logged and dropped instead +of re-queued. Without this distinction, one permanently-doomed batch sits at +the front of the FIFO buffer and blocks every event queued behind it — +including brand new, perfectly valid ones — for the full backoff ceiling, over +and over, since a batch that can never succeed is retried exactly like one +that might. + +Note that "permanent" is a specific list of statuses, NOT the whole 4xx range — +see `_PERMANENT_STATUSES`. """ from __future__ import annotations @@ -17,6 +29,42 @@ from collections.abc import Callable from typing import Any +from .exceptions import LagoApiError + +# Statuses where re-sending the SAME batch can never succeed: the request itself +# is the problem (malformed body, bad credentials, a transaction_id Lago has +# already accepted). Deliberately an explicit list rather than the 400-499 range, +# because two 4xx statuses mean "try again, later": 429 (rate limited) and 408 +# (request timeout). Treating those as permanent dropped billable events AND fanned +# one throttled batch out into up to `max_batch_size` extra requests aimed at the +# server that had just asked us to slow down. +# +# 413/402/415 are in the set for the OPPOSITE reason to 429: re-sending the same batch +# provably cannot succeed (too large, payment required, wrong media type), so treating +# them as transient re-prepended the identical batch at the head of the FIFO and backed +# off to 60s forever, blocking every event behind it until the buffer overflowed. Being +# "permanent" here routes them to `_send_individually`, which SPLITS the batch and +# delivers what is deliverable — so a 413 on a 100-event batch becomes 100 single-event +# sends rather than a stalled queue. That is the behaviour we want, and the batch that +# most needs splitting was the one that never reached it. 405/410 stay transient: they +# usually indicate a misrouted or retired endpoint, which a deploy can fix. +_PERMANENT_STATUSES = frozenset({400, 401, 402, 403, 404, 409, 413, 415, 422}) + + +def _is_permanent_failure(exc: Exception) -> bool: + """True when re-sending this exact batch can never succeed. + + A validation 4xx (bad request, duplicate transaction_id, revoked key) will + fail identically forever, so it is isolated and dropped. Everything else — + 5xx, a network-level exception (timeout, connection error, no LagoApiError at + all), and the throttling 4xxs 429/408 — might succeed later and stays + retryable. An unrecognized 4xx is treated as transient: waiting on an event + that would have been dropped costs a delay, dropping one that would have been + accepted costs revenue. + """ + return isinstance(exc, LagoApiError) and exc.status in _PERMANENT_STATUSES + + logger = logging.getLogger("lago_agent_sdk.queue") @@ -47,6 +95,8 @@ def __init__( self._stopping = threading.Event() self._backoff_seconds = 0.0 self._http_calls = 0 # for tests + # Per-thread "already reporting an overflow" flag — see push(). + self._reporting = threading.local() self._thread = threading.Thread(target=self._run, name="lago-queue", daemon=True) self._thread.start() @@ -66,6 +116,7 @@ def _after_in_child(self) -> None: self._buffer = deque() # don't replay parent's events from the child self._backoff_seconds = 0.0 self._http_calls = 0 + self._reporting = threading.local() # Note: the PricingProvider self-heals on fork via a PID check inside # lookup()/maybe_refresh(); we deliberately do NOT call into it from this # fork handler (touching it here changes thread timing enough to trip @@ -73,13 +124,55 @@ def _after_in_child(self) -> None: self._thread = threading.Thread(target=self._run, name="lago-queue", daemon=True) self._thread.start() + def wake(self) -> None: + """Nudge the background thread to run its tick (drain + pricing + `maybe_refresh()`) right now instead of waiting up to + `flush_interval` seconds for its next scheduled tick. Just sets an + in-memory flag — never blocks, never does I/O on the caller's + thread.""" + self._wake.set() + def push(self, event: dict[str, Any]) -> None: with self._lock: - if len(self._buffer) >= self._max_buffer_size: + overflowed = len(self._buffer) >= self._max_buffer_size + if overflowed: self._buffer.popleft() - logger.warning("lago queue overflow at %d events; dropping oldest", self._max_buffer_size) self._buffer.append(event) should_wake = len(self._buffer) >= self._max_batch_size + # Both of these run with the lock RELEASED, the same shape `should_wake` + # already used. Reporting inside the lock deadlocked the caller: `_lock` is a + # plain Lock, not an RLock, so a customer `on_error` that touched the SDK at + # all — emitting a diagnostic, forcing a flush — blocked forever on a lock its + # own thread already held. Overflow happens under sustained load, which is + # exactly when such a hook fires, and the failure is worse than the drop it + # reports: an unnoticed dropped event costs one event, a hung producer thread + # costs the application. The surrounding try/except cannot help, because a + # deadlock is not an exception. + # + # Keeping the callback out of the lock also stops a full buffer from running + # the hook plus a log write synchronously on the customer's LLM-call thread + # while holding the lock every producer and the drain thread need. + if overflowed and not getattr(self._reporting, "active", False): + # Re-entrancy guard, per thread. Moving the report out of the lock fixed + # the deadlock but exposed the other half: the buffer is full again by the + # time the hook runs, so a hook that calls `push()` overflows again and + # re-enters without bound. Suppressing the nested report breaks the cycle + # while still letting the hook's own event be buffered. Per-thread so one + # producer's hook can never silence another producer's report. + self._reporting.active = True + try: + logger.warning("lago queue overflow at %d events; dropping oldest", self._max_buffer_size) + # Also through on_error: an overflow drops BILLABLE events, and a + # customer watching only the error hook — the documented channel for + # every other billing gap — never learned revenue had been lost. + self._report_error( + RuntimeError( + f"queue overflow at {self._max_buffer_size} events; dropped the oldest event" + ), + "overflow", + ) + finally: + self._reporting.active = False if should_wake: self._wake.set() @@ -118,6 +211,51 @@ def _replay_failed(self, batch: list[dict[str, Any]]) -> None: with self._lock: self._buffer.extendleft(reversed(batch)) + def _report_error(self, exc: Exception, where: str = "send_batch") -> None: + """Best-effort `on_error` callback — a customer's own callback must + never be allowed to break the queue's send/retry loop.""" + if self._on_error: + try: + self._on_error(exc, where) + except Exception: # noqa: BLE001 + pass + + def _send_individually(self, batch: list[dict[str, Any]], batch_exc: Exception) -> None: + """Recovery path for a batch that failed with a permanent (4xx) error. + + Each event is sent alone: one that individually 4xxs (e.g. its own + transaction_id really is a duplicate) is logged and dropped for good; + one that succeeds alone is done; one that hits a TRANSIENT error while + isolated is re-queued for the normal backoff-and-retry path, same as + any other event. Reports once via on_error for the batch as a whole + (the original exception) so a caller isn't flooded with N callbacks + for what's really one root cause. + """ + self._report_error(batch_exc) + # Collected and re-queued ONCE at the end, not per event. `_replay_failed` + # prepends, so calling it inside the loop reversed the survivors' relative + # order: a 413 batch of a,b,c,d,e whose b,c,d fail transiently while isolated + # came back as d,c,b. FIFO is the queue's contract — it is what makes the + # oldest-dropped-first overflow policy and Lago's own event ordering + # meaningful — so a recovery path must not silently invert it. + retry: list[dict[str, Any]] = [] + for event in batch: + try: + self._http_calls += 1 + self._sender([event]) + except Exception as exc: # noqa: BLE001 + if _is_permanent_failure(exc): + logger.warning( + "lago dropping event (permanent failure, will not retry): transaction_id=%s: %s", + event.get("transaction_id"), + exc, + ) + else: + logger.warning("lago send failed for isolated event, will retry: %s", exc) + retry.append(event) + if retry: + self._replay_failed(retry) + def _run(self) -> None: while not self._stopping.is_set(): self._wake.wait(timeout=self._flush_interval) @@ -143,12 +281,19 @@ def _run(self) -> None: self._sender(batch) self._backoff_seconds = 0.0 except Exception as exc: # noqa: BLE001 + if _is_permanent_failure(exc): + # Lago's batch endpoint is all-or-nothing: a single bad + # transaction_id fails the WHOLE batch, even if the rest + # are perfectly valid — re-queuing the batch as-is would + # retry (and re-fail) forever, but dropping it outright + # would silently lose those valid events too. Isolate by + # falling back to one-by-one for this batch only; only + # the events that individually 4xx get dropped. + self._send_individually(batch, exc) + self._backoff_seconds = 0.0 + continue self._replay_failed(batch) - if self._on_error: - try: - self._on_error(exc, "send_batch") - except Exception: # noqa: BLE001 - pass + self._report_error(exc) logger.warning("lago send_batch failed: %s", exc) self._backoff_seconds = ( 1.0 @@ -156,10 +301,39 @@ def _run(self) -> None: else min(self._backoff_seconds * 2, self._max_retry_seconds) ) break - # drain on exit - batch = self._take_batch() - if batch: + # Drain on exit — keep sending until the buffer is truly empty, not + # just one batch's worth (a buffer holding more than max_batch_size + # events at shutdown previously left the rest never even attempted). + # No more retries are possible once this thread exits, so unlike the + # main loop, a transient failure here is ALSO final: it must be + # logged, never silently swallowed the way a bare `except: pass` + # previously did — that's what actually lost events, not the network + # blip itself, which by itself is recoverable if it's just reported. + # `_send_individually` re-queues transient sub-failures for retry — + # appropriate for the main loop, which lives on, but during this exit + # drain that could spin forever against a persistently-down network. + # Bound the whole drain by wall-clock time; whatever's still in the + # buffer once the budget is spent is logged as lost, not retried + # forever in an exiting daemon thread. + drain_deadline = time.monotonic() + min(self._max_retry_seconds, 10.0) + while time.monotonic() < drain_deadline: + batch = self._take_batch() + if not batch: + break try: self._sender(batch) - except Exception: # noqa: BLE001 - pass + except Exception as exc: # noqa: BLE001 + if _is_permanent_failure(exc): + self._send_individually(batch, exc) + else: + self._report_error(exc) + logger.warning( + "lago: %d event(s) LOST on shutdown — final drain failed with no more " + "retries possible: %s", + len(batch), + exc, + ) + with self._lock: + stranded = len(self._buffer) + if stranded: + logger.warning("lago: %d event(s) LOST on shutdown — drain time budget exhausted", stranded) diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index dfe8d06..c9edfbd 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -6,6 +6,7 @@ import logging import time import uuid +from collections.abc import Iterable from typing import Any from .canonical import CanonicalUsage @@ -13,7 +14,16 @@ from .detector import detect_client_kind from .exceptions import PricingUnavailableError, UnknownClientError from .lago_client import LagoClient -from .pricing import PricingProvider, coerce_markup, compute_cost +from .pricing import ( + CostBreakdown, + PricingProvider, + apply_markup, + coerce_markup, + compute_cost, + compute_precomputed_cost, + deoverlapped_token_total, + money_str_to_cents, +) from .queue import EventQueue logger = logging.getLogger("lago_agent_sdk") @@ -27,26 +37,51 @@ class LagoSDK: def __init__( self, api_key: str, - api_url: str = "https://api.getlago.com/api/v1", + api_url: str | None = None, default_subscription_id: str | None = None, config: LagoConfig | None = None, + verify_ssl: bool | None = None, ) -> None: - self.config = config or LagoConfig( - api_key=api_key, - api_url=api_url, - default_subscription_id=default_subscription_id, - ) - # explicit args win over `config` + """Explicit args win over anything set on ``config``; ``config`` supplies + every field they don't mention. + + ``api_url`` defaults to None, NOT to the production URL. That distinction + is load-bearing: with a truthy default, ``if api_url:`` always fired and + overwrote ``config.api_url``, so + ``LagoSDK(api_key=k, config=LagoConfig(api_url="http://localhost:3000/api/v1"))`` + silently sent every event to PRODUCTION Lago. That is the shortest path to + the bug, too — a custom ``api_url`` and ``verify_ssl=False`` go together in + exactly the local-dev-Lago setup ``verify_ssl`` exists to serve. + + ``verify_ssl`` is accepted directly so that setup needs no ``LagoConfig`` + at all: a local instance behind a self-signed cert (Traefik's default) is + reachable with ``LagoSDK(api_key=..., api_url=..., verify_ssl=False)``. + """ + self.config = config or LagoConfig(api_key=api_key) + # explicit args win over `config` — guarded on "was it actually passed?" + # rather than on truthiness, so a config value survives when it wasn't. self.config.api_key = api_key or self.config.api_key + # `api_url` is the one exception: an EMPTY string must not win either. The bug + # this guard was written for was a *truthy default* overwriting config, so + # accepting "" swapped one silent misroute for a worse one — + # `api_url=os.environ.get("LAGO_API_URL", "")` with the var unset used to keep + # the production URL and instead wrote "". Downstream that is unrecoverable: + # `requests` raises MissingSchema, which is not a LagoApiError, so the queue + # classifies it transient, re-prepends the batch and retries at the 60s ceiling + # forever. All billing stops, nothing is dropped or escalated, and the only + # symptom is a growing buffer. if api_url: self.config.api_url = api_url if default_subscription_id is not None: self.config.default_subscription_id = default_subscription_id + if verify_ssl is not None: + self.config.verify_ssl = verify_ssl self._lago_client = LagoClient( api_key=self.config.api_key, api_url=self.config.api_url, timeout=self.config.request_timeout_seconds, + verify_ssl=self.config.verify_ssl, ) # Pricing provider (price mode). Default does no network until a # price-mode lookup flags a source stale; refreshes run on the queue @@ -55,6 +90,9 @@ def __init__( ttl_seconds=self.config.pricing_ttl_seconds, default_region=self.config.bedrock_default_region, on_error=self.config.on_error, + cloudflare_account_id=self.config.cloudflare_account_id, + cloudflare_api_token=self.config.cloudflare_api_token, + mistral_api_key=self.config.mistral_api_key, ) if self.config.pricing_mode == "price": self._pricing.prime() # eager warm when price mode is the global default @@ -80,6 +118,63 @@ def reset_subscription(self, token: contextvars.Token[str | None]) -> None: def _resolve_subscription(self, override: str | None) -> str | None: return override or _subscription_var.get() or self.config.default_subscription_id + def _auto_prime_pricing_for(self, kind: str, client: Any) -> None: + """Best-effort, automatic, non-blocking warm-up for the two + credential-gated pricing sources — triggered by `wrap()` itself, + which the customer already calls, so there's no separate function to + remember. `wrap()` almost always happens some real time before the + customer's first actual completion call (building the prompt, + setting up messages, etc.), so kicking the fetch off here — instead + of waiting for that first completion call to flag it stale — gives + it a real head start: often enough to be warm before that first call + even lands, not just for every call after it. + + Only runs when `pricing_mode == "price"` is the global default — + otherwise there's nothing to warm for. `prime()`/`wake()` are both + pure in-memory (no I/O on this thread); the actual HTTP fetch still + happens on the queue's background thread, never here. + """ + if self.config.pricing_mode != "price": + return + provider: str | None = None + if kind == "mistral": + # The client being wrapped already carries the exact Mistral API + # key needed to call Mistral's own /v1/models for alias + # resolution — no separate LagoConfig.mistral_api_key required. + key = self._extract_mistral_api_key(client) + if key: + self._pricing.learn_mistral_api_key(key) + provider = "mistral" + elif kind == "openai": + # A generic OpenAI-shaped client can point at real OpenAI OR, via + # Cloudflare's `.../compat` endpoint, at Workers AI — the client + # kind alone can't tell them apart. `base_url` is the one signal + # that can, without waiting for a response to resolve a model + # string. Defensive: some client variants may not expose it. + try: + base_url = str(getattr(client, "base_url", "") or "") + except Exception: # noqa: BLE001 + base_url = "" + if "gateway.ai.cloudflare.com" in base_url: + provider = "workers-ai" + if provider: + self._pricing.prime([provider]) + self._queue.wake() + + @staticmethod + def _extract_mistral_api_key(client: Any) -> str | None: + """The mistralai SDK stores the constructor's `api_key=...` at + `client.sdk_configuration.security.api_key` (verified against a real + client instance). Defensive: an SDK version change to this internal + path degrades to "no key learned" (falls back to + `LagoConfig.mistral_api_key` if set, else the existing lazy-miss + behavior) rather than raising.""" + try: + key = client.sdk_configuration.security.api_key + except Exception: # noqa: BLE001 + return None + return key if isinstance(key, str) and key else None + # ------------------------------------------------------------------ # Wrap() # ------------------------------------------------------------------ @@ -87,6 +182,7 @@ def wrap( self, client: Any, dimensions: dict[str, Any] | None = None, subscription: str | None = None ) -> Any: kind = detect_client_kind(client) + self._auto_prime_pricing_for(kind, client) if kind == "bedrock": from .wrappers.boto3_bedrock import wrap_boto3_bedrock_client @@ -138,6 +234,8 @@ def emit( dimensions: dict[str, Any] | None = None, mode: str | None = None, markup: float | None = None, + usd_cost: float | None = None, + event_id: str | None = None, ) -> None: """Emit usage to Lago. @@ -145,26 +243,77 @@ def emit( In ``price`` mode, pushes a single dollar-cost event; if no price is available it falls back to token events and reports via on_error. Precedence for mode/markup: per-call arg > config default. + + ``usd_cost``: skip this SDK's own OpenRouter/Bedrock price lookup and + bill this exact amount instead. For a gateway that reports its own + real, metered cost per call (e.g. Cloudflare AI Gateway's `cost` + field on a log entry), that number is more accurate than anything we'd + compute ourselves — this is the connector's one-call entrypoint rather + than hand-building a `precise_total_amount_cents` event. Only consulted + when the effective mode is "price"; ignored in token mode. + + ``event_id``: use this as Lago's idempotency key (`transaction_id`) + instead of a random UUID — pass the source log entry's own id when + replaying/backfilling from a gateway's logs, so re-running against the + same window never double-bills. A live, one-shot call has no natural + id to reuse and should leave this as None. + + Both multi-event paths suffix per field so they don't collide with each + other, and they use DIFFERENT namespaces so they can't collide across + modes either: + + * token events ``f"{event_id}_tok_{field_name}"`` + * split cost events ``f"{event_id}_cost_{field_name}"`` + * single cost event ``event_id`` (one event, nothing to disambiguate) + + The namespaces are load-bearing. Both paths are reachable for the SAME + `event_id`: a price lookup that misses falls back to token events, and + the same window re-run once the table is warm takes the cost path. Under + one shared namespace the second run re-sent `{event_id}_input` under a + different metric code, Lago rejected it as a duplicate — and because + `/events/batch` is all-or-nothing, that rejection failed every other + event in the batch too. Net effect: the dollar amounts for that window + were never billed, only the raw token counts, and nothing surfaced it. """ try: sub = self._resolve_subscription(subscription) if not sub: - logger.error( - "lago: dropping events for model=%s — no resolvable subscription", - usage.model, + # `_report_error` is the single channel: it invokes on_error AND + # logs. An extra logger.error here emitted the same drop twice under + # two different levels, so a customer grepping logs counted one lost + # call as two — and the JS port logged nothing at all, so the two + # repos reported 2 lines vs 0 for the same event. + self._report_error( + ValueError( + f"no resolvable subscription for model={usage.model!r}; events dropped. " + f"Pass subscription=..., use with_subscription(), or set " + f"LagoConfig.default_subscription_id." + ), + "emit", ) return effective_mode = mode or self.config.pricing_mode if effective_mode != "price": - self._emit_token_events(usage, sub, dimensions) - return - - price = self._pricing.lookup(usage.provider, usage.model, usage.api) - if price is None: - # Don't silently under-bill: fall back to token events + report. - self._report_error(PricingUnavailableError(usage.provider, usage.model, usage.api), "pricing") - self._emit_token_events(usage, sub, dimensions) + if usd_cost is not None: + # A caller who went to the trouble of supplying a real metered + # cost gets told it was dropped, rather than discovering later + # that a whole backfill billed token counts only. Reported per + # occurrence, deliberately not deduped: the number of discarded + # costs is exactly what a caller reconciling on `on_error` + # needs, and the documented backfill pattern passes an explicit + # `mode="price"`, so reaching this at volume means a real + # misconfiguration rather than normal operation. + self._report_error( + ValueError( + f"usd_cost={usd_cost!r} ignored: effective pricing mode is " + f"{effective_mode!r}, not 'price' — emitting token counts " + f"instead. Pass mode='price' per call, or set " + f"LagoConfig.pricing_mode='price'." + ), + "pricing", + ) + self._emit_token_events(usage, sub, dimensions, event_id) return markup_value, ok = coerce_markup(markup if markup is not None else self.config.markup) @@ -175,12 +324,40 @@ def emit( ), "pricing", ) - self._emit_cost_event(usage, price, markup_value, sub, dimensions) + + if usd_cost is not None: + breakdown = compute_precomputed_cost(usd_cost, markup_value) + else: + price = self._pricing.lookup(usage.provider, usage.model, usage.api) + if price is None: + # Don't silently under-bill: fall back to token events + report. + self._report_error( + PricingUnavailableError(usage.provider, usage.model, usage.api), "pricing" + ) + self._emit_token_events(usage, sub, dimensions, event_id) + return + breakdown = compute_cost(usage, price, markup_value) + + self._push_cost_event(usage, breakdown, sub, dimensions, event_id) except Exception as exc: # noqa: BLE001 — never raise from emit self._report_error(exc, "emit") - def _emit_token_events(self, usage: CanonicalUsage, sub: str, dimensions: dict[str, Any] | None) -> None: + def _emit_token_events( + self, usage: CanonicalUsage, sub: str, dimensions: dict[str, Any] | None, event_id: str | None = None + ) -> None: nonzero = usage.nonzero_numeric() + # A negative count is silently unbillable — Lago would otherwise sum it into + # a negative quantity. It was the only drop path in the SDK that never + # reached on_error, so a caller who built a CanonicalUsage with a bad delta + # saw nothing at all. Reported before the empty-check, because an event whose + # only fields were negative leaves `nonzero` empty and would return below + # without a word. + negatives = usage.negative_numeric() + if negatives: + self._report_error( + ValueError(f"dropped negative token counts for model={usage.model!r}: {negatives}"), + "negative_tokens", + ) if not nonzero: # Mistral legacy / empty — nothing to bill return @@ -190,7 +367,11 @@ def _emit_token_events(self, usage: CanonicalUsage, sub: str, dimensions: dict[s if not code: continue event = { - "transaction_id": str(uuid.uuid4()), + # `_tok_` namespace: the cost path suffixes with the same field + # vocabulary, and both are reachable for one `event_id` (price + # miss -> token fallback, then the cost path once the table is + # warm). See emit()'s docstring for what a shared namespace cost. + "transaction_id": f"{event_id}_tok_{field_name}" if event_id else str(uuid.uuid4()), "external_subscription_id": sub, "code": code, "timestamp": now, @@ -204,49 +385,102 @@ def _emit_token_events(self, usage: CanonicalUsage, sub: str, dimensions: dict[s } self._queue.push(event) - def _emit_cost_event( + def _push_cost_event( self, usage: CanonicalUsage, - price: Any, - markup: Any, + breakdown: CostBreakdown, sub: str, dimensions: dict[str, Any] | None, + event_id: str | None = None, ) -> None: - breakdown = compute_cost(usage, price, markup) - # `unit` = total tokens for the call — the quantity the sum-aggregation - # billable metric sums (the dynamic charge's fee comes from - # precise_total_amount_cents; unit is the displayed usage quantity). - # Sum the *billed* per-field counts from the breakdown, which compute_cost - # has already de-overlapped (e.g. cache_read carved out of input), so - # subset fields aren't double-counted in the displayed total. - unit = sum(int(parts["tokens"]) for parts in breakdown.fields.values()) - properties: dict[str, Any] = { - "unit": str(unit), - "value": breakdown.total, - "base_cost": breakdown.base, - "markup": breakdown.markup, + """Push one llm_cost event — or several, one per token_type, when a + real per-field breakdown exists. + + `breakdown.fields` only exists when we priced via our own per-token + table (`compute_cost`): the live wrap() path, where an OpenRouter/ + Bedrock unit price is available for input/output/cache/reasoning + separately. There, billing is split one event per field, each tagged + `token_type`, so Lago's `grouped_by: ["model", "token_type"]` charge + can break llm_cost down by both dimensions. + + A precomputed breakdown (`usd_cost` — e.g. Cloudflare AI Gateway's own + already-metered `cost` per call) has no such split: the gateway gives + one lump sum, not "$X of this was input tokens" — inventing a + proportional split would substitute our own guess for the real number + we specifically avoided guessing at. That path bills a single event, + grouped by model only; no `token_type` at all rather than a fabricated + one. + """ + now = int(time.time()) + # Caller dimensions are spread LAST in each `properties` below, not here — + # they must win over every SDK-computed key, exactly as they already do in + # `_emit_token_events`. Spreading them into `base_properties` put them + # *before* `unit`/`value`/`base_cost`/`unit_price`, so those four silently + # overwrote a caller's same-named dimension on this path while honouring it + # on the token path — one customer config, two different outcomes. + base_properties: dict[str, Any] = { "model": usage.model, "provider": usage.provider, "api": usage.api, "price_source": breakdown.source, + "markup": breakdown.markup, } + + if not breakdown.fields: + properties = { + **base_properties, + # Same basis as the split path below (which reports the + # de-overlapped per-field `parts["tokens"]`), so the two branches + # can't report different quantities for one call. `input + output` + # dropped `reasoning` and `cache_write` entirely — on a real + # captured Gemini row with 9 in / 21 out / 852 reasoning it + # published unit="30" for a call that consumed 882 — and counted a + # cache-inclusive provider's cached tokens at full weight. + "unit": str(deoverlapped_token_total(usage)), + "value": breakdown.total, + "base_cost": breakdown.base, + **(dimensions or {}), + } + self._queue.push( + { + # Unsuffixed: this branch pushes exactly ONE event, so there is + # nothing to disambiguate. It cannot collide with the namespaced + # multi-event ids below or in _emit_token_events. + "transaction_id": event_id or str(uuid.uuid4()), + "external_subscription_id": sub, + "code": self.config.cost_metric_code, + "timestamp": now, + "precise_total_amount_cents": breakdown.total_cents, + "properties": properties, + } + ) + return + for field_name, parts in breakdown.fields.items(): - properties[f"{field_name}_tokens"] = parts["tokens"] - properties[f"{field_name}_unit_price"] = parts["unit_price"] - properties[f"{field_name}_cost"] = parts["cost"] - properties.update(dimensions or {}) - self._queue.push( - { - "transaction_id": str(uuid.uuid4()), - "external_subscription_id": sub, - "code": self.config.cost_metric_code, - "timestamp": int(time.time()), - # Top-level amount (in cents) for Lago's dynamic charge model — - # the charge sums these into a single fee. - "precise_total_amount_cents": breakdown.total_cents, - "properties": properties, + # parts["cost"] is PRE-markup (compute_cost only applies markup to + # the summed total) — apply it here or a markup != 1.0 silently + # vanishes from every split event. + billed_cost = apply_markup(parts["cost"], breakdown.markup) + properties = { + **base_properties, + "token_type": field_name, + "unit": parts["tokens"], + "value": billed_cost, + "base_cost": parts["cost"], + "unit_price": parts["unit_price"], + **(dimensions or {}), } - ) + self._queue.push( + { + # `_cost_` namespace — see the `_tok_` note in _emit_token_events. + "transaction_id": f"{event_id}_cost_{field_name}" if event_id else str(uuid.uuid4()), + "external_subscription_id": sub, + "code": self.config.cost_metric_code, + "timestamp": now, + "precise_total_amount_cents": money_str_to_cents(billed_cost), + "properties": properties, + } + ) def _report_error(self, exc: Exception, where: str) -> None: if self.config.on_error: @@ -256,6 +490,40 @@ def _report_error(self, exc: Exception, where: str) -> None: pass logger.warning("lago %s failed: %s", where, exc) + def warm_pricing(self, providers: Iterable[str] = ()) -> None: + """Block until the given price table(s) are fetched, instead of + waiting for the queue's background thread to pick them up on its + next tick (up to `flush_interval` seconds later, by default ~1s). + + A call made immediately after construction — the common shape in a + script, notebook, or one-shot job, as opposed to a long-running server + where the first real call naturally lands well after that first tick + — races a still-cold cache. `emit()` never silently under-bills, so a + miss falls back to token events; but with no token-metric charge + configured at all (a single `llm_cost`-only billing setup), there is + nowhere left to fall back to and the event is lost. Call this once, + right after constructing the SDK with `pricing_mode="price"`, to close + that window deterministically for OpenRouter — the table nearly every + native provider prices against — which is always warmed regardless + of `providers`. + + Cloudflare Workers AI and Mistral alias resolution are NOT warmed by + default: both are credential-gated and provider-specific, and + eagerly hitting either's API at construction time regardless of + whether that provider is ever actually called would be pure waste + for the common case. Left alone, they stay reactive — the first real + call to that provider triggers the fetch, and every call after that + (even the one a moment later) is cached — so only a session's first + Workers AI or Mistral call can race a cold cache. + + If you already know you're about to call one or both this session, + say so and skip that one-time cost too: `providers=["mistral"]` + and/or `["workers-ai"]`. A no-op for any source that isn't stale + (e.g. the SDK isn't in price mode, was already warmed, or the + provider name wasn't recognized).""" + self._pricing.prime(providers) + self._pricing.maybe_refresh() + def flush(self, timeout: float = 5.0) -> bool: return self._queue.flush(timeout=timeout) diff --git a/src/lago_agent_sdk/wrappers/anthropic.py b/src/lago_agent_sdk/wrappers/anthropic.py index ded2252..c864a82 100644 --- a/src/lago_agent_sdk/wrappers/anthropic.py +++ b/src/lago_agent_sdk/wrappers/anthropic.py @@ -9,6 +9,18 @@ - AsyncMessages.create(...) — async non-streaming and stream=True - AsyncMessages.stream(...) — async context-manager helper +Gateway cache-hit detection (non-streaming .create(...) only): + Non-streaming calls go through `.with_raw_response.create(...)` instead of + `.create(...)` so we can see response headers before parsing the body. If a + gateway in front of the provider (e.g. Cloudflare AI Gateway) marks the + response `cf-aig-cache-status: HIT`, the provider served it from cache at zero + cost to the customer — we skip billing it. `.parse()` on the raw response + returns the exact same object `.create()` would, so nothing downstream changes. + This is a no-op when there's no gateway in the path: the header is simply + absent. `.stream()` is NOT covered — Anthropic recommends + `.with_streaming_response` for that, which behaves differently and hasn't been + verified end-to-end. + Per-call override: pop `extra_lago={"subscription": ..., "dimensions": ...}` from kwargs before forwarding so Anthropic's strict validation doesn't reject it. """ @@ -46,7 +58,21 @@ def _is_message_like(obj: Any) -> bool: return False -def _merge_stream_usage(accumulated: dict[str, Any], payload: Any) -> None: +def _is_cache_hit(raw_response: Any) -> bool: + """True if a gateway in front of the provider served this from cache. + + A cache hit (Cloudflare AI Gateway: `cf-aig-cache-status: HIT`) costs the + provider — and the customer — nothing. Billing it would overcharge for a + call that never actually happened. Safe no-op with no gateway in the path: + `.headers.get(...)` simply returns None. + """ + try: + return bool(raw_response.headers.get("cf-aig-cache-status") == "HIT") + except Exception: # noqa: BLE001 + return False + + +def _merge_stream_usage(accumulated: dict[str, Any], payload: Any) -> str | None: """Fold one streaming event's usage into the running accumulator. Anthropic splits authoritative usage across two events: @@ -60,19 +86,30 @@ def _merge_stream_usage(accumulated: dict[str, Any], payload: Any) -> None: would bill ``input_tokens=0``. Merge both locations; ``dict.update`` lets the more complete / more recent values win while preserving the input counts from ``message_start`` when a delta omits them. + + Returns the model this event reported, if any. ``message_start`` carries the + RESOLVED snapshot under ``message.model`` — e.g. "claude-sonnet-4-5-20250929" + for a requested "claude-sonnet-4-5" — and discarding it made every streaming + call attribute (and price) the alias instead, the same bug the non-streaming + path was fixed for. The caller keeps the first one it sees. """ if not isinstance(payload, dict): - return - # message_start: input/cache live under message.usage + return None + found: str | None = None + # message_start: input/cache live under message.usage, resolved model alongside message = payload.get("message") if isinstance(message, dict): nested = message.get("usage") if isinstance(nested, dict): accumulated.update(nested) + model = message.get("model") + if isinstance(model, str) and model: + found = model # message_delta (and others): cumulative usage at the top level top = payload.get("usage") if isinstance(top, dict): accumulated.update(top) + return found def wrap_anthropic_client( @@ -95,6 +132,7 @@ def wrap_anthropic_client( return client original_create = getattr(messages, "create", None) + raw_create = getattr(getattr(messages, "with_raw_response", None), "create", None) original_stream = getattr(messages, "stream", None) is_async = type(client).__name__.startswith("Async") @@ -121,6 +159,18 @@ def _create(*args: Any, **kwargs: Any) -> Any: lago_opts = _pop_lago_kwarg(kwargs) model_id = kwargs.get("model", "") opts = _resolve_opts(lago_opts) + + if not kwargs.get("stream") and raw_create is not None: + # Non-streaming with `.with_raw_response` available: see gateway + # headers before parsing — `.parse()` returns the identical object + # `.create()` would have, so the customer sees no difference. + raw = raw_create(*args, **kwargs) + response = raw.parse() + if _is_message_like(response) and not _is_cache_hit(raw): + _emit_from(response, model_id, opts) + return response + + # Streaming, or `.with_raw_response` unavailable (older/custom client). response = original_create(*args, **kwargs) if _is_message_like(response): @@ -131,14 +181,15 @@ def _create(*args: Any, **kwargs: Any) -> Any: # (input/cache) and message_delta (cumulative output) before emitting. def _wrap_stream(src: Iterator[Any]) -> Iterator[Any]: accumulated: dict[str, Any] = {} + resolved_model: str | None = None try: for event in src: payload = event.model_dump() if hasattr(event, "model_dump") else event - _merge_stream_usage(accumulated, payload) + resolved_model = _merge_stream_usage(accumulated, payload) or resolved_model yield event finally: if accumulated: - _emit_from({"usage": accumulated}, model_id, opts) + _emit_from({"usage": accumulated, "model": resolved_model}, model_id, opts) return _wrap_stream(response) @@ -150,6 +201,14 @@ async def _create_async(*args: Any, **kwargs: Any) -> Any: lago_opts = _pop_lago_kwarg(kwargs) model_id = kwargs.get("model", "") opts = _resolve_opts(lago_opts) + + if not kwargs.get("stream") and raw_create is not None: + raw = await raw_create(*args, **kwargs) + response = raw.parse() + if _is_message_like(response) and not _is_cache_hit(raw): + _emit_from(response, model_id, opts) + return response + response = await original_create(*args, **kwargs) if _is_message_like(response): @@ -158,14 +217,15 @@ async def _create_async(*args: Any, **kwargs: Any) -> Any: async def _wrap_async_stream(src: AsyncIterator[Any]) -> AsyncIterator[Any]: accumulated: dict[str, Any] = {} + resolved_model: str | None = None try: async for event in src: payload = event.model_dump() if hasattr(event, "model_dump") else event - _merge_stream_usage(accumulated, payload) + resolved_model = _merge_stream_usage(accumulated, payload) or resolved_model yield event finally: if accumulated: - _emit_from({"usage": accumulated}, model_id, opts) + _emit_from({"usage": accumulated, "model": resolved_model}, model_id, opts) return _wrap_async_stream(response) diff --git a/src/lago_agent_sdk/wrappers/gemini.py b/src/lago_agent_sdk/wrappers/gemini.py index ccf4144..a620beb 100644 --- a/src/lago_agent_sdk/wrappers/gemini.py +++ b/src/lago_agent_sdk/wrappers/gemini.py @@ -91,11 +91,33 @@ def _stream(*args: Any, **kwargs: Any) -> Iterator[Any]: def _iter() -> Iterator[Any]: last_with_usage: Any = None + resolved_model: str | None = None try: for chunk in src: payload = chunk.model_dump() if hasattr(chunk, "model_dump") else chunk - if isinstance(payload, dict) and payload.get("usage_metadata"): - last_with_usage = {"usage_metadata": payload["usage_metadata"]} + if isinstance(payload, dict): + # `model_version` must PERSIST across chunks, not be read + # off whichever chunk happens to carry usage. Gemini + # hot-swaps "-latest" aliases server-side and announces + # the resolved version on an EARLY chunk, while usage + # arrives on the last one — so reading it from the + # usage-bearing chunk alone found nothing and silently + # reverted to the requested alias. Measured on identical + # input: this emitted "gemini-flash-latest" where the JS + # port, which already persisted it, emitted + # "gemini-2.5-flash-002". Pricing keys off the resolved + # version, so the two ports priced the same call + # differently. Both spellings are accepted because + # `model_dump()` yields snake_case while a raw REST dict + # is camelCase. + mv = payload.get("model_version") or payload.get("modelVersion") + if isinstance(mv, str) and mv: + resolved_model = mv + if payload.get("usage_metadata"): + last_with_usage = { + "usage_metadata": payload["usage_metadata"], + "model_version": resolved_model, + } yield chunk finally: if last_with_usage is not None: @@ -114,11 +136,33 @@ async def _stream_async(*args: Any, **kwargs: Any) -> AsyncIterator[Any]: async def _aiter() -> AsyncIterator[Any]: last_with_usage: Any = None + resolved_model: str | None = None try: async for chunk in src: payload = chunk.model_dump() if hasattr(chunk, "model_dump") else chunk - if isinstance(payload, dict) and payload.get("usage_metadata"): - last_with_usage = {"usage_metadata": payload["usage_metadata"]} + if isinstance(payload, dict): + # `model_version` must PERSIST across chunks, not be read + # off whichever chunk happens to carry usage. Gemini + # hot-swaps "-latest" aliases server-side and announces + # the resolved version on an EARLY chunk, while usage + # arrives on the last one — so reading it from the + # usage-bearing chunk alone found nothing and silently + # reverted to the requested alias. Measured on identical + # input: this emitted "gemini-flash-latest" where the JS + # port, which already persisted it, emitted + # "gemini-2.5-flash-002". Pricing keys off the resolved + # version, so the two ports priced the same call + # differently. Both spellings are accepted because + # `model_dump()` yields snake_case while a raw REST dict + # is camelCase. + mv = payload.get("model_version") or payload.get("modelVersion") + if isinstance(mv, str) and mv: + resolved_model = mv + if payload.get("usage_metadata"): + last_with_usage = { + "usage_metadata": payload["usage_metadata"], + "model_version": resolved_model, + } yield chunk finally: if last_with_usage is not None: diff --git a/src/lago_agent_sdk/wrappers/openai.py b/src/lago_agent_sdk/wrappers/openai.py index 90ccb2f..1796b87 100644 --- a/src/lago_agent_sdk/wrappers/openai.py +++ b/src/lago_agent_sdk/wrappers/openai.py @@ -14,6 +14,18 @@ usage payload we need to bill. Without that flag, OpenAI's stream emits no usage data and the customer gets silent under-billing. +Gateway cache-hit detection (non-streaming only): + Non-streaming calls go through `.with_raw_response.create(...)` instead of + `.create(...)` so we can see response headers before parsing the body. If a + gateway in front of the provider (e.g. Cloudflare AI Gateway) marks the + response `cf-aig-cache-status: HIT`, the provider served it from cache at zero + cost to the customer — we skip billing it. `.parse()` on the raw response + returns the exact same object `.create()` would, so nothing downstream changes. + This is a no-op when there's no gateway in the path: the header is simply + absent. Streaming calls are NOT covered — OpenAI recommends + `.with_streaming_response` for that, which behaves differently and hasn't been + verified end-to-end, so streaming keeps using the plain `.create()` path. + Per-call override: pop `extra_lago={"subscription": ..., "dimensions": ...}` from kwargs before forwarding so OpenAI's strict validation doesn't reject it. """ @@ -68,6 +80,20 @@ def _is_response_like(obj: Any) -> bool: return False +def _is_cache_hit(raw_response: Any) -> bool: + """True if a gateway in front of the provider served this from cache. + + A cache hit (Cloudflare AI Gateway: `cf-aig-cache-status: HIT`) costs the + provider — and the customer — nothing. Billing it would overcharge for a + call that never actually happened. Safe no-op with no gateway in the path: + `.headers.get(...)` simply returns None. + """ + try: + return bool(raw_response.headers.get("cf-aig-cache-status") == "HIT") + except Exception: # noqa: BLE001 + return False + + def wrap_openai_client( sdk: Any, client: Any, @@ -106,21 +132,31 @@ def _extract_stream_usage(payload: Any) -> dict[str, Any] | None: Responses API: usage sits under `event.response.usage` on the terminal `response.completed` event (`{"type": "response.completed", "response": {"usage": {...}}}`). + + Carries the chunk's own `model` through alongside the usage. Rebuilding a + usage-ONLY payload made `resolve_model` fall back to the requested alias + on every streaming call, which is precisely the attribution bug the + non-streaming path was fixed for: a streamed `gpt-5-chat-latest` stayed + `gpt-5-chat-latest` instead of resolving to the dated snapshot OpenRouter + lists, so price mode missed and silently degraded to token events. It + matters most on a gateway, where the resolved name is what decides which + price table the call is even looked up in. """ if not isinstance(payload, dict): return None usage = payload.get("usage") if isinstance(usage, dict) and usage: - return {"usage": usage} - # Responses API stream events nest usage under `.response.usage` + return {"usage": usage, "model": payload.get("model")} + # Responses API stream events nest usage under `.response.usage` — and the + # resolved model under `.response.model`, not at the event's top level. response = payload.get("response") if isinstance(response, dict): nested = response.get("usage") if isinstance(nested, dict) and nested: - return {"usage": nested} + return {"usage": nested, "model": response.get("model")} return None - def _make_sync_create(original: Any, is_responses_api: bool = False) -> Any: + def _make_sync_create(original: Any, raw_create: Any | None, is_responses_api: bool = False) -> Any: def _create(*args: Any, **kwargs: Any) -> Any: lago_opts = _pop_lago_kwarg(kwargs) # `stream_options.include_usage` is a Chat-Completions-only knob. @@ -129,13 +165,25 @@ def _create(*args: Any, **kwargs: Any) -> Any: _ensure_stream_options_include_usage(kwargs) model_id = kwargs.get("model", "") opts = _resolve_opts(lago_opts) + + if not kwargs.get("stream") and raw_create is not None: + # Non-streaming with `.with_raw_response` available: see gateway + # headers before parsing — `.parse()` returns the identical object + # `.create()` would have, so the customer sees no difference. + raw = raw_create(*args, **kwargs) + response = raw.parse() + if _is_response_like(response) and not _is_cache_hit(raw): + _emit_from(response, model_id, opts) + return response + + # Streaming, or `.with_raw_response` unavailable (older/custom client) + # — plain `.create()`. No cache-hit detection possible on this path. response = original(*args, **kwargs) if _is_response_like(response): _emit_from(response, model_id, opts) return response - # Streaming — wrap the iterator to capture the final usage on close. def _wrap_stream(src: Iterator[Any]) -> Iterator[Any]: last_usage: dict[str, Any] | None = None try: @@ -153,13 +201,21 @@ def _wrap_stream(src: Iterator[Any]) -> Iterator[Any]: return _create - def _make_async_create(original: Any, is_responses_api: bool = False) -> Any: + def _make_async_create(original: Any, raw_create: Any | None, is_responses_api: bool = False) -> Any: async def _create_async(*args: Any, **kwargs: Any) -> Any: lago_opts = _pop_lago_kwarg(kwargs) if not is_responses_api: _ensure_stream_options_include_usage(kwargs) model_id = kwargs.get("model", "") opts = _resolve_opts(lago_opts) + + if not kwargs.get("stream") and raw_create is not None: + raw = await raw_create(*args, **kwargs) + response = raw.parse() + if _is_response_like(response) and not _is_cache_hit(raw): + _emit_from(response, model_id, opts) + return response + response = await original(*args, **kwargs) if _is_response_like(response): @@ -190,11 +246,12 @@ async def _wrap_async_stream(src: AsyncIterator[Any]) -> AsyncIterator[Any]: completions = getattr(chat, "completions", None) if chat is not None else None if completions is not None: original_chat_create = getattr(completions, "create", None) + raw_chat_create = getattr(getattr(completions, "with_raw_response", None), "create", None) if original_chat_create is not None: completions.create = ( - _make_async_create(original_chat_create, is_responses_api=False) + _make_async_create(original_chat_create, raw_chat_create, is_responses_api=False) if is_async - else _make_sync_create(original_chat_create, is_responses_api=False) + else _make_sync_create(original_chat_create, raw_chat_create, is_responses_api=False) ) # ------------------------------------------------------------------ @@ -203,11 +260,14 @@ async def _wrap_async_stream(src: AsyncIterator[Any]) -> AsyncIterator[Any]: responses_namespace = getattr(client, "responses", None) if responses_namespace is not None: original_responses_create = getattr(responses_namespace, "create", None) + raw_responses_create = getattr( + getattr(responses_namespace, "with_raw_response", None), "create", None + ) if original_responses_create is not None: responses_namespace.create = ( - _make_async_create(original_responses_create, is_responses_api=True) + _make_async_create(original_responses_create, raw_responses_create, is_responses_api=True) if is_async - else _make_sync_create(original_responses_create, is_responses_api=True) + else _make_sync_create(original_responses_create, raw_responses_create, is_responses_api=True) ) setattr(client, _INSTRUMENTED_ATTR, True) diff --git a/tests/integration/test_lago_reconciliation.py b/tests/integration/test_lago_reconciliation.py deleted file mode 100644 index 46edb09..0000000 --- a/tests/integration/test_lago_reconciliation.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Live Lago reconciliation — emit N events, poll current_usage, verify exact match. - -Skipped unless LAGO_API_URL, LAGO_API_KEY, and LAGO_EXTERNAL_SUBSCRIPTION_ID -are set. Requires `truststore` if Lago is on a self-signed dev cert. -""" - -from __future__ import annotations - -import os -import time - -import pytest -import requests - -from lago_agent_sdk import CanonicalUsage, LagoSDK - -try: - import truststore - - truststore.inject_into_ssl() -except Exception: # noqa: BLE001 - pass - -API_URL = (os.environ.get("LAGO_API_URL") or "").rstrip("/") -API_KEY = os.environ.get("LAGO_API_KEY") or "" -SUB_ID = os.environ.get("LAGO_EXTERNAL_SUBSCRIPTION_ID") or "" -CUST_ID = os.environ.get("LAGO_EXTERNAL_CUSTOMER_ID") or "cust_demo" - -pytestmark = pytest.mark.skipif( - not (API_URL and API_KEY and SUB_ID), - reason="LAGO_API_URL / LAGO_API_KEY / LAGO_EXTERNAL_SUBSCRIPTION_ID not set", -) - - -def _read_usage() -> dict[str, float]: - r = requests.get( - f"{API_URL}/customers/{CUST_ID}/current_usage", - params={"external_subscription_id": SUB_ID}, - headers={"Authorization": f"Bearer {API_KEY}"}, - timeout=15, - ) - r.raise_for_status() - out: dict[str, float] = {} - for c in r.json().get("customer_usage", {}).get("charges_usage", []) or []: - code = c.get("billable_metric", {}).get("code", "") - out[code] = float(c.get("units", 0) or 0) - return out - - -def test_emit_then_reconcile_with_live_lago(): - """Send 5 known-shape events; assert input/output totals incremented correctly.""" - sdk = LagoSDK(api_key=API_KEY, api_url=API_URL, default_subscription_id=SUB_ID) - - before = _read_usage() - in_before = before.get("llm_input_tokens", 0.0) - out_before = before.get("llm_output_tokens", 0.0) - - # Emit 5 events with stable values for arithmetic - for _ in range(5): - sdk.emit( - CanonicalUsage( - input=100, - output=200, - model="claude-sonnet-4-6", - provider="anthropic", - api="bedrock_invoke", - ) - ) - - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=3.0) - - # Lago is async — poll for up to 30s - deadline = time.time() + 30 - after = before - while time.time() < deadline: - after = _read_usage() - in_delta = after.get("llm_input_tokens", 0.0) - in_before - out_delta = after.get("llm_output_tokens", 0.0) - out_before - if in_delta >= 500 and out_delta >= 1000: - break - time.sleep(1.0) - - in_delta = after.get("llm_input_tokens", 0.0) - in_before - out_delta = after.get("llm_output_tokens", 0.0) - out_before - assert in_delta == 500, f"input delta {in_delta} != 500 — events lost or duplicated" - assert out_delta == 1000, f"output delta {out_delta} != 1000 — events lost or duplicated" diff --git a/tests/integration/test_live_anthropic.py b/tests/integration/test_live_anthropic.py deleted file mode 100644 index 73c4e35..0000000 --- a/tests/integration/test_live_anthropic.py +++ /dev/null @@ -1,146 +0,0 @@ -"""End-to-end Anthropic integration test — live API + mocked Lago. - -Skipped unless ANTHROPIC_API_KEY is set. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest - -from lago_agent_sdk import LagoSDK - -pytestmark = pytest.mark.skipif( - not os.environ.get("ANTHROPIC_API_KEY"), - reason="ANTHROPIC_API_KEY not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def test_live_anthropic_messages_create_emits_to_lago() -> None: - from anthropic import Anthropic - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])) - client.messages.create( - model="claude-haiku-4-5-20251001", - max_tokens=20, - messages=[{"role": "user", "content": "Say hi"}], - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "native" - assert e["properties"]["provider"] == "anthropic" - finally: - server.shutdown() - - -def test_live_anthropic_streaming_emits_from_final_delta() -> None: - from anthropic import Anthropic - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])) - for _ in client.messages.create( - model="claude-haiku-4-5-20251001", - max_tokens=20, - messages=[{"role": "user", "content": "Say hi"}], - stream=True, - ): - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -def test_live_anthropic_messages_stream_context_manager() -> None: - from anthropic import Anthropic - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])) - with client.messages.stream( - model="claude-haiku-4-5-20251001", - max_tokens=20, - messages=[{"role": "user", "content": "Say hi"}], - ) as stream: - for _ in stream.text_stream: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -@pytest.mark.asyncio -async def test_live_async_anthropic_messages_stream_context_manager_emits() -> None: - """Live regression test for the async messages.stream(...) context manager. - - Bug: __aexit__ called the sync _emit_final, which invoked - get_final_message() without await. On AsyncMessageStream that method is - a coroutine, so the un-awaited object fell through to the adapter as {} - → zero usage emitted, plus a "coroutine was never awaited" RuntimeWarning. - """ - from anthropic import AsyncAnthropic - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])) - async with client.messages.stream( - model="claude-haiku-4-5-20251001", - max_tokens=20, - messages=[{"role": "user", "content": "Say hi"}], - ) as stream: - async for _ in stream.text_stream: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() diff --git a/tests/integration/test_live_bedrock.py b/tests/integration/test_live_bedrock.py deleted file mode 100644 index e972324..0000000 --- a/tests/integration/test_live_bedrock.py +++ /dev/null @@ -1,91 +0,0 @@ -"""End-to-end integration test — live Bedrock REST + mocked Lago endpoint. - -Skipped unless `AWS_BEARER_TOKEN_BEDROCK` is set. Mocks Lago so no real -events are sent. Verifies that wrapping the bearer-token REST flow -produces correctly-shaped events at the Lago HTTP boundary. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest -import requests - -from lago_agent_sdk import LagoSDK -from lago_agent_sdk.adapters import extract_bedrock_converse - -REGION = "eu-west-1" -PROMPT = "One sentence about dolphins." - -pytestmark = pytest.mark.skipif( - not os.environ.get("AWS_BEARER_TOKEN_BEDROCK"), - reason="AWS_BEARER_TOKEN_BEDROCK not set — skipping live Bedrock integration", -) - - -class _MockLagoHandler(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - length = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(length).decode("utf-8") - self.server.received_payloads.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): # silence - return - - -def _start_mock_lago() -> tuple[HTTPServer, str]: - server = HTTPServer(("127.0.0.1", 0), _MockLagoHandler) - server.received_payloads = [] # type: ignore[attr-defined] - t = threading.Thread(target=server.serve_forever, daemon=True) - t.start() - return server, f"http://127.0.0.1:{server.server_port}" - - -def _bearer_call_converse(api_key: str, model_id: str) -> dict: - url = f"https://bedrock-runtime.{REGION}.amazonaws.com/model/{model_id}/converse" - body = { - "messages": [{"role": "user", "content": [{"text": PROMPT}]}], - "inferenceConfig": {"maxTokens": 50}, - } - r = requests.post( - url, - headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, - json=body, - timeout=60, - ) - r.raise_for_status() - return r.json() - - -def test_live_converse_to_mocked_lago(): - api_key = os.environ["AWS_BEARER_TOKEN_BEDROCK"] - server, base_url = _start_mock_lago() - try: - sdk = LagoSDK(api_key="lago_dummy", api_url=base_url, default_subscription_id="sub_int") - model_id = "eu.amazon.nova-lite-v1:0" - # Use the bearer-token REST surface (works without IAM creds in env) - resp = _bearer_call_converse(api_key, model_id) - usage = extract_bedrock_converse(resp, model_id=model_id) - sdk.emit(usage) - assert sdk.flush(timeout=5.0) - sdk.shutdown(timeout=2.0) - - assert len(server.received_payloads) >= 1 # type: ignore[attr-defined] - events = [e for p in server.received_payloads for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["external_subscription_id"] == "sub_int" - assert e["properties"]["api"] == "bedrock_converse" - assert e["properties"]["provider"] == "amazon" - finally: - server.shutdown() diff --git a/tests/integration/test_live_gemini.py b/tests/integration/test_live_gemini.py deleted file mode 100644 index 4ac5de6..0000000 --- a/tests/integration/test_live_gemini.py +++ /dev/null @@ -1,154 +0,0 @@ -"""End-to-end Gemini integration test — live API + mocked Lago. - -Skipped unless GEMINI_API_KEY is set. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest - -from lago_agent_sdk import LagoSDK - -pytestmark = pytest.mark.skipif( - not os.environ.get("GEMINI_API_KEY"), - reason="GEMINI_API_KEY not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def _collect_events(server) -> list[dict]: - return [e for p in server.received for e in p["events"]] - - -def _codes(events) -> set[str]: - return {e["code"] for e in events} - - -def test_live_gemini_generate_content_emits_to_lago() -> None: - from google import genai - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(genai.Client(api_key=os.environ["GEMINI_API_KEY"])) - client.models.generate_content( - model="gemini-2.5-flash", - contents="Say hi", - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "native" - assert e["properties"]["provider"] == "gemini" - finally: - server.shutdown() - - -def test_live_gemini_streaming_captures_usage_from_final_chunk() -> None: - from google import genai - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(genai.Client(api_key=os.environ["GEMINI_API_KEY"])) - for _ in client.models.generate_content_stream( - model="gemini-2.5-flash", - contents="Count from 1 to 3.", - ): - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -def test_live_gemini_thinking_emits_reasoning() -> None: - """Gemini 2.5 emits thoughts_token_count → llm_reasoning_tokens event.""" - from google import genai - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(genai.Client(api_key=os.environ["GEMINI_API_KEY"])) - client.models.generate_content( - model="gemini-2.5-flash", - contents="What is 17 * 23? Show your reasoning step by step.", - ) - assert sdk.flush(timeout=15.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - # Gemini 2.5 reasons even without explicit thinking_config - assert "llm_reasoning_tokens" in codes - finally: - server.shutdown() - - -def test_live_gemini_tool_use_emits_tool_calls() -> None: - from google import genai - from google.genai import types as genai_types - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(genai.Client(api_key=os.environ["GEMINI_API_KEY"])) - weather_fn = genai_types.FunctionDeclaration( - name="get_weather", - description="Get the current weather for a city.", - parameters=genai_types.Schema( - type="OBJECT", - properties={"city": genai_types.Schema(type="STRING")}, - required=["city"], - ), - ) - client.models.generate_content( - model="gemini-2.5-flash", - contents="What's the weather in Tokyo?", - config=genai_types.GenerateContentConfig( - tools=[genai_types.Tool(function_declarations=[weather_fn])], - tool_config=genai_types.ToolConfig( - function_calling_config=genai_types.FunctionCallingConfig(mode="ANY"), - ), - ), - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - assert "llm_tool_calls" in _codes(events) - finally: - server.shutdown() diff --git a/tests/integration/test_live_mistral.py b/tests/integration/test_live_mistral.py deleted file mode 100644 index 72fe916..0000000 --- a/tests/integration/test_live_mistral.py +++ /dev/null @@ -1,120 +0,0 @@ -"""End-to-end Mistral integration test — live API + mocked Lago. - -Skipped unless MISTRAL_API_KEY is set. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest - -from lago_agent_sdk import LagoSDK - -pytestmark = pytest.mark.skipif( - not os.environ.get("MISTRAL_API_KEY"), - reason="MISTRAL_API_KEY not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def test_live_mistral_chat_complete_emits_to_lago(): - from mistralai.client import Mistral - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Mistral(api_key=os.environ["MISTRAL_API_KEY"])) - client.chat.complete( - model="mistral-small-latest", - messages=[{"role": "user", "content": "Say hi"}], - max_tokens=20, - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "native" - assert e["properties"]["provider"] == "mistral" - finally: - server.shutdown() - - -def test_live_mistral_chat_stream_emits_to_lago(): - from mistralai.client import Mistral - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Mistral(api_key=os.environ["MISTRAL_API_KEY"])) - for _ in client.chat.stream( - model="mistral-small-latest", - messages=[{"role": "user", "content": "Say hi"}], - max_tokens=20, - ): - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -@pytest.mark.asyncio -async def test_live_mistral_chat_stream_async_emits_to_lago() -> None: - """Live regression test for chat.stream_async. - - Bug: the wrapper iterated `original_stream_async(*args, **kwargs)` without - awaiting it. In mistralai v2 this method is `async def`, so calling it - returns a coroutine — `async for` raises "got coroutine" TypeError. - """ - from mistralai.client import Mistral - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Mistral(api_key=os.environ["MISTRAL_API_KEY"])) - stream = await client.chat.stream_async( - model="mistral-small-latest", - messages=[{"role": "user", "content": "Say hi"}], - max_tokens=20, - ) - async for _ in stream: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() diff --git a/tests/integration/test_live_openai.py b/tests/integration/test_live_openai.py deleted file mode 100644 index c79f345..0000000 --- a/tests/integration/test_live_openai.py +++ /dev/null @@ -1,225 +0,0 @@ -"""End-to-end OpenAI integration test — live API + mocked Lago. - -Skipped unless OPENAI_API_KEY is set. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest - -from lago_agent_sdk import LagoSDK - -pytestmark = pytest.mark.skipif( - not os.environ.get("OPENAI_API_KEY"), - reason="OPENAI_API_KEY not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def _collect_events(server) -> list[dict]: - return [e for p in server.received for e in p["events"]] - - -def _codes(events) -> set[str]: - return {e["code"] for e in events} - - -# -------------------------------------------------------------------------- -# Chat Completions -# -------------------------------------------------------------------------- -def test_live_openai_chat_completions_create_emits_to_lago() -> None: - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - client.chat.completions.create( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Say hi"}], - max_completion_tokens=20, - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "chat_completions" - assert e["properties"]["provider"] == "openai" - finally: - server.shutdown() - - -def test_live_openai_chat_completions_streaming_emits_from_final_chunk() -> None: - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - # Note: stream_options.include_usage is auto-injected by the wrapper - for _ in client.chat.completions.create( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Say hi"}], - max_completion_tokens=20, - stream=True, - ): - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -def test_live_openai_chat_completions_tool_use_emits_tool_calls() -> None: - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - client.chat.completions.create( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "What's the weather in Tokyo?"}], - tools=[ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather for a city.", - "parameters": { - "type": "object", - "properties": {"city": {"type": "string"}}, - "required": ["city"], - }, - }, - } - ], - tool_choice={"type": "function", "function": {"name": "get_weather"}}, - max_completion_tokens=200, - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - assert "llm_tool_calls" in _codes(events) - finally: - server.shutdown() - - -def test_live_openai_reasoning_model_emits_reasoning_tokens() -> None: - """o-series models populate completion_tokens_details.reasoning_tokens. - First provider to actually expose this metric.""" - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - client.chat.completions.create( - model="o4-mini", - messages=[{"role": "user", "content": "What is 17 * 23? Just the number."}], - max_completion_tokens=2000, - ) - assert sdk.flush(timeout=30.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - assert "llm_reasoning_tokens" in codes # ← the key win for OpenAI - finally: - server.shutdown() - - -# -------------------------------------------------------------------------- -# Responses API -# -------------------------------------------------------------------------- -def test_live_openai_responses_create_emits_to_lago() -> None: - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - client.responses.create( - model="gpt-4o-mini", - input="Say hi", - max_output_tokens=20, - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "responses" - assert e["properties"]["provider"] == "openai" - finally: - server.shutdown() - - -def test_live_openai_responses_create_with_stream_emits_to_lago() -> None: - """Live regression test for two bugs in the Responses API streaming path: - - 1. The wrapper must NOT inject `stream_options.include_usage` — Responses - rejects that param and the call would fail with HTTP 400. - 2. The wrapper must extract usage from `event.response.usage` on the - terminal `response.completed` event (not from a top-level `event.usage`). - """ - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - stream = client.responses.create( - model="gpt-4o-mini", - input="Say hi", - max_output_tokens=20, - stream=True, - ) - # Drain — also verifies the customer's call wasn't broken by injection. - for _ in stream: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "responses" - finally: - server.shutdown() diff --git a/tests/integration/test_live_pricing.py b/tests/integration/test_live_pricing.py deleted file mode 100644 index c03b104..0000000 --- a/tests/integration/test_live_pricing.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Live pricing test — hits the real OpenRouter + AWS Bedrock bulk APIs. - -Skipped unless LAGO_LIVE_PRICING=1 (it makes real network calls, no keys needed -since both sources are public). Validates that the real fetchers build tables -and that known models resolve to sane USD-per-token prices — in particular it -exercises the AWS Bedrock offer-file parser against the live schema. -""" - -from __future__ import annotations - -import os -from decimal import Decimal - -import pytest - -from lago_agent_sdk.pricing import HttpPricingFetcher, lookup_bedrock, lookup_openrouter - -pytestmark = pytest.mark.skipif( - os.environ.get("LAGO_LIVE_PRICING") != "1", - reason="LAGO_LIVE_PRICING != 1 (live network test)", -) - - -def test_openrouter_live_table_and_known_models() -> None: - table = HttpPricingFetcher(timeout=30).fetch_openrouter() - exact = table["exact"] - assert len(exact) > 50, "expected a substantial OpenRouter model list" - - # A few well-known models should resolve with a positive input price. - resolved = 0 - for provider, model in [ - ("openai", "gpt-4o"), - ("anthropic", "claude-3.5-sonnet"), - ("google", "gemini-2.5-flash"), - ]: - mp = lookup_openrouter(table, provider, model) - if mp is not None and mp.input is not None and mp.input >= Decimal(0): - resolved += 1 - assert resolved >= 1, "expected at least one well-known OpenRouter model to resolve" - - -def test_bedrock_live_table_builds_and_resolves() -> None: - region = "us-east-1" - table = HttpPricingFetcher(timeout=30).fetch_bedrock(region) - # The parser should extract at least some priced models from the live offer. - assert table, "AWS Bedrock offer parsed to an empty table — schema may have changed" - priced = [mp for mp in table.values() if mp.input is not None or mp.output is not None] - assert priced, "no Bedrock models had input/output token prices" - - # A common Bedrock model should resolve (best-effort; logs the key on miss). - for model in [ - "anthropic.claude-3-5-sonnet-20240620-v1:0", - "anthropic.claude-3-haiku-20240307-v1:0", - ]: - mp = lookup_bedrock(table, model) - if mp is not None and (mp.input or mp.output): - return - pytest.skip(f"no probed Bedrock model matched; {len(table)} keys built — refine matcher if needed") diff --git a/tests/integration/test_live_streaming.py b/tests/integration/test_live_streaming.py deleted file mode 100644 index 0556c8d..0000000 --- a/tests/integration/test_live_streaming.py +++ /dev/null @@ -1,101 +0,0 @@ -"""Live streaming end-to-end against real Bedrock + mock Lago. - -Skipped unless AWS_BEARER_TOKEN_BEDROCK is set. Drives real -`converse_stream` and `invoke_model_with_response_stream` via the bearer -REST surface, reshaped into the same flow our wrapper drains. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import boto3 -import pytest - -from lago_agent_sdk import LagoSDK - -REGION = "eu-west-1" -PROMPT = "One sentence about dolphins." - -pytestmark = pytest.mark.skipif( - not os.environ.get("AWS_BEARER_TOKEN_BEDROCK"), - reason="AWS_BEARER_TOKEN_BEDROCK not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): # silence - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def _fresh_client(sdk: LagoSDK): - return sdk.wrap(boto3.client("bedrock-runtime", region_name=REGION)) - - -def test_live_converse_stream_emits_events(): - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = _fresh_client(sdk) - resp = client.converse_stream( - modelId="eu.amazon.nova-lite-v1:0", - messages=[{"role": "user", "content": [{"text": PROMPT}]}], - inferenceConfig={"maxTokens": 30}, - ) - # Drain — wrapper extracts usage from the metadata event - for _event in resp["stream"]: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes and "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "bedrock_converse" - finally: - server.shutdown() - - -def test_live_invoke_model_stream_emits_events(): - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = _fresh_client(sdk) - body = json.dumps( - { - "anthropic_version": "bedrock-2023-05-31", - "max_tokens": 40, - "messages": [{"role": "user", "content": PROMPT}], - } - ) - resp = client.invoke_model_with_response_stream(modelId="eu.anthropic.claude-sonnet-4-6", body=body) - for _event in resp["body"]: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes and "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "bedrock_invoke" - finally: - server.shutdown() diff --git a/tests/integration/test_outage_replay.py b/tests/integration/test_outage_replay.py deleted file mode 100644 index 0b6e0af..0000000 --- a/tests/integration/test_outage_replay.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Outage replay — Lago fails for N seconds; events buffer and arrive in order on recovery.""" - -from __future__ import annotations - -import json -import threading -import time -from http.server import BaseHTTPRequestHandler, HTTPServer - -from lago_agent_sdk import CanonicalUsage, LagoSDK - - -class _ToggleableLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - if self.server.failing: # type: ignore[attr-defined] - self.send_response(503) - self.end_headers() - return - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn(): - s = HTTPServer(("127.0.0.1", 0), _ToggleableLago) - s.received = [] # type: ignore[attr-defined] - s.failing = False # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def test_outage_replay_preserves_order_and_count(): - server, url = _spawn() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_test") - # Cap backoff low so the test doesn't take a minute - sdk._queue._max_retry_seconds = 1.0 # type: ignore[attr-defined] - - # 1. Lago is down — push 200 events - server.failing = True # type: ignore[attr-defined] - for i in range(200): - sdk.emit( - CanonicalUsage(input=1, model=f"m{i:03d}", provider="p", api="bedrock_invoke"), - ) - - # Give the queue worker a few attempts during the outage - time.sleep(2.0) - - # 2. Lago comes back - server.failing = False # type: ignore[attr-defined] - assert sdk.flush(timeout=15.0), "queue did not drain after recovery" - sdk.shutdown(timeout=2.0) - finally: - server.shutdown() - - flat = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - assert len(flat) == 200, f"expected 200 events, got {len(flat)}" - - # Order preserved — model field is m000, m001, ..., m199 - models = [e["properties"]["model"] for e in flat] - assert models == [f"m{i:03d}" for i in range(200)] - - -def test_long_outage_at_buffer_cap_drops_oldest_then_drains(): - """Outage long enough to overflow the (small) buffer — oldest dropped, rest drain.""" - server, url = _spawn() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_test") - # Tiny buffer + tiny backoff so the test runs quickly - sdk._queue._max_buffer_size = 30 # type: ignore[attr-defined] - sdk._queue._max_retry_seconds = 0.5 # type: ignore[attr-defined] - - server.failing = True # type: ignore[attr-defined] - # Push 50 — buffer caps at 30, so 20 oldest get dropped (model='m00'..'m19') - for i in range(50): - sdk.emit(CanonicalUsage(input=1, model=f"m{i:02d}", provider="p", api="bedrock_invoke")) - time.sleep(0.5) - - server.failing = False # type: ignore[attr-defined] - assert sdk.flush(timeout=15.0) - sdk.shutdown(timeout=2.0) - finally: - server.shutdown() - - flat = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - # Expect exactly 30 events, the most recent ones (m20..m49) - assert len(flat) == 30 - models = sorted({e["properties"]["model"] for e in flat}) - assert models == [f"m{i:02d}" for i in range(20, 50)] diff --git a/tests/unit/adapters/test_anthropic_native.py b/tests/unit/adapters/test_anthropic_native.py index 13bff69..f1e6c99 100644 --- a/tests/unit/adapters/test_anthropic_native.py +++ b/tests/unit/adapters/test_anthropic_native.py @@ -98,6 +98,35 @@ def test_unknown_top_usage_field_lands_in_extras() -> None: assert "server_tool_use" in u.extras +# -------------------------------------------------------------------------- +# Model attribution — bill on what answered, not what was requested +# -------------------------------------------------------------------------- +def test_model_resolves_to_response_value_not_request_alias() -> None: + """Anthropic resolves a short alias to a dated snapshot in the response. + + Reproduced live against the real API (no gateway involved): requesting + "claude-sonnet-4-5" answered as "claude-sonnet-4-5-20250929". Pricing/ + attribution must key off what actually answered, or every alias-based call + gets billed under the wrong model. + """ + resp = { + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "hi"}], + "usage": {"input_tokens": 16, "output_tokens": 7}, + } + u = extract_anthropic_native(resp, model_id="claude-sonnet-4-5") + assert u.model == "claude-sonnet-4-5-20250929" + + +def test_model_falls_back_to_request_when_response_is_silent() -> None: + """The synthetic usage blob the streaming wrapper builds carries no top-level + `model` — fall back to the requested model rather than emitting an empty string.""" + u = extract_anthropic_native( + {"usage": {"input_tokens": 1, "output_tokens": 1}}, model_id="claude-sonnet-4-5" + ) + assert u.model == "claude-sonnet-4-5" + + # -------------------------------------------------------------------------- # Synthetic # -------------------------------------------------------------------------- diff --git a/tests/unit/adapters/test_gemini_native.py b/tests/unit/adapters/test_gemini_native.py index ffaab30..be88705 100644 --- a/tests/unit/adapters/test_gemini_native.py +++ b/tests/unit/adapters/test_gemini_native.py @@ -213,3 +213,32 @@ def test_traffic_type_lands_in_known_fields_not_extras() -> None: } u = extract_gemini_native(resp, model_id="gemini-2.5-flash") assert "traffic_type" not in u.extras + + +# -------------------------------------------------------------------------- +# Model attribution — bill on what answered, not the alias that was requested +# -------------------------------------------------------------------------- +def test_model_resolves_to_response_value_not_request_alias() -> None: + """Gemini hot-swaps "-latest" aliases (e.g. "gemini-flash-latest") to a + dated snapshot server-side and reports the resolved id in + `model_version`. Pricing must key off that, not the alias requested — + same failure mode as OpenAI's alias resolution, and previously mishandled + here: the field was available in every response but ignored in favor of + the requested string.""" + resp = { + "model_version": "gemini-flash-latest-002", + "usage_metadata": {"prompt_token_count": 10, "candidates_token_count": 20}, + } + u = extract_gemini_native(resp, model_id="gemini-flash-latest") + assert u.model == "gemini-flash-latest-002" + + +def test_model_falls_back_to_request_when_response_is_silent() -> None: + """The synthetic usage blob the streaming wrapper builds when no final + chunk carries `model_version` — fall back to the requested model rather + than emitting an empty string.""" + u = extract_gemini_native( + {"usage_metadata": {"prompt_token_count": 10, "candidates_token_count": 20}}, + model_id="gemini-2.5-flash", + ) + assert u.model == "gemini-2.5-flash" diff --git a/tests/unit/adapters/test_openai_native.py b/tests/unit/adapters/test_openai_native.py index 71ffb16..f710a8a 100644 --- a/tests/unit/adapters/test_openai_native.py +++ b/tests/unit/adapters/test_openai_native.py @@ -144,6 +144,30 @@ def test_responses_api_shape_detected() -> None: assert u.api == "responses" +# -------------------------------------------------------------------------- +# Model attribution — bill on what answered, not what was requested +# -------------------------------------------------------------------------- +def test_model_resolves_to_response_value_not_request_alias() -> None: + """OpenAI resolves a short alias to a dated snapshot in the response. + + Every non-streaming fixture in this suite shows this exact mismatch — e.g. + `model_id="gpt-4o-mini"` was requested, but the response reports + "gpt-4o-mini-2024-07-18". Pricing/attribution must key off what actually + answered, or every alias-based call gets billed under the wrong model. + """ + model_id, resp = _load("01_plain_chat.json") + assert model_id == "gpt-4o-mini" # sanity: the alias that was requested + u = extract_openai_native(resp, model_id=model_id) + assert u.model == "gpt-4o-mini-2024-07-18" # the resolved model that actually answered + + +def test_model_falls_back_to_request_when_response_is_silent() -> None: + """The synthetic usage blob the streaming wrapper builds carries no top-level + `model` — fall back to the requested model rather than emitting an empty string.""" + u = extract_openai_native({"usage": {"prompt_tokens": 1, "completion_tokens": 1}}, model_id="gpt-4o-mini") + assert u.model == "gpt-4o-mini" + + # -------------------------------------------------------------------------- # Robustness # -------------------------------------------------------------------------- @@ -226,3 +250,27 @@ def test_audio_output_mapped_from_completion_details() -> None: u = extract_openai_native(resp, model_id="gpt-4o-audio") assert u.audio_input == 0 assert u.audio_output == 33 + + +def test_workers_ai_model_via_openai_sdk_infers_correct_provider() -> None: + """Real shape: the openai SDK pointed at Cloudflare's `.../compat` endpoint, + routed to a Workers AI model. The SDK shape looks identical to a real + OpenAI response — "provider" can only be told apart by the resolved model + string itself. Getting this wrong made Workers AI calls permanently + unpriceable in price mode (stamped "openai", which has no Workers AI + entries in its price table) — this is what fixed it.""" + resp = { + "model": "@cf/meta/llama-3.3-70b-instruct-fp8-fast", + "usage": {"prompt_tokens": 38, "completion_tokens": 2}, + } + u = extract_openai_native(resp, model_id="workers-ai/@cf/meta/llama-3.3-70b-instruct-fp8-fast") + assert u.provider == "workers-ai" + assert u.model == "@cf/meta/llama-3.3-70b-instruct-fp8-fast" + + +def test_real_openai_model_still_gets_openai_provider() -> None: + """The inference rule must not become over-eager — a genuine OpenAI model + (no "@cf/" prefix) still gets "openai", unchanged.""" + resp = {"model": "gpt-4o-mini-2024-07-18", "usage": {"prompt_tokens": 10, "completion_tokens": 5}} + u = extract_openai_native(resp, model_id="gpt-4o-mini") + assert u.provider == "openai" diff --git a/tests/unit/fixtures/pricing/money_golden.json b/tests/unit/fixtures/pricing/money_golden.json index 5ac4612..3865382 100644 --- a/tests/unit/fixtures/pricing/money_golden.json +++ b/tests/unit/fixtures/pricing/money_golden.json @@ -1,5 +1,5 @@ { - "_comment": "Cross-repo golden money cases. Python (Decimal) and JS (scaled BigInt) must produce identical base/total/total_cents strings. Money is floored to 12 decimal places (ROUND_DOWN). Prices are USD per token (strings); markup is a string multiplier. total_cents = total USD x 100 (Lago dynamic charge precise_total_amount_cents).", + "_comment": "Cross-repo golden money cases. Python (Decimal) and JS (scaled BigInt) must produce identical base/total/total_cents strings. Money is floored to 12 decimal places (ROUND_DOWN). Prices are USD per token (strings); markup is a string multiplier. total_cents = total USD x 100 (Lago dynamic charge precise_total_amount_cents). `cases` drive compute_cost (optional `provider` selects the provider's token semantics, default 'p'); `precomputed_cases` drive compute_precomputed_cost, where usd_cost is a gateway-reported lump sum and may be a JSON number in exponential notation.", "cases": [ { "name": "input+output, no markup", @@ -54,6 +54,87 @@ "base": "0.001", "total": "0.001333333333", "total_cents": "0.1333333333" + }, + { + "name": "workers-ai: cache_read is a SUBSET of input, billed once", + "_note": "Real counts from a live OpenAI-shaped cached call (prompt 23233, cached 23168) at live Cloudflare @cf/moonshotai/kimi-k2.6 rates. Only 23233-23168=65 tokens may be billed at the input rate; billing all 23233 double-charges the cached portion.", + "provider": "workers-ai", + "prices": { "input": "0.00000095", "cache_read": "0.00000016" }, + "counts": { "input": 23233, "cache_read": 23168 }, + "markup": "1", + "base": "0.00376863", + "total": "0.00376863", + "total_cents": "0.376863" + }, + { + "name": "anthropic: cache_read is ADDITIVE, input not reduced", + "_note": "Same counts and rates as the workers-ai case above; the only difference is the provider's token semantics. Anthropic reports input exclusive of cache, so all 23233 input tokens are billed.", + "provider": "anthropic", + "prices": { "input": "0.00000095", "cache_read": "0.00000016" }, + "counts": { "input": 23233, "cache_read": 23168 }, + "markup": "1", + "base": "0.02577823", + "total": "0.02577823", + "total_cents": "2.577823" + } + ], + "precomputed_cases": [ + { + "name": "real Cloudflare gateway cost below 1e-6 (JSON number -> exponential)", + "_note": "Verbatim `cost` from a real AI Gateway log entry for @cf/meta/llama-3.2-1b-instruct (14 in / 3 out). String(n) in JS renders this as '9.807224944233895e-7', which a decimal-only parser rejected and then billed as zero.", + "usd_cost": 9.807224944233895e-7, + "markup": "1", + "base": "0.000000980722", + "total": "0.000000980722", + "total_cents": "0.0000980722" + }, + { + "name": "same sub-1e-6 cost with 1.5x markup", + "usd_cost": 9.807224944233895e-7, + "markup": "1.5", + "base": "0.000000980722", + "total": "0.000001471083", + "total_cents": "0.0001471083" + }, + { + "name": "real sub-1e-6 cost, second sample", + "usd_cost": 8.91e-7, + "markup": "1", + "base": "0.000000891", + "total": "0.000000891", + "total_cents": "0.0000891" + }, + { + "name": "exponential supplied as a STRING", + "usd_cost": "9.78e-07", + "markup": "1", + "base": "0.000000978", + "total": "0.000000978", + "total_cents": "0.0000978" + }, + { + "name": "plain decimal notation still works (real anthropic gateway cost)", + "usd_cost": 0.002839, + "markup": "1", + "base": "0.002839", + "total": "0.002839", + "total_cents": "0.2839" + }, + { + "name": "below the 12dp floor -> zero", + "usd_cost": 1e-13, + "markup": "1", + "base": "0", + "total": "0", + "total_cents": "0" + }, + { + "name": "negative is rejected -> zero", + "usd_cost": -5, + "markup": "1", + "base": "0", + "total": "0", + "total_cents": "0" } ] } diff --git a/tests/integration/__init__.py b/tests/unit/gateway/__init__.py similarity index 100% rename from tests/integration/__init__.py rename to tests/unit/gateway/__init__.py diff --git a/tests/unit/gateway/adapters/__init__.py b/tests/unit/gateway/adapters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/01_real_anthropic_call.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/01_real_anthropic_call.json new file mode 100644 index 0000000..1ae48bb --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/01_real_anthropic_call.json @@ -0,0 +1,24 @@ +{ + "id": "01KZ3Y993DV0Z5CAQCA4CJ3GRD", + "created_at": "2026-08-03T13:51:26.857Z", + "provider": "anthropic", + "model": "claude-sonnet-4-5-20250929", + "model_type": "", + "path": "v1/messages", + "status_code": 200, + "success": true, + "cached": false, + "tokens_in": 16, + "tokens_out": 7, + "metadata": {"lago_subscription": "cf_gateway_test_sub"}, + "step": 0, + "cost": 0.000153, + "custom_cost": false, + "usage_metadata": { + "input_tokens": 16, + "output_tokens": 7, + "total_tokens": 23, + "input_cache_creation_tokens": 0, + "input_cached_tokens": 0 + } +} diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/02_real_wholesale_credits_failure.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/02_real_wholesale_credits_failure.json new file mode 100644 index 0000000..80b2a2d --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/02_real_wholesale_credits_failure.json @@ -0,0 +1,18 @@ +{ + "id": "01KZ3VXTW2YVNFPVQ7QV6V8HPB", + "created_at": "2026-08-03T13:10:06.565Z", + "provider": "anthropic", + "model": "claude-sonnet-4-5", + "model_type": "", + "path": "v1/messages", + "status_code": 402, + "success": false, + "cached": false, + "tokens_in": 0, + "tokens_out": 0, + "metadata": null, + "step": 0, + "cost": 0, + "custom_cost": false, + "usage_metadata": null +} diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/03_real_workers_ai_failed.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/03_real_workers_ai_failed.json new file mode 100644 index 0000000..d47fd23 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/03_real_workers_ai_failed.json @@ -0,0 +1,18 @@ +{ + "id": "01KZ3VR781DDRC2Z1BK0FX43ND", + "created_at": "2026-08-03T13:07:02.287Z", + "provider": "workers-ai", + "model": "@cf/moonshotai/kimi-k2.7-code", + "model_type": "", + "path": "chat/completions", + "status_code": 403, + "success": false, + "cached": false, + "tokens_in": 0, + "tokens_out": 0, + "metadata": null, + "step": 0, + "cost": 0, + "custom_cost": false, + "usage_metadata": null +} diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/06_real_native_binding_with_metadata.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/06_real_native_binding_with_metadata.json new file mode 100644 index 0000000..376cb50 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/06_real_native_binding_with_metadata.json @@ -0,0 +1,52 @@ +{ + "id": "01KZ8GBQ2FK36MSXZXT8GE04Z7", + "created_at": "2026-08-05T08:24:10.327Z", + "updated_at": "2026-08-05 08:24:10", + "event_id": "5e4ef893-b905-4f0f-8af3-51ed48366314", + "provider": "workers-ai", + "model": "@cf/meta/llama-3.2-1b-instruct", + "model_type": "text-generation", + "path": "@cf/meta/llama-3.2-1b-instruct", + "duration": 244, + "request_type": "workers-ai", + "request_content_type": "application/json", + "status_code": 200, + "response_content_type": "application/json", + "success": true, + "cached": false, + "tokens_in": 19, + "tokens_out": 35, + "metadata": { + "lago_subscription": "cf_gateway_test_sub" + }, + "step": 0, + "timings": { + "total": 243.16022199999998, + "latency": 239.53808700000002 + }, + "location": { + "region": "", + "colo": "" + }, + "cost": 7.513e-06, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "cloudflare-worker", + "usage_metadata": { + "input_tokens": 19, + "output_tokens": 35, + "total_tokens": 54, + "input_cached_tokens": 0, + "neurons": 0.6854994362220168 + }, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/07_real_rest_anthropic_402.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/07_real_rest_anthropic_402.json new file mode 100644 index 0000000..ca544fb --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/07_real_rest_anthropic_402.json @@ -0,0 +1,44 @@ +{ + "id": "6aac7b704ccccd862902ca6eae7c28aa9222c4dae2d5696036d834e0c89e1696", + "created_at": "2026-08-05T08:01:32.561Z", + "updated_at": "2026-08-05 08:01:32", + "event_id": "", + "provider": "anthropic", + "model": "anthropic/claude-opus-4.8", + "model_type": "", + "path": "/run", + "duration": 428, + "request_type": "run", + "request_content_type": "", + "status_code": 402, + "response_content_type": "application/json", + "success": false, + "cached": false, + "tokens_in": 0, + "tokens_out": 0, + "metadata": null, + "step": 0, + "timings": { + "total": 164.2084809988737, + "latency": 0 + }, + "location": { + "region": "", + "colo": "" + }, + "cost": 0, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "curl/8.7.1", + "usage_metadata": null, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/08_real_unified_compat_success.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/08_real_unified_compat_success.json new file mode 100644 index 0000000..b27b1fe --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/08_real_unified_compat_success.json @@ -0,0 +1,50 @@ +{ + "id": "01KZ8F1DX5WPEDD85DBXD1ED9E", + "created_at": "2026-08-05T08:01:05.726Z", + "updated_at": "2026-08-05 08:01:06", + "event_id": "681991e6-9062-45b6-a6ca-5ca26ea6a6e0", + "provider": "workers-ai", + "model": "@cf/meta/llama-3.3-70b-instruct-fp8-fast", + "model_type": "text-generation", + "path": "chat/completions", + "duration": 1272, + "request_type": "provider", + "request_content_type": "application/json", + "status_code": 200, + "response_content_type": "application/json", + "success": true, + "cached": false, + "tokens_in": 43, + "tokens_out": 41, + "metadata": null, + "step": 0, + "timings": { + "total": 1271.8018139973283, + "latency": 1267.7829030007124 + }, + "location": { + "region": "\u00cele-de-France", + "colo": "CDG" + }, + "cost": 0.00010472, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "OpenAI/Python 2.38.0", + "usage_metadata": { + "input_tokens": 43, + "output_tokens": 41, + "total_tokens": 84, + "input_cached_tokens": 0, + "neurons": 9.54372787475586 + }, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/09_real_rest_bare_model_success.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/09_real_rest_bare_model_success.json new file mode 100644 index 0000000..a445aba --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/09_real_rest_bare_model_success.json @@ -0,0 +1,50 @@ +{ + "id": "01KZ8DWFHB5BEEB4VMKF25YCYX", + "created_at": "2026-08-05T07:40:54.076Z", + "updated_at": "2026-08-05 07:40:54", + "event_id": "94b6195b-5bef-430f-9784-6fec3b9e5783", + "provider": "workers-ai", + "model": "@cf/meta/llama-3.2-3b-instruct", + "model_type": "text-generation", + "path": "@cf/meta/llama-3.2-3b-instruct", + "duration": 331, + "request_type": "workers-ai", + "request_content_type": "application/json", + "status_code": 200, + "response_content_type": "application/json", + "success": true, + "cached": false, + "tokens_in": 38, + "tokens_out": 8, + "metadata": null, + "step": 0, + "timings": { + "total": 330.33441799879074, + "latency": 325.5145410001278 + }, + "location": { + "region": "", + "colo": "" + }, + "cost": 4.658e-06, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "curl/8.7.1", + "usage_metadata": { + "input_tokens": 38, + "output_tokens": 8, + "total_tokens": 46, + "input_cached_tokens": 0, + "neurons": 0.41953223943710327 + }, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/10_real_llama_guard_moderation.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/10_real_llama_guard_moderation.json new file mode 100644 index 0000000..78dedfa --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/10_real_llama_guard_moderation.json @@ -0,0 +1,50 @@ +{ + "id": "01KZ8GQQ25SC32YT1PXJBXQS61", + "created_at": "2026-08-05T08:30:43.414Z", + "updated_at": "2026-08-05 08:30:43", + "event_id": "a2acb8fb-063d-4831-8215-91f74a984386", + "provider": "workers-ai", + "model": "@cf/meta/llama-guard-3-8b", + "model_type": "text-generation", + "path": "@cf/meta/llama-guard-3-8b", + "duration": 96, + "request_type": "workers-ai", + "request_content_type": "application/json", + "status_code": 200, + "response_content_type": "application/json", + "success": true, + "cached": false, + "tokens_in": 203, + "tokens_out": 3, + "metadata": null, + "step": 0, + "timings": { + "total": 95.55636099912226, + "latency": 92.99995299987495 + }, + "location": { + "region": "", + "colo": "" + }, + "cost": 9.753e-05, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "python-requests/2.34.2", + "usage_metadata": { + "input_tokens": 203, + "output_tokens": 3, + "total_tokens": 206, + "input_cached_tokens": 0, + "neurons": 8.940808348928998 + }, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/11_real_paid_plan_required_403.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/11_real_paid_plan_required_403.json new file mode 100644 index 0000000..cbfc3e9 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/11_real_paid_plan_required_403.json @@ -0,0 +1,44 @@ +{ + "id": "01KZ8GQTV6Y2VCM1MVFT4WV056", + "created_at": "2026-08-05T08:30:47.182Z", + "updated_at": "2026-08-05 08:30:47", + "event_id": "c8003445-e7d3-4f99-9b12-1b2d92f61df9", + "provider": "workers-ai", + "model": "@cf/moonshotai/kimi-k2.7-code", + "model_type": "", + "path": "@cf/moonshotai/kimi-k2.7-code", + "duration": 38, + "request_type": "workers-ai", + "request_content_type": "application/json", + "status_code": 403, + "response_content_type": "", + "success": false, + "cached": false, + "tokens_in": 0, + "tokens_out": 0, + "metadata": null, + "step": 0, + "timings": { + "total": 37.37525200005621, + "latency": 33.722323999973014 + }, + "location": { + "region": "", + "colo": "" + }, + "cost": 0, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "python-requests/2.34.2", + "usage_metadata": null, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/12_real_cache_read.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/12_real_cache_read.json new file mode 100644 index 0000000..ba16cd8 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/12_real_cache_read.json @@ -0,0 +1,50 @@ +{ + "id": "01KZ8HCAGMGZWBBDAYR4WR9J8P", + "created_at": "2026-08-05T08:42:01.279Z", + "updated_at": "2026-08-05 08:42:01", + "event_id": "2a80105e-33f9-41fa-8fcf-c56b48cf05d4", + "provider": "anthropic", + "model": "claude-sonnet-4-5-20250929", + "model_type": "", + "path": "v1/messages", + "duration": 2707, + "request_type": "provider", + "request_content_type": "application/json", + "status_code": 200, + "response_content_type": "application/json", + "success": true, + "cached": false, + "tokens_in": 10, + "tokens_out": 4, + "metadata": null, + "step": 0, + "timings": { + "total": 2706.3450969997793, + "latency": 2704.368099000305 + }, + "location": { + "region": "\u00cele-de-France", + "colo": "CDG" + }, + "cost": 0.0011187, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "Anthropic/Python 0.103.1", + "usage_metadata": { + "input_tokens": 10, + "output_tokens": 4, + "total_tokens": 14, + "input_cache_creation_tokens": 0, + "input_cached_tokens": 3429 + }, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/13_real_cache_write.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/13_real_cache_write.json new file mode 100644 index 0000000..aeebcb5 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/13_real_cache_write.json @@ -0,0 +1,50 @@ +{ + "id": "01KZ8HC6EJB8QRTBYWSDC8TTAH", + "created_at": "2026-08-05T08:41:56.585Z", + "updated_at": "2026-08-05 08:41:57", + "event_id": "b157d276-2829-49ea-919c-aaa8db496936", + "provider": "anthropic", + "model": "claude-sonnet-4-5-20250929", + "model_type": "", + "path": "v1/messages", + "duration": 2125, + "request_type": "provider", + "request_content_type": "application/json", + "status_code": 200, + "response_content_type": "application/json", + "success": true, + "cached": false, + "tokens_in": 9, + "tokens_out": 5, + "metadata": null, + "step": 0, + "timings": { + "total": 2124.7226979993284, + "latency": 2117.614604000002 + }, + "location": { + "region": "\u00cele-de-France", + "colo": "CDG" + }, + "cost": 0.01296075, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "Anthropic/Python 0.103.1", + "usage_metadata": { + "input_tokens": 9, + "output_tokens": 5, + "total_tokens": 14, + "input_cache_creation_tokens": 3429, + "input_cached_tokens": 0 + }, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/14_real_gateway_cache_hit.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/14_real_gateway_cache_hit.json new file mode 100644 index 0000000..5432131 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/14_real_gateway_cache_hit.json @@ -0,0 +1,44 @@ +{ + "id": "01KZ8HF00A57XFC18M957BQVRW", + "created_at": "2026-08-05T08:43:26.094Z", + "updated_at": "2026-08-05 08:43:26", + "event_id": "c9101ff1-00d6-4276-b18e-0123e3d6f9da", + "provider": "workers-ai", + "model": "@cf/meta/llama-3.2-1b-instruct", + "model_type": "text-generation", + "path": "@cf/meta/llama-3.2-1b-instruct", + "duration": 8, + "request_type": "workers-ai", + "request_content_type": "application/json", + "status_code": 200, + "response_content_type": "application/json", + "success": true, + "cached": true, + "tokens_in": 0, + "tokens_out": 0, + "metadata": null, + "step": 0, + "timings": { + "total": 7.906569000333548, + "latency": 0 + }, + "location": { + "region": "", + "colo": "" + }, + "cost": 0, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "curl/8.7.1", + "usage_metadata": null, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/15_real_mistral_via_dedicated_endpoint.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/15_real_mistral_via_dedicated_endpoint.json new file mode 100644 index 0000000..f337ded --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/15_real_mistral_via_dedicated_endpoint.json @@ -0,0 +1,49 @@ +{ + "id": "01KZ8KG4JMCZNH73Q8ENDS1MJ3", + "created_at": "2026-08-05T09:19:01.892Z", + "updated_at": "2026-08-05 09:19:02", + "event_id": "73af6f65-e9fc-4f56-a50e-daa5092e55e8", + "provider": "mistral", + "model": "mistral-small-latest", + "model_type": "text-generation", + "path": "v1/chat/completions", + "duration": 1069, + "request_type": "provider", + "request_content_type": "application/json", + "status_code": 200, + "response_content_type": "application/json", + "success": true, + "cached": false, + "tokens_in": 23, + "tokens_out": 30, + "metadata": null, + "step": 0, + "timings": { + "total": 1068.4983109980822, + "latency": 1065.597853999585 + }, + "location": { + "region": "\u00cele-de-France", + "colo": "CDG" + }, + "cost": 2.145e-05, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "mistral-client-python/2.4.5", + "usage_metadata": { + "input_tokens": 23, + "output_tokens": 30, + "total_tokens": 53, + "input_cached_tokens": 0 + }, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/16_real_gemini_via_dedicated_endpoint.json b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/16_real_gemini_via_dedicated_endpoint.json new file mode 100644 index 0000000..a6a71bf --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/cloudflare_gateway/16_real_gemini_via_dedicated_endpoint.json @@ -0,0 +1,50 @@ +{ + "id": "01KZ8KE7ST217Z969EPKT12SCS", + "created_at": "2026-08-05T09:18:03.260Z", + "updated_at": "2026-08-05 09:18:03", + "event_id": "26231144-7ac3-4f20-89d5-53364a7f10c0", + "provider": "google-ai-studio", + "model": "gemini-2.5-flash", + "model_type": "text-generation", + "path": "v1beta/models/gemini-2.5-flash:generateContent", + "duration": 4774, + "request_type": "provider", + "request_content_type": "application/json", + "status_code": 200, + "response_content_type": "application/json; charset=utf-8", + "success": true, + "cached": false, + "tokens_in": 9, + "tokens_out": 21, + "metadata": null, + "step": 0, + "timings": { + "total": 4773.519632000476, + "latency": 4770.581113997847 + }, + "location": { + "region": "\u00cele-de-France", + "colo": "CDG" + }, + "cost": 5.52e-05, + "custom_cost": false, + "feedback": null, + "score": null, + "request": "", + "response": "", + "prompts": null, + "guardrails": null, + "authentication": true, + "wholesale": false, + "byok": null, + "user_agent": "google-genai-sdk/2.7.0 gl-python/3.11.15", + "usage_metadata": { + "input_tokens": 9, + "output_tokens": 21, + "total_tokens": 882, + "reasoningTokens": 852, + "input_text_tokens": 9 + }, + "dlp_action": null, + "dlp_profiles": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/test_cloudflare_gateway.py b/tests/unit/gateway/adapters/test_cloudflare_gateway.py new file mode 100644 index 0000000..1e6c8cc --- /dev/null +++ b/tests/unit/gateway/adapters/test_cloudflare_gateway.py @@ -0,0 +1,464 @@ +"""Cloudflare AI Gateway log adapter — verified against a real captured log entry.""" + +from __future__ import annotations + +import json +import pathlib +from decimal import Decimal + +import pytest + +from lago_agent_sdk import CanonicalUsage +from lago_agent_sdk.gateway.adapters import extract_cloudflare_log, resolve_subscription +from lago_agent_sdk.pricing import compute_cost, lookup_openrouter, parse_openrouter + +FIX = pathlib.Path(__file__).parent / "fixtures" / "cloudflare_gateway" + + +def _load(name: str) -> dict: + return json.loads((FIX / name).read_text()) + + +# -------------------------------------------------------------------------- +# Real fixtures +# -------------------------------------------------------------------------- +def test_real_anthropic_call() -> None: + """The exact log entry captured against a live Cloudflare account + real + Anthropic call. These numbers were independently confirmed to roll up + correctly in a real Lago instance (16.0 / 7.0 units billed, exact match).""" + entry = _load("01_real_anthropic_call.json") + u = extract_cloudflare_log(entry) + assert u.input == 16 + assert u.output == 7 + assert u.cache_read == 0 + assert u.cache_write == 0 + assert u.model == "claude-sonnet-4-5-20250929" + assert u.provider == "anthropic" + assert u.api == "cloudflare_gateway" + assert u.extras["cached"] is False + assert u.extras["step"] == 0 + assert u.extras["log_id"] == "01KZ3Y993DV0Z5CAQCA4CJ3GRD" + assert resolve_subscription(entry) == "cf_gateway_test_sub" + + +def test_real_wholesale_credits_failure_has_zero_usage() -> None: + """A 402 (Unified Billing out of credits) never reaches the provider — + tokens_in/out are 0 and usage_metadata is null. Must not raise, must not + fabricate nonzero usage.""" + entry = _load("02_real_wholesale_credits_failure.json") + u = extract_cloudflare_log(entry) + assert u.input == 0 + assert u.output == 0 + assert not u.nonzero_numeric() + assert resolve_subscription(entry) is None # metadata is null on this entry + + +def test_real_workers_ai_provider_and_model_pass_through() -> None: + """A different provider entirely — confirms the mapping isn't Anthropic/OpenAI- + specific; provider/model pass through verbatim regardless of which one it is.""" + entry = _load("03_real_workers_ai_failed.json") + u = extract_cloudflare_log(entry) + assert u.provider == "workers-ai" + assert u.model == "@cf/moonshotai/kimi-k2.7-code" + assert u.input == 0 + assert u.output == 0 + + +# -------------------------------------------------------------------------- +# Real fixtures — three separate ingress methods into the same gateway. +# extract_cloudflare_log() never sees how the call was made (curl, the real +# OpenAI SDK, or a Workers AI binding) — only Cloudflare's own normalized log +# entry. These four fixtures prove that holds across every ingress method. +# -------------------------------------------------------------------------- +def test_real_rest_api_bare_model_success() -> None: + """REST API (`POST /accounts/{account}/ai/run`), a bare Workers AI model + string with no provider prefix. Real call, real success, no BYOK needed — + Workers AI is billed directly by Cloudflare.""" + entry = _load("09_real_rest_bare_model_success.json") + u = extract_cloudflare_log(entry) + assert u.input == 38 + assert u.output == 8 + assert u.provider == "workers-ai" + assert u.model == "@cf/meta/llama-3.2-3b-instruct" + assert resolve_subscription(entry) is None # no metadata sent on this call + + +def test_real_rest_api_anthropic_402_is_provider_agnostic() -> None: + """Same funding failure as 02_real_wholesale_credits_failure.json, but for + Anthropic instead of OpenAI — confirms the "no BYOK/wholesale credits" + failure mode isn't specific to one provider, and still extracts as zero + usage regardless of which provider was requested.""" + entry = _load("07_real_rest_anthropic_402.json") + u = extract_cloudflare_log(entry) + assert u.input == 0 + assert u.output == 0 + assert u.provider == "anthropic" + assert u.model == "anthropic/claude-opus-4.8" + assert resolve_subscription(entry) is None + + +def test_real_unified_compat_success() -> None: + """Unified API (`.../compat/chat/completions`), called with the real `openai` + Python client pointed at Cloudflare's compat endpoint, routed to a Workers AI + model. `path` and `user_agent` on the raw log entry ("OpenAI/Python 2.38.0") + confirm this came from a real SDK call, not a raw curl — proves the log + schema is identical regardless of which client library made the request.""" + entry = _load("08_real_unified_compat_success.json") + assert entry["path"] == "chat/completions" + u = extract_cloudflare_log(entry) + assert u.input == 43 + assert u.output == 41 + assert u.provider == "workers-ai" + assert u.model == "@cf/meta/llama-3.3-70b-instruct-fp8-fast" + + +def test_real_llama_guard_moderation_model_unusual_token_shape() -> None: + """A moderation/classifier model (llama-guard), not a chat model — input is + dominated by the full conversation-plus-policy being classified (203 tokens) + against a tiny 3-token verdict output. Confirms extraction doesn't assume a + "normal" chat-shaped input/output ratio; captured from a real sweep across + 22 distinct Workers AI models with zero extraction failures.""" + entry = _load("10_real_llama_guard_moderation.json") + u = extract_cloudflare_log(entry) + assert u.input == 203 + assert u.output == 3 + assert u.provider == "workers-ai" + assert u.model == "@cf/meta/llama-guard-3-8b" + + +def test_real_paid_plan_required_403_is_distinct_from_funding_402() -> None: + """A different real failure mode: 403 "requires a Workers Paid plan", not the + 402 "insufficient balance" case covered elsewhere. Different status code, + same shape otherwise — still extracts as zero usage, no attribution.""" + entry = _load("11_real_paid_plan_required_403.json") + assert entry["status_code"] == 403 + u = extract_cloudflare_log(entry) + assert u.input == 0 + assert u.output == 0 + assert not u.nonzero_numeric() + assert resolve_subscription(entry) is None + + +def test_real_mistral_via_dedicated_endpoint() -> None: + """Real `mistralai` SDK client, wrapped via `wrap_mistral_client`, pointed at + Cloudflare's dedicated `.../mistral` passthrough (not the Unified/compat + endpoint) — proves Path A generalizes to a fourth native SDK, using a real + customer-supplied Mistral key rather than Cloudflare-side BYOK/credits.""" + entry = _load("15_real_mistral_via_dedicated_endpoint.json") + u = extract_cloudflare_log(entry) + assert u.input == 23 + assert u.output == 30 + assert u.provider == "mistral" + assert u.model == "mistral-small-latest" + + +def test_real_gemini_reasoning_tokens_mapped_from_camelcase_field() -> None: + """Real `google-genai` SDK client, wrapped via `wrap_gemini_client`, through + Cloudflare's dedicated `.../google-ai-studio` passthrough. + + This is the fixture that caught a real gap: Cloudflare's log for this call + has `usage_metadata.reasoningTokens: 852` (camelCase, unlike Anthropic's + snake_case `input_cached_tokens`) — `extract_cloudflare_log()` didn't map it + until this fixture surfaced it. `tokens_out` itself is only 21 (just the + visible completion); the 852 reasoning tokens exist ONLY in usage_metadata.""" + entry = _load("16_real_gemini_via_dedicated_endpoint.json") + assert entry["usage_metadata"]["reasoningTokens"] == 852 + u = extract_cloudflare_log(entry) + assert u.input == 9 + assert u.output == 21 + assert u.reasoning == 852 + # Cloudflare logs this as "google-ai-studio"; the SDK's own vocabulary calls + # it "gemini", which is what the price and token-semantics tables key off. + assert entry["provider"] == "google-ai-studio" + assert u.provider == "gemini" + assert u.model == "gemini-2.5-flash" + + +def test_real_native_binding_with_metadata_resolves_subscription() -> None: + """Native/binding method (`env.AI.run(model, input, {gateway: {id, metadata}})`), + only reachable from inside a deployed Cloudflare Worker — `user_agent` on the + raw entry is literally "cloudflare-worker". The binding's `gateway.metadata` + option maps to the same `metadata` field as the `cf-aig-metadata` header used + by the other two methods, so attribution resolves identically.""" + entry = _load("06_real_native_binding_with_metadata.json") + assert entry["user_agent"] == "cloudflare-worker" + u = extract_cloudflare_log(entry) + assert u.input == 19 + assert u.output == 35 + assert u.provider == "workers-ai" + assert u.model == "@cf/meta/llama-3.2-1b-instruct" + assert resolve_subscription(entry) == "cf_gateway_test_sub" + + +# -------------------------------------------------------------------------- +# Two DIFFERENT "cache" concepts, both verified live — don't conflate them: +# 1. Gateway-level response cache (`cached` boolean) — the entire call was +# served from Cloudflare's own cache, costing the provider (and customer) +# nothing at all. +# 2. Provider-level PROMPT cache (`usage_metadata.input_cache_creation_tokens` +# / `input_cached_tokens`) — a real, separately-priced Anthropic feature +# (`cache_control` on a content block) for reusing part of a long prompt +# across calls that still fully execute. +# -------------------------------------------------------------------------- +def test_real_gateway_cache_hit_has_zero_tokens() -> None: + """Captured by sending the exact same request twice with cf-aig-cache-ttl + set; the second call came back in 8ms (vs 296ms) with `cached: true`. + + Correction from an earlier assumption: a gateway cache HIT does NOT report + the token counts the call "would have" cost — Cloudflare's own log already + reports tokens_in/tokens_out as 0. Billing policy doesn't need to branch on + `cached` at all; a real cache hit already extracts as zero usage.""" + entry = _load("14_real_gateway_cache_hit.json") + u = extract_cloudflare_log(entry) + assert u.input == 0 + assert u.output == 0 + assert u.extras["cached"] is True + + +def test_real_cache_write_then_read_from_anthropic_prompt_cache() -> None: + """Two real, back-to-back Anthropic calls through the gateway with the same + long (>1024 token) `cache_control: {"type": "ephemeral"}` system block. + + Call 1 (cache miss, writes the cache): Anthropic's own response reported + cache_creation_input_tokens=3429, cache_read_input_tokens=0 — Cloudflare's + log matches those exact numbers under different field names. + Call 2 (cache hit, reads it back): the numbers flip — Anthropic reported + cache_creation_input_tokens=0, cache_read_input_tokens=3429 — again an + exact match in the gateway log. Unlike the gateway-level cache above, this + call still executes and still bills the non-cached tokens normally.""" + write_entry = _load("13_real_cache_write.json") + read_entry = _load("12_real_cache_read.json") + + w = extract_cloudflare_log(write_entry) + assert w.input == 9 + assert w.output == 5 + assert w.cache_write == 3429 + assert w.cache_read == 0 + + r = extract_cloudflare_log(read_entry) + assert r.input == 10 + assert r.output == 4 + assert r.cache_write == 0 + assert r.cache_read == 3429 + + +# -------------------------------------------------------------------------- +# Attribution +# -------------------------------------------------------------------------- +def test_resolve_subscription_missing_metadata_key_returns_none() -> None: + entry = {"metadata": {"some_other_key": "x"}} + assert resolve_subscription(entry) is None + + +def test_resolve_subscription_empty_string_returns_none() -> None: + """An empty string is falsy attribution, not a real subscription id.""" + entry = {"metadata": {"lago_subscription": ""}} + assert resolve_subscription(entry) is None + + +def test_resolve_subscription_non_dict_metadata_returns_none() -> None: + assert resolve_subscription({"metadata": "not-a-dict"}) is None + assert resolve_subscription({"metadata": None}) is None + assert resolve_subscription({}) is None + + +# -------------------------------------------------------------------------- +# Robustness — a poller processes entries in a batch; one malformed entry +# must not take down the whole run. +# -------------------------------------------------------------------------- +def test_survives_missing_fields() -> None: + u = extract_cloudflare_log({}) + assert u.input == 0 + assert u.output == 0 + assert u.model == "" + assert u.provider == "" + assert not u.nonzero_numeric() + + +def test_survives_non_dict_usage_metadata() -> None: + u = extract_cloudflare_log({"tokens_in": 5, "tokens_out": 3, "usage_metadata": "bogus"}) + assert u.input == 5 + assert u.output == 3 + assert u.cache_read == 0 + assert u.cache_write == 0 + + +def test_survives_non_string_model_and_provider() -> None: + u = extract_cloudflare_log({"model": 123, "provider": None, "tokens_in": 1, "tokens_out": 1}) + assert u.model == "" + assert u.provider == "" + + +def test_survives_negative_and_non_numeric_tokens() -> None: + assert extract_cloudflare_log({"tokens_in": -5}).input == 0 + assert extract_cloudflare_log({"tokens_out": "bogus"}).output == 0 + + +# -------------------------------------------------------------------------- +# Provider vocabulary — Cloudflare's names are not the SDK's names +# -------------------------------------------------------------------------- +def test_provider_aliases_map_onto_sdk_vocabulary() -> None: + """Cloudflare's log vocabulary differs from the names the pricing tables and + the token-semantics sets key off. Verified live: `lookup_openrouter` with + provider="google-ai-studio" missed against the real 400-model OpenRouter + table and hit as "gemini".""" + for raw, expected in [ + ("google-ai-studio", "gemini"), + ("google-vertex-ai", "gemini"), + ("vertex", "gemini"), + ("azure-openai", "openai"), + ("azureopenai", "openai"), + ("workersai", "workers-ai"), + ]: + u = extract_cloudflare_log({"provider": raw, "tokens_in": 1}) + assert u.provider == expected, f"{raw} -> {u.provider}, expected {expected}" + + +def test_provider_passthrough_for_names_we_already_agree_on() -> None: + for raw in ("anthropic", "openai", "mistral", "workers-ai"): + assert extract_cloudflare_log({"provider": raw, "tokens_in": 1}).provider == raw + + +def test_unknown_provider_passes_through_untouched() -> None: + """An unrecognized provider is one we have no price table for; a clean miss + falls back to token events, which beats inventing a mapping.""" + assert extract_cloudflare_log({"provider": "perplexity", "tokens_in": 1}).provider == "perplexity" + # AWS Bedrock is deliberately NOT aliased — its prices key off the `api` + # field, which this connector always sets to "cloudflare_gateway". + assert extract_cloudflare_log({"provider": "bedrock", "tokens_in": 1}).provider == "bedrock" + + +def test_normalized_gemini_provider_prices_and_bills_cache_correctly() -> None: + """The two downstream consequences of the alias, end to end: the Gemini + price is now findable, and cache_read is treated as a SUBSET of input + (Gemini's semantics) instead of being billed on top of it.""" + entry = _load("16_real_gemini_via_dedicated_endpoint.json") + u = extract_cloudflare_log(entry) + table = parse_openrouter( + { + "data": [ + { + "id": "google/gemini-2.5-flash", + "pricing": { + "prompt": "0.0000003", + "completion": "0.0000025", + "input_cache_read": "0.000000075", + }, + } + ] + } + ) + price = lookup_openrouter(table, u.provider, u.model) + assert price is not None, "gemini price must resolve after normalization" + + cached = CanonicalUsage(model=u.model, provider=u.provider, api=u.api, input=1000, cache_read=800) + b = compute_cost(cached, price, Decimal("1")) + assert b.fields["input"]["tokens"] == "200" # 1000 - 800, not 1000 + assert b.fields["cache_read"]["tokens"] == "800" + + +def test_lookup_openrouter_strips_a_redundant_vendor_prefix() -> None: + """Real fixture 07 reports model="anthropic/claude-opus-4.8" alongside + provider="anthropic"; unstripped that built "anthropic/anthropic/..." and + never matched.""" + table = parse_openrouter( + {"data": [{"id": "anthropic/claude-opus-4.8", "pricing": {"prompt": "0.000005"}}]} + ) + assert lookup_openrouter(table, "anthropic", "anthropic/claude-opus-4.8") is not None + assert lookup_openrouter(table, "anthropic", "claude-opus-4.8") is not None + # Still vendor-gated: a model claiming a different vendor must not match. + assert lookup_openrouter(table, "openai", "anthropic/claude-opus-4.8") is None + + +# ---------------------------------------------------------------------- +# Cache-key casing. The gateway forwards some provider keys unnormalized — the +# real Gemini fixture carries camelCase `reasoningTokens` — and a missed cache +# key does not merely lose a field: `gemini` is in _INPUT_INCLUDES_CACHE_READ, +# so compute_cost needs `cache_read` populated to SUBTRACT the cached portion +# out of `input`. A silent 0 bills those tokens at the full prompt rate. +# ---------------------------------------------------------------------- +@pytest.mark.parametrize( + "key", + ["input_cached_tokens", "inputCachedTokens", "cachedContentTokenCount"], +) +def test_cache_read_is_read_under_every_plausible_spelling(key: str) -> None: + u = extract_cloudflare_log( + {"tokens_in": 100, "tokens_out": 10, "provider": "google-ai-studio", "usage_metadata": {key: 90}} + ) + assert u.cache_read == 90, f"{key} must resolve" + + +@pytest.mark.parametrize( + "key", + ["input_cache_creation_tokens", "inputCacheCreationTokens", "cache_creation_input_tokens"], +) +def test_cache_write_is_read_under_every_plausible_spelling(key: str) -> None: + u = extract_cloudflare_log( + {"tokens_in": 100, "tokens_out": 10, "provider": "anthropic", "usage_metadata": {key: 40}} + ) + assert u.cache_write == 40, f"{key} must resolve" + + +def test_a_zeroed_alias_falls_through_to_the_real_count() -> None: + """Fallthrough is on a falsy value, not just a missing key. A provider sending + both its own name and the gateway's, with one zeroed, must resolve to the real + count — this is where JS's `??` diverged from Python's `or`.""" + u = extract_cloudflare_log( + { + "tokens_in": 100, + "tokens_out": 10, + "provider": "google-ai-studio", + "usage_metadata": {"input_cached_tokens": 0, "cachedContentTokenCount": 77}, + } + ) + assert u.cache_read == 77 + + +def test_cache_read_still_zero_when_genuinely_absent() -> None: + u = extract_cloudflare_log( + {"tokens_in": 100, "tokens_out": 10, "provider": "anthropic", "usage_metadata": {}} + ) + assert u.cache_read == 0 + assert u.cache_write == 0 + + +def test_provider_native_cache_and_reasoning_spellings_are_accepted() -> None: + """SYNTHETIC entries — no provider-native key appears in ANY of the 14 captured + fixtures (they carry only Cloudflare's own vocabulary). These pin the unobserved + insurance spellings so the fallthrough list cannot be trimmed by accident. + + The direction of the harm differs by provider, which is why both matter: + Anthropic's cache_read is ADDITIVE, so a missed key means those tokens are never + billed (under-bill); Gemini's is SUBTRACTIVE, so a missed key bills them at the + full prompt rate (over-bill). + """ + anthropic_native = extract_cloudflare_log( + { + "tokens_in": 100, + "tokens_out": 10, + "provider": "anthropic", + "usage_metadata": {"cache_read_input_tokens": 4242}, + } + ) + assert anthropic_native.cache_read == 4242 + + gemini_native = extract_cloudflare_log( + { + "tokens_in": 100, + "tokens_out": 10, + "provider": "google-ai-studio", + "usage_metadata": {"thoughtsTokenCount": 852}, + } + ) + assert gemini_native.reasoning == 852 + + # Cloudflare's own spelling still wins when both are present + both = extract_cloudflare_log( + { + "tokens_in": 100, + "tokens_out": 10, + "provider": "anthropic", + "usage_metadata": {"input_cached_tokens": 11, "cache_read_input_tokens": 4242}, + } + ) + assert both.cache_read == 11 diff --git a/tests/unit/test_auto_prime_pricing.py b/tests/unit/test_auto_prime_pricing.py new file mode 100644 index 0000000..5d47e6e --- /dev/null +++ b/tests/unit/test_auto_prime_pricing.py @@ -0,0 +1,189 @@ +"""wrap()-triggered automatic, non-blocking pricing warm-up. + +Covers `LagoSDK._auto_prime_pricing_for`/`_extract_mistral_api_key`: the +customer calls `sdk.wrap(client)` (already part of their normal flow, no new +function to remember) and that alone should be enough for the session's +FIRST Mistral/Workers AI call to have a real shot at pricing correctly, +without ever declaring `LagoConfig.mistral_api_key` separately — the client +being wrapped already carries the exact credential needed. +""" + +from __future__ import annotations + +import time +from decimal import Decimal + +from lago_agent_sdk import LagoConfig, LagoSDK, ModelPrice +from lago_agent_sdk.pricing import HttpPricingFetcher, PricingProvider, parse_mistral_aliases + + +def _wait_until(predicate, timeout: float = 2.0) -> bool: + """`LagoSDK.wrap()` wakes the REAL background queue thread (see + `EventQueue.wake()`), which races any direct `provider.maybe_refresh()` + call in the test's own thread — both are legitimate, concurrent + triggers. Poll instead of asserting immediately after one call.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.01) + return predicate() + + +class _FakeSecurity: + def __init__(self, api_key: str): + self.api_key = api_key + + +class _FakeSdkConfiguration: + def __init__(self, api_key: str): + self.security = _FakeSecurity(api_key) + + +class FakeMistralClient: + """Mimics the real shape verified against mistralai.client.Mistral: + `client.sdk_configuration.security.api_key`.""" + + __module__ = "mistralai.client.sdk" + + def __init__(self, api_key: str): + self.sdk_configuration = _FakeSdkConfiguration(api_key) + + +class FakeOpenAIClient: + """Mimics openai.OpenAI's `base_url` attribute (a plain str is enough — + real usage is an httpx.URL, but only `str(...)` on it is ever read).""" + + __module__ = "openai.client" + + def __init__(self, base_url: str): + self.base_url = base_url + + +_MISTRAL_ALIASES = parse_mistral_aliases( + {"data": [{"id": "mistral-small-2603", "aliases": ["mistral-small-latest"]}]} +) +_OPENROUTER = { + "exact": {}, + "norm": { + # matches parse_openrouter's own convention: (vendor, full suffix after "vendor/") + ("mistralai", "mistral-small-2603"): ModelPrice( + source="openrouter", input=Decimal("0.00000015"), output=Decimal("0.0000006") + ) + }, +} + + +class _CloudflareCallCountingFetcher(HttpPricingFetcher): + """Shared by both wrap()-vs-Cloudflare-base_url tests below.""" + + def __init__(self): + super().__init__() + self.cloudflare_calls = 0 + + def fetch_cloudflare_workers_ai(self): + self.cloudflare_calls += 1 + return {} + + +def _sdk_with_provider(provider: PricingProvider) -> LagoSDK: + cfg = LagoConfig( + api_key="dummy", default_subscription_id="sub_test", pricing_mode="price", pricing_provider=provider + ) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: None # type: ignore[attr-defined] + return sdk + + +def test_extract_mistral_api_key_reads_the_real_attribute_path(): + client = FakeMistralClient(api_key="sk-from-client") + assert LagoSDK._extract_mistral_api_key(client) == "sk-from-client" + + +def test_extract_mistral_api_key_returns_none_when_attribute_missing(): + class Empty: + pass + + assert LagoSDK._extract_mistral_api_key(Empty()) is None + + +def test_wrap_mistral_learns_key_and_primes_without_config_key(): + """The whole point: no LagoConfig.mistral_api_key anywhere, and the + session's first Mistral lookup still resolves correctly because wrap() + learned the key from the client and kicked off the fetch.""" + + class _StubFetcher(HttpPricingFetcher): + def __init__(self): + super().__init__() + self.seen_keys: list[str | None] = [] + + def fetch_mistral_aliases(self, api_key=None): + self.seen_keys.append(api_key) + return _MISTRAL_ALIASES + + def fetch_openrouter(self): + return _OPENROUTER + + fetcher = _StubFetcher() + provider = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + sdk = _sdk_with_provider(provider) + + client = FakeMistralClient(api_key="sk-from-client") + sdk.wrap(client) # <-- the only thing the customer does + + assert _wait_until(lambda: fetcher.seen_keys == ["sk-from-client"]) + mp = provider.lookup("mistral", "mistral-small-latest", "native") + assert mp is not None + assert mp.input == Decimal("0.00000015") + + +def test_wrap_openai_pointed_at_cloudflare_gateway_primes_workers_ai(): + fetcher = _CloudflareCallCountingFetcher() + provider = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + sdk = _sdk_with_provider(provider) + + client = FakeOpenAIClient(base_url="https://gateway.ai.cloudflare.com/v1/acct/gw/compat") + sdk.wrap(client) + + assert _wait_until(lambda: fetcher.cloudflare_calls == 1) + + +def test_wrap_openai_pointed_at_real_openai_does_not_prime_workers_ai(): + """A generic OpenAI client NOT pointed at Cloudflare must not trigger the + Workers AI fetch — only the base_url signal should do that.""" + fetcher = _CloudflareCallCountingFetcher() + provider = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + sdk = _sdk_with_provider(provider) + + client = FakeOpenAIClient(base_url="https://api.openai.com/v1") + sdk.wrap(client) + provider.maybe_refresh() + + assert fetcher.cloudflare_calls == 0 + + +def test_auto_prime_is_a_noop_in_token_mode(): + """No point flagging anything stale for a customer who never opted into + price mode — the credential-gated sources should stay completely untouched.""" + + class _StubFetcher(HttpPricingFetcher): + def __init__(self): + super().__init__() + self.mistral_calls = 0 + + def fetch_mistral_aliases(self, api_key=None): + self.mistral_calls += 1 + return {} + + fetcher = _StubFetcher() + provider = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + cfg = LagoConfig( + api_key="dummy", default_subscription_id="sub_test", pricing_provider=provider + ) # tokens (default) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: None # type: ignore[attr-defined] + + sdk.wrap(FakeMistralClient(api_key="sk-from-client")) + provider.maybe_refresh() + + assert fetcher.mistral_calls == 0 diff --git a/tests/unit/test_buffer_overflow.py b/tests/unit/test_buffer_overflow.py index d1c9907..a791393 100644 --- a/tests/unit/test_buffer_overflow.py +++ b/tests/unit/test_buffer_overflow.py @@ -14,10 +14,18 @@ def test_overflow_drops_oldest_at_exact_boundary(): def slow_sender(batch): paused.wait(timeout=30.0) + # max_batch_size > max_buffer_size keeps the background worker from ever being + # woken by push (the buffer can't reach max_batch_size), the same determinism + # fix test_repeated_overflow_keeps_window_sliding already carries. Without it + # push() signals the worker on the 10_000th event and the worker then races the + # assertions below for `_lock` — it is entitled to drain at any point once woken, + # so "the buffer still holds what I just pushed" is not a property the test can + # rely on. Nothing here needs the worker to run; every assertion is about buffer + # contents, and the finally block only unpauses so shutdown can finish. q = EventQueue( sender=slow_sender, flush_interval=10.0, # never timer-flush during the test - max_batch_size=10_000, # match buffer so worker takes everything once unpaused + max_batch_size=20_000, max_buffer_size=10_000, ) try: diff --git a/tests/unit/test_lago_client.py b/tests/unit/test_lago_client.py new file mode 100644 index 0000000..8fbec26 --- /dev/null +++ b/tests/unit/test_lago_client.py @@ -0,0 +1,50 @@ +"""LagoClient — verify_ssl passthrough. + +A local dev Lago instance behind a self-signed certificate is a real, common +setup; the only alternative without this flag is routing every request +through a public tunnel purely to get a browser-trusted cert. +""" + +from __future__ import annotations + +from unittest.mock import patch + +from lago_agent_sdk.config import LagoConfig +from lago_agent_sdk.lago_client import LagoClient +from lago_agent_sdk.sdk import LagoSDK + + +def test_verify_ssl_defaults_to_true() -> None: + client = LagoClient(api_key="k", api_url="https://api.getlago.com/api/v1") + assert client.verify_ssl is True + with patch("requests.post") as mock_post: + mock_post.return_value.status_code = 200 + client.send_batch([{"transaction_id": "t1"}]) + assert mock_post.call_args.kwargs["verify"] is True + + +def test_verify_ssl_false_is_passed_through_to_requests() -> None: + client = LagoClient(api_key="k", api_url="https://api.lago.dev/api/v1", verify_ssl=False) + assert client.verify_ssl is False + with patch("requests.post") as mock_post: + mock_post.return_value.status_code = 200 + client.send_batch([{"transaction_id": "t1"}]) + assert mock_post.call_args.kwargs["verify"] is False + + +def test_lago_config_verify_ssl_defaults_to_true() -> None: + assert LagoConfig(api_key="k").verify_ssl is True + + +def test_sdk_threads_verify_ssl_from_config_to_its_internal_client() -> None: + sdk = LagoSDK(api_key="k", config=LagoConfig(api_key="k", verify_ssl=False)) + try: + assert sdk._lago_client.verify_ssl is False + finally: + sdk.shutdown(timeout=1.0) + + sdk2 = LagoSDK(api_key="k") # default config — verify_ssl stays True + try: + assert sdk2._lago_client.verify_ssl is True + finally: + sdk2.shutdown(timeout=1.0) diff --git a/tests/unit/test_pricing.py b/tests/unit/test_pricing.py index 8b83926..c66079b 100644 --- a/tests/unit/test_pricing.py +++ b/tests/unit/test_pricing.py @@ -4,21 +4,34 @@ import json import pathlib +import re +import uuid from decimal import Decimal from typing import Any import pytest from lago_agent_sdk import CanonicalUsage, LagoConfig, LagoSDK, ModelPrice +from lago_agent_sdk.adapters.openai_native import extract_openai_native +from lago_agent_sdk.canonical import WORKERS_AI_COMPAT_PREFIX from lago_agent_sdk.pricing import ( + HttpPricingFetcher, PricingProvider, + _parse_price, + _pick_mistral_canonical, + apply_markup, bedrock_model_key, coerce_markup, compute_cost, + compute_precomputed_cost, + deoverlapped_token_total, lookup_bedrock, + lookup_cloudflare_workers_ai, lookup_openrouter, parse_bedrock_offer, parse_bedrock_region, + parse_cloudflare_workers_ai, + parse_mistral_aliases, parse_openrouter, ) @@ -29,11 +42,22 @@ # Stub fetcher (no network) — mirrors the queue's injectable sender pattern # ---------------------------------------------------------------------- class StubFetcher: - def __init__(self, openrouter: dict | None = None, bedrock: dict | None = None) -> None: + def __init__( + self, + openrouter: dict | None = None, + bedrock: dict | None = None, + cloudflare_workers_ai: dict[str, ModelPrice] | None = None, + mistral_aliases: dict[str, str] | None = None, + ) -> None: self._openrouter = openrouter or {"exact": {}, "norm": {}} self._bedrock = bedrock or {} + self._cloudflare_workers_ai = cloudflare_workers_ai or {} + self._mistral_aliases = mistral_aliases or {} self.openrouter_calls = 0 self.bedrock_calls: list[str] = [] + self.cloudflare_workers_ai_calls = 0 + self.mistral_aliases_calls = 0 + self.last_mistral_api_key: str | None = None def fetch_openrouter(self) -> dict[str, Any]: self.openrouter_calls += 1 @@ -43,6 +67,15 @@ def fetch_bedrock(self, region: str) -> dict[str, ModelPrice]: self.bedrock_calls.append(region) return self._bedrock.get(region, {}) + def fetch_cloudflare_workers_ai(self) -> dict[str, ModelPrice]: + self.cloudflare_workers_ai_calls += 1 + return self._cloudflare_workers_ai + + def fetch_mistral_aliases(self, api_key: str | None = None) -> dict[str, str]: + self.mistral_aliases_calls += 1 + self.last_mistral_api_key = api_key + return self._mistral_aliases + _OPENROUTER_RAW = { "data": [ @@ -66,6 +99,16 @@ def fetch_bedrock(self, region: str) -> dict[str, ModelPrice]: }, }, {"id": "mistralai/mistral-large", "pricing": {"prompt": "0.000002", "completion": "0.000006"}}, + # Real case: OpenRouter lists the dated snapshot, never the "-latest" + # alias a customer actually requests. + { + "id": "mistralai/mistral-small-2603", + "pricing": { + "prompt": "0.00000015", + "completion": "0.0000006", + "input_cache_read": "0.000000015", + }, + }, { "id": "google/gemini-2.5-flash", "pricing": { @@ -78,6 +121,97 @@ def fetch_bedrock(self, region: str) -> dict[str, ModelPrice]: ] } +# Real data, captured live from /accounts/{id}/ai/models/search — this exact +# shape (including the non-token unit types and the no-price model) is what's +# actually in the catalog, not a synthetic guess at its structure. +_CLOUDFLARE_MODELS_RAW = [ + { + "name": "@cf/meta/llama-3.3-70b-instruct-fp8-fast", + "properties": [ + {"property_id": "context_window", "value": "24000"}, + { + "property_id": "price", + "value": [ + {"unit": "per M input tokens", "price": 0.293, "currency": "USD"}, + {"unit": "per M output tokens", "price": 2.253, "currency": "USD"}, + ], + }, + ], + }, + { + "name": "@cf/moonshotai/kimi-k2.7-code", + "properties": [ + { + "property_id": "price", + "value": [ + {"unit": "per M input tokens", "price": 0.95, "currency": "USD"}, + {"unit": "per M output tokens", "price": 4, "currency": "USD"}, + {"unit": "per M cached input tokens", "price": 0.19, "currency": "USD"}, + ], + }, + ], + }, + { + # Real non-token-priced model — must be skipped entirely, not stored + # with a bogus/zero token price. + "name": "@cf/pipecat-ai/smart-turn-v2", + "properties": [ + { + "property_id": "price", + "value": [{"unit": "per audio minute", "price": 0.000338, "currency": "USD"}], + }, + ], + }, + { + # Real case: some models have no `price` property at all. + "name": "@cf/some/unpriced-model", + "properties": [{"property_id": "context_window", "value": "8192"}], + }, +] + +# Real data, captured live from Mistral's own /v1/models — "mistral-small-2603" +# is the dated snapshot that actually answers; "mistral-small-latest" (what a +# customer requests) is one of several aliases pointing at it. +_MISTRAL_MODELS_RAW = { + "data": [ + { + "id": "mistral-small-2603", + "aliases": ["mistral-small-latest", "mistral-vibe-cli-fast", "magistral-small-latest"], + }, + {"id": "mistral-large-2411", "aliases": ["mistral-large-latest"]}, + {"id": "codestral-2508", "aliases": []}, + ] +} + +# Real data, captured live — the messy shape that actually broke this feature +# in production. Mistral's real /v1/models does NOT have one clean canonical +# entry with pure aliases: "mistral-small-2603", "mistral-small-latest", AND +# "magistral-small-latest" each appear as their OWN top-level `id`, each +# listing the other two as `aliases`. A naive "map each alias -> this +# entry's id" parser resolves "mistral-small-latest" to whichever of these +# three entries happens to be processed last — here, "magistral-small-latest" +# (index 13, after "mistral-small-latest" at index 11) — instead of the real +# dated snapshot OpenRouter lists. +_MISTRAL_MODELS_RAW_MUTUAL_ALIASING = { + "data": [ + { + "id": "mistral-small-2603", + "aliases": ["mistral-small-latest", "mistral-vibe-cli-fast", "magistral-small-latest"], + }, + { + "id": "mistral-small-latest", + "aliases": ["mistral-small-2603", "mistral-vibe-cli-fast", "magistral-small-latest"], + }, + {"id": "mistral-vibe-cli-fast", "aliases": ["mistral-small-2603"]}, + { + "id": "magistral-small-latest", + "aliases": ["mistral-small-2603", "mistral-small-latest", "mistral-vibe-cli-fast"], + }, + {"id": "voxtral-small-2507", "aliases": ["voxtral-small-latest"]}, + {"id": "voxtral-small-latest", "aliases": ["voxtral-small-2507"]}, + ] +} + # ---------------------------------------------------------------------- # OpenRouter parsing + matching @@ -112,6 +246,60 @@ def test_openrouter_date_version_stripped_match() -> None: assert mp.input == Decimal("0.000001") +@pytest.mark.parametrize("alias_first", [True, False]) +def test_openrouter_moving_alias_never_overwrites_a_real_listing(alias_first: bool) -> None: + """A "~" alias and a real listing can collide on the same key. Which one wins + must not depend on catalog order — with plain assignment it did, and the moving + alias's rate (0.009) could replace the real listing's (0.001) purely by + arriving later in the response.""" + real = {"id": "google/gemini-flash-latest", "pricing": {"prompt": "0.001"}} + alias = {"id": "~google/gemini-flash-latest", "pricing": {"prompt": "0.009"}} + data = [alias, real] if alias_first else [real, alias] + table = parse_openrouter({"data": data}) + mp = lookup_openrouter(table, "gemini", "gemini-flash-latest") + assert mp is not None + assert mp.input == Decimal("0.001"), ( + f"real listing must win regardless of order (alias_first={alias_first})" + ) + # the "~"-spelled id still resolves to its own entry + assert table["exact"]["~google/gemini-flash-latest"].input == Decimal("0.009") + + +def test_three_digit_revision_is_stripped_for_openrouter_only() -> None: + """The "-002" arm exists for Gemini's revision, which only OpenRouter omits. + + It must NOT reach the AWS/Bedrock key builder: there a shortened key does not + merely miss, it collapses two models onto one key whose per-direction prices + are assigned in place, so one silently overwrites the other's rate. + """ + table = parse_openrouter({"data": [{"id": "google/gemini-2.5-flash", "pricing": {"prompt": "0.001"}}]}) + assert lookup_openrouter(table, "gemini", "gemini-2.5-flash-002") is not None + + # shared helper keeps a 3-digit tail, so distinct Bedrock ids stay distinct + assert bedrock_model_key("amazon.titan-text-001") == "titantext001" + assert bedrock_model_key("amazon.titan-text-002") == "titantext002" + assert bedrock_model_key("amazon.titan-text-001") != bedrock_model_key("amazon.titan-text-002") + # real dated/versioned ids are unaffected + assert bedrock_model_key("anthropic.claude-haiku-4-5-20251001-v1:0") == "claudehaiku45" + assert bedrock_model_key("eu.anthropic.claude-sonnet-4-6") == "claudesonnet46" + + +def test_unparseable_markup_keeps_the_cost_instead_of_zeroing_it() -> None: + """Defence in depth, and cross-port parity — `coerce_markup` means neither + branch is reachable through `emit()` (see the coerce test further down). + + The two bad inputs are not interchangeable: a bad COST leaves nothing to bill, + but a bad MARKUP only loses the multiplier, and returning 0 for it would + discard a good cost. JS already fell back to 1.0 here; Python returned "0", so + the ports would have billed differently had anything reached it. + """ + assert apply_markup("0.0042", "1.5") == "0.0063" + for bad in ("abc", "", "1,5", "None"): + assert apply_markup("0.0042", bad) == "0.0042", f"markup={bad!r} must not zero the cost" + # an unparseable COST is different: there is nothing to bill + assert apply_markup("abc", "1.5") == "0" + + def test_openrouter_miss_returns_none() -> None: table = parse_openrouter(_OPENROUTER_RAW) assert lookup_openrouter(table, "anthropic", "totally-made-up-model") is None @@ -119,6 +307,542 @@ def test_openrouter_miss_returns_none() -> None: assert lookup_openrouter(table, "openai", "claude-opus-4-8") is None +# ---------------------------------------------------------------------- +# Cloudflare Workers AI parsing + matching +# ---------------------------------------------------------------------- +# ---------------------------------------------------------------------- +# OpenRouter's "~" moving-alias marker. Measured live: 11 ids across 6 vendors, +# every one a "-latest" moniker with real token pricing, and every one +# unpriceable before this — the vendor parsed as "~anthropic"/"~openai"/"~google", +# which match nothing in _VENDOR_MAP. +# ---------------------------------------------------------------------- +_TILDE_RAW = { + "data": [ + { + "id": "~anthropic/claude-sonnet-latest", + "pricing": {"prompt": "0.000002", "completion": "0.00001"}, + }, + {"id": "~openai/gpt-latest", "pricing": {"prompt": "0.0000025", "completion": "0.000015"}}, + { + "id": "~google/gemini-flash-latest", + "pricing": {"prompt": "0.000000375", "completion": "0.000001875"}, + }, + ] +} + + +@pytest.mark.parametrize( + ("provider", "model"), + [ + ("anthropic", "claude-sonnet-latest"), + ("openai", "gpt-latest"), + ("gemini", "gemini-flash-latest"), + ], +) +def test_moving_alias_ids_are_priceable(provider: str, model: str) -> None: + """A "-latest" alias a customer plausibly requests must resolve. Billing + nothing at all is the outcome in an llm_cost-only setup.""" + t = parse_openrouter(_TILDE_RAW) + assert lookup_openrouter(t, provider, model) is not None + + +def test_moving_alias_still_indexed_under_its_verbatim_id() -> None: + """Stripping the marker must ADD a key, not replace one — the raw id stays + resolvable so nothing that already worked breaks.""" + t = parse_openrouter(_TILDE_RAW) + assert "~openai/gpt-latest" in t["exact"] + assert "openai/gpt-latest" in t["exact"] + + +def test_three_digit_revision_suffix_strips_to_a_hit() -> None: + """Gemini's `model_version` can report a "-002" revision where OpenRouter + lists only the bare name. Verified against the live catalog that no real id's + model part ends in exactly three digits, so this arm is safe.""" + t = parse_openrouter(_TILDE_RAW) + assert lookup_openrouter(t, "gemini", "gemini-flash-latest-002") is not None + + +def test_cloudflare_parses_real_price_shape() -> None: + table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + mp = lookup_cloudflare_workers_ai(table, "@cf/meta/llama-3.3-70b-instruct-fp8-fast") + assert mp is not None + assert mp.source == "cloudflare_workers_ai" + # $0.293/M input -> $0.000000293/token; $2.253/M output -> $0.000002253/token + assert mp.input == Decimal("0.000000293") + assert mp.output == Decimal("0.000002253") + assert mp.cache_read is None # this model has no cached-input price + + +def test_cloudflare_maps_cached_input_tokens_to_cache_read() -> None: + table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + mp = lookup_cloudflare_workers_ai(table, "@cf/moonshotai/kimi-k2.7-code") + assert mp is not None + assert mp.input == Decimal("0.00000095") + assert mp.output == Decimal("0.000004") + assert mp.cache_read == Decimal("0.00000019") + + +def test_cloudflare_skips_non_token_priced_model() -> None: + """A real model priced only in "per audio minute" — not a canonical priced + field — must be absent from the table entirely, not stored with a bogus + zero/None price that could be mistaken for "free".""" + table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + assert "@cf/pipecat-ai/smart-turn-v2" not in table + + +def test_cloudflare_skips_model_with_no_price_property() -> None: + table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + assert "@cf/some/unpriced-model" not in table + + +def test_cloudflare_lookup_miss_returns_none() -> None: + table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + assert lookup_cloudflare_workers_ai(table, "@cf/totally/made-up-model") is None + + +def test_cloudflare_lookup_version_suffix_fallback() -> None: + """Real drift we've observed: a live response naming a model with a + trailing "-v2" the catalog itself doesn't have listed separately.""" + table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + mp = lookup_cloudflare_workers_ai(table, "@cf/meta/llama-3.3-70b-instruct-fp8-fast-v2") + assert mp is not None + assert mp.input == Decimal("0.000000293") + + +@pytest.mark.parametrize( + "requested", + [ + "@cf/meta/llama-3.3-70b-instruct-fp8-fast", + "workers-ai/@cf/meta/llama-3.3-70b-instruct-fp8-fast", + # The routing prefix and the version-suffix drift, together. + "workers-ai/@cf/meta/llama-3.3-70b-instruct-fp8-fast-v2", + ], +) +def test_cloudflare_lookup_accepts_the_compat_routing_prefix(requested: str) -> None: + """Cloudflare's catalog lists bare "@cf/..." names, but reaching a model + through the gateway's OpenAI-compatible `/compat` endpoint requires the + "workers-ai/" prefix — the form the README prescribes and the only form a + streaming call reports. Both must price to the same rate.""" + table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + mp = lookup_cloudflare_workers_ai(table, requested) + assert mp is not None, f"{requested} should have priced" + assert mp.input == Decimal("0.000000293") + + +def test_cloudflare_lookup_miss_is_still_a_miss_with_the_prefix() -> None: + """The prefix strip must not turn an unknown model into a false hit.""" + table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + assert lookup_cloudflare_workers_ai(table, "workers-ai/@cf/nope/not-a-model") is None + + +@pytest.mark.parametrize( + "requested", + [ + "@cf/meta/llama-3.3-70b-instruct-fp8-fast", + "workers-ai/@cf/meta/llama-3.3-70b-instruct-fp8-fast", + ], +) +def test_workers_ai_provider_inferred_from_both_spellings(requested: str) -> None: + """A streaming Workers AI call carries no response model, so the requested + string — which the docs give in prefixed form — is all `_infer_provider` + has. Stamping "openai" there priced it against OpenRouter, missed, and + silently degraded to token events.""" + u = extract_openai_native({"usage": {"prompt_tokens": 10, "completion_tokens": 5}}, model_id=requested) + assert u.provider == "workers-ai" + # The model keeps the spelling the customer used — the strip happens at lookup, + # so reporting stays faithful to the request. + assert u.model == requested + + +@pytest.mark.parametrize( + ("usage", "expected", "why"), + [ + # Ancor's cited case: a real captured Gemini row. `input + output` dropped + # 852 additive reasoning tokens and published unit="30" for 882 consumed. + ( + CanonicalUsage(input=9, output=21, reasoning=852, provider="gemini", api="x", model="m"), + 882, + "gemini reasoning is additive", + ), + # Cache-inclusive provider: cache_read sits INSIDE input, so counting both + # would double it. + ( + CanonicalUsage(input=10000, output=100, cache_read=9000, provider="openai", api="x", model="m"), + 10100, + "openai cache_read is a subset of input", + ), + # Additive provider: cache_read/cache_write are real extra consumption, and + # the old basis under-reported this by 9.6x. + ( + CanonicalUsage( + input=1000, + output=100, + cache_read=9000, + cache_write=500, + provider="anthropic", + api="x", + model="m", + ), + 10600, + "anthropic cache is additive", + ), + # reasoning ⊆ output for openai — must not be added on top. + ( + CanonicalUsage(input=10, output=100, reasoning=80, provider="openai", api="x", model="m"), + 110, + "openai reasoning is a subset of output", + ), + # tool_calls is a CALL COUNT, not tokens, so it must never land in a token total. + ( + CanonicalUsage(input=10, output=20, tool_calls=3, provider="openai", api="x", model="m"), + 30, + "tool_calls excluded", + ), + ], +) +def test_deoverlapped_token_total(usage: CanonicalUsage, expected: int, why: str) -> None: + assert deoverlapped_token_total(usage) == expected, why + + +@pytest.mark.parametrize("provider", ["openai", "workers-ai"]) +def test_openai_shaped_providers_treat_reasoning_as_a_subset(provider: str) -> None: + """workers-ai is reached ONLY through Cloudflare's OpenAI-compatible endpoint, so + reasoning is a subset of output there exactly as it is for real OpenAI. Omitting it + from _OUTPUT_INCLUDES_REASONING counted the subset twice — 1900 against 1100.""" + u = CanonicalUsage(input=100, output=1000, reasoning=800, model="m", provider=provider, api="chat") + assert deoverlapped_token_total(u) == 1100 + + +def test_precomputed_unit_matches_the_split_path_basis() -> None: + """The two cost branches must report the same quantity for one call — that was + the actual complaint: `unit` on the single-event path used a different basis + from `parts["tokens"]` on the split path.""" + received: list = [] + provider = _warm_provider() + sdk, got = _price_sdk(provider) + u = CanonicalUsage( + input=1000, output=100, cache_read=900, model="claude-opus-4-8", provider="anthropic", api="native" + ) + # Split path (real per-field breakdown). + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + split = [e for batch in got for e in batch] + split_total = sum(int(e["properties"]["unit"]) for e in split) + + # Single-event path (precomputed cost). + sdk2, got2 = _price_sdk(_warm_provider()) + sdk2.emit(u, usd_cost=0.05) + assert sdk2.flush(timeout=2.0) + sdk2.shutdown(timeout=1.0) + single = [e for batch in got2 for e in batch] + assert len(single) == 1 + assert int(single[0]["properties"]["unit"]) == split_total, ( + f"single-event unit {single[0]['properties']['unit']} != split total {split_total}" + ) + _ = received + + +def test_cloudflare_entry_with_null_properties_does_not_unprice_everything() -> None: + """`.get("properties", [])` only defaults when the key is ABSENT — an explicit + JSON null returns None and `for p in None` raised TypeError out of this + function into maybe_refresh's handler, leaving the whole table None. One + malformed entry would unprice EVERY Workers AI model, not just its own.""" + raw = [ + {"name": "@cf/broken/model", "properties": None}, + { + "name": "@cf/good/model", + "properties": [ + { + "property_id": "price", + "value": [{"unit": "per M input tokens", "price": 1.0, "currency": "USD"}], + } + ], + }, + ] + table = parse_cloudflare_workers_ai(raw) + assert "@cf/good/model" in table, "a sibling entry must survive a malformed one" + assert "@cf/broken/model" not in table + + +def _cf_pages(*counts: int, total_count: int | None = None) -> list[dict]: + """Fake paged responses: `counts[i]` models on page i+1.""" + pages = [] + for n in counts: + info: dict = {"page": len(pages) + 1, "per_page": 50, "count": n} + if total_count is not None: + info["total_count"] = total_count + pages.append( + { + "result": [ + { + "name": f"@cf/m/p{len(pages) + 1}-{i}", + "properties": [ + { + "property_id": "price", + "value": [{"unit": "per M input tokens", "price": 1.0, "currency": "USD"}], + } + ], + } + for i in range(n) + ], + "result_info": info, + } + ) + return pages + + +def _run_cf_fetch(pages: list[dict]) -> tuple[int, list[int]]: + """Drive fetch_cloudflare_workers_ai against faked pages; return (models, pages hit).""" + import requests as _rq + + seen: list[int] = [] + orig = _rq.get + + class _Resp: + def __init__(self, body): + self._b = body + + def raise_for_status(self): + pass + + def json(self): + return self._b + + def fake(url, **kw): + page = int((kw.get("params") or {}).get("page", 1)) + seen.append(page) + return _Resp(pages[page - 1] if page - 1 < len(pages) else {"result": [], "result_info": {}}) + + _rq.get = fake + try: + f = HttpPricingFetcher(cloudflare_account_id="acct", cloudflare_api_token="tok") + table = f.fetch_cloudflare_workers_ai() + finally: + _rq.get = orig + return len(table), seen + + +def test_cloudflare_pagination_walks_until_a_short_page() -> None: + """Matches the real endpoint, which serves 50 then 14 then 0.""" + n, seen = _run_cf_fetch(_cf_pages(50, 14, total_count=291)) + assert seen == [1, 2], f"should stop after the short page, hit {seen}" + assert n == 64 + + +def test_cloudflare_pagination_survives_a_missing_total_count() -> None: + """The bug: `total_count` defaulting to len(models) made an absent count break + after page one, silently keeping 50 of the 64 available.""" + n, seen = _run_cf_fetch(_cf_pages(50, 14, total_count=None)) + assert seen == [1, 2], f"a missing total_count must not stop paging, hit {seen}" + assert n == 64 + + +def test_cloudflare_pagination_ignores_a_wrong_total_count() -> None: + """Measured live: the endpoint reports total_count=291 while serving 64, so a + `len(models) >= total` test can never be the terminator.""" + n, _ = _run_cf_fetch(_cf_pages(50, 14, total_count=291)) + assert n == 64 + + +def test_cloudflare_pagination_is_bounded() -> None: + """This runs on the queue's flush tick ahead of the drain, so an endpoint that + always returns a full page must not stall event delivery.""" + n, seen = _run_cf_fetch(_cf_pages(*([50] * 60), total_count=100000)) + assert len(seen) <= 40, f"loop must be bounded, hit {len(seen)} pages" + + +def test_cloudflare_fetcher_returns_empty_without_credentials() -> None: + """No account id / token set — Workers AI pricing is simply unavailable, + not an error; the fetch never even makes a request.""" + fetcher = HttpPricingFetcher() + assert fetcher.fetch_cloudflare_workers_ai() == {} + + +# ---------------------------------------------------------------------- +# Mistral alias resolution +# ---------------------------------------------------------------------- +def test_mistral_parses_real_alias_shape() -> None: + aliases = parse_mistral_aliases(_MISTRAL_MODELS_RAW) + assert aliases["mistral-small-latest"] == "mistral-small-2603" + assert aliases["mistral-vibe-cli-fast"] == "mistral-small-2603" + assert aliases["magistral-small-latest"] == "mistral-small-2603" + assert aliases["mistral-large-latest"] == "mistral-large-2411" + + +def test_mistral_model_with_no_aliases_contributes_nothing() -> None: + aliases = parse_mistral_aliases(_MISTRAL_MODELS_RAW) + assert "codestral-2508" not in aliases # it's an id, never requested as an alias + + +def test_mistral_alias_resolves_to_a_real_openrouter_listing() -> None: + """The whole point: the resolved id isn't a dead end — OpenRouter lists it.""" + aliases = parse_mistral_aliases(_MISTRAL_MODELS_RAW) + table = parse_openrouter(_OPENROUTER_RAW) + resolved = aliases["mistral-small-latest"] + mp = lookup_openrouter(table, "mistral", resolved) + assert mp is not None + assert mp.input == Decimal("0.00000015") + assert mp.output == Decimal("0.0000006") + assert mp.cache_read == Decimal("0.000000015") + + +def test_mistral_fetcher_returns_empty_without_credentials() -> None: + """No API key set — alias resolution is simply skipped, not an error; the + fetch never even makes a request.""" + fetcher = HttpPricingFetcher() + assert fetcher.fetch_mistral_aliases() == {} + + +def test_mistral_fetcher_accepts_a_key_passed_at_call_time() -> None: + """The key learned from a wrapped client (see PricingProvider.learn_mistral_api_key) + is passed per-call, not baked into the fetcher at construction — no + explicit config key is required for this path to work.""" + fetcher = HttpPricingFetcher() + calls = [] + + class _FakeResp: + def raise_for_status(self): + pass + + def json(self): + return {"data": []} + + def _fake_get(url, headers=None, timeout=None): + calls.append(headers) + return _FakeResp() + + import requests as _requests + + orig = _requests.get + _requests.get = _fake_get + try: + fetcher.fetch_mistral_aliases(api_key="learned-key-123") + finally: + _requests.get = orig + assert calls == [{"Authorization": "Bearer learned-key-123"}] + + +def test_mistral_fetcher_explicit_config_key_wins_over_learned_key() -> None: + """A key deliberately set via LagoConfig.mistral_api_key must not be + silently shadowed by one auto-detected from a wrapped client.""" + fetcher = HttpPricingFetcher(mistral_api_key="configured-key") + calls = [] + + class _FakeResp: + def raise_for_status(self): + pass + + def json(self): + return {"data": []} + + def _fake_get(url, headers=None, timeout=None): + calls.append(headers) + return _FakeResp() + + import requests as _requests + + orig = _requests.get + _requests.get = _fake_get + try: + fetcher.fetch_mistral_aliases(api_key="learned-key-123") + finally: + _requests.get = orig + assert calls == [{"Authorization": "Bearer configured-key"}] + + +def test_mistral_mutual_aliasing_resolves_to_the_dated_snapshot_not_another_alias() -> None: + """Real bug, found live: naively mapping "each alias -> this entry's id" + is order-dependent when Mistral lists a "-latest" moniker as its OWN + top-level `id` too (it does, for every alias in this real shape) — it + resolved "mistral-small-latest" to "magistral-small-latest" (whichever + entry got processed last), not "mistral-small-2603". OpenRouter lists + the dated snapshot, never the sibling alias, so that resolution was a + dead end in production. Every one of the 4 mutually-aliasing names must + land on the single dated snapshot, regardless of which entry mentions + which or what order they're processed in.""" + aliases = parse_mistral_aliases(_MISTRAL_MODELS_RAW_MUTUAL_ALIASING) + assert aliases["mistral-small-latest"] == "mistral-small-2603" + assert aliases["mistral-vibe-cli-fast"] == "mistral-small-2603" + assert aliases["magistral-small-latest"] == "mistral-small-2603" + # The canonical name itself is never a key — nothing should "resolve" it + # to something else. + assert "mistral-small-2603" not in aliases + + +def test_mistral_mutual_aliasing_reversed_input_order_gives_same_result() -> None: + """The result must not depend on which entry the source API happens to + list first — that's exactly the bug this replaced (last-write-wins).""" + reversed_data = {"data": list(reversed(_MISTRAL_MODELS_RAW_MUTUAL_ALIASING["data"]))} + aliases = parse_mistral_aliases(reversed_data) + assert aliases["mistral-small-latest"] == "mistral-small-2603" + assert aliases["magistral-small-latest"] == "mistral-small-2603" + + +def test_mistral_two_way_aliasing_still_resolves() -> None: + """The simplest mutual case — just id A and id B each listing the + other — must also converge on one canonical (the dated one), not stay + as a symmetric pair or resolve backwards.""" + aliases = parse_mistral_aliases(_MISTRAL_MODELS_RAW_MUTUAL_ALIASING) + assert aliases["voxtral-small-latest"] == "voxtral-small-2507" + assert "voxtral-small-2507" not in aliases + + +def test_mistral_family_resolves_to_the_NEWEST_dated_snapshot() -> None: + """Regression: the tie-break used to resolve on the date ASCENDING. + + Every dated id in one family is the same length, so `(len(n), n)` fell + through to the alphabetical term — which for `-2402` / `-2407` / `-2411` is + the date, oldest first. The whole family collapsed onto `mistral-large-2402` + and got priced at a two-year-old rate. + """ + family = [ + "mistral-large-2402", + "mistral-large-2407", + "mistral-large-2411", + "mistral-large-latest", + ] + data = {"data": [{"id": n, "aliases": [x for x in family if x != n]} for n in family]} + aliases = parse_mistral_aliases(data) + assert aliases["mistral-large-latest"] == "mistral-large-2411" + + +def test_mistral_explicit_dated_snapshot_is_never_remapped() -> None: + """An exact snapshot request is already the id OpenRouter lists, so it must + pass through untouched. Remapping it onto the group's canonical priced it at + a sibling's rate — a mispricing, not a miss.""" + family = ["mistral-large-2402", "mistral-large-2411", "mistral-large-latest"] + data = {"data": [{"id": n, "aliases": [x for x in family if x != n]} for n in family]} + aliases = parse_mistral_aliases(data) + assert "mistral-large-2402" not in aliases + assert "mistral-large-2411" not in aliases + assert aliases["mistral-large-latest"] == "mistral-large-2411" + + +@pytest.mark.parametrize( + ("names", "expected"), + [ + # Mistral's own 4-digit YYMM convention. + (["m-2402", "m-2411", "m-latest"], "m-2411"), + # Mixed widths: "20250929" sorts BELOW "2411" as a raw string, so the + # normalization to one scale is what makes this come out right. + (["m-2411", "m-20250929", "m-latest"], "m-20250929"), + # No dated candidate at all — deterministic shortest-then-code-point. + (["mm-latest", "m-latest"], "m-latest"), + ], +) +def test_mistral_canonical_picks_newest_across_suffix_shapes(names: list[str], expected: str) -> None: + assert _pick_mistral_canonical(names) == expected + + +def test_mistral_canonical_orders_by_code_point_not_locale() -> None: + """Cross-repo parity: the JS port must not use `localeCompare`, which is + ICU/locale-dependent. Both repos must pick the same canonical for a group + whose members differ only by case/separator — and the pick has to be the one + that still normalizes onto a name OpenRouter lists.""" + assert _pick_mistral_canonical(["mistral-small-2603", "Mistral-Small-2603"]) == "Mistral-Small-2603" + + # ---------------------------------------------------------------------- # Bedrock region + key + offer parsing # ---------------------------------------------------------------------- @@ -238,18 +962,127 @@ def test_compute_cost_only_unpriced_fields_yields_zero() -> None: assert b.fields == {} +def test_compute_precomputed_cost_matches_gateway_reported_amount() -> None: + """Cloudflare AI Gateway reports its own real cost per call (e.g. the + `cost` field on a log entry, in USD) — this must bill that exact amount, + not something recomputed from a per-token table.""" + b = compute_precomputed_cost(0.00010472, Decimal("1")) + assert b.total == "0.00010472" + assert b.total_cents == "0.010472" + assert b.base == "0.00010472" + assert b.source == "precomputed" + assert b.fields == {} # no per-field breakdown — Cloudflare gives one lump sum + + +def test_compute_precomputed_cost_applies_markup() -> None: + b = compute_precomputed_cost(0.0001, Decimal("2")) + assert b.base == "0.0001" + assert b.total == "0.0002" + assert b.total_cents == "0.02" + + +def test_compute_precomputed_cost_negative_floors_to_zero() -> None: + b = compute_precomputed_cost(-5, Decimal("1")) + assert b.total == "0" + assert b.base == "0" + + def test_money_golden_cases() -> None: cases = json.loads((FIXTURES / "money_golden.json").read_text())["cases"] for c in cases: prices = {k: Decimal(v) for k, v in c["prices"].items()} price = ModelPrice(source="openrouter", **prices) - usage = CanonicalUsage(model="m", provider="p", api="native", **c["counts"]) + # `provider` is optional and defaults to a name in no _INCLUDES_ set, so + # the pre-existing cases keep their original semantics; cases that pin + # per-provider token semantics set it explicitly. + usage = CanonicalUsage(model="m", provider=c.get("provider", "p"), api="native", **c["counts"]) b = compute_cost(usage, price, Decimal(c["markup"])) assert b.base == c["base"], f"{c['name']}: base {b.base} != {c['base']}" assert b.total == c["total"], f"{c['name']}: total {b.total} != {c['total']}" assert b.total_cents == c["total_cents"], f"{c['name']}: cents {b.total_cents} != {c['total_cents']}" +def test_money_golden_precomputed_cases() -> None: + """The gateway path: a lump sum the caller already knows. + + Several of these are verbatim `cost` values from real Cloudflare AI Gateway + log entries. JS renders any number below 1e-6 in exponential notation, so + these are the cases where the two repos silently disagreed on real money. + """ + cases = json.loads((FIXTURES / "money_golden.json").read_text())["precomputed_cases"] + for c in cases: + b = compute_precomputed_cost(c["usd_cost"], Decimal(c["markup"])) + assert b.base == c["base"], f"{c['name']}: base {b.base} != {c['base']}" + assert b.total == c["total"], f"{c['name']}: total {b.total} != {c['total']}" + assert b.total_cents == c["total_cents"], f"{c['name']}: cents {b.total_cents} != {c['total_cents']}" + + +def test_workers_ai_cache_read_is_subtracted_from_input() -> None: + """Regression: Workers AI is reached only through Cloudflare's OpenAI-COMPATIBLE + endpoint, so its `prompt_tokens` already includes the cached tokens. Counts and + rates here are real — a live cached call reported prompt=23233/cached=23168, and + @cf/moonshotai/kimi-k2.6 lists input $0.95/M with cached input $0.16/M. Billing + all 23233 at the input rate charged the cached portion twice (+583%). + """ + price = ModelPrice( + source="cloudflare_workers_ai", + input=Decimal("0.00000095"), + cache_read=Decimal("0.00000016"), + ) + usage = CanonicalUsage( + model="@cf/moonshotai/kimi-k2.6", + provider="workers-ai", + api="chat.completions", + input=23233, + cache_read=23168, + ) + b = compute_cost(usage, price, Decimal("1")) + # only the 65 uncached tokens are billed at the input rate + assert b.fields["input"]["tokens"] == "65" + assert b.fields["cache_read"]["tokens"] == "23168" + assert b.total == "0.00376863" + + +def test_anthropic_cache_read_stays_additive() -> None: + """The other side of the same rule: Anthropic reports input EXCLUSIVE of cache, + so nothing may be subtracted. Same counts/rates as the workers-ai case above.""" + price = ModelPrice( + source="openrouter", + input=Decimal("0.00000095"), + cache_read=Decimal("0.00000016"), + ) + usage = CanonicalUsage( + model="claude-x", provider="anthropic", api="native", input=23233, cache_read=23168 + ) + b = compute_cost(usage, price, Decimal("1")) + assert b.fields["input"]["tokens"] == "23233" + assert b.total == "0.02577823" + + +def test_parse_price_accepts_exponential_notation() -> None: + """Real gateway costs below 1e-6 arrive in exponential form. Python has always + handled these; the golden fixture pins JS to the same values.""" + assert _parse_price(9.807224944233895e-07) == Decimal("0.000000980722") + assert _parse_price("8.91e-7") == Decimal("0.000000891") + assert _parse_price("9.78e-07") == Decimal("0.000000978") + # below the 12dp floor -> zero, not None (a real but unbillably small amount) + assert _parse_price(1e-13) == Decimal(0) + + +def test_parse_price_returns_none_instead_of_raising_on_huge_values() -> None: + """Regression: `.quantize()` sat outside the try, so any value >= 1e16 raised + InvalidOperation straight out of this function — past every caller that relies + on the documented None, and out of compute_precomputed_cost into emit()'s + catch-all, where the event was dropped as an unknown error rather than taking + the normal 'no price' path.""" + assert _parse_price("1e15") == Decimal("1000000000000000") + assert _parse_price("1e16") is None + assert _parse_price("1e30") is None + assert _parse_price("1e999999999") is None + # and the tiny end stays a real zero, not a None + assert _parse_price("1e-999999999") == Decimal(0) + + def test_coerce_markup() -> None: assert coerce_markup(1.2) == (Decimal("1.2"), True) assert coerce_markup("2") == (Decimal("2"), True) @@ -283,6 +1116,191 @@ def test_provider_token_mode_does_no_fetch() -> None: assert fetcher.openrouter_calls == 0 +def test_provider_cloudflare_workers_ai_cold_then_warm() -> None: + cf_table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + fetcher = StubFetcher(cloudflare_workers_ai=cf_table) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + # cold: no table yet -> None, and flags it for refresh + assert p.lookup("workers-ai", "@cf/meta/llama-3.3-70b-instruct-fp8-fast", "cloudflare_gateway") is None + assert fetcher.cloudflare_workers_ai_calls == 0 + p.maybe_refresh() + assert fetcher.cloudflare_workers_ai_calls == 1 + mp = p.lookup("workers-ai", "@cf/meta/llama-3.3-70b-instruct-fp8-fast", "cloudflare_gateway") + assert mp is not None and mp.input == Decimal("0.000000293") + + +def test_provider_cloudflare_workers_ai_only_fetched_for_workers_ai_provider() -> None: + """A lookup for a totally different provider must not flag the Cloudflare + source stale — each source only ever fetches for the traffic that needs it.""" + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.lookup("anthropic", "claude-opus-4-8", "native") + p.maybe_refresh() + assert fetcher.cloudflare_workers_ai_calls == 0 + + +def test_provider_mistral_alias_cold_miss_then_warm_resolves() -> None: + """Cold: the alias table hasn't been fetched yet, so the raw alias string + is looked up against OpenRouter directly and misses safely — never worse + than before this resolution step existed. Warm: it resolves and hits.""" + fetcher = StubFetcher( + openrouter=parse_openrouter(_OPENROUTER_RAW), + mistral_aliases=parse_mistral_aliases(_MISTRAL_MODELS_RAW), + ) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + # cold: openrouter table is ALSO cold here, so this exercises both misses + # at once — the important thing is it's a clean None, not an exception. + assert p.lookup("mistral", "mistral-small-latest", "native") is None + p.maybe_refresh() + assert fetcher.mistral_aliases_calls == 1 + assert fetcher.openrouter_calls == 1 + mp = p.lookup("mistral", "mistral-small-latest", "native") + assert mp is not None + assert mp.input == Decimal("0.00000015") + assert mp.output == Decimal("0.0000006") + + +def test_learn_mistral_api_key_is_used_on_next_fetch() -> None: + """No LagoConfig.mistral_api_key was ever configured — the key is + learned instead (e.g. from a wrapped client) and still reaches the + fetcher on the next refresh.""" + fetcher = StubFetcher(mistral_aliases=parse_mistral_aliases(_MISTRAL_MODELS_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.learn_mistral_api_key("learned-from-client-key") + p.prime(providers=["mistral"]) + p.maybe_refresh() + assert fetcher.mistral_aliases_calls == 1 + assert fetcher.last_mistral_api_key == "learned-from-client-key" + + +def test_learn_mistral_api_key_does_not_overwrite_an_already_learned_key() -> None: + """First-learned key wins — a second call (e.g. wrap() invoked again for + a second client) doesn't clobber it.""" + fetcher = StubFetcher(mistral_aliases=parse_mistral_aliases(_MISTRAL_MODELS_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.learn_mistral_api_key("first-key") + p.learn_mistral_api_key("second-key") + p.prime(providers=["mistral"]) + p.maybe_refresh() + assert fetcher.last_mistral_api_key == "first-key" + + +def test_learn_mistral_api_key_ignores_empty_string() -> None: + fetcher = StubFetcher(mistral_aliases=parse_mistral_aliases(_MISTRAL_MODELS_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.learn_mistral_api_key("") + p.prime(providers=["mistral"]) + p.maybe_refresh() + assert fetcher.last_mistral_api_key is None + + +def test_provider_mistral_lookup_without_credentials_falls_back_to_raw_model() -> None: + """No Mistral API key configured -> fetch_mistral_aliases returns {} -> + the alias string is looked up as-is against OpenRouter, same behavior as + before this feature existed (a safe miss for an alias, a hit for a + non-aliased model like "mistral-large" that's already the exact id).""" + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) # mistral_aliases defaults to {} + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.lookup("mistral", "mistral-large", "native") + p.maybe_refresh() + mp = p.lookup("mistral", "mistral-large", "native") + assert mp is not None and mp.input == Decimal("0.000002") + + +def test_provider_mistral_alias_only_fetched_for_mistral_provider() -> None: + """A lookup for a totally different provider must not flag the Mistral + alias source stale — each source only ever fetches for the traffic that + needs it.""" + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.lookup("anthropic", "claude-opus-4-8", "native") + p.maybe_refresh() + assert fetcher.mistral_aliases_calls == 0 + assert fetcher.openrouter_calls == 1 + + +def test_prime_only_eagerly_warms_openrouter_not_cloudflare_or_mistral() -> None: + """prime() (called automatically when pricing_mode="price" is the global + default, and by warm_pricing()) must not force-fetch Cloudflare/Mistral — + both are credential-gated and provider-specific, and most price-mode + customers never call either. Eagerly hitting their APIs at construction + time regardless of actual usage would be pure waste. Only a real lookup + for that specific provider should ever trigger their fetch.""" + fetcher = StubFetcher( + openrouter=parse_openrouter(_OPENROUTER_RAW), + cloudflare_workers_ai=parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW), + mistral_aliases=parse_mistral_aliases(_MISTRAL_MODELS_RAW), + ) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.prime() + p.maybe_refresh() + assert fetcher.openrouter_calls == 1 + assert fetcher.cloudflare_workers_ai_calls == 0 + assert fetcher.mistral_aliases_calls == 0 + # Confirms it's not just "hasn't fetched yet" — a real lookup for either + # provider afterward still works, fetching lazily on its own trigger. + assert p.lookup("workers-ai", "@cf/meta/llama-3.3-70b-instruct-fp8-fast", "cloudflare_gateway") is None + p.maybe_refresh() + assert fetcher.cloudflare_workers_ai_calls == 1 + mp = p.lookup("workers-ai", "@cf/meta/llama-3.3-70b-instruct-fp8-fast", "cloudflare_gateway") + assert mp is not None + + +def test_prime_with_providers_eagerly_warms_the_named_ones_too() -> None: + """Opt-in escape hatch: a caller who already knows they're about to call + Mistral and/or Workers AI this session can say so up front and skip the + one-time lazy cold-start cost for THAT provider's first call too — + without going back to unconditionally warming both for every customer.""" + fetcher = StubFetcher( + openrouter=parse_openrouter(_OPENROUTER_RAW), + cloudflare_workers_ai=parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW), + mistral_aliases=parse_mistral_aliases(_MISTRAL_MODELS_RAW), + ) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.prime(providers=["mistral", "workers-ai"]) + p.maybe_refresh() + assert fetcher.openrouter_calls == 1 + assert fetcher.cloudflare_workers_ai_calls == 1 + assert fetcher.mistral_aliases_calls == 1 + # Both now resolve correctly on their very first real lookup — no cold miss. + assert p.lookup("mistral", "mistral-small-latest", "native") is not None + assert ( + p.lookup("workers-ai", "@cf/meta/llama-3.3-70b-instruct-fp8-fast", "cloudflare_gateway") is not None + ) + + +def test_prime_with_unknown_provider_name_is_ignored_not_an_error() -> None: + """This is a hint, not a contract — a typo'd or unrecognized provider + name is silently ignored rather than raising.""" + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.prime(providers=["totally-made-up-provider"]) + p.maybe_refresh() + assert fetcher.openrouter_calls == 1 + assert fetcher.cloudflare_workers_ai_calls == 0 + assert fetcher.mistral_aliases_calls == 0 + + +def test_warm_pricing_with_providers_threads_through_from_sdk() -> None: + """Same opt-in escape hatch, exercised through LagoSDK.warm_pricing() + rather than the PricingProvider directly.""" + fetcher = StubFetcher( + openrouter=parse_openrouter(_OPENROUTER_RAW), + mistral_aliases=parse_mistral_aliases(_MISTRAL_MODELS_RAW), + ) + provider = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + cfg = LagoConfig( + api_key="dummy", + default_subscription_id="sub_default", + pricing_mode="price", + pricing_provider=provider, + ) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk.warm_pricing(providers=["mistral"]) + assert fetcher.mistral_aliases_calls == 1 + assert provider.lookup("mistral", "mistral-small-latest", "native") is not None + + def test_provider_bedrock_region_routing() -> None: bedrock_table = parse_bedrock_offer( { @@ -326,6 +1344,17 @@ def _warm_provider() -> PricingProvider: return p +def _warm_cloudflare_provider() -> PricingProvider: + cf_table = parse_cloudflare_workers_ai(_CLOUDFLARE_MODELS_RAW) + p = PricingProvider(fetcher=StubFetcher(cloudflare_workers_ai=cf_table), ttl_seconds=3600) + # prime() no longer eagerly warms Cloudflare (it's credential-gated and + # provider-specific — see prime()'s docstring) — a real first lookup for + # this provider is what flags it stale, same as production. + p.lookup("workers-ai", "@cf/meta/llama-3.3-70b-instruct-fp8-fast", "cloudflare_gateway") + p.maybe_refresh() + return p + + def _price_sdk( provider: PricingProvider, default_sub: str = "sub_default", on_error=None, markup: float = 1.0 ): @@ -343,40 +1372,144 @@ def _price_sdk( return sdk, received -def test_price_mode_emits_single_cost_event() -> None: +def _by_token_type(received: list) -> dict[str, dict]: + flat = [e for batch in received for e in batch] + assert all(e["code"] == "llm_cost" for e in flat) + return {e["properties"]["token_type"]: e for e in flat} + + +def test_warm_pricing_closes_the_cold_start_race() -> None: + """Without warm_pricing(), a call made immediately after construction hits + a cold table and emit() falls back to token events (or, with no token + metric configured, loses the event entirely). warm_pricing() blocks until + the table is fetched, so the very first call in price mode prices + correctly instead of racing the background thread's first tick.""" + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + provider = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + cfg = LagoConfig( + api_key="dummy", + default_subscription_id="sub_default", + pricing_mode="price", + pricing_provider=provider, + ) + sdk = LagoSDK(api_key="dummy", config=cfg) + received: list = [] + sdk._queue._sender = lambda b: received.append(list(b)) # type: ignore[attr-defined] + try: + assert provider.lookup("anthropic", "claude-opus-4-8", "native") is None # genuinely cold + + sdk.warm_pricing() + + assert provider.lookup("anthropic", "claude-opus-4-8", "native") is not None # now warm + u = CanonicalUsage( + input=1000, output=500, model="claude-opus-4-8", provider="anthropic", api="native" + ) + sdk.emit(u) + assert sdk.flush(timeout=2.0) + finally: + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert all(e["code"] == "llm_cost" for e in flat) # priced, not a token-event fallback + + +def test_bad_markup_is_coerced_to_one_reported_and_still_bills_the_cost() -> None: + """`markup` is customer input (`extra_lago={"markup": ...}`), so a comma decimal + like "1,5" genuinely arrives. `coerce_markup` is the guard that catches it; this + pins the end-to-end consequence, which no test covered: the cost is still billed + at 1.0 rather than zeroed, and the lost markup reaches on_error.""" + seen: list[tuple[Exception, str]] = [] + sdk, received = _price_sdk(_warm_provider(), on_error=lambda e, c: seen.append((e, c))) + try: + u = CanonicalUsage( + input=1000, output=500, model="claude-opus-4-8", provider="anthropic", api="native" + ) + # annotated `float | None`, but it arrives from untyped customer input + sdk.emit(u, markup="1,5") # type: ignore[arg-type] + assert sdk.flush(timeout=2.0) + finally: + sdk.shutdown(timeout=1.0) + + events = _by_token_type(received) + assert events, "a bad markup must not drop the cost events" + for token_type, ev in events.items(): + props = ev["properties"] + assert props["markup"] == "1", f"{token_type}: bad markup should coerce to 1.0" + assert props["value"] == props["base_cost"], ( + f"{token_type}: should bill the un-marked-up cost, not {props['value']!r}" + ) + assert Decimal(props["value"]) > 0, f"{token_type}: a bad markup must not zero the bill" + + contexts = [c for _, c in seen] + assert "pricing" in contexts, f"the invalid markup must reach on_error; got {contexts}" + assert any("markup" in str(e) and "1,5" in str(e) for e, _ in seen) + + +def test_price_mode_emits_one_event_per_token_type() -> None: + """A real per-field breakdown (OpenRouter has both input/output prices for + this model) splits into one llm_cost event per token_type, so Lago's + `grouped_by: ["model", "token_type"]` charge can break it down by both — + not one summed event that hides the split.""" sdk, received = _price_sdk(_warm_provider()) u = CanonicalUsage(input=1000, output=500, model="claude-opus-4-8", provider="anthropic", api="native") sdk.emit(u) assert sdk.flush(timeout=2.0) sdk.shutdown(timeout=1.0) - flat = [e for batch in received for e in batch] - assert len(flat) == 1 - ev = flat[0] - assert ev["code"] == "llm_cost" - # Lago dynamic charge: top-level cents amount = 0.0175 USD * 100 = 1.75 - assert ev["precise_total_amount_cents"] == "1.75" - props = ev["properties"] - # `unit` = total tokens (1000 + 500) — the sum-aggregation quantity - assert props["unit"] == "1500" - # 1000*0.000005 + 500*0.000025 = 0.005 + 0.0125 = 0.0175 - assert props["value"] == "0.0175" - assert props["base_cost"] == "0.0175" - assert props["price_source"] == "openrouter" - assert props["input_tokens"] == "1000" - assert props["input_unit_price"] == "0.000005" - assert props["output_cost"] == "0.0125" - - -def test_price_mode_markup_scales_value() -> None: + by_type = _by_token_type(received) + assert set(by_type) == {"input", "output"} + + inp = by_type["input"] + assert inp["properties"]["unit"] == "1000" + assert inp["properties"]["value"] == "0.005" # 1000 * 0.000005 + assert inp["properties"]["unit_price"] == "0.000005" + assert inp["properties"]["model"] == "claude-opus-4-8" + assert inp["properties"]["price_source"] == "openrouter" + # Lago dynamic charge cents = 0.005 USD * 100 = 0.5 + assert inp["precise_total_amount_cents"] == "0.5" + + out = by_type["output"] + assert out["properties"]["unit"] == "500" + assert out["properties"]["value"] == "0.0125" # 500 * 0.000025 + assert out["precise_total_amount_cents"] == "1.25" + + # Same call's split transaction ids don't collide with each other. + assert inp["transaction_id"] != out["transaction_id"] + + +def test_price_mode_workers_ai_uses_cloudflare_catalog_not_openrouter() -> None: + """Real captured shape: 38 input / 2 output tokens through + "@cf/meta/llama-3.3-70b-instruct-fp8-fast" — same call this catalog price + was verified against live (predicted $0.00001564 vs Cloudflare's own + real-charged $0.00001552; the ~0.8% gap is the catalog's own displayed + rate rounding to 3dp, not our computation).""" + sdk, received = _price_sdk(_warm_cloudflare_provider()) + u = CanonicalUsage( + input=38, + output=2, + model="@cf/meta/llama-3.3-70b-instruct-fp8-fast", + provider="workers-ai", + api="cloudflare_gateway", + ) + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + by_type = _by_token_type(received) + assert by_type["input"]["properties"]["price_source"] == "cloudflare_workers_ai" + # 38 * 0.000000293 + 2 * 0.000002253 = 0.000011134 + 0.000004506 = 0.00001564 + assert by_type["input"]["properties"]["value"] == "0.000011134" + assert by_type["output"]["properties"]["value"] == "0.000004506" + + +def test_price_mode_markup_scales_each_token_type_event() -> None: sdk, received = _price_sdk(_warm_provider(), markup=2.0) u = CanonicalUsage(input=1000, output=500, model="claude-opus-4-8", provider="anthropic", api="native") sdk.emit(u) assert sdk.flush(timeout=2.0) sdk.shutdown(timeout=1.0) - ev = [e for batch in received for e in batch][0] - assert ev["properties"]["base_cost"] == "0.0175" - assert ev["properties"]["value"] == "0.035" # 0.0175 * 2 - assert ev["properties"]["markup"] == "2" + by_type = _by_token_type(received) + assert by_type["input"]["properties"]["base_cost"] == "0.005" + assert by_type["input"]["properties"]["value"] == "0.01" # 0.005 * 2 + assert by_type["input"]["properties"]["markup"] == "2" + assert by_type["output"]["properties"]["value"] == "0.025" # 0.0125 * 2 def test_per_call_markup_overrides_global() -> None: @@ -385,8 +1518,9 @@ def test_per_call_markup_overrides_global() -> None: sdk.emit(u, markup=3.0) assert sdk.flush(timeout=2.0) sdk.shutdown(timeout=1.0) - ev = [e for batch in received for e in batch][0] - assert ev["properties"]["value"] == "0.0525" # 0.0175 * 3 + by_type = _by_token_type(received) + assert by_type["input"]["properties"]["value"] == "0.015" # 0.005 * 3 + assert by_type["output"]["properties"]["value"] == "0.0375" # 0.0125 * 3 # ---------------------------------------------------------------------- @@ -403,15 +1537,15 @@ def test_price_mode_openai_cache_read_subset_not_double_billed() -> None: sdk.emit(u) assert sdk.flush(timeout=2.0) sdk.shutdown(timeout=1.0) - props = [e for batch in received for e in batch][0]["properties"] + by_type = _by_token_type(received) + assert set(by_type) == {"input", "cache_read", "output"} # input billed for only the non-cached portion (1000 - 800); cache billed at cache rate - assert props["input_tokens"] == "200" - assert props["cache_read_tokens"] == "800" - # 200*0.0000025 + 800*0.00000125 + 500*0.00001 = 0.0005 + 0.001 + 0.005 = 0.0065 - # (the bug would bill input at full 1000 -> 0.0085) - assert props["value"] == "0.0065" - # unit = billed tokens 200 + 800 + 500 = 1500 = prompt(1000) + completion(500) - assert props["unit"] == "1500" + assert by_type["input"]["properties"]["unit"] == "200" + assert by_type["cache_read"]["properties"]["unit"] == "800" + # 200*0.0000025=0.0005, 800*0.00000125=0.001, 500*0.00001=0.005 (the bug would bill input at full 1000) + assert by_type["input"]["properties"]["value"] == "0.0005" + assert by_type["cache_read"]["properties"]["value"] == "0.001" + assert by_type["output"]["properties"]["value"] == "0.005" def test_price_mode_gemini_cache_subset_and_reasoning_additive() -> None: @@ -429,15 +1563,17 @@ def test_price_mode_gemini_cache_subset_and_reasoning_additive() -> None: sdk.emit(u) assert sdk.flush(timeout=2.0) sdk.shutdown(timeout=1.0) - props = [e for batch in received for e in batch][0]["properties"] - assert props["input_tokens"] == "700" # 1000 - 300 cached - assert props["cache_read_tokens"] == "300" - assert props["output_tokens"] == "400" - assert props["reasoning_tokens"] == "100" # billed separately (additive for Gemini) - # 700*3e-7 + 300*7.5e-8 + 400*2.5e-6 + 100*2.5e-6 = 0.00021+0.0000225+0.001+0.00025 = 0.0014825 - assert props["value"] == "0.0014825" - # unit = 700+300+400+100 = 1500 = prompt(1000)+candidates(400)+thoughts(100) - assert props["unit"] == "1500" + by_type = _by_token_type(received) + assert set(by_type) == {"input", "cache_read", "output", "reasoning"} + assert by_type["input"]["properties"]["unit"] == "700" # 1000 - 300 cached + assert by_type["cache_read"]["properties"]["unit"] == "300" + assert by_type["output"]["properties"]["unit"] == "400" + assert by_type["reasoning"]["properties"]["unit"] == "100" # billed separately (additive for Gemini) + # 700*3e-7=0.00021, 300*7.5e-8=0.0000225, 400*2.5e-6=0.001, 100*2.5e-6=0.00025 + assert by_type["input"]["properties"]["value"] == "0.00021" + assert by_type["cache_read"]["properties"]["value"] == "0.0000225" + assert by_type["output"]["properties"]["value"] == "0.001" + assert by_type["reasoning"]["properties"]["value"] == "0.00025" def test_price_mode_openai_reasoning_in_output_not_double_billed() -> None: @@ -447,13 +1583,13 @@ def test_price_mode_openai_reasoning_in_output_not_double_billed() -> None: sdk.emit(u) assert sdk.flush(timeout=2.0) sdk.shutdown(timeout=1.0) - props = [e for batch in received for e in batch][0]["properties"] - # reasoning folded into output — no separate reasoning line, output billed in full - assert "reasoning_tokens" not in props - assert props["output_tokens"] == "500" - # 100*0.0000025 + 500*0.00001 = 0.00025 + 0.005 = 0.00525 (bug would add 200*1e-5=0.002) - assert props["value"] == "0.00525" - assert props["unit"] == "600" # 100 + 500; reasoning not double-counted + by_type = _by_token_type(received) + # reasoning folded into output — no separate reasoning event, output billed in full + assert set(by_type) == {"input", "output"} + assert by_type["output"]["properties"]["unit"] == "500" + # 100*0.0000025=0.00025, 500*0.00001=0.005 (bug would add a separate 200*1e-5=0.002 reasoning event) + assert by_type["input"]["properties"]["value"] == "0.00025" + assert by_type["output"]["properties"]["value"] == "0.005" def test_price_mode_anthropic_cache_is_additive() -> None: @@ -471,13 +1607,154 @@ def test_price_mode_anthropic_cache_is_additive() -> None: sdk.emit(u) assert sdk.flush(timeout=2.0) sdk.shutdown(timeout=1.0) - props = [e for batch in received for e in batch][0]["properties"] - assert props["input_tokens"] == "1000" # unchanged — additive provider - assert props["cache_read_tokens"] == "400" - assert props["cache_write_tokens"] == "200" - # 1000*5e-6 + 500*25e-6 + 400*5e-7 + 200*6.25e-6 = 0.005+0.0125+0.0002+0.00125 = 0.01895 - assert props["value"] == "0.01895" - assert props["unit"] == "2100" # 1000+500+400+200, all additive + by_type = _by_token_type(received) + assert set(by_type) == {"input", "output", "cache_read", "cache_write"} + assert by_type["input"]["properties"]["unit"] == "1000" # unchanged — additive provider + assert by_type["cache_read"]["properties"]["unit"] == "400" + assert by_type["cache_write"]["properties"]["unit"] == "200" + # 1000*5e-6=0.005, 500*25e-6=0.0125, 400*5e-7=0.0002, 200*6.25e-6=0.00125 + assert by_type["input"]["properties"]["value"] == "0.005" + assert by_type["output"]["properties"]["value"] == "0.0125" + assert by_type["cache_read"]["properties"]["value"] == "0.0002" + assert by_type["cache_write"]["properties"]["value"] == "0.00125" + + +# ---------------------------------------------------------------------- +# usd_cost — the gateway-connector entrypoint: skip our own price lookup +# entirely and bill the caller's already-known real cost. +# ---------------------------------------------------------------------- +def test_usd_cost_skips_pricing_lookup_entirely() -> None: + """A COLD, never-warmed provider — if this passed, `emit` would have had + to fall back to token events (no price available). It doesn't: usd_cost + bypasses `_pricing.lookup` altogether, so a cold provider is irrelevant.""" + cold_provider = PricingProvider(fetcher=StubFetcher(openrouter={}), ttl_seconds=3600) + sdk, received = _price_sdk(cold_provider) + u = CanonicalUsage( + input=38, output=41, model="@cf/meta/llama-3.3-70b", provider="workers-ai", api="cloudflare_gateway" + ) + sdk.emit(u, usd_cost=0.00010472) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert len(flat) == 1 + ev = flat[0] + assert ev["code"] == "llm_cost" + assert ev["precise_total_amount_cents"] == "0.010472" + props = ev["properties"] + assert props["price_source"] == "precomputed" + assert props["value"] == "0.00010472" + # No per-field breakdown available — unit falls back to raw input+output. + assert props["unit"] == "79" + assert "input_tokens" not in props + + +def test_usd_cost_applies_markup_same_as_looked_up_price() -> None: + cold_provider = PricingProvider(fetcher=StubFetcher(openrouter={}), ttl_seconds=3600) + sdk, received = _price_sdk(cold_provider, markup=1.5) + u = CanonicalUsage(input=10, output=5, model="m", provider="workers-ai", api="cloudflare_gateway") + sdk.emit(u, usd_cost=0.0001) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + ev = [e for batch in received for e in batch][0] + assert ev["properties"]["base_cost"] == "0.0001" + assert ev["properties"]["value"] == "0.00015" + + +def test_usd_cost_ignored_in_token_mode() -> None: + """usd_cost is a price-mode-only override — in the default token mode it + must not do anything; the call still emits ordinary token events.""" + cold_provider = PricingProvider(fetcher=StubFetcher(openrouter={}), ttl_seconds=3600) + received: list = [] + cfg = LagoConfig(api_key="dummy", default_subscription_id="sub_default", pricing_provider=cold_provider) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: received.append(list(b)) # type: ignore[attr-defined] + u = CanonicalUsage(input=10, output=5, model="m", provider="workers-ai", api="cloudflare_gateway") + sdk.emit(u, usd_cost=0.0001) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + codes = {e["code"] for e in flat} + assert codes == {"llm_input_tokens", "llm_output_tokens"} + + +def test_event_id_used_as_transaction_id_in_price_mode() -> None: + """The connector's idempotency key: pass the source log entry's own id so + re-running a backfill over the same window doesn't double-bill.""" + cold_provider = PricingProvider(fetcher=StubFetcher(openrouter={}), ttl_seconds=3600) + sdk, received = _price_sdk(cold_provider) + u = CanonicalUsage(input=10, output=5, model="m", provider="workers-ai", api="cloudflare_gateway") + sdk.emit(u, usd_cost=0.0001, event_id="backfill_01ABC") + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + ev = [e for batch in received for e in batch][0] + assert ev["transaction_id"] == "backfill_01ABC" + + +def test_event_id_suffixed_per_field_in_token_mode() -> None: + """Token mode can push several events from one call (input, output, ...); + reusing the same event_id verbatim for all of them would collide, so each + field gets its own suffix off the same base id — in the `_tok_` namespace, + which keeps it distinct from the cost path's suffix for the same field.""" + received: list = [] + cfg = LagoConfig(api_key="dummy", default_subscription_id="sub_default") + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: received.append(list(b)) # type: ignore[attr-defined] + u = CanonicalUsage(input=10, output=5, model="m", provider="workers-ai", api="cloudflare_gateway") + sdk.emit(u, event_id="backfill_01ABC") + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + ids = {e["transaction_id"] for e in flat} + assert ids == {"backfill_01ABC_tok_input", "backfill_01ABC_tok_output"} + + +def test_token_fallback_and_cost_ids_never_collide_for_one_event_id() -> None: + """The bug this namespacing exists for. + + A price miss falls back to token events; the SAME window re-run once the + table is warm takes the cost path. Under one shared namespace both emitted + `{event_id}_input`, so Lago rejected the second as a duplicate — and since + `/events/batch` is all-or-nothing, that rejection failed every other event + in the batch too. The dollar amounts for that window were never billed, + only the raw token counts, and nothing surfaced it. + """ + u = CanonicalUsage(input=10, output=5, model="claude-opus-4-8", provider="anthropic", api="native") + + # Run 1: cold table -> price miss -> token fallback, same event_id. + cold = PricingProvider(fetcher=StubFetcher(openrouter={}), ttl_seconds=3600) + sdk_cold, got_cold = _price_sdk(cold) + sdk_cold.emit(u, event_id="backfill_01ABC") + assert sdk_cold.flush(timeout=2.0) + sdk_cold.shutdown(timeout=1.0) + cold_ids = {e["transaction_id"] for batch in got_cold for e in batch} + + # Run 2: warm table -> real per-field cost events, same event_id. + sdk_warm, got_warm = _price_sdk(_warm_provider()) + sdk_warm.emit(u, event_id="backfill_01ABC") + assert sdk_warm.flush(timeout=2.0) + sdk_warm.shutdown(timeout=1.0) + warm_ids = {e["transaction_id"] for batch in got_warm for e in batch} + + assert cold_ids, "cold run should have emitted token events" + assert warm_ids, "warm run should have emitted cost events" + assert not (cold_ids & warm_ids), ( + f"token-fallback and cost ids must not collide; overlap={cold_ids & warm_ids}" + ) + assert all("_tok_" in i for i in cold_ids) + assert all("_cost_" in i for i in warm_ids) + + +def test_no_event_id_still_falls_back_to_random_uuid() -> None: + """A live, one-shot call has no natural id to reuse — must still work + exactly as before this option existed.""" + cold_provider = PricingProvider(fetcher=StubFetcher(openrouter={}), ttl_seconds=3600) + sdk, received = _price_sdk(cold_provider) + u = CanonicalUsage(input=10, output=5, model="m", provider="workers-ai", api="cloudflare_gateway") + sdk.emit(u, usd_cost=0.0001) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + ev = [e for batch in received for e in batch][0] + uuid.UUID(ev["transaction_id"]) # raises if not a valid UUID def test_price_unavailable_falls_back_to_token_events_and_reports() -> None: @@ -508,7 +1785,7 @@ def test_per_call_price_mode_overrides_global_tokens() -> None: assert sdk.flush(timeout=2.0) sdk.shutdown(timeout=1.0) flat = [e for batch in received for e in batch] - assert len(flat) == 1 and flat[0]["code"] == "llm_cost" + assert len(flat) == 2 and all(e["code"] == "llm_cost" for e in flat) # one per token_type: input, output def test_default_mode_is_tokens_unchanged() -> None: @@ -523,3 +1800,22 @@ def test_default_mode_is_tokens_unchanged() -> None: sdk.shutdown(timeout=1.0) flat = [e for batch in received for e in batch] assert {e["code"] for e in flat} == {"llm_input_tokens", "llm_output_tokens"} + + +def test_workers_ai_compat_prefix_is_defined_exactly_once() -> None: + """Two unrelated layers must agree on this string: `adapters/openai_native` + decides the PROVIDER from it, `pricing` strips it before a catalog lookup. They + must never import each other, so it lives in `canonical`. A drift between two + copies is a silently unpriced call, not a crash — which is why this is asserted + rather than left to review.""" + src = pathlib.Path(__file__).resolve().parents[2] / "src" / "lago_agent_sdk" + definitions = [ + f"{path.relative_to(src)}:{i}" + for path in src.rglob("*.py") + for i, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1) + if re.match(r"\s*_?WORKERS_AI_COMPAT_PREFIX\s*=", line) + ] + assert definitions == ["canonical.py:19"] or len(definitions) == 1, ( + f"expected one definition, found {definitions}" + ) + assert WORKERS_AI_COMPAT_PREFIX == "workers-ai/" diff --git a/tests/unit/test_queue.py b/tests/unit/test_queue.py index 1728dfa..0423f6b 100644 --- a/tests/unit/test_queue.py +++ b/tests/unit/test_queue.py @@ -2,9 +2,15 @@ from __future__ import annotations +import pathlib +import subprocess +import sys import threading import time +import pytest + +from lago_agent_sdk.exceptions import LagoApiError from lago_agent_sdk.queue import EventQueue @@ -79,6 +85,364 @@ def test_flush_returns_true_when_drained(): q.shutdown(timeout=1.0) +# ---------------------------------------------------------------------- +# Permanent (4xx) vs transient failures. A duplicate transaction_id from +# replaying/backfilling the same window twice will NEVER succeed by retrying +# the same batch — it must be isolated and dropped, not block real events +# queued behind it forever the same way a genuine transient failure would. +# ---------------------------------------------------------------------- +def test_permanent_failure_isolates_bad_events_from_good_ones_in_same_batch(): + """Lago's batch endpoint is all-or-nothing: one duplicate transaction_id + fails the WHOLE batch even though the other events are perfectly valid. + Naively dropping the batch would silently lose those valid events too — + the queue must fall back to one-by-one to tell them apart.""" + sent_individually = [] + + def sender(batch): + if len(batch) > 1: + raise LagoApiError(422, '{"error_details":{"transaction_id":["value_already_exist"]}}') + event = batch[0] + sent_individually.append(event["id"]) + if event["id"] in ("dup_1", "dup_2"): + raise LagoApiError(422, '{"error_details":{"transaction_id":["value_already_exist"]}}') + # "good_*" events succeed alone. + + errors: list[tuple[Exception, str]] = [] + q = EventQueue( + sender=sender, + flush_interval=0.05, + max_batch_size=10, + on_error=lambda exc, where: errors.append((exc, where)), + ) + try: + for eid in ["dup_1", "good_1", "dup_2", "good_2"]: + q.push({"id": eid}) + assert q.flush(timeout=2.0) + finally: + q.shutdown(timeout=1.0) + + # All four were tried individually — the two "good" ones weren't silently + # dropped along with the two duplicates just because they shared a batch. + assert set(sent_individually) == {"dup_1", "good_1", "dup_2", "good_2"} + # on_error fires once for the batch-level failure, not once per dropped item. + assert len(errors) == 1 + assert errors[0][1] == "send_batch" + + +def test_permanent_failure_does_not_apply_backoff(): + """Retrying a permanently-doomed batch with exponential backoff is + pointless — the isolate-and-drop path must not slow down subsequent + genuinely-transient failures by leaving a stale backoff in place.""" + calls = {"n": 0} + + def sender(batch): + calls["n"] += 1 + if len(batch) > 1: + raise LagoApiError(422, "duplicate") + raise LagoApiError(422, "duplicate") # every isolated event is also a dup here + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=0.5) + try: + q.push({"id": "dup_1"}) + q.push({"id": "dup_2"}) + assert q.flush(timeout=2.0) # drains fast — no backoff wait, unlike a transient failure + assert q._backoff_seconds == 0.0 + finally: + q.shutdown(timeout=1.0) + + +def test_transient_failure_during_isolation_still_gets_retried(): + """An event that hits a network-level (non-4xx) error while being sent + individually is a real transient failure — it must still go through the + normal re-queue-and-retry path, not get treated as permanent.""" + attempts = {"flaky": 0} + + def sender(batch): + if len(batch) > 1: + raise LagoApiError(422, "duplicate") # forces the isolate-one-by-one path + event = batch[0] + if event["id"] == "flaky": + attempts["flaky"] += 1 + if attempts["flaky"] == 1: + raise RuntimeError("transient network blip") # not a LagoApiError at all + return # succeeds on the retried attempt + if event["id"] == "dup": + raise LagoApiError(422, "duplicate") + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=0.5) + try: + q.push({"id": "dup"}) + q.push({"id": "flaky"}) + deadline = time.monotonic() + 3.0 + while time.monotonic() < deadline and attempts["flaky"] < 2: + time.sleep(0.05) + assert attempts["flaky"] >= 2, "the transient failure should have been retried, not dropped" + finally: + q.shutdown(timeout=2.0) + + +_REENTRANT_OVERFLOW_PROGRAM = """ +import sys, threading +sys.path.insert(0, {src!r}) +from lago_agent_sdk.queue import EventQueue + +calls = {{"n": 0}} +def on_error(exc, where): + calls["n"] += 1 + if calls["n"] > 200: # runaway guard so this exits rather than spinning + raise SystemExit(3) + q.push({{"diagnostic": True}}) # re-enters push() from inside the hook + +q = EventQueue(sender=lambda b: None, flush_interval=10.0, max_batch_size=1000, + max_buffer_size=1, on_error=on_error) +for i in range(3): + q.push({{"i": i}}) +q.shutdown(timeout=1.0) +print("OK", calls["n"]) +""" + + +def _run_reentrant_overflow(timeout: float = 20.0) -> subprocess.CompletedProcess: + """Run the re-entrant-overflow scenario in a SUBPROCESS. + + It has to be a subprocess: once the deadlock happens, the wedged producer holds + `_lock` forever, and `EventQueue.__init__`'s `atexit` shutdown then blocks on that + same lock at interpreter exit. The process is poisoned, so an in-process test + would hang the whole session instead of reporting a failure. Out-of-process, a + hang is just a timeout we can assert on. + """ + src = str(pathlib.Path(__file__).resolve().parents[2] / "src") + return subprocess.run( + [sys.executable, "-c", _REENTRANT_OVERFLOW_PROGRAM.format(src=src)], + capture_output=True, + text=True, + timeout=timeout, + ) + + +def test_overflow_report_does_not_deadlock_a_reentrant_callback(): + """The report must run with the lock RELEASED, and must not re-enter unboundedly. + + Two failure modes, one scenario. `_lock` is a plain Lock, not an RLock, so + reporting inside it meant a customer `on_error` that touched the SDK at all — + emitting a diagnostic, forcing a flush — blocked forever on a lock its own thread + already held; overflow happens under sustained load, which is exactly when such a + hook fires. Moving the report out of the lock then exposed the other half: the + buffer is full again by the time the hook runs, so a hook that pushes overflows + again and re-enters without bound. + """ + try: + proc = _run_reentrant_overflow() + except subprocess.TimeoutExpired: + raise AssertionError( + "re-entrant on_error during overflow hung the process — the report is holding the lock" + ) from None + assert proc.returncode == 0, f"exit {proc.returncode}: " + ( + "runaway re-entrant reporting" if proc.returncode == 3 else proc.stderr[-400:] + ) + assert proc.stdout.startswith("OK"), proc.stdout + reports = int(proc.stdout.split()[1]) + assert reports <= 10, f"expected a bounded number of overflow reports, got {reports}" + + +def test_overflow_is_reported_through_on_error(): + """An overflow drops BILLABLE events. It was logger.warning only, so a customer + watching on_error never learned revenue had been lost; the JS port already + reported it.""" + errors: list = [] + q = EventQueue( + sender=lambda b: None, + flush_interval=10.0, # keep the worker idle so the buffer really fills + max_batch_size=1000, + max_buffer_size=2, + on_error=lambda exc, where: errors.append((str(exc), where)), + ) + try: + for i in range(5): + q.push({"id": i}) + assert errors, "overflow must reach on_error" + assert any(w == "overflow" for _, w in errors) + assert any("overflow" in m for m, _ in errors) + finally: + q.shutdown(timeout=1.0) + + +# ---------------------------------------------------------------------- +# The throttling 4xxs. 429 and 408 sit inside the 400-499 range but mean "try +# again, later" — classifying them as permanent dropped billable events and +# aimed `max_batch_size` extra requests at a server that had just asked us to +# slow down. +# ---------------------------------------------------------------------- +@pytest.mark.parametrize("status", [413, 402, 415]) +def test_batch_only_4xx_is_split_not_head_of_line_blocked(status: int): + """For these the SAME batch can never succeed, but its events can individually. + + Treating them as transient re-prepended the identical batch at the head of the FIFO + and backed off to 60s forever, blocking everything behind it. Routing them to + `_send_individually` splits the batch and delivers what is deliverable — which is + what the isolation path was built for, and it was unreachable for exactly the batch + that most needed it. + """ + sent_individually: list = [] + + def sender(batch): + if len(batch) > 1: + raise LagoApiError(status, "batch too large / unacceptable as-is") + sent_individually.append(batch[0]["id"]) # each event succeeds alone + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=0.5) + try: + for i in range(4): + q.push({"id": i}) + assert q.flush(timeout=3.0), "queue should drain, not head-of-line block" + assert sorted(sent_individually) == [0, 1, 2, 3], ( + f"every event should have been delivered individually, got {sent_individually}" + ) + assert q._backoff_seconds == 0.0, "splitting must not leave a stale backoff" + finally: + q.shutdown(timeout=1.0) + + +@pytest.mark.parametrize("status", [429, 408]) +def test_throttling_4xx_is_retried_not_dropped(status: int): + """A rate-limited or timed-out batch must reach Lago eventually. Dropping + it loses revenue, and isolating it one-by-one multiplies the load on a + server that is already shedding it.""" + attempts = {"n": 0} + delivered: list = [] + + def sender(batch): + attempts["n"] += 1 + if attempts["n"] == 1: + raise LagoApiError(status, '{"error":"too many requests"}') + delivered.extend(batch) # succeeds once the throttle lifts + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=0.5) + try: + q.push({"id": "a"}) + q.push({"id": "b"}) + deadline = time.monotonic() + 3.0 + while time.monotonic() < deadline and not delivered: + time.sleep(0.05) + assert [e["id"] for e in delivered] == ["a", "b"], "throttled events must still be delivered" + # Delivered as one batch, i.e. never fanned out into per-event requests. + assert attempts["n"] == 2 + finally: + q.shutdown(timeout=2.0) + + +@pytest.mark.parametrize("status", [429, 408]) +def test_throttling_4xx_applies_backoff(status: int): + """The inverse of test_permanent_failure_does_not_apply_backoff: a + throttling failure is transient, so it MUST leave a backoff in place — + that pause is the whole point of respecting a rate limit.""" + + def sender(batch): + raise LagoApiError(status, "slow down") + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=0.5) + try: + q.push({"id": "a"}) + deadline = time.monotonic() + 2.0 + while time.monotonic() < deadline and q._backoff_seconds == 0.0: + time.sleep(0.05) + assert q._backoff_seconds > 0.0, "a throttling 4xx must back off, not isolate-and-drop" + finally: + q.shutdown(timeout=1.0) + + +def test_unrecognized_4xx_is_treated_as_transient(): + """Only the enumerated validation statuses are permanent. An unfamiliar 4xx + errs toward retrying: a needless delay costs latency, a wrong drop costs + revenue.""" + attempts = {"n": 0} + delivered: list = [] + + def sender(batch): + attempts["n"] += 1 + if attempts["n"] == 1: + raise LagoApiError(418, "i am a teapot") + delivered.extend(batch) + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=0.5) + try: + q.push({"id": "a"}) + deadline = time.monotonic() + 3.0 + while time.monotonic() < deadline and not delivered: + time.sleep(0.05) + assert [e["id"] for e in delivered] == ["a"] + finally: + q.shutdown(timeout=2.0) + + +@pytest.mark.parametrize("status", [400, 401, 403, 404, 409, 422]) +def test_validation_4xx_still_isolates_and_drops(status: int): + """The statuses that genuinely cannot succeed on a re-send keep the + isolate-one-by-one behaviour, so a single bad transaction_id still doesn't + take the rest of its batch down with it.""" + sent_individually: list = [] + + def sender(batch): + if len(batch) > 1: + raise LagoApiError(status, "batch rejected") + sent_individually.append(batch[0]["id"]) + if batch[0]["id"].startswith("bad"): + raise LagoApiError(status, "this one really is invalid") + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=0.5) + try: + q.push({"id": "bad_1"}) + q.push({"id": "good_1"}) + assert q.flush(timeout=2.0) + assert set(sent_individually) == {"bad_1", "good_1"} + assert q._backoff_seconds == 0.0 + finally: + q.shutdown(timeout=1.0) + + +# ---------------------------------------------------------------------- +# Shutdown's final drain. Previously: `except Exception: pass` on a single +# attempt at a single batch — any failure at all was silently swallowed, and +# a buffer holding more than one batch's worth of events at shutdown time +# left the rest never even attempted. +# ---------------------------------------------------------------------- +def test_shutdown_drains_more_than_one_batch(): + """Buffer holds 3 batches' worth of events right as shutdown starts — + every one of them must be attempted, not just the first.""" + sent = [] + q = EventQueue(sender=lambda b: sent.extend(b), flush_interval=10.0, max_batch_size=5) + try: + for i in range(15): # 3 full batches of 5, worker hasn't had a flush tick yet + q.push({"i": i}) + finally: + q.shutdown(timeout=2.0) + assert len(sent) == 15 + + +def test_shutdown_reports_transient_failure_instead_of_silently_swallowing(): + """A persistently-failing sender at shutdown time must surface via + on_error — not vanish behind a bare `except: pass` the way it used to.""" + errors: list[tuple[Exception, str]] = [] + + def always_fails(batch): + raise RuntimeError("network still down") + + q = EventQueue( + sender=always_fails, + flush_interval=10.0, + max_batch_size=10, + max_retry_seconds=1.0, + on_error=lambda exc, where: errors.append((exc, where)), + ) + try: + q.push({"i": 1}) + finally: + q.shutdown(timeout=3.0) + assert len(errors) >= 1 + assert errors[0][1] == "send_batch" + assert "network still down" in str(errors[0][0]) + + def test_flush_returns_false_on_timeout(): blocking = threading.Event() @@ -96,3 +460,23 @@ def slow(batch): finally: blocking.set() q.shutdown(timeout=2.0) + + +def test_isolated_retries_keep_their_fifo_order() -> None: + """`_replay_failed` PREPENDS, so calling it once per event inside the isolation + loop reversed the survivors: a 413 batch of a,b,c,d,e whose b,c,d fail + transiently while isolated came back as d,c,b. FIFO is the queue's contract — + it is what makes oldest-dropped-first overflow and Lago's own event ordering + mean anything — so a recovery path must not silently invert it.""" + + def sender(batch): + if batch[0]["id"] in ("b", "c", "d"): + raise LagoApiError(503, "transient while isolated") + + q = EventQueue(sender=sender, flush_interval=60.0, max_batch_size=10, max_buffer_size=100) + try: + q._send_individually([{"id": i} for i in "abcde"], LagoApiError(413, "too large")) + with q._lock: + assert [e["id"] for e in q._buffer] == ["b", "c", "d"] + finally: + q.shutdown(timeout=1.0) diff --git a/tests/unit/test_sdk.py b/tests/unit/test_sdk.py index ed86c45..64711c1 100644 --- a/tests/unit/test_sdk.py +++ b/tests/unit/test_sdk.py @@ -2,9 +2,11 @@ from __future__ import annotations +import logging + import pytest -from lago_agent_sdk import CanonicalUsage, LagoSDK +from lago_agent_sdk import CanonicalUsage, LagoConfig, LagoSDK from lago_agent_sdk.exceptions import UnknownClientError @@ -83,6 +85,153 @@ def test_wrap_unknown_client_raises_at_wrap_time(): sdk.shutdown(timeout=1.0) +# ---------------------------------------------------------------------- +# Constructor precedence. `api_url`'s default used to be the production URL, so +# `if api_url:` always fired and clobbered a config-supplied one — sending a +# local-dev customer's events to production Lago. +# ---------------------------------------------------------------------- +def test_config_only_api_url_survives(): + """The bug: a customer who configures ONLY via LagoConfig must not have their + events redirected to production.""" + sdk = LagoSDK(api_key="k", config=LagoConfig(api_url="http://localhost:3000/api/v1")) + try: + assert sdk.config.api_url == "http://localhost:3000/api/v1" + finally: + sdk.shutdown(timeout=1.0) + + +def test_explicit_api_url_still_wins_over_config(): + """The documented rule — explicit args beat config — must still hold.""" + sdk = LagoSDK( + api_key="k", + api_url="http://explicit:3000/api/v1", + config=LagoConfig(api_url="http://fromconfig:3000/api/v1"), + ) + try: + assert sdk.config.api_url == "http://explicit:3000/api/v1" + finally: + sdk.shutdown(timeout=1.0) + + +@pytest.mark.parametrize("empty", ["", None]) +def test_empty_or_absent_api_url_keeps_the_production_default(empty): + """`api_url=os.environ.get("LAGO_API_URL", "")` with the var unset must NOT write + "". Downstream that is unrecoverable: requests raises MissingSchema, which is not a + LagoApiError, so the queue treats it as transient and retries at the 60s ceiling + forever — all billing stops with only a growing buffer as the symptom.""" + sdk = LagoSDK(api_key="k", api_url=empty) + try: + assert sdk.config.api_url == "https://api.getlago.com/api/v1" + finally: + sdk.shutdown(timeout=1.0) + + +def test_default_api_url_is_still_production_when_nothing_is_passed(): + """Changing the parameter default to None must not change this.""" + sdk = LagoSDK(api_key="k") + try: + assert sdk.config.api_url == "https://api.getlago.com/api/v1" + finally: + sdk.shutdown(timeout=1.0) + + +def test_verify_ssl_needs_no_config_object(): + """A local Lago on a self-signed cert is reachable without building a + LagoConfig — which is what pushed callers toward the clobber in the first + place, since a custom api_url and verify_ssl=False go together.""" + sdk = LagoSDK(api_key="k", api_url="https://api.lago.dev/api/v1", verify_ssl=False) + try: + assert sdk.config.verify_ssl is False + assert sdk._lago_client.verify_ssl is False + finally: + sdk.shutdown(timeout=1.0) + + +def test_explicit_verify_ssl_wins_over_config(): + sdk = LagoSDK(api_key="k", verify_ssl=True, config=LagoConfig(verify_ssl=False)) + try: + assert sdk.config.verify_ssl is True + finally: + sdk.shutdown(timeout=1.0) + + +def test_ignored_usd_cost_is_reported_not_silently_dropped(): + """A caller who supplies a real metered cost while the effective mode isn't + 'price' had it discarded with no log and no on_error — so a hand-rolled + backfill could bill token counts only and look successful.""" + errors: list = [] + received: list = [] + cfg = LagoConfig( + api_key="dummy", + default_subscription_id="sub", + pricing_mode="tokens", + on_error=lambda exc, where: errors.append((str(exc), where)), + ) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: received.extend(b) # type: ignore[attr-defined] + u = CanonicalUsage(input=10, output=5, model="m", provider="anthropic", api="native") + sdk.emit(u, usd_cost=0.0123) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + + assert errors, "an ignored usd_cost must reach on_error" + msg, where = errors[0] + assert "usd_cost" in msg and "0.0123" in msg + assert where == "pricing" + # And the call is still billed as token counts — reporting must not drop events. + assert {e["code"] for e in received} == {"llm_input_tokens", "llm_output_tokens"} + + +def test_no_usd_cost_in_token_mode_reports_nothing(): + """The common case must stay silent — only an explicitly supplied cost that + gets discarded is worth reporting.""" + errors: list = [] + cfg = LagoConfig( + api_key="dummy", + default_subscription_id="sub", + pricing_mode="tokens", + on_error=lambda exc, where: errors.append((str(exc), where)), + ) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: None # type: ignore[attr-defined] + sdk.emit(CanonicalUsage(input=10, output=5, model="m", provider="anthropic", api="native")) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert errors == [] + + +def test_no_resolvable_subscription_is_reported_not_just_logged(): + """Dropping a call for lack of a subscription loses its billing entirely, so it + must reach on_error — the documented channel for every other billing gap. It + was logger.error only, while the JS port already reported it.""" + errors: list = [] + cfg = LagoConfig( + api_key="dummy", + default_subscription_id=None, + on_error=lambda exc, where: errors.append((str(exc), where)), + ) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: None # type: ignore[attr-defined] + sdk.emit(CanonicalUsage(input=10, model="m", provider="p", api="x")) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert errors, "a dropped call must reach on_error" + assert "subscription" in errors[0][0] + + +def test_negative_counts_are_never_emitted(): + """`nonzero_numeric` filtered on truthiness, so a negative survived and was + emitted verbatim as value="-100" — a negative billable quantity. JS already + filtered on > 0.""" + sdk, received = _new_sdk(default_sub="sub") + sdk.emit(CanonicalUsage(input=-100, output=5, model="m", provider="p", api="x")) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert {e["code"] for e in flat} == {"llm_output_tokens"} + assert all(float(e["properties"]["value"]) > 0 for e in flat) + + def test_dimensions_merge_into_event_properties(): sdk, received = _new_sdk(default_sub="sub") u = CanonicalUsage(input=1, model="m", provider="p", api="bedrock_invoke") @@ -92,3 +241,115 @@ def test_dimensions_merge_into_event_properties(): flat = [e for batch in received for e in batch] assert flat[0]["properties"]["project"] == "demo" assert flat[0]["properties"]["tenant"] == "acme" + + +def test_caller_dimensions_win_on_a_collision_on_both_emitters(): + """One rule across both paths: a caller dimension overrides every + SDK-computed property of the same name. + + The cost path used to spread dimensions into `base_properties`, i.e. BEFORE + `unit`/`value`/`base_cost`/`unit_price`, so those four silently overwrote a + same-named caller dimension there while the token path honoured it. Same + customer config, two different outcomes depending on the mode. + """ + dims = {"unit": "seat", "value": "CUSTOM", "model": "my-label", "team": "platform"} + u = CanonicalUsage(input=100, output=50, model="claude-sonnet-4-5", provider="anthropic", api="native") + + # Token path. + sdk_tok, got_tok = _new_sdk(default_sub="sub") + sdk_tok.emit(u, dimensions=dims, mode="tokens") + assert sdk_tok.flush(timeout=2.0) + sdk_tok.shutdown(timeout=1.0) + tok = [e for batch in got_tok for e in batch] + + # Cost path (precomputed, so no price table needed). + sdk_cost, got_cost = _new_sdk(default_sub="sub") + sdk_cost.emit(u, dimensions=dims, mode="price", usd_cost=0.01) + assert sdk_cost.flush(timeout=2.0) + sdk_cost.shutdown(timeout=1.0) + cost = [e for batch in got_cost for e in batch] + + assert tok and cost + for label, events in (("token", tok), ("cost", cost)): + for e in events: + p = e["properties"] + assert p["unit"] == "seat", f"{label}: caller `unit` must win" + assert p["value"] == "CUSTOM", f"{label}: caller `value` must win" + assert p["model"] == "my-label", f"{label}: caller `model` must win" + assert p["team"] == "platform" + + # The accepted consequence of that rule, pinned deliberately: a dimension + # named `value` overrides the reported quantity. It is NOT able to touch the + # charged amount on a cost event, because `precise_total_amount_cents` is a + # sibling of `properties`, not a member of it. + assert cost[0]["precise_total_amount_cents"] == "1" + + +def test_negative_token_counts_are_reported_not_just_dropped(caplog) -> None: + """`CanonicalUsage` is exported and `emit()` takes one directly — the documented + way to backfill usage the SDK did not intercept — so a caller computing a delta + wrongly really can hand us a negative. Dropping it was correct (Lago would sum a + negative billable quantity) but it was the only drop path that never reached + on_error.""" + seen: list[tuple[Exception, str]] = [] + received: list = [] + cfg = LagoConfig(api_key="k", default_subscription_id="sub", on_error=lambda e, c: seen.append((e, c))) + sdk = LagoSDK(api_key="k", config=cfg) + sdk._queue._sender = lambda b: received.append(list(b)) # type: ignore[attr-defined] + try: + sdk.emit(CanonicalUsage(input=-100, output=50, model="m", provider="anthropic", api="native")) + assert sdk.flush(timeout=2.0) + finally: + sdk.shutdown(timeout=1.0) + + flat = [e for batch in received for e in batch] + values = [e["properties"]["value"] for e in flat] + assert all(not str(v).startswith("-") for v in values), f"a negative was billed: {values}" + assert "negative_tokens" in [c for _, c in seen], f"drop must reach on_error; got {seen}" + assert "input" in str(next(e for e, c in seen if c == "negative_tokens")) + + +def test_a_dropped_event_logs_exactly_once(caplog) -> None: + """`_report_error` invokes on_error AND logs. An extra logger.error alongside it + emitted the same drop twice under two levels, so a customer grepping logs counted + one lost call as two — while the JS port logged nothing at all for it.""" + cfg = LagoConfig(api_key="k") # no default subscription -> the drop path + sdk = LagoSDK(api_key="k", config=cfg) + try: + with caplog.at_level(logging.DEBUG, logger="lago_agent_sdk"): + sdk.emit(CanonicalUsage(input=10, output=5, model="m", provider="anthropic", api="native")) + drop_lines = [r for r in caplog.records if "subscription" in r.getMessage()] + assert len(drop_lines) == 1, f"expected one line per drop, got {[r.getMessage() for r in drop_lines]}" + finally: + sdk.shutdown(timeout=1.0) + + +def test_verify_ssl_false_survives_a_broken_urllib3(monkeypatch) -> None: + """Suppressing the InsecureRequestWarning is an optional convenience; it must + never be able to fail construction. This sits on an advertised path — `verify_ssl` + is a first-class constructor arg the docstring recommends for local dev — so an + ImportError/AttributeError here would take down `LagoSDK()` for exactly the setup + the flag exists to serve. The old code reached through `requests.packages`, a + legacy shim with no guarantee of existing.""" + import builtins + + real_import = builtins.__import__ + + def boom(name, *args, **kwargs): + if name == "urllib3": + raise ImportError("no urllib3 for you") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", boom) + # Also remove the legacy shim the old code reached through, so this test fails + # against that version instead of silently passing: `requests.packages` is a + # compatibility alias, not API, and nothing guarantees it is present. + import requests + + monkeypatch.delattr(requests, "packages", raising=False) + + sdk = LagoSDK(api_key="k", api_url="https://example.invalid/api/v1", verify_ssl=False) + try: + assert sdk.config.verify_ssl is False + finally: + sdk.shutdown(timeout=1.0) diff --git a/tests/unit/test_wrapper_anthropic.py b/tests/unit/test_wrapper_anthropic.py index 85804cf..1292f93 100644 --- a/tests/unit/test_wrapper_anthropic.py +++ b/tests/unit/test_wrapper_anthropic.py @@ -20,6 +20,12 @@ def model_dump(self) -> dict[str, Any]: return self._payload +# What Anthropic resolves the requested "claude-sonnet-4-6" alias to. Only +# `message_start` reports it, so the wrapper has to keep it across the whole +# stream or pricing looks up an alias OpenRouter doesn't list. +_RESOLVED_STREAM_MODEL = "claude-sonnet-4-6-20260214" + + class FakeStreamEvent: """Mimics one of Anthropic's MessageStreamEvent objects (MessageDelta/Start/etc.).""" @@ -30,10 +36,52 @@ def model_dump(self) -> dict[str, Any]: return self._payload +class FakeRawResponse: + """Mimics the return value of `.with_raw_response.create(...)`: `.headers` + `.parse()`.""" + + def __init__(self, parsed: Any, headers: dict[str, str] | None = None) -> None: + self._parsed = parsed + self.headers = headers or {} + + def parse(self) -> Any: + return self._parsed + + +class _RawResponseProxy: + """Mimics `.with_raw_response` — delegates to the owner's `.create()`, wraps the + result with whatever headers the test configured on `owner.raw_response_headers`. + + Captures the owner's `.create` bound method at construction time (i.e. before + `sdk.wrap()` can monkey-patch it) — looking it up dynamically via + `self._owner.create` at call time would resolve to the *wrapped* method once + `sdk.wrap()` reassigns it, causing infinite recursion. + """ + + def __init__(self, owner: Any) -> None: + self._owner = owner + self._original_create = owner.create + + def create(self, **kwargs: Any) -> FakeRawResponse: + parsed = self._original_create(**kwargs) + return FakeRawResponse(parsed, self._owner.raw_response_headers) + + +class _AsyncRawResponseProxy: + def __init__(self, owner: Any) -> None: + self._owner = owner + self._original_create = owner.create + + async def create(self, **kwargs: Any) -> FakeRawResponse: + parsed = await self._original_create(**kwargs) + return FakeRawResponse(parsed, self._owner.raw_response_headers) + + class FakeMessages: def __init__(self) -> None: self.create_calls = 0 self.stream_calls = 0 + self.raw_response_headers: dict[str, str] = {} + self.with_raw_response = _RawResponseProxy(self) def create(self, **kwargs: Any) -> Any: self.create_calls += 1 @@ -49,12 +97,15 @@ def create(self, **kwargs: Any) -> Any: { "type": "message_start", "message": { + # message_start is also where the RESOLVED snapshot + # arrives — the requested alias never appears again. + "model": _RESOLVED_STREAM_MODEL, "usage": { "input_tokens": 12, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "output_tokens": 1, - } + }, }, } ), @@ -182,6 +233,20 @@ def test_wrap_double_wrap_is_idempotent() -> None: assert fake.messages.create_calls == 1 +def test_stream_attributes_the_resolved_model_not_the_requested_alias() -> None: + """Only `message_start` carries the resolved snapshot, and the wrapper + accumulates usage across several events before emitting — so the model has + to survive the whole stream. Rebuilding a usage-only payload reverted the + attribution to the requested alias, which OpenRouter doesn't list.""" + sdk, received = _new_sdk() + client = sdk.wrap(FakeAnthropic()) + list(client.messages.create(model="claude-sonnet-4-6", messages=[], stream=True)) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + models = {e["properties"]["model"] for e in received} + assert models == {_RESOLVED_STREAM_MODEL}, f"expected the resolved snapshot, got {models}" + + def test_wrap_create_with_stream_merges_message_start_and_delta() -> None: """Regression: input/cache come from message_start, output from message_delta. @@ -216,6 +281,47 @@ def test_wrap_messages_stream_context_manager_emits_on_close() -> None: assert by_code["llm_output_tokens"] == 11 +# -------------------------------------------------------------------------- +# Gateway cache-hit detection (non-streaming only) +# -------------------------------------------------------------------------- +def test_wrap_cache_miss_still_bills_normally() -> None: + """No gateway, or a MISS: bills exactly as before — .with_raw_response is the + new code path, but must be behaviorally invisible with no cache header set.""" + sdk, received = _new_sdk() + fake = FakeAnthropic() + client = sdk.wrap(fake) + client.messages.create(model="claude-sonnet-4-6", messages=[]) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + by_code = {e["code"]: int(float(e["properties"]["value"])) for e in received} + assert by_code["llm_input_tokens"] == 8 + assert by_code["llm_output_tokens"] == 16 + + +def test_wrap_cache_hit_skips_billing() -> None: + """A gateway-served cache HIT cost the customer nothing — bill nothing for it.""" + sdk, received = _new_sdk() + fake = FakeAnthropic() + fake.messages.raw_response_headers = {"cf-aig-cache-status": "HIT"} + client = sdk.wrap(fake) + resp = client.messages.create(model="claude-sonnet-4-6", messages=[]) + assert resp.usage["input_tokens"] == 8 # customer still gets the real response + sdk.shutdown(timeout=1.0) + assert received == [] + + +def test_wrap_cache_status_other_than_hit_still_bills() -> None: + """Only an exact "HIT" suppresses billing — "MISS", "EXPIRED", or anything else bills.""" + sdk, received = _new_sdk() + fake = FakeAnthropic() + fake.messages.raw_response_headers = {"cf-aig-cache-status": "MISS"} + client = sdk.wrap(fake) + client.messages.create(model="claude-sonnet-4-6", messages=[]) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert len(received) == 2 + + def test_instrumentation_failure_does_not_break_call() -> None: sdk, _ = _new_sdk() @@ -262,6 +368,8 @@ def __init__(self) -> None: self.create_calls = 0 self.stream_calls = 0 self.final_message_awaited = False # tracks whether the async path was actually awaited + self.raw_response_headers: dict[str, str] = {} + self.with_raw_response = _AsyncRawResponseProxy(self) async def create(self, **kwargs: Any) -> Any: self.create_calls += 1 @@ -374,6 +482,18 @@ async def test_async_wrap_messages_create_emits() -> None: assert by_code["llm_output_tokens"] == 16 +@pytest.mark.asyncio +async def test_async_wrap_cache_hit_skips_billing() -> None: + sdk, received = _new_sdk() + fake = FakeAsyncAnthropic() + fake.messages.raw_response_headers = {"cf-aig-cache-status": "HIT"} + client = sdk.wrap(fake) + resp = await client.messages.create(model="claude-sonnet-4-6", messages=[]) + assert resp.usage["input_tokens"] == 8 + sdk.shutdown(timeout=1.0) + assert received == [] + + @pytest.mark.asyncio async def test_async_wrap_messages_create_stream_captures_usage() -> None: """Async iteration of `messages.create(stream=True)` — wraps an async generator.""" diff --git a/tests/unit/test_wrapper_gemini.py b/tests/unit/test_wrapper_gemini.py index 615cc44..ba8fd65 100644 --- a/tests/unit/test_wrapper_gemini.py +++ b/tests/unit/test_wrapper_gemini.py @@ -18,6 +18,12 @@ def model_dump(self) -> dict: return self._payload +# What Gemini resolves the requested "gemini-2.5-flash" alias to. Google +# hot-swaps these server-side, so the chunk's own version is what OpenRouter +# lists and what pricing must key off. +_RESOLVED_STREAM_MODEL = "gemini-2.5-flash-002" + + class FakeStreamChunk: def __init__(self, payload: dict): self._payload = payload @@ -55,11 +61,15 @@ def generate_content_stream(self, **kwargs: Any) -> Any: { "candidates": [{"content": {"parts": [{"text": "hi"}]}}], "usage_metadata": None, # intermediate chunks don't carry usage + # Gemini hot-swaps "-latest" aliases, so every chunk reports + # the version that actually answered. + "model_version": _RESOLVED_STREAM_MODEL, } ), FakeStreamChunk( { "candidates": [{"content": {"parts": [{"text": "."}]}, "finish_reason": "STOP"}], + "model_version": _RESOLVED_STREAM_MODEL, "usage_metadata": { "prompt_token_count": 9, "candidates_token_count": 4, @@ -168,6 +178,79 @@ def generate_content(self, *_a, **_k): # pragma: no cover - never called sdk.shutdown(timeout=1.0) +def test_stream_attributes_the_resolved_model_not_the_requested_alias() -> None: + """Gemini resolves "-latest" and short aliases server-side and reports the + real version as `model_version`. The stream wrapper rebuilt a usage-only + payload and dropped it, reverting attribution to the requested alias.""" + sdk, received = _make_sdk() + client = sdk.wrap(FakeGeminiClient()) + list(client.models.generate_content_stream(model="gemini-2.5-flash", contents="hi")) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + models = {e["properties"]["model"] for e in flat} + assert models == {_RESOLVED_STREAM_MODEL}, f"expected the resolved version, got {models}" + + +class _VersionOnlyOnEarlyChunk: + """A SYNTHETIC stream: the resolved version arrives on an early chunk, usage on + the last one, with no version on it. + + This is deliberately NOT what real Gemini does — verified live 2026-08-20 that + every streaming chunk carries BOTH `model_version` and `usage_metadata`, which + is why `FakeGeminiClient` puts the version on both and why this hazard is + invisible there. The point of this case is the robustness property, not a + captured shape: `model_version` must be remembered across chunks rather than + read off whichever chunk happens to carry usage. Python read it from the + usage-bearing chunk alone, so on this input it reverted to the requested alias + while the JS port (which already persisted it) reported the resolved version — + the two repos priced the same call differently. + """ + + __module__ = "google.genai.client" # so the detector routes it to the gemini wrapper + + def __init__(self) -> None: + self.models = self + + def generate_content_stream(self, **kwargs: Any) -> Any: + return iter( + [ + FakeStreamChunk( + { + "candidates": [{"content": {"parts": [{"text": "hi"}]}}], + "model_version": _RESOLVED_STREAM_MODEL, + "usage_metadata": None, + } + ), + FakeStreamChunk( + { + "candidates": [{"content": {"parts": [{"text": "."}]}, "finish_reason": "STOP"}], + # no model_version here + "usage_metadata": { + "prompt_token_count": 9, + "candidates_token_count": 4, + "thoughts_token_count": 0, + "total_token_count": 13, + }, + } + ), + ] + ) + + +def test_stream_remembers_the_resolved_version_from_an_earlier_chunk() -> None: + sdk, received = _make_sdk() + client = sdk.wrap(_VersionOnlyOnEarlyChunk()) + list(client.models.generate_content_stream(model="gemini-flash-latest", contents="hi")) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + models = {e["properties"]["model"] for e in flat} + assert models == {_RESOLVED_STREAM_MODEL}, ( + f"the version from the earlier chunk must survive to the usage chunk; got {models}" + ) + + def test_wrap_generate_content_stream_captures_usage_from_final_chunk() -> None: sdk, received = _make_sdk() fake = FakeGeminiClient() @@ -277,11 +360,13 @@ async def _aiter(): { "candidates": [{"content": {"parts": [{"text": "hi"}]}}], "usage_metadata": None, + "model_version": _RESOLVED_STREAM_MODEL, } ) yield FakeStreamChunk( { "candidates": [{"content": {"parts": [{"text": "."}]}, "finish_reason": "STOP"}], + "model_version": _RESOLVED_STREAM_MODEL, "usage_metadata": { "prompt_token_count": 9, "candidates_token_count": 4, diff --git a/tests/unit/test_wrapper_openai.py b/tests/unit/test_wrapper_openai.py index 44e3b19..dda6126 100644 --- a/tests/unit/test_wrapper_openai.py +++ b/tests/unit/test_wrapper_openai.py @@ -30,6 +30,12 @@ def model_dump(self) -> dict[str, Any]: return self._payload +# What OpenAI resolves the requested "gpt-4o-mini" alias to. Streaming chunks +# report it on every frame; the wrapper must carry it through to the event, or +# pricing looks up an alias OpenRouter doesn't list. +_RESOLVED_STREAM_MODEL = "gpt-4o-mini-2024-07-18" + + class FakeStreamChunk: """Mimics a ChatCompletionChunk.""" @@ -40,10 +46,52 @@ def model_dump(self) -> dict[str, Any]: return self._payload +class FakeRawResponse: + """Mimics the return value of `.with_raw_response.create(...)`: `.headers` + `.parse()`.""" + + def __init__(self, parsed: Any, headers: dict[str, str] | None = None) -> None: + self._parsed = parsed + self.headers = headers or {} + + def parse(self) -> Any: + return self._parsed + + +class _RawResponseProxy: + """Mimics `.with_raw_response` — delegates to the owner's `.create()`, wraps the + result with whatever headers the test configured on `owner.raw_response_headers`. + + Captures the owner's `.create` bound method at construction time (i.e. before + `sdk.wrap()` can monkey-patch it) — looking it up dynamically via + `self._owner.create` at call time would resolve to the *wrapped* method once + `sdk.wrap()` reassigns it, causing infinite recursion. + """ + + def __init__(self, owner: Any) -> None: + self._owner = owner + self._original_create = owner.create + + def create(self, **kwargs: Any) -> FakeRawResponse: + parsed = self._original_create(**kwargs) + return FakeRawResponse(parsed, self._owner.raw_response_headers) + + +class _AsyncRawResponseProxy: + def __init__(self, owner: Any) -> None: + self._owner = owner + self._original_create = owner.create + + async def create(self, **kwargs: Any) -> FakeRawResponse: + parsed = await self._original_create(**kwargs) + return FakeRawResponse(parsed, self._owner.raw_response_headers) + + class FakeCompletions: def __init__(self) -> None: self.create_calls = 0 self.last_kwargs: dict[str, Any] | None = None + self.raw_response_headers: dict[str, str] = {} + self.with_raw_response = _RawResponseProxy(self) def create(self, **kwargs: Any) -> Any: self.create_calls += 1 @@ -54,13 +102,20 @@ def create(self, **kwargs: Any) -> Any: if kwargs.get("stream") is True: # Stream yields several chunks; the LAST one carries usage # (because the wrapper auto-injects stream_options.include_usage). + # Every real chunk carries the RESOLVED model — a short alias like + # "gpt-4o-mini" comes back as a dated snapshot. Pricing keys off it. chunks = [ FakeStreamChunk( - {"choices": [{"delta": {"content": "hi"}}], "usage": None}, + { + "choices": [{"delta": {"content": "hi"}}], + "usage": None, + "model": _RESOLVED_STREAM_MODEL, + }, ), FakeStreamChunk( { "choices": [], + "model": _RESOLVED_STREAM_MODEL, "usage": { "prompt_tokens": 12, "completion_tokens": 22, @@ -98,6 +153,8 @@ def __init__(self) -> None: class FakeResponsesNamespace: def __init__(self) -> None: self.create_calls = 0 + self.raw_response_headers: dict[str, str] = {} + self.with_raw_response = _RawResponseProxy(self) def create(self, **kwargs: Any) -> Any: self.create_calls += 1 @@ -196,6 +253,25 @@ def test_wrap_create_with_stream_captures_usage_from_final_chunk() -> None: assert by_code["llm_output_tokens"] == 22 +def test_stream_attributes_the_resolved_model_not_the_requested_alias() -> None: + """The model-attribution fix has to reach the streaming path too. + + The wrapper rebuilds a synthetic usage payload from the chunks, and dropping + the chunk's own `model` made `resolve_model` fall back to the requested alias + — so a streamed call was attributed (and priced) as "gpt-4o-mini" while the + identical non-streaming call correctly resolved to the dated snapshot. In + price mode that means the OpenRouter lookup misses and silently degrades to + token events. + """ + sdk, received = _new_sdk() + client = sdk.wrap(FakeOpenAI()) + list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True)) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + models = {e["properties"]["model"] for e in received} + assert models == {_RESOLVED_STREAM_MODEL}, f"expected the resolved snapshot, got {models}" + + def test_wrap_auto_injects_stream_options_include_usage() -> None: """Customer passes stream=True without stream_options — wrapper injects include_usage:True.""" sdk, _ = _new_sdk() @@ -260,6 +336,58 @@ def test_wrap_responses_create_emits_input_output_and_tool_calls() -> None: assert by_code["llm_tool_calls"] == 1 +# -------------------------------------------------------------------------- +# Gateway cache-hit detection (non-streaming only) +# -------------------------------------------------------------------------- +def test_wrap_cache_miss_still_bills_normally() -> None: + """No gateway, or a MISS: bills exactly as before — .with_raw_response is the + new code path, but must be behaviorally invisible with no cache header set.""" + sdk, received = _new_sdk() + fake = FakeOpenAI() + client = sdk.wrap(fake) + client.chat.completions.create(model="gpt-4o-mini", messages=[]) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + by_code = {e["code"]: int(float(e["properties"]["value"])) for e in received} + assert by_code["llm_input_tokens"] == 8 + assert by_code["llm_output_tokens"] == 16 + + +def test_wrap_cache_hit_skips_billing_chat_completions() -> None: + """A gateway-served cache HIT cost the customer nothing — bill nothing for it.""" + sdk, received = _new_sdk() + fake = FakeOpenAI() + fake.chat.completions.raw_response_headers = {"cf-aig-cache-status": "HIT"} + client = sdk.wrap(fake) + resp = client.chat.completions.create(model="gpt-4o-mini", messages=[]) + assert resp.usage["prompt_tokens"] == 8 # customer still gets the real response + sdk.shutdown(timeout=1.0) + assert received == [] + + +def test_wrap_cache_hit_skips_billing_responses_api() -> None: + sdk, received = _new_sdk() + fake = FakeOpenAI() + fake.responses.raw_response_headers = {"cf-aig-cache-status": "HIT"} + client = sdk.wrap(fake) + resp = client.responses.create(model="gpt-4o-mini", input="hi") + assert resp.usage["input_tokens"] == 53 + sdk.shutdown(timeout=1.0) + assert received == [] + + +def test_wrap_cache_status_other_than_hit_still_bills() -> None: + """Only an exact "HIT" suppresses billing — "MISS", "EXPIRED", or anything else bills.""" + sdk, received = _new_sdk() + fake = FakeOpenAI() + fake.chat.completions.raw_response_headers = {"cf-aig-cache-status": "MISS"} + client = sdk.wrap(fake) + client.chat.completions.create(model="gpt-4o-mini", messages=[]) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert len(received) == 2 + + # -------------------------------------------------------------------------- # Failure isolation # -------------------------------------------------------------------------- @@ -308,6 +436,8 @@ class FakeAsyncCompletions: def __init__(self) -> None: self.create_calls = 0 self.last_kwargs: dict[str, Any] | None = None + self.raw_response_headers: dict[str, str] = {} + self.with_raw_response = _AsyncRawResponseProxy(self) async def create(self, **kwargs: Any) -> Any: self.create_calls += 1 @@ -358,6 +488,8 @@ class FakeAsyncResponsesNamespace: def __init__(self) -> None: self.create_calls = 0 self.last_kwargs: dict[str, Any] | None = None + self.raw_response_headers: dict[str, str] = {} + self.with_raw_response = _AsyncRawResponseProxy(self) async def create(self, **kwargs: Any) -> Any: self.create_calls += 1 @@ -418,6 +550,18 @@ async def test_async_wrap_chat_completions_emits() -> None: assert by_code["llm_output_tokens"] == 16 +@pytest.mark.asyncio +async def test_async_wrap_cache_hit_skips_billing() -> None: + sdk, received = _new_sdk() + fake = FakeAsyncOpenAI() + fake.chat.completions.raw_response_headers = {"cf-aig-cache-status": "HIT"} + client = sdk.wrap(fake) + resp = await client.chat.completions.create(model="gpt-4o-mini", messages=[]) + assert resp.usage["prompt_tokens"] == 8 + sdk.shutdown(timeout=1.0) + assert received == [] + + @pytest.mark.asyncio async def test_async_wrap_chat_completions_stream_captures_usage() -> None: sdk, received = _new_sdk()