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..450db6a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,42 @@ 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 +- **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. +- **`_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/README.md b/README.md index 70562a2..2389c66 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) | - -**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`. +| image_input | `llm_image_input_tokens` | ✗ | ✗ | ✗ | ✗ | ✓ (multimodal IMAGE) | -**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. +**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`). -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 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/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..21ae9d2 100644 --- a/src/lago_agent_sdk/adapters/openai_native.py +++ b/src/lago_agent_sdk/adapters/openai_native.py @@ -40,6 +40,7 @@ from typing import Any, cast from ..canonical import CanonicalUsage +from ._common import resolve_model # Top-level usage fields we recognize across BOTH chat completions and responses APIs. _KNOWN_USAGE_FIELDS = { @@ -101,6 +102,21 @@ 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.""" + if resolved_model.startswith("@cf/"): + 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 +158,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 +167,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/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..5ab8655 --- /dev/null +++ b/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py @@ -0,0 +1,126 @@ +"""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_tokens → reasoning + model, provider → passed straight through + +`usage_metadata`'s exact key casing is NOT normalized by Cloudflare — it passes +through whatever convention the underlying provider's own usage object used +(Anthropic/OpenAI: snake_case `input_cached_tokens`; a real captured Gemini +entry: camelCase `reasoningTokens`). Both cases are checked for every field +we map; this is observed behavior across two providers, not a documented +guarantee, so a third provider could use a convention we haven't seen yet. + +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 "" + + +# 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")), + cache_read=_safe_int(usage_meta.get("input_cached_tokens")), + cache_write=_safe_int(usage_meta.get("input_cache_creation_tokens")), + reasoning=_safe_int(usage_meta.get("reasoningTokens") or usage_meta.get("reasoning_tokens")), + 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..55ff810 100644 --- a/src/lago_agent_sdk/lago_client.py +++ b/src/lago_agent_sdk/lago_client.py @@ -11,10 +11,20 @@ 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. + # Access urllib3 via requests' own re-export — it's only a + # transitive dependency for us, not one we declare directly. + requests.packages.urllib3.disable_warnings( # type: ignore[attr-defined] + requests.packages.urllib3.exceptions.InsecureRequestWarning # type: ignore[attr-defined] + ) def __repr__(self) -> str: if not self.api_key: @@ -23,7 +33,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 +47,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..1504676 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,7 +46,7 @@ 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 @@ -40,6 +58,8 @@ 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,7 +69,16 @@ # 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 @@ -75,6 +104,19 @@ "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", +} + # Bedrock cross-region inference prefix -> a representative AWS region. _BEDROCK_REGION_PREFIX = { "us": "us-east-1", @@ -112,7 +154,34 @@ 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.""" + return _fmt_money((Decimal(usd) * Decimal(markup)).quantize(_Q, rounding=ROUND_DOWN)) def _fmt_money(d: Decimal) -> str: @@ -205,19 +274,41 @@ 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 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) @@ -258,9 +349,102 @@ def parse_openrouter(data: Any) -> dict[str, Any]: 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 _pick_mistral_canonical(names: list[str]) -> str: + """Prefer a dated snapshot id (what OpenRouter actually lists models + under) over a "-latest"-style moniker. Falls back to shortest-then- + alphabetical so the choice is always deterministic even with no dated + candidate in the group.""" + dated = [n for n in names if _MISTRAL_DATED_ID.search(n)] + pool = dated or names + return sorted(pool, 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: + 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 @@ -278,6 +462,68 @@ def lookup_openrouter(table: dict[str, Any], provider: str, model: str) -> Model 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 + price_prop = next( + (p for p in m.get("properties", []) 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".""" + hit = table.get(model) + if hit is not None: + return hit + return table.get(_strip_version(model)) + + # ---------------------------------------------------------------------- # Bedrock parsing + matching # @@ -416,13 +662,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 +714,43 @@ 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": 50, "page": page}, timeout=self._timeout + ) + resp.raise_for_status() + body = resp.json() + batch = body.get("result") or [] + models.extend(batch) + total = body.get("result_info", {}).get("total_count", len(models)) + if len(batch) < 50 or len(models) >= total: + 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 +762,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 +784,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 +808,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 +875,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 +912,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 +946,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..5415f13 100644 --- a/src/lago_agent_sdk/queue.py +++ b/src/lago_agent_sdk/queue.py @@ -1,9 +1,17 @@ """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 (Lago 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. """ from __future__ import annotations @@ -17,6 +25,17 @@ from collections.abc import Callable from typing import Any +from .exceptions import LagoApiError + + +def _is_permanent_failure(exc: Exception) -> bool: + """A Lago 4xx (bad request, validation error, duplicate transaction_id, + ...) will never succeed by retrying the exact same batch. A 5xx or a + network-level exception (timeout, connection error, no LagoApiError at + all) might — those stay retryable.""" + return isinstance(exc, LagoApiError) and 400 <= exc.status < 500 + + logger = logging.getLogger("lago_agent_sdk.queue") @@ -73,6 +92,14 @@ 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: @@ -118,6 +145,42 @@ 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) + 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) + self._replay_failed([event]) + def _run(self) -> None: while not self._stopping.is_set(): self._wake.wait(timeout=self._flush_interval) @@ -143,12 +206,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 +226,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..168fbe5 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,15 @@ 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, + money_str_to_cents, +) from .queue import EventQueue logger = logging.getLogger("lago_agent_sdk") @@ -47,6 +56,7 @@ def __init__( 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 +65,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 +93,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 +157,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 +209,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,6 +218,22 @@ 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. In token mode, which can + push several events from one call, each field's event is suffixed + (``f"{event_id}_{field_name}"``) so they don't collide with each other. """ try: sub = self._resolve_subscription(subscription) @@ -157,14 +246,7 @@ def emit( 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) + 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,11 +257,27 @@ 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() if not nonzero: # Mistral legacy / empty — nothing to bill @@ -190,7 +288,7 @@ def _emit_token_events(self, usage: CanonicalUsage, sub: str, dimensions: dict[s if not code: continue event = { - "transaction_id": str(uuid.uuid4()), + "transaction_id": f"{event_id}_{field_name}" if event_id else str(uuid.uuid4()), "external_subscription_id": sub, "code": code, "timestamp": now, @@ -204,49 +302,84 @@ 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()) + base_properties: dict[str, Any] = { "model": usage.model, "provider": usage.provider, "api": usage.api, "price_source": breakdown.source, + "markup": breakdown.markup, + **(dimensions or {}), } + + if not breakdown.fields: + properties = { + **base_properties, + "unit": str(usage.input + usage.output), + "value": breakdown.total, + "base_cost": breakdown.base, + } + self._queue.push( + { + "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"], } - ) + self._queue.push( + { + "transaction_id": f"{event_id}_{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 +389,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..7d26b58 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,6 +58,20 @@ def _is_message_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 _merge_stream_usage(accumulated: dict[str, Any], payload: Any) -> None: """Fold one streaming event's usage into the running accumulator. @@ -95,6 +121,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 +148,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): @@ -150,6 +189,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): diff --git a/src/lago_agent_sdk/wrappers/openai.py b/src/lago_agent_sdk/wrappers/openai.py index 90ccb2f..106d8f9 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, @@ -120,7 +146,7 @@ def _extract_stream_usage(payload: Any) -> dict[str, Any] | None: return {"usage": nested} 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 +155,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 +191,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 +236,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 +250,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/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/unit/gateway/__init__.py b/tests/unit/gateway/__init__.py new file mode 100644 index 0000000..e69de29 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..503d81a --- /dev/null +++ b/tests/unit/gateway/adapters/test_cloudflare_gateway.py @@ -0,0 +1,368 @@ +"""Cloudflare AI Gateway log adapter — verified against a real captured log entry.""" + +from __future__ import annotations + +import json +import pathlib +from decimal import Decimal + +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 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_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..6787e3d 100644 --- a/tests/unit/test_pricing.py +++ b/tests/unit/test_pricing.py @@ -4,6 +4,7 @@ import json import pathlib +import uuid from decimal import Decimal from typing import Any @@ -11,14 +12,20 @@ from lago_agent_sdk import CanonicalUsage, LagoConfig, LagoSDK, ModelPrice from lago_agent_sdk.pricing import ( + HttpPricingFetcher, PricingProvider, + _parse_price, bedrock_model_key, coerce_markup, compute_cost, + compute_precomputed_cost, 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 +36,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 +61,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 +93,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 +115,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 @@ -119,6 +247,192 @@ def test_openrouter_miss_returns_none() -> None: assert lookup_openrouter(table, "openai", "claude-opus-4-8") is None +# ---------------------------------------------------------------------- +# Cloudflare Workers AI parsing + matching +# ---------------------------------------------------------------------- +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") + + +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 + + # ---------------------------------------------------------------------- # Bedrock region + key + offer parsing # ---------------------------------------------------------------------- @@ -238,18 +552,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 +706,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 +934,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 +962,112 @@ 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_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 +1076,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 +1095,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 +1121,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 +1141,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 +1165,117 @@ 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.""" + 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_input", "backfill_01ABC_output"} + + +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 +1306,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: diff --git a/tests/unit/test_queue.py b/tests/unit/test_queue.py index 1728dfa..8bd22b2 100644 --- a/tests/unit/test_queue.py +++ b/tests/unit/test_queue.py @@ -5,6 +5,7 @@ import threading import time +from lago_agent_sdk.exceptions import LagoApiError from lago_agent_sdk.queue import EventQueue @@ -79,6 +80,145 @@ 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) + + +# ---------------------------------------------------------------------- +# 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() diff --git a/tests/unit/test_wrapper_anthropic.py b/tests/unit/test_wrapper_anthropic.py index 85804cf..ac8b013 100644 --- a/tests/unit/test_wrapper_anthropic.py +++ b/tests/unit/test_wrapper_anthropic.py @@ -30,10 +30,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 @@ -216,6 +258,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 +345,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 +459,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_openai.py b/tests/unit/test_wrapper_openai.py index 44e3b19..cfa88af 100644 --- a/tests/unit/test_wrapper_openai.py +++ b/tests/unit/test_wrapper_openai.py @@ -40,10 +40,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 @@ -98,6 +140,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 @@ -260,6 +304,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 +404,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 +456,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 +518,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()