diff --git a/CHANGELOG.md b/CHANGELOG.md index 450db6a..a1cc959 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,16 +4,23 @@ 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. +### Changed + +- **A Databricks-hosted model no longer reports a price failure it can never avoid.** In price mode, `provider="databricks"` is deliberately unmatchable (see below), so `emit()` used to log `lago pricing failed: no price for provider='databricks' model='meta-llama-4-maverick-040225'` on **every single call** and route it to `on_error`. That description is wrong: nothing failed. Databricks bills hosted models in DBUs at a per-model rate that exists on an HTML page and in no system table — verified across every column of all 88 of them — so token counts are the *complete* answer for them, not a degraded fallback, and no refresh could ever supply the missing rate. New `TOKEN_BILLED_PROVIDERS` in `pricing.py` (exported from the package) names the providers this applies to; `emit()` skips the lookup for them, emits token counts, and states the reason **once per model** at info level instead of warning once per call. + - **Deliberately a narrow exception to invariant "never silently under-bill".** That invariant exists so a price miss can't pass unnoticed, and it still holds for every miss a customer could act on — a cold table, an unmatched model name, a mistyped provider all still raise `PricingUnavailableError` through `on_error`. This exception covers only the case where the miss is *structural and permanent*. The reason to make it is that an alarm which always fires is one nobody reads: leaving it in place taught the reader to ignore `on_error`, which is precisely how a real miss gets missed. + - Keys on the **provider**, so it covers Databricks-*hosted* traffic only. BYOK through the same gateway is stamped `openai`/`anthropic` and keeps pricing normally — still verified exact against Databricks' own metered spend on 38 of 38 buckets. ### Fixed + +- **Prompt-cached Mistral calls were over-billed by up to 6.15x in price mode.** `mistral` was missing from `_INPUT_INCLUDES_CACHE_READ`, so the cached portion of a prompt was billed twice: once at the full input rate because `input` was never reduced, and again at the cache-read rate. Mistral's API is OpenAI-shaped and reports `prompt_tokens_details.cached_tokens` as a SUBSET of `prompt_tokens` — its own documented example (`prompt_tokens=1013`, `cached_tokens=1008`, `total_tokens=1043=prompt+completion`) only reconciles if the cached tokens sit inside the prompt count, and Mistral bills them at 10% of the input rate. Measured 6.15x over-bill on that exact payload. 13 of 18 Mistral models on OpenRouter publish a cache-read rate, so the wrong path was reachable for most of them, including Mistral routed through a Cloudflare gateway (the gateway adapter leaves `provider="mistral"` unmapped). Token mode was unaffected — only the price computation was wrong. This is the second provider missing from that set (after `workers-ai`); `money_golden.json` gains a `mistral` case, but the set is still hand-maintained and a completeness check remains the real fix. + +- **A comment in the Databricks adapter documented a correction that would under-bill 13%.** It stated that a computed-cost fallback for this table "must key off `api == \"databricks_gateway\"`, **never** the vendor name". Keying off `api` alone is exactly the mistake: it correctly separates a table row from a live call, but `compute_cost` already subtracts `cache_read` for providers in `_INPUT_INCLUDES_CACHE_READ`, so an OpenAI row corrected that way is double-subtracted — measured at `$0.00354` against a true `$0.004065`. The correction needs both keys. Comment only; no code path reads it today, which is why the error survived review. + +- **Price mode silently missed every current OpenAI model.** `_strip_version` only matched a COMPACT trailing date (`-YYYYMMDD`), which is Anthropic's convention (`claude-sonnet-4-5-20250929`). OpenAI stamps a **hyphenated** one (`gpt-5-2025-08-07`, `gpt-4.1-2025-04-14`, `o3-2025-04-16`), and OpenRouter lists the **bare** id (`openai/gpt-5`) — so a name we could not strip back to bare never matched. Because `resolve_model` prefers the response's own `model` over the requested one, `create(model="gpt-5")` resolves to `gpt-5-2025-08-07` and misses: verified against the live 400-model OpenRouter table with the repo's own `lookup_openrouter` that `gpt-4.1`, `gpt-4.1-mini`, `gpt-5`, `gpt-5-mini`, `o3` and `o4-mini` all fell through to token events, i.e. anyone in price mode on a current OpenAI model was getting no cost at all. `gpt-4o` looked fine only by luck — OpenRouter happens to list `openai/gpt-4o-2024-08-06` verbatim. The pattern now accepts both shapes; all six resolve, the Anthropic compact cases still pass, and Workers AI ids (`@cf/meta/llama-3.3-70b-instruct-fp8-fast`) are left untouched. Found while validating a Databricks AI Gateway backfill, where 5 of 39 real calls could not be priced without it. +- **Tokens reported in neither named bucket were silently dropped.** For genuine OpenAI, `total_tokens` always equals `prompt_tokens + completion_tokens` — reasoning is a SUBSET of completion, never additive — verified across every captured real response with zero deltas. Behind an OpenAI-**compatible** proxy fronting a thinking model that invariant breaks: measured against Gemini through Google's own compat layer, `prompt_tokens: 57`, `completion_tokens: 47`, `total_tokens: 1253`, with the 1,149 thinking tokens reported nowhere and no `completion_tokens_details` to recover them from. Billing prompt+completion dropped **92% of the call**, at the output rate — a silent under-bill, which the "never silently under-bill" invariant exists to prevent. A positive delta is now folded into `output` and recorded as `extras["unaccounted_output_tokens"]`. Deliberately **not** assigned to `reasoning`: `compute_cost` zeroes reasoning whenever the provider is in `_OUTPUT_INCLUDES_REASONING`, and an OpenAI-shaped payload is stamped `provider="openai"` by definition, so that would set the field and immediately discard it — measured as recovering exactly $0. Applies to any such proxy (Cloudflare `/compat`, OpenRouter, LiteLLM, Bifrost); a no-op for real OpenAI, where the delta is always 0. +- **The drift contract did not hold one level down, and a real field vanished because of it.** `extras` swept only *top-level* usage keys, but `prompt_tokens_details` is itself a KNOWN top-level key — so nothing nested inside it was ever inspected. A live `gpt-5.6-sol` response carries `prompt_tokens_details.cache_write_tokens: 3022`, and those 3,022 tokens were discarded with no error and no `on_error`. Every drift test passed, because none of them looked inside a details object. The sweep now recurses into `prompt_tokens_details` / `completion_tokens_details` / `input_tokens_details` / `output_tokens_details`, surfacing anything unmapped under a dotted key — which also finally delivers the Predicted Outputs counts (`accepted_prediction_tokens`, `rejected_prediction_tokens`) that this module's own docstring already promised customers could read from `extras`. + - **Deliberately NOT mapped to `CanonicalUsage.cache_write`, because that would have been the worse bug.** For OpenAI these tokens sit *inside* `prompt_tokens` (measured: `prompt_tokens=3025` with `cache_write_tokens=3022`) and bill at the plain input rate — cross-checked against Databricks' own `system.ai_gateway.external_model_spend`, which charged **$0.015245**, exactly what the SDK already produced by billing all 3,025 as input. But OpenRouter publishes a separate `input_cache_write` rate ($6.25/M) for the model, so mapping the field would have charged those 3,022 tokens twice: **$0.0341 against a true $0.0152, a 2.24× over-bill**. Anthropic is the opposite case — its `cache_creation_input_tokens` sits *outside* `input_tokens`, which is why mapping is correct there and wrong here. No `_INPUT_INCLUDES_CACHE_WRITE` table is warranted: it would only be needed if some provider reported cache-write tokens outside its input count, and none observed does. + - **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. @@ -28,11 +35,35 @@ All notable changes to this project will be documented here. Format follows [Kee - **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 + +- **Databricks AI Gateway connector** — `lago_agent_sdk.gateway.adapters.databricks_gateway`. `extract_databricks_log()` maps a `system.ai_gateway.usage` row to `CanonicalUsage`; `resolve_databricks_subscription()` reads Lago attribution from the caller's `Databricks-Ai-Gateway-Request-Tags` header. Second entry in the `gateway/` namespace, alongside Cloudflare. Verified against real rows read from a live workspace over the SQL Statement Execution API — 226 rows, all 36 columns (the public docs undercount at ~28, omitting `service_type`/`service_id`/`service_name`/`service_tags`/`mcp_metadata`/`routing_information`/`invocation_metadata`), with 24 captured fixtures covering both destination types, cache read/write, reasoning, embeddings, and all three failure shapes. + - **Databricks has no unified endpoint, unlike Cloudflare's `/compat`.** Each provider is reachable only through its own native surface, confirmed by trying the BYOK services against the mlflow endpoint: `400 INVALID_PARAMETER_VALUE: Unsupported native_api_type for OpenAI v1 surface`. Two of the four surfaces use the *same* `openai.OpenAI` class but need different price tables, so `base_url` discrimination is load-bearing rather than cosmetic — hence the new `provider_hint` parameter on `extract_openai_native`, supplied by the wrapper from the client's `base_url` (the response body cannot reveal it: a hosted call echoes a served-entity name like `meta-llama-4-maverick-040225` with no distinguishing marker). `/ai-gateway/mlflow/` stamps `provider="databricks"`; the BYOK surfaces keep their real vendor. + - **`provider="databricks"` is deliberately unmatchable.** It hits no vendor in `_VENDOR_MAP`, so a hosted call cannot reach a price table and `emit()` emits token counts instead (see the `### Changed` entry above — no error is reported, because no rate exists to miss). Databricks bills these in DBUs against a rate card published only as HTML and present in no system table (verified by searching every column of all 88 system tables), while OpenRouter *does* list bare `openai/gpt-oss-20b` and `meta-llama/llama-4-maverick` at 0.2-0.4x of Databricks' real rate — so being stamped `"openai"` would silently under-bill 2.5-5x the moment a served-entity rename let `_strip_version` match it. Same trap as Workers AI, made impossible by construction. + - **BYOK pricing verified exact against Databricks' own billing: 38 of 38 priceable buckets, zero divergences.** Each `system.ai_gateway.external_model_spend` bucket was joined to its usage rows on Databricks' own grouping key `(hour, provider, model, request_tags)` and priced through the real pipeline at live OpenRouter rates. 13 models across both BYOK providers, both cache conventions, four reasoning models, costs spanning $0.0000036 to $0.015245. The four non-matches are all one model, `gpt-5.6`: the spend table records the requested alias while OpenRouter lists only `gpt-5.6-sol`, so it prices live and misses on backfill — a miss falling back to token events, not a mispricing. + - **Two mapping quirks that a docs-only reading gets wrong**, both caught by real rows. `destination_name` is the model for hosted rows but the *provider service* (a Unity Catalog credential name) for BYOK, so a single "model, falling back to name" rule bills a credential as the model. And `destination_model` is unstable for hosted models — the same `destination_name` was observed reporting both `llama-4-maverick` and the display label `Llama 4 Maverick`. Model resolution is therefore keyed off `destination_type`, and `provider` off `api_type`'s leading segment, which already *is* this SDK's provider vocabulary. + - **This table's `input_tokens` INCLUDES `cache_read` and `cache_write`** — the inverse of the providers' own response bodies, confirmed per row (`input=1825, cache_read=1812` for a call whose body reported `input_tokens: 13`). The adapter extracts faithfully and does not subtract, because the intended billing path takes Databricks' own metered USD and never touches token counts. Documented in the module because computing from these tokens instead over-bills 3.04x with no correction — and the correction is per-provider, not uniform: `compute_cost` already subtracts `cache_read` for providers in `_INPUT_INCLUDES_CACHE_READ`, so an OpenAI row must pass through while an Anthropic row must be pre-subtracted. Getting that uniform under-bills 13% one way and over-bills 3x the other. + - Failed calls (403/404, and every Gemini call while that connection is broken) are recorded with NULL token counts and extract to all-zero, so nothing is emitted — the same way a Cloudflare cache hit does. Gemini itself is out of scope: its Databricks connection returns an unhandled `500` with an empty body to Databricks' own documented code sample, which their KB attributes to using a Google AI Studio key with the Vertex-typed provider. + - **`gateway/databricks.py` — the one piece of gateway code that does I/O, deliberately.** `DatabricksSource.read_usage(window)` returns rows already shaped for `emit()`, and `LagoSDK.backfill_databricks(source, "7 days")` bills a whole window in one call, returning `{"cost": n, "tokens": n, "skipped": n}`. Cloudflare's read is one paginated GET and rightly stays in the example notebook; Databricks needs a SQL warehouse, the Statement Execution API, columnar-to-dict zipping, chunked result fetching, and two tables reconciled against each other — ~100 lines in which three money-losing mistakes are easy, all three of which the first hand-rolled version of the demo notebook actually made. **Silent truncation:** only chunk 0 arrives inline, so a window wide enough to span `manifest.total_chunk_count > 1` bills a fraction of itself with no error. **Double billing:** a BYOK call appears in *both* `ai_gateway.usage` and `external_model_spend`. **Unscoped idempotency keys:** `transaction_id` is unique account-wide, so a row id built from the source row alone blocks that row from ever reaching a second subscription — and because an untagged row is billed to the caller's default rather than to its own (absent) tag, the key has to be built from the subscription actually billed, which is what `DatabricksUsageRow.event_id_for()` exists for. Uses `requests` (Python) / global `fetch` (JS), already present, so nothing is added to the install; `databricks-sql-connector` remains the better choice for interactive analysis and the pure adapter still accepts its rows. Verified live: 107 billable rows over a 7-day window → 148 events, byte-identical `transaction_id`s across a re-run. + - **The window is validated, not escaped.** It reaches SQL by interpolation, so `read_usage("1 day; DROP TABLE …")` is refused outright — only a bare count plus unit, or a `datetime`, is ever accepted. + - **Token counts are summed per spend bucket.** `external_model_spend` aggregates per `(hour, model, provider, request_tags)`, so N calls in one hour collapse to one dollar row while `ai_gateway.usage` still holds N token rows; reporting only the first would understate the tokens behind a cost the customer can see in their own console. + - **Every backfilled event carries the grouping key of the Databricks surface it came from**, which is what makes the connector checkable rather than merely correct: `endpoint_name` for hosted rows (how the AI Gateway usage page groups) and `bucket`, the hour, for BYOK rows (`external_model_spend`'s own aggregation key). Without it a side-by-side comparison fails on naming alone — the SDK's `model` is normalized (`qwen35-122b-a10b`) where the gateway page shows `system.ai.qwen35-122b-a10b` or even a display label (`GPT OSS 20B`). `backfill_databricks()` gained a `dimensions=` argument for the caller's own keys, applied after the automatic ones so an explicit key wins rather than being silently overwritten. Deliberately NOT emitted: `invocation_id`/`request_id`/`status_code` — one Lago group per request is a list, not a comparison, and on an hourly aggregate they state one sampled request's value as if it described the whole hour. Verified live: 148 events over a 7-day window, 88 hosted across 13 endpoints, 60 BYOK across 4 hours, none without a key. + - **Hosted models bill as token counts on BOTH paths, and the earlier claim that backfill yields a dollar cost from `system.billing.usage` × `list_prices` was wrong** — nothing in the source tree ever queried those tables, and the README table row contradicted the paragraph two lines below it. The dollars are genuinely available (`list_prices`, or `account_prices` for an account's contract rate), so "hosted USD is impossible" was also wrong; that applies only to the tokens→DBU rate, which exists on an HTML page and in no table. They are not billed from because they come from a *different Databricks screen* than the gateway view: `custom_tags` is `{}` on every `billing.usage` row, so per-subscription splits would be ours rather than Databricks', and the table lags the gateway by roughly a day — measured at `max(usage_start_time) = 2026-08-10T17:00` against `max(event_time) = 2026-08-11T10:09`. Emitting only what a Databricks *gateway* page also shows is the property being protected. + - **A spend bucket no longer reports one sampled request's fields as its own.** `_merge_usage` copied `extras` from the first row of the hourly bucket, so a merged BYOK row carried an arbitrary request's `invocation_id`/`status_code` as though it described the hour. Harmless while `extras` was never emitted; a live mis-statement the moment the reconciliation dimensions above read from it. Bucket representatives now keep only the endpoint-describing keys (`endpoint_name`, `endpoint_id`, `destination_type`, `destination_name`, `api_type`), and the filter is applied to single-request buckets too — otherwise `status_code` survived on quiet hours and vanished on busy ones. + - **Hosted model names shed a second prefix.** Most hosted entities are named `system.ai.databricks-`, not `system.ai.` — 38 of 48 distinct hosted `destination_name`s on a live workspace carry that inner `databricks-`, which is a serving-endpoint artefact rather than part of the model id. Left in, it emitted `databricks-qwen35-122b-a10b` as the model, splitting one model into two rows in Lago against the live path's own name for it. It is **not** stripped unconditionally, because Databricks also publishes models genuinely named that way (`databricks-dbrx-instruct`, `databricks-dolly-v2`) and no string inspection tells the two apart: `destination_model` is the tie-breaker, and a disagreement keeps the raw name rather than guessing. + - `examples/databricks_gateway_demo.ipynb` (Python) demonstrates both halves end to end and was re-run against a live workspace after being rewritten onto the helper: 107 billable rows over 7 days → 60 dollar-cost events plus 88 token events, all `transaction_id`s unique. + +- **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. + - **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. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 368a059..92713c6 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -70,6 +70,7 @@ uv lock --upgrade-package X # bump a single package - `src/lago_agent_sdk/canonical.py` — the normalized usage shape sent to Lago - `src/lago_agent_sdk/queue.py` — async event queue with backoff - `src/lago_agent_sdk/lago_client.py` — thin HTTP client to `/events/batch` +- `src/lago_agent_sdk/gateway/` — second front door: gateway usage logs → `CanonicalUsage`, for backfill - `tests/unit/` — unit tests, organized to mirror `src/` - `tests/unit/adapters/fixtures/` — captured real provider responses, used by adapter tests - `tests/integration/` — live tests, gated on credential env vars @@ -84,6 +85,83 @@ uv lock --upgrade-package X # bump a single package 6. Add unit tests against the captured fixtures. 7. Add a live integration test gated on the provider's API key env var. +## Adding a gateway + +`gateway/` is a **second front door** into the same kernel, separate from the provider-native +`adapters/` used by `wrap()`. A gateway connector reads a gateway's own usage log and maps it into +`CanonicalUsage` for backfill; there is no client to patch. Two exist: Cloudflare and Databricks. + +1. Capture real rows/entries from a live gateway into + `tests/unit/gateway/adapters/fixtures//`, one file per scenario. Cover both success and + every failure shape you can produce — failed calls are where the surprises live. +2. Write `src/lago_agent_sdk/gateway/adapters/.py` exporting + `extract__log(entry) -> CanonicalUsage` and `resolve__subscription(entry) -> str | None`. + Keep it a **pure function**: no HTTP, no SDK state. +3. Export both from `gateway/adapters/__init__.py` under explicitly gateway-scoped names, so no + gateway is the implicit default. +4. Add `tests/unit/gateway/adapters/test_.py` against the captured fixtures. +5. Add a `## AI Gateway` README section and a `CHANGELOG.md` entry. +6. Add `examples/_gateway_demo.ipynb` showing backfill and live calls. + +### A connector is only as good as the comparison + +The reason the Cloudflare connector reads well is that you can put the gateway's own +dashboard beside Lago and see the same numbers. Two rules protect that, and both were +learned the hard way on Databricks: + +- **Emit the gateway's own grouping key as a dimension.** Our `model` is normalized; the + gateway's page is not. Group Lago by one and the dashboard by the other and the + comparison fails on naming alone, before any number is even wrong. Attach the key the + gateway's surface aggregates by, and only keys that are true of the whole row — a + per-request field on an hourly aggregate is one sampled value dressed up as a property + of the bucket. +- **Never bill from a surface the gateway UI doesn't show.** Databricks does expose exact + dollars for its hosted models, in `system.billing.usage` x `list_prices` — on a + different screen, with no attribution tags, about a day behind. Billing from it would + produce a number the customer cannot find anywhere, which costs more trust than the + feature adds. Hosted therefore bills token counts, matching the page they do look at. + +### Does the read itself belong in the SDK? + +Default: **no.** The adapters stay pure and the fetching lives in the example notebook, as Cloudflare's +does — its whole read is one paginated GET, and an SDK wrapper around that would be indirection for +nothing. + +Databricks earned the exception, in `gateway/databricks.py` (a sibling module, so the adapter stays +pure). The bar it cleared, and the one to hold a third gateway to: the read is long enough that a +customer will reimplement it wrong, and the ways it goes wrong lose money silently. Databricks needs a +SQL warehouse, the Statement Execution API, columnar-to-dict zipping, chunked result fetching and two +tables reconciled against each other — and the first hand-rolled version in the demo notebook truncated +at chunk 0, which bills a fraction of a wide window with no error at all. If a gateway's read is a loop +over one endpoint, leave it in the notebook. + +When it does clear the bar: name it `gateway/.py`, expose a `Source` with an explicit +window and a `read_usage()` that yields rows already shaped for `emit()`, add the `backfill_()` +one-liner to `LagoSDK`, and use a dependency that is already core (`requests` / `undici`). No scheduler, +no cursor store, no credential store — that is the poller, and it stays out of the SDK. + +### Things both existing connectors had to get right + +These are the traps, and every one of them cost real debugging: + +- **Which cost is authoritative.** Gateway traffic bills from the *gateway's* metered cost, not one we + compute — it keeps Lago reconcilable against the dashboard the customer looks at. Note the gateway + may under-report: Cloudflare's `cost` omits additive reasoning tokens, measured at 22.8x low on a + real call. +- **Token semantics are per-gateway, not per-vendor.** Cloudflare passes Anthropic's cache counts + through *additively*; Databricks' table folds them *into* `input_tokens`. Same provider, opposite + conventions. Never assume the vendor's own convention survives the gateway. +- **`provider` must be unmatchable when you cannot price it.** If a gateway bills on its own rate card, + stamp a provider that hits nothing in `_VENDOR_MAP` so the lookup misses honestly. Stamping a real + vendor name lets a near-miss model string match at 2.5-5x the wrong rate, silently. +- **Idempotency keys must be subscription-scoped.** `transaction_id` is unique org-wide, so + `f"{prefix}_{subscription}_{row_id}"` — an unscoped id silently blocks a row from ever reaching a + second subscription. +- **Failed calls appear in the log.** Extract them to all-zero so `nonzero_numeric()` is empty and + nothing is emitted, rather than billing zeros. +- **Drift.** An unrecognized field must reach `extras`, including one level down inside nested + `*_details` objects. `test_drift.py` pins it. + ## Pull request checklist - [ ] Unit tests cover the change diff --git a/README.md b/README.md index 2389c66..424c8fb 100644 --- a/README.md +++ b/README.md @@ -166,6 +166,95 @@ sdk.flush() See [`examples/cloudflare_gateway_demo.ipynb`](examples/cloudflare_gateway_demo.ipynb) for a runnable end-to-end version of both. +**Gateway-routed calls are billed at the gateway's metered cost.** Cloudflare reports its own `cost` per log entry and the backfill passes that straight through, so Lago reconciles against the dashboard you actually look at. One measured consequence to be aware of: that field excludes additive *reasoning* tokens, so a thinking-heavy Gemini call bills about 4% of what Google charges (verified live at 22.8x on one call, 39.6x on another — the ratio tracks each prompt's thinking-to-output ratio). Cloudflare is exact on input, output, cache-read and cache-write. + +**If you hand-roll a poller, don't use `urllib`.** `gateway.ai.cloudflare.com` returns `403` with body `error code: 1010` to `Python-urllib` — its bot-signature check. Any other User-Agent passes, and `requests` (which this SDK uses) is fine. The failure looks like an auth error because the body is otherwise empty. + +## Databricks AI Gateway + +Unlike Cloudflare, Databricks has **no unified endpoint** — each provider is reachable only through its own native surface, and two of them use the same `openai.OpenAI` class. Which `base_url` you point at decides how the call is priced. + +**Databricks-hosted foundation models** (`system.ai.*`) — billed by Databricks in DBUs: + +```python +from openai import OpenAI +from lago_agent_sdk import LagoSDK + +sdk = LagoSDK(api_key="...", default_subscription_id="sub_acme") +client = sdk.wrap(OpenAI( + api_key=DATABRICKS_TOKEN, + base_url=f"{DATABRICKS_HOST}/ai-gateway/mlflow/v1", + default_headers={"Databricks-Ai-Gateway-Request-Tags": json.dumps({"lago_subscription": "sub_acme"})}, +)) +client.chat.completions.create(model="system.ai.llama-4-maverick", messages=[{"role": "user", "content": "Hi"}]) +``` + +**Your own vendor key (BYOK)** — Anthropic via its native passthrough, note `api_key="unused"` because the real credential goes in `Authorization`, and the Unity Catalog connection holding your Anthropic key is named in `Databricks-Model-Provider-Service`: + +```python +from anthropic import Anthropic +client = sdk.wrap(Anthropic( + api_key="unused", + base_url=f"{DATABRICKS_HOST}/ai-gateway/anthropic", + default_headers={ + "Authorization": f"Bearer {DATABRICKS_TOKEN}", + "Databricks-Model-Provider-Service": "workspace.default.anthropickey", + }, +)) +``` + +OpenAI BYOK is the same `OpenAI` class as the hosted example, against `/ai-gateway/openai/v1` with its own `Databricks-Model-Provider-Service`. + +### What gets billed + +| Path | Live `wrap()` | Backfill | +|---|---|---| +| BYOK (OpenAI / Anthropic) | **dollar cost**, priced from the vendor's published rates | dollar cost from Databricks' own `external_model_spend` | +| Hosted (`system.ai.*`) | **token counts** | **token counts** | + +BYOK prices live because you pay the vendor directly, so the vendor's rate *is* your cost — verified against Databricks' own metered spend on 38 of 38 real buckets, exactly. Hosted models bill in DBUs against a rate card published only as HTML and present in no system table, so there is no rate to look up: those calls emit token counts instead of a dollar cost. That is the complete answer for them, not a degraded one, so it is **not** reported as an error — `TOKEN_BILLED_PROVIDERS` lists the providers this applies to, and the SDK notes it once per model at info level rather than warning on every call. A genuine price miss — a cold table, an unmatched model name — still reports through `on_error` as before. + +**Hosted dollars exist, and are deliberately not billed from.** `system.billing.usage` × `list_prices` (or `account_prices` for your contract rate) does yield exact USD per hour and endpoint. It is not used because it comes from a *different Databricks screen* than the gateway view: it carries no `request_tags`, so per-subscription splits would be ours rather than Databricks', and it lags the gateway by roughly a day — measured at ~19h on a live workspace. Every number this connector sends is one you can find on a Databricks **gateway** page, which is the property that makes it checkable. + +**Grouping matches the Databricks page.** Each backfilled event carries the grouping key of the surface it came from — `endpoint_name` for hosted, `bucket` (the hour) for BYOK. Group Lago by `endpoint_name` and you get the AI Gateway → Usage table row for row. Pass `dimensions={...}` to add your own keys; yours win on a name collision. + +**Don't run the live path and the backfill over the same hosted traffic.** Both emit token events, with different `transaction_id`s, so Lago accepts both and the counts double. Pick one per traffic stream: `wrap()` for real time, the backfill for completeness. + +`Databricks-Ai-Gateway-Request-Tags` is what makes attribution work. It lands in `request_tags` on `system.ai_gateway.usage` **and** is a first-class aggregation dimension on `external_model_spend`, so tagging `lago_subscription` means BYOK cost arrives already split per subscription — no apportioning needed. + +### Backfill — give it a window, it does the rest + +```python +from lago_agent_sdk.gateway.databricks import DatabricksSource + +source = DatabricksSource.from_env() # DATABRICKS_HOST / _TOKEN / _WAREHOUSE_ID +print(sdk.backfill_databricks(source, "7 days", default_subscription="sub_default")) +sdk.flush() +# {'cost': 60, 'tokens': 47, 'skipped': 0} +``` + +Pass a `datetime` instead of `"7 days"` for an exact lower bound, and `unified=True` to bill the whole window to `default_subscription` regardless of per-call tags. + +Unlike Cloudflare's single paginated GET, this one is worth having in the SDK — hand-rolling it is ~100 lines with three money-losing traps in them. The Statement Execution API returns only **chunk 0** inline, so a wide window silently truncates and bills a fraction of it with no error. A BYOK call appears in **both** `ai_gateway.usage` and `external_model_spend`, so billing both charges twice. And `transaction_id` is unique account-wide, so an unscoped row id blocks that row from ever reaching a second subscription. + +To inspect a window before billing it, or to route rows yourself, read them directly — each row is already shaped for `emit()`: + +```python +for row in source.read_usage("7 days"): + print(row.usage.model, row.subscription, row.usd_cost) # usd_cost is None for hosted +``` + +Reading the system tables needs a PAT with the **`sql`** scope plus a SQL warehouse — the live calls above need neither. `examples/databricks_gateway_demo.ipynb` is a complete worked example of both halves. The pure `extract_databricks_log(row)` / `resolve_databricks_subscription(row)` functions stay available from `lago_agent_sdk.gateway.adapters` if you already have rows from `databricks-sql-connector` or your own warehouse job. + +**One cost note:** a SQL warehouse is a real cost centre. Measured on a test workspace, the warehouse queries cost roughly 1,500× the model-serving usage they were reporting on. Run the backfill as one query over a wide window, never as a tight polling loop. + +### Gotchas worth knowing + +- **`gpt-oss` models inflate input by ~100 tokens** from a server-injected preamble — a 2-character prompt bills 102. Not an SDK error. +- **`claude-opus-4-5` does not cache through this gateway**: reproducibly `cache_read`/`cache_write` of 0 with the full prompt billed as input, on a request shape where `claude-sonnet-4-5` caches fine. An opus customer silently gets no cache discount. +- **Hosted models report three different name strings.** `system.ai.llama-4-maverick` and `databricks-llama-4-maverick` both work as requests, and the response echoes a third (`meta-llama-4-maverick-040225`). Pricing keys off the resolved name, so reconciling by requested id will not line up. +- **Embeddings** work on `/ai-gateway/mlflow/v1/embeddings` and report input only — no `completion_tokens` at all. + ## Multi-tenant — pick a subscription per call Three ways to set the `external_subscription_id`, in priority order: diff --git a/examples/.env.example b/examples/.env.example index 247a32b..2aa1925 100644 --- a/examples/.env.example +++ b/examples/.env.example @@ -1,6 +1,9 @@ # Copy this file to examples/.env and fill in real values. # examples/.env is gitignored — never commit real credentials. +# --------------------------------------------------------------------------- +# cloudflare_gateway_demo.ipynb +# --------------------------------------------------------------------------- CF_ACCOUNT_ID= CF_GATEWAY_ID= CF_LOGS_TOKEN= @@ -21,3 +24,22 @@ ANTHROPIC_API_KEY= # resolution too (see LagoConfig.mistral_api_key) — no separate credential # needed for that. MISTRAL_API_KEY= + +# --------------------------------------------------------------------------- +# databricks_gateway_demo.ipynb +# --------------------------------------------------------------------------- +DATABRICKS_HOST=https://dbc-xxxxxxxx-xxxx.cloud.databricks.com +# Part 2 (live calls) works with any token that has gateway inference access. +# Part 1 (backfill) additionally needs the `sql` scope to read +# system.ai_gateway.usage — without it every warehouse route returns +# 403 "does not have required scopes: sql", including the Thrift path the +# databricks-sql-connector uses, so changing client library does not help. +DATABRICKS_TOKEN= +# Backfill only. SQL Warehouses -> your warehouse -> Connection details. +# Just the id here (the notebook builds the rest), e.g. a292ad231ac2d202. +DATABRICKS_WAREHOUSE_ID= +# Unity Catalog connections holding your own vendor keys (BYOK). Only needed for +# whichever provider you actually call in Part 2. Three-level UC names, e.g. +# workspace.default.anthropickey. +DATABRICKS_PROVIDER_SERVICE_ANTHROPIC= +DATABRICKS_PROVIDER_SERVICE_OPENAI= diff --git a/examples/databricks_gateway_demo.ipynb b/examples/databricks_gateway_demo.ipynb new file mode 100644 index 0000000..44393ce --- /dev/null +++ b/examples/databricks_gateway_demo.ipynb @@ -0,0 +1,382 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "db00setup", + "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.\n", + "\n", + "import json\n", + "import os\n", + "import sys\n", + "\n", + "sys.path.insert(0, \"../src\") # run this notebook from examples/, or adjust to your install\n", + "\n", + "\n", + "def _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", + "\n", + "from lago_agent_sdk import LagoSDK # noqa: E402\n", + "from lago_agent_sdk.config import LagoConfig # noqa: E402\n", + "from lago_agent_sdk.gateway.databricks import DatabricksSource # noqa: E402\n", + "\n", + "_REQUIRED = [\"DATABRICKS_HOST\", \"DATABRICKS_TOKEN\", \"LAGO_API_KEY\"]\n", + "_missing = [name for name in _REQUIRED if not os.environ.get(name)]\n", + "if _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 .env.example.\"\n", + " )\n", + "\n", + "DBX_HOST = os.environ[\"DATABRICKS_HOST\"].rstrip(\"/\")\n", + "DBX_TOKEN = os.environ[\"DATABRICKS_TOKEN\"]\n", + "# Backfill only — see Part 1 for the scope it needs.\n", + "DBX_WAREHOUSE_ID = os.environ.get(\"DATABRICKS_WAREHOUSE_ID\", \"\")\n", + "# Unity Catalog connections holding your own vendor keys (BYOK). Only needed for\n", + "# the provider you actually call in Part 2.\n", + "SVC_ANTHROPIC = os.environ.get(\"DATABRICKS_PROVIDER_SERVICE_ANTHROPIC\", \"\")\n", + "SVC_OPENAI = os.environ.get(\"DATABRICKS_PROVIDER_SERVICE_OPENAI\", \"\")\n", + "\n", + "LAGO_API_KEY = os.environ[\"LAGO_API_KEY\"]\n", + "LAGO_API_URL = os.environ.get(\"LAGO_API_URL\", \"https://api.getlago.com/api/v1\")\n", + "LAGO_SUBSCRIPTION_ID = os.environ.get(\"LAGO_SUBSCRIPTION_ID\", \"databricks_gateway_demo_sub\")\n", + "LAGO_VERIFY_SSL = os.environ.get(\"LAGO_VERIFY_SSL\", \"true\").lower() != \"false\"\n", + "\n", + "sdk = 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,\n", + " api_url=LAGO_API_URL,\n", + " pricing_mode=\"price\",\n", + " verify_ssl=LAGO_VERIFY_SSL,\n", + " ),\n", + ")\n", + "# Blocks until OpenRouter's table is fetched, closing the cold-start race for the\n", + "# very first live call. That table is what prices the BYOK paths in Part 2 —\n", + "# verified exact against Databricks' own metered spend on 39 of 39 real calls.\n", + "# Databricks-HOSTED models are deliberately NOT priceable from it: they bill in\n", + "# DBUs against Databricks' own rate card, so they fall back to token events by\n", + "# design rather than being priced at some other vendor's rate for the same\n", + "# open-weight model.\n", + "sdk.warm_pricing()\n", + "print(\"SDK ready — billing to\", LAGO_SUBSCRIPTION_ID)\n", + "print(\"backfill enabled:\", bool(DBX_WAREHOUSE_ID))\n" + ] + }, + { + "cell_type": "markdown", + "id": "db01md1", + "metadata": {}, + "source": [ + "## Part 1 — Backfill historic usage from Databricks\n", + "\n", + "Databricks has no REST logs API. Usage lands in Unity Catalog system tables read\n", + "over SQL, and cost lives in a *different* table from the token counts:\n", + "\n", + "| | table | unit |\n", + "|---|---|---|\n", + "| tokens, per request, with your attribution tags | `system.ai_gateway.usage` | counts |\n", + "| cost for BYOK providers, already attributed | `system.ai_gateway.external_model_spend` | **USD** |\n", + "\n", + "`DatabricksSource` reads both and reconciles them, so the whole backfill is a\n", + "window plus one call. Doing it by hand is about a hundred lines with three\n", + "money-losing traps in them: the Statement Execution API returns only **chunk 0**\n", + "inline, so a wide window silently truncates; a BYOK call appears in **both** tables\n", + "and billing both charges twice; and `transaction_id` is unique account-wide, so an\n", + "unscoped row id blocks that row from ever reaching a second subscription.\n", + "\n", + "The billing rule is the same one the Cloudflare connector follows: **the gateway is\n", + "the metering authority**, so a BYOK row bills from Databricks' own `usage_quantity`\n", + "via `emit(usd_cost=...)` rather than from a price we compute ourselves.\n", + "Databricks-hosted models have no per-request USD anywhere in Databricks' system\n", + "tables, so they bill as token events.\n", + "\n", + "`request_tags` is a first-class aggregation dimension on the spend table, so tagging\n", + "`lago_subscription` on the call means cost arrives **already split per\n", + "subscription** — no apportioning by token share.\n", + "\n", + "One caveat that is easy to miss: a SQL warehouse is a real cost centre. Measured on\n", + "the test workspace behind this notebook, warehouse queries cost roughly **1,500x**\n", + "the model-serving usage they were reporting on. Read one wide window per run; never\n", + "poll in a tight loop.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db02query", + "metadata": {}, + "outputs": [], + "source": [ + "# Needs a PAT carrying the `sql` scope plus a SQL warehouse — the live calls in\n", + "# Part 2 need neither. A token without `sql` fails every warehouse route with\n", + "# 403 \"does not have required scopes: sql\", including the Thrift path the\n", + "# databricks-sql-connector uses, so switching client libraries does not help.\n", + "if not DBX_WAREHOUSE_ID:\n", + " raise SystemExit(\"DATABRICKS_WAREHOUSE_ID is unset — backfill needs a SQL warehouse.\")\n", + "\n", + "source = DatabricksSource(host=DBX_HOST, token=DBX_TOKEN, warehouse_id=DBX_WAREHOUSE_ID)\n", + "# ...or DatabricksSource.from_env(), which reads the same three variables.\n", + "\n", + "WINDOW = \"7 days\" # or a datetime, for an exact lower bound\n", + "\n", + "rows = list(source.read_usage(WINDOW))\n", + "byok = [r for r in rows if r.is_byok]\n", + "hosted = [r for r in rows if not r.is_byok]\n", + "\n", + "print(f\"{len(rows)} billable rows in the last {WINDOW}\")\n", + "print(f\" {len(byok):>3} BYOK ${sum(r.usd_cost for r in byok):.6f} metered by Databricks\")\n", + "print(f\" {len(hosted):>3} hosted token counts only — no per-request USD exists\")\n", + "for r in rows[:5]:\n", + " print(f\" {r.usage.provider:<10} {r.usage.model:<24} {r.subscription or '(untagged)'}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db03emit", + "metadata": {}, + "outputs": [], + "source": [ + "# One call: resolve each row's subscription, pick cost-vs-tokens per row, and emit.\n", + "# unified=True bills everything to one subscription, ignoring per-call tags — right\n", + "# when this gateway's traffic all belongs to one customer. Set it False to respect\n", + "# real per-call attribution and fall back to the default only for untagged rows.\n", + "# `rows` from the cell above is passed straight in, so the window is read ONCE.\n", + "# Handing `source` + WINDOW here instead would re-run both warehouse queries — and a\n", + "# warehouse costs ~1,500x the model-serving usage it reports on, so that doubles the\n", + "# expensive half of this notebook. It would also let rows land between the two reads,\n", + "# making the summary printed above disagree with what was billed.\n", + "counts = sdk.backfill_databricks(\n", + " rows,\n", + " default_subscription=LAGO_SUBSCRIPTION_ID,\n", + " unified=True,\n", + ")\n", + "assert sdk.flush(timeout=30.0), \"queue did not flush in time\"\n", + "print(counts)\n", + "\n", + "# Re-run this cell: every event id is derived from the source row and scoped by\n", + "# subscription, so Lago rejects the duplicates instead of billing the window twice.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Compare it against Databricks\n", + "\n", + "Below is what Databricks' own tables say for this window, next to what was sent to Lago.\n", + "\n", + "They are equal **by construction**, not by luck — a BYOK row bills `usage_quantity`\n", + "verbatim via `emit(usd_cost=...)` with no price lookup on our side, and a hosted row\n", + "bills the table's own token counts. So read this as a visible restatement, not as an\n", + "independent audit.\n", + "\n", + "What makes it *checkable* is the last block: every event carries the grouping key of the\n", + "Databricks surface it came from — `endpoint_name` for hosted, `bucket` (the hour) for\n", + "BYOK. Group Lago by `endpoint_name` and you get the table below, row for row, next to\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Everything here comes from the `rows` already read above — no extra query, so this\n", + "# cell costs nothing. (A SQL warehouse is expensive relative to the traffic it reports\n", + "# on: measured at ~1,500x the model-serving usage in this workspace.)\n", + "from collections import defaultdict\n", + "\n", + "byok_usd = sum(r.usd_cost for r in byok)\n", + "hosted_in = sum(r.usage.input for r in hosted)\n", + "hosted_out = sum(r.usage.output for r in hosted)\n", + "\n", + "n_cost = counts[\"cost\"]\n", + "n_in = sum(1 for r in hosted if r.usage.input)\n", + "n_out = sum(1 for r in hosted if r.usage.output)\n", + "\n", + "print(\"Databricks says -> sent to Lago\")\n", + "print(f\"BYOK external_model_spend ${byok_usd:>8.6f} -> llm_cost \"\n", + " f\"${byok_usd:>8.6f} {n_cost} events\")\n", + "print(f\"hosted ai_gateway.usage in {hosted_in:>9,} -> llm_input_tokens \"\n", + " f\"{hosted_in:>10,} {n_in} events\")\n", + "print(f\"{'':31}out {hosted_out:>8,} -> llm_output_tokens \"\n", + " f\"{hosted_out:>10,} {n_out} events\")\n", + "\n", + "# Keyed by endpoint_name — the same column the AI Gateway usage page groups by, and the\n", + "# dimension now on every hosted event.\n", + "per_endpoint = defaultdict(lambda: [0, 0])\n", + "for r in hosted:\n", + " key = r.usage.extras.get(\"endpoint_name\") or r.usage.model\n", + " per_endpoint[key][0] += r.usage.input\n", + " per_endpoint[key][1] += r.usage.output\n", + "\n", + "print(\"\\nper endpoint — read against the AI Gateway usage page:\")\n", + "for endpoint, (tin, tout) in sorted(per_endpoint.items(), key=lambda kv: -sum(kv[1])):\n", + " print(f\" {endpoint:<40} in {tin:>8,} out {tout:>8,}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "db04md2", + "metadata": {}, + "source": [ + "## Part 2 — Live calls through the gateway\n", + "\n", + "Each provider is reachable **only** through its own native surface — there is no\n", + "single OpenAI-compatible front door the way Cloudflare offers `/compat`. So the\n", + "same `openai.OpenAI` client means two different things depending on `base_url`:\n", + "`/ai-gateway/mlflow/v1` is a Databricks-hosted model billed in DBUs, while\n", + "`/ai-gateway/openai/v1` is your own OpenAI account.\n", + "\n", + "`Databricks-Ai-Gateway-Request-Tags` carries the Lago attribution and is what\n", + "makes cost arrive pre-split per subscription in Part 1.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db05prompt", + "metadata": {}, + "outputs": [], + "source": [ + "PROMPT = \"Tell me about getLago, the billing company - give as many details as you can find\"\n", + "TAGS = json.dumps({\"lago_subscription\": LAGO_SUBSCRIPTION_ID, \"team\": \"lago-demo\"})\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db06hosted", + "metadata": {}, + "outputs": [], + "source": [ + "# Databricks-HOSTED foundation model, via the unified mlflow surface.\n", + "# Prices in DBUs against Databricks' own rate card, which OpenRouter does not\n", + "# carry — so this bills as token events, deliberately, rather than being matched\n", + "# to some other vendor's price for the same open-weight model.\n", + "from openai import OpenAI\n", + "\n", + "client = sdk.wrap(OpenAI(\n", + " api_key=DBX_TOKEN,\n", + " base_url=f\"{DBX_HOST}/ai-gateway/mlflow/v1\",\n", + " default_headers={\"Databricks-Ai-Gateway-Request-Tags\": TAGS},\n", + "))\n", + "resp = client.chat.completions.create(\n", + " model=\"system.ai.llama-4-maverick\",\n", + " messages=[{\"role\": \"user\", \"content\": PROMPT}],\n", + " max_tokens=400,\n", + ")\n", + "text = resp.choices[0].message.content\n", + "print(\"resolved model:\", resp.model, \"| usage:\", resp.usage)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db07print", + "metadata": {}, + "outputs": [], + "source": [ + "print(text)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db08anthropic", + "metadata": {}, + "outputs": [], + "source": [ + "# Anthropic BYOK, via the native passthrough. Two quirks: the Anthropic SDK wants\n", + "# an api_key, so it gets a placeholder and the real credential goes in\n", + "# Authorization; and the Unity Catalog connection holding your Anthropic key is\n", + "# named in Databricks-Model-Provider-Service.\n", + "from anthropic import Anthropic\n", + "\n", + "client = sdk.wrap(Anthropic(\n", + " api_key=\"unused\",\n", + " base_url=f\"{DBX_HOST}/ai-gateway/anthropic\",\n", + " default_headers={\n", + " \"Authorization\": f\"Bearer {DBX_TOKEN}\",\n", + " \"Databricks-Model-Provider-Service\": SVC_ANTHROPIC,\n", + " \"Databricks-Ai-Gateway-Request-Tags\": TAGS,\n", + " },\n", + "))\n", + "resp = client.messages.create(\n", + " model=\"claude-sonnet-4-5\",\n", + " max_tokens=400,\n", + " messages=[{\"role\": \"user\", \"content\": PROMPT}],\n", + ")\n", + "print(\"resolved model:\", resp.model, \"| usage:\", resp.usage)\n", + "print(resp.content[0].text)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db09openai", + "metadata": {}, + "outputs": [], + "source": [ + "# OpenAI BYOK, via the native OpenAI surface. Same client class as the hosted cell\n", + "# above — only base_url differs, and that difference decides which price table\n", + "# applies. This path is priced from OpenRouter and matched Databricks' own metered\n", + "# spend to the digit on every real call tested.\n", + "from openai import OpenAI\n", + "\n", + "client = sdk.wrap(OpenAI(\n", + " api_key=DBX_TOKEN,\n", + " base_url=f\"{DBX_HOST}/ai-gateway/openai/v1\",\n", + " default_headers={\n", + " \"Databricks-Model-Provider-Service\": SVC_OPENAI,\n", + " \"Databricks-Ai-Gateway-Request-Tags\": TAGS,\n", + " },\n", + "))\n", + "resp = client.chat.completions.create(\n", + " model=\"gpt-4o\",\n", + " messages=[{\"role\": \"user\", \"content\": PROMPT}],\n", + " max_tokens=400,\n", + ")\n", + "print(\"resolved model:\", resp.model, \"| usage:\", resp.usage)\n", + "print(resp.choices[0].message.content)\n", + "\n", + "assert sdk.flush(timeout=30.0), \"queue did not flush in time\"\n", + "print(\"\\nflushed — check Lago for llm_cost / token events\")\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/lago_agent_sdk/__init__.py b/src/lago_agent_sdk/__init__.py index de2a037..c3548e7 100644 --- a/src/lago_agent_sdk/__init__.py +++ b/src/lago_agent_sdk/__init__.py @@ -9,7 +9,13 @@ PricingUnavailableError, UnknownClientError, ) -from .pricing import HttpPricingFetcher, ModelPrice, PricingProvider, compute_cost +from .pricing import ( + TOKEN_BILLED_PROVIDERS, + HttpPricingFetcher, + ModelPrice, + PricingProvider, + compute_cost, +) from .sdk import LagoSDK __all__ = [ @@ -27,5 +33,6 @@ "HttpPricingFetcher", "ModelPrice", "compute_cost", + "TOKEN_BILLED_PROVIDERS", ] __version__ = "0.2.0" diff --git a/src/lago_agent_sdk/adapters/openai_native.py b/src/lago_agent_sdk/adapters/openai_native.py index 21ae9d2..f614b23 100644 --- a/src/lago_agent_sdk/adapters/openai_native.py +++ b/src/lago_agent_sdk/adapters/openai_native.py @@ -57,6 +57,38 @@ "output_tokens_details", } +# Nested keys inside the *_tokens_details sub-objects that we actually MAP onto a +# CanonicalUsage field. Anything nested that isn't listed here is drift and gets +# surfaced in `extras` under a dotted key. +# +# Sweeping only top-level keys was a real hole: `prompt_tokens_details` is itself +# a KNOWN top-level key, so nothing inside it was ever inspected. A live +# gpt-5.6-sol response carries `prompt_tokens_details.cache_write_tokens: 3022` +# and those 3022 tokens vanished with no error — a silent violation of the drift +# contract test_drift.py exists to pin, which passed only because it never looked +# one level down. +# +# NOTE the billing subtlety: cache_write_tokens must NOT be mapped to +# CanonicalUsage.cache_write. For OpenAI it sits INSIDE prompt_tokens (measured: +# prompt_tokens=3025 with cache_write_tokens=3022) and bills at the plain input +# rate — Databricks charged exactly what billing all 3025 as input produces. But +# OpenRouter does publish a separate cache_write rate for the model, so mapping it +# would charge those tokens twice: $0.0341 against a true $0.0152, a 2.24x +# over-bill. Anthropic is the opposite — its cache_creation_input_tokens sits +# OUTSIDE input_tokens, which is why mapping is correct there and wrong here. +# Surfacing in extras keeps the field visible without touching the money. +_MAPPED_DETAIL_FIELDS = { + "prompt_tokens_details": {"cached_tokens", "audio_tokens"}, + "input_tokens_details": {"cached_tokens", "audio_tokens"}, + "completion_tokens_details": {"reasoning_tokens", "audio_tokens"}, + # NOTE `output_tokens_details` deliberately omits `audio_tokens`: the Responses + # branch hardcodes `audio_output = 0` because the API does not expose it today, so + # listing it here would exclude a real, unmapped count from `extras` — 500 audio + # tokens vanishing with no error, which is the exact hole this table closes. Add it + # back only together with a Responses branch that reads it. + "output_tokens_details": {"reasoning_tokens"}, +} + def _safe_dict(v: Any) -> dict[str, Any]: return v if isinstance(v, dict) else {} @@ -117,11 +149,19 @@ def _infer_provider(resolved_model: str) -> str: return "openai" -def extract_openai_native(response: Any, model_id: str = "") -> CanonicalUsage: +def extract_openai_native(response: Any, model_id: str = "", provider_hint: str = "") -> CanonicalUsage: """Translate an OpenAI response (chat completion or responses API) → CanonicalUsage. Accepts the SDK's pydantic objects, dicts (e.g. captured fixtures), or the synthetic `{"usage": {...}}` blob produced by the streaming wrapper. + + `provider_hint` overrides the model-string inference below. Only the wrapper + can supply it, because the only reliable signal for some gateways is the + client's `base_url` — which the response never carries. Databricks is the + case that forced it: a Databricks-HOSTED model answers on + `/ai-gateway/mlflow/v1` but echoes a served-entity name + ("meta-llama-4-maverick-040225") with no marker of its own, so no rule based + on the model string can identify it. See `wrappers/openai.py`. """ resp = _to_dict(response) if not isinstance(response, dict) else response usage = _safe_dict(resp.get("usage")) @@ -158,6 +198,51 @@ def extract_openai_native(response: Any, model_id: str = "") -> CanonicalUsage: if k not in _KNOWN_USAGE_FIELDS: extras[k] = v + # Drift sweep one level down, into the *_tokens_details sub-objects. Without + # this, an unrecognized nested field is silently dropped (see + # _MAPPED_DETAIL_FIELDS) because its container is a known top-level key. + for container, mapped in _MAPPED_DETAIL_FIELDS.items(): + for k, v in _safe_dict(usage.get(container)).items(): + if k not in mapped: + extras[f"{container}.{k}"] = v + + # Consistency guard: for genuine OpenAI, total_tokens always equals + # prompt + completion (reasoning is a SUBSET of completion, never additive). + # Verified across every captured real OpenAI-shaped response — zero deltas. + # So a POSITIVE delta means tokens exist that neither named bucket accounts + # for, which only happens behind an OpenAI-COMPATIBLE proxy that under-reports. + # + # Measured on Gemini through Google's own OpenAI-compat layer: + # prompt_tokens=57, completion_tokens=47, total_tokens=1253 — 1149 real + # thinking tokens reported nowhere, and no completion_tokens_details to + # recover them from. Billing prompt+completion drops 92% of the call, at the + # output rate. Folding the remainder into `output` is the honest read: the + # provider's own total proves those tokens were generated. + # + # Deliberately NOT assigned to `reasoning`: compute_cost zeroes reasoning for + # providers in _OUTPUT_INCLUDES_REASONING, so for real OpenAI that would set the + # field and immediately discard it, recovering nothing. + # + # `reasoning` is subtracted from the accounted total, and that subtraction is + # load-bearing rather than cosmetic. This adapter no longer only ever emits + # provider="openai" — it also emits "workers-ai" (Cloudflare `/compat`) and + # "databricks" (via provider_hint), and for those compute_cost bills reasoning + # ADDITIVELY. A payload reporting both `reasoning_tokens` and an inflated + # `total_tokens` would then be charged for them twice: once inside the grown + # `output` and again as a separate reasoning line. Subtracting first means a + # provider that already broke reasoning out gets no second bill, while the case + # this guard exists for — a thinking model behind a proxy that reports NO + # breakdown at all (measured: prompt 57, completion 47, total 1253) — still + # recovers its 1,149 tokens, because reasoning is 0 there. + # + # A no-op for real OpenAI either way: total always equals prompt + completion. + declared_total = _safe_int(usage.get("total_tokens")) + if declared_total: + unaccounted = declared_total - (input_tokens + output_tokens + reasoning) + if unaccounted > 0: + output_tokens += unaccounted + extras["unaccounted_output_tokens"] = unaccounted + resolved_model = resolve_model(resp.get("model"), model_id) return CanonicalUsage( input=input_tokens, @@ -168,7 +253,7 @@ def extract_openai_native(response: Any, model_id: str = "") -> CanonicalUsage: audio_output=audio_output, tool_calls=tool_calls, model=resolved_model, - provider=_infer_provider(resolved_model), + provider=provider_hint or _infer_provider(resolved_model), api=api, extras=extras, ) diff --git a/src/lago_agent_sdk/gateway/adapters/__init__.py b/src/lago_agent_sdk/gateway/adapters/__init__.py index 1e0177b..7f1aabb 100644 --- a/src/lago_agent_sdk/gateway/adapters/__init__.py +++ b/src/lago_agent_sdk/gateway/adapters/__init__.py @@ -1,6 +1,15 @@ from .cloudflare_gateway import extract_cloudflare_log, resolve_subscription +from .databricks_gateway import extract_databricks_log, resolve_databricks_subscription + +# `resolve_subscription` predates the second gateway and reads Cloudflare's +# `cf-aig-metadata`. Exported under an explicit name too, so the two gateways read +# symmetrically at the call site and neither is the implicit default. +resolve_cloudflare_subscription = resolve_subscription __all__ = [ "extract_cloudflare_log", + "extract_databricks_log", + "resolve_cloudflare_subscription", + "resolve_databricks_subscription", "resolve_subscription", ] diff --git a/src/lago_agent_sdk/gateway/adapters/databricks_gateway.py b/src/lago_agent_sdk/gateway/adapters/databricks_gateway.py new file mode 100644 index 0000000..8488b20 --- /dev/null +++ b/src/lago_agent_sdk/gateway/adapters/databricks_gateway.py @@ -0,0 +1,230 @@ +"""Databricks AI Gateway usage adapter — maps a `system.ai_gateway.usage` row to CanonicalUsage. + +Verified against real rows read from a live workspace over the SQL Statement +Execution API (226 rows, all 36 columns; the public docs undercount at ~28 and +omit `service_*`, `mcp_metadata`, `routing_information`, `invocation_metadata`). + +Unlike Cloudflare, Databricks exposes no REST logs API — usage lands in a Unity +Catalog Delta table queried over SQL. The row reaches this function as a plain +dict: `databricks-sql-connector` yields `Row` objects with `.asDict()`, the +Node driver yields column-keyed objects natively, and the raw Statement +Execution API returns columnar `data_array` the caller zips. All three end up +here as `{column_name: value}`. + +Field mapping (`system.ai_gateway.usage`): + input_tokens → input + output_tokens → output + token_details.cache_read_input_tokens → cache_read + token_details.cache_creation_input_tokens → cache_write + token_details.output_reasoning_tokens → reasoning + destination_type + destination_name/_model → model, provider (see below) + api → hardcoded "databricks_gateway" + extras → routing/identity columns + +`total_tokens` is deliberately NOT mapped: it is derived from the others and +mapping it would double-count. Same reason the Cloudflare adapter skips +`usage_metadata.total_tokens`. + +TWO MEASURED QUIRKS drive the shapes below. Both were wrong in an earlier draft +of this connector that reasoned from the docs alone. + +1. `destination_name` means DIFFERENT things per destination type. For a hosted + model it is the model (`system.ai.llama-4-maverick`); for BYOK it is the + PROVIDER SERVICE (`workspace.default.anthropickey`) — a credential name, not + a model. So a single "model, falling back to name" rule yields a credential + as the model for every BYOK row. + +2. `destination_model` is unstable for hosted models. The same + `destination_name` was observed reporting both `llama-4-maverick` and + `Llama 4 Maverick` — a human display label with spaces and capitals — and + likewise `gpt-oss-20b` / `GPT OSS 20B`. It is clean and stable for BYOK + (`claude-sonnet-4-5`, `gpt-4o`), so it is authoritative there and unusable + for hosted. + +BILLING HAZARD, documented because it is the inverse of every other adapter +here: this table's `input_tokens` INCLUDES both cache_read and cache_write, +where the providers' own response bodies EXCLUDE them. Measured per row — +`input=1825, cache_read=1812` for a call whose response body reported +`input_tokens: 13`. Only one of cache_read/cache_write is ever non-zero per +row, so `input - cache_read - cache_write` recovers the true non-cached input +exactly. This adapter extracts the row FAITHFULLY and does not subtract: the +intended billing path takes Databricks' own metered USD from +`system.ai_gateway.external_model_spend` via `emit(usd_cost=...)`, which never +touches token counts. Computing cost from these tokens instead would over-bill +3.04x with no subtraction, or 1.40x subtracting only cache_read. + +If a computed fallback is ever added, the correction needs BOTH keys, not one. +`api == "databricks_gateway"` alone distinguishes a table row from a live call +(a `provider="anthropic"` row from this table needs correcting; a live +`provider="anthropic"` call must not) — but it is not sufficient, because +`compute_cost` ALREADY subtracts cache_read for providers in +_INPUT_INCLUDES_CACHE_READ. So an openai/gemini row must pass through untouched +while an anthropic row must be pre-subtracted. Measured by getting it wrong: +correcting an openai row double-subtracts and billed $0.00354 against a true +$0.004065, a 13% UNDER-bill. + +Failed calls (403/404, and every Gemini call while that connection is broken) +are recorded with NULL token counts. They extract to all-zero, so +`nonzero_numeric()` is empty and the caller emits nothing — the same way a +Cloudflare cache hit extracts to zero. +""" + +from __future__ import annotations + +import json +from typing import Any + +from ...canonical import CanonicalUsage + +# Databricks' own name for a first-party pay-per-token foundation model. Any other +# destination type (observed: "EXTERNAL_FOUNDATION_MODEL", or NULL on rows rejected +# before routing) is BYOK — the customer's own vendor credential behind a Unity +# Catalog connection. +_HOSTED_DESTINATION_TYPE = "PAY_PER_TOKEN_FOUNDATION_MODEL" + +# Unity Catalog prefix on every hosted model's `destination_name`. +_HOSTED_NAME_PREFIX = "system.ai." + +# A second, INNER prefix that most hosted entities also carry: +# `system.ai.databricks-claude-sonnet-4-5`, `system.ai.databricks-qwen35-122b-a10b`. +# Measured on a live workspace: 38 of 48 distinct hosted `destination_name`s have it +# and 10 do not (`system.ai.gpt-oss-20b`, `system.ai.llama-4-maverick`, ...). It is a +# serving-endpoint naming artefact, not part of the model id — leaving it in emits +# `databricks-qwen35-122b-a10b` as the model, which both reads as a vendor prefix and +# splits one model into two rows in Lago against the live path's own name. +# +# It is NOT safe to strip unconditionally: Databricks also publishes models whose own +# names begin the same way (`databricks-dbrx-instruct`, `databricks-dolly-v2`), and no +# amount of string inspection tells the two apart. `destination_model` does — it was +# the clean name on all 38 prefixed rows — so the prefix comes off only when the two +# columns agree that it is an artefact. Disagreement keeps the raw name: a model +# emitted under a slightly ugly id is recoverable, a silently renamed one is not. +_HOSTED_ENDPOINT_PREFIX = "databricks-" + + +def _safe_dict(v: Any) -> dict[str, Any]: + """Coerce a STRUCT/MAP column to a dict, accepting either shape it arrives in. + + The SQL drivers hand back real dicts (pyarrow-backed), but the raw Statement + Execution API serializes STRUCT and MAP columns as JSON STRINGS — measured: + `token_details` arrives as '{"cache_read_input_tokens":1812}'. Tolerating both + means the adapter works whichever access path the caller chose, rather than + silently reading zeros from a string it never parsed. + """ + if isinstance(v, dict): + return v + if isinstance(v, str) and v.strip().startswith("{"): + try: + parsed = json.loads(v) + except ValueError: + return {} + return parsed if isinstance(parsed, dict) else {} + return {} + + +def _safe_int(v: Any) -> int: + """Coerce to a non-negative int. Token columns arrive as STRINGS over the REST + API ("1825") and as NULL on failed calls; both must land on 0 rather than raise.""" + try: + return max(0, int(v or 0)) + except (TypeError, ValueError): + return 0 + + +def _safe_str(v: Any) -> str: + return v if isinstance(v, str) else "" + + +def _model_and_provider(row: dict[str, Any]) -> tuple[str, str]: + """Resolve (model, provider) — type-dependent, for the reasons in the module docstring.""" + destination_type = _safe_str(row.get("destination_type")) + destination_name = _safe_str(row.get("destination_name")) + + if destination_type == _HOSTED_DESTINATION_TYPE: + # `destination_name` is the stable id here; `destination_model` flips + # between a slug and a display label for the very same model — measured, + # `system.ai.gpt-oss-20b` reports both "gpt-oss-20b" and "GPT OSS 20B". + model = destination_name + if model.startswith(_HOSTED_NAME_PREFIX): + model = model[len(_HOSTED_NAME_PREFIX) :] + if model.startswith(_HOSTED_ENDPOINT_PREFIX): + shed = model[len(_HOSTED_ENDPOINT_PREFIX) :] + if shed == _safe_str(row.get("destination_model")): + model = shed + # Deliberately "databricks", which matches no vendor in pricing's + # _VENDOR_MAP — so a price lookup CANNOT hit and emit() falls back to + # token events (see TOKEN_BILLED_PROVIDERS — no error, since no rate + # exists to miss), rather than silently + # pricing a DBU-billed model at some other vendor's rate. OpenRouter does + # list bare `openai/gpt-oss-20b` etc. at 0.2-0.4x of Databricks' own rate, + # so an accidental match here would under-bill 2.5-5x. + return model, "databricks" + + # BYOK: `destination_model` is the clean requested alias, and `api_type` names + # the native surface the call went through — "anthropic/v1/messages", + # "openai/v1/chat/completions", "gemini/v1/generateContent". Its leading + # segment already IS this SDK's provider vocabulary, so no alias table is + # needed. "unmanaged" (an unrecognized path) yields "unmanaged", which no + # vendor matches — an honest miss, and those rows carry no usage anyway. + provider = _safe_str(row.get("api_type")).split("/")[0] + return _safe_str(row.get("destination_model")), provider + + +def extract_databricks_log(row: dict[str, Any]) -> CanonicalUsage: + """Translate one `system.ai_gateway.usage` row → CanonicalUsage. + + Missing/malformed fields degrade to zero/empty rather than raising, matching + the defensive style of the other adapters — a backfill processing a batch of + rows must not have one malformed row take down the whole run. + """ + details = _safe_dict(row.get("token_details")) + model, provider = _model_and_provider(row) + + return CanonicalUsage( + input=_safe_int(row.get("input_tokens")), + output=_safe_int(row.get("output_tokens")), + cache_read=_safe_int(details.get("cache_read_input_tokens")), + cache_write=_safe_int(details.get("cache_creation_input_tokens")), + reasoning=_safe_int(details.get("output_reasoning_tokens")), + model=model, + provider=provider, + api="databricks_gateway", + extras={ + # `invocation_id` is per individual inference call while `request_id` + # is per request — one request with a fallback produces several + # invocations, the same distinction Cloudflare's `step` marks. Keep + # both; `invocation_id` is the row's natural idempotency key. + "request_id": row.get("request_id"), + "invocation_id": row.get("invocation_id"), + # A THIRD naming variant: `endpoint_name` is the requested form + # (`databricks-llama-4-maverick`, `system.ai.gemma-3-12b`) where + # `destination_name` is the resolved entity (`system.ai.gemma-3-12b-it`). + # Kept for reconciliation; never price off it. + "endpoint_name": row.get("endpoint_name"), + "endpoint_id": row.get("endpoint_id"), + "destination_type": row.get("destination_type"), + "destination_name": row.get("destination_name"), + "api_type": row.get("api_type"), + "status_code": row.get("status_code"), + }, + ) + + +def resolve_databricks_subscription(row: dict[str, Any]) -> str | None: + """Pull the Lago subscription id from the caller's `request_tags`. + + Customers set these with the `Databricks-Ai-Gateway-Request-Tags` header (a + JSON object of string→string), the direct analogue of Cloudflare's + `cf-aig-metadata`. Note they are also a first-class AGGREGATION DIMENSION on + `system.ai_gateway.external_model_spend`, so tagging `lago_subscription` + yields cost already attributed per subscription — no token-share + apportioning needed for BYOK. + + Returns None if the caller never set `lago_subscription` — untagged calls do + produce rows, with `request_tags` empty. The caller decides what to do with + an unattributed row (drop it, route to a default, warn); this function only + reports whether attribution is present. + """ + tags = _safe_dict(row.get("request_tags")) + value = tags.get("lago_subscription") + return value if isinstance(value, str) and value else None diff --git a/src/lago_agent_sdk/gateway/databricks.py b/src/lago_agent_sdk/gateway/databricks.py new file mode 100644 index 0000000..2522816 --- /dev/null +++ b/src/lago_agent_sdk/gateway/databricks.py @@ -0,0 +1,519 @@ +"""Databricks AI Gateway usage reader — the I/O half of the connector. + +`gateway/adapters/databricks_gateway.py` stays a pure function with no I/O. This +module is its sibling: it does the reading, and it exists because reading usage +out of Databricks is genuinely hard in a way Cloudflare's is not. + +Cloudflare is one paginated GET, about twelve lines. Databricks needs a SQL +warehouse, the Statement Execution API, columnar-to-dict zipping, chunked result +fetching, and TWO different tables whose rows must not be billed twice. Hand-rolled +that is ~100 lines in which several money-losing mistakes are easy: + + * **Silent truncation.** The Statement Execution API returns only chunk 0 inline; + `manifest.total_chunk_count` can be higher and the rest need separate fetches. + A naive reader works on a small window and quietly bills a fraction of a large + one, with no error. + * **Double billing.** A BYOK call appears in BOTH `ai_gateway.usage` (tokens) and + `ai_gateway.external_model_spend` (USD). Bill both and you charge twice. + * **Unscoped idempotency keys.** `transaction_id` is unique account-wide, so an + unscoped row id silently blocks that row from ever reaching a second + subscription. + +Deliberately NOT here: scheduler, cursor store, credential store. You pass an +explicit window and this returns what it finds; it does not remember where it got +to. That is the poller, and it stays a separate concern — as the Cloudflare +connector's changelog already states. + +Uses `requests`, already a core dependency, so this adds nothing to the install. +`databricks-sql-connector` would also work and is the better choice for +interactive analysis, but it is a heavy extra to require for a batch read. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import re +from collections.abc import Iterator +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any + +from ..canonical import CanonicalUsage +from .adapters.databricks_gateway import ( + _safe_str, + extract_databricks_log, + resolve_databricks_subscription, +) + +logger = logging.getLogger("lago_agent_sdk.gateway.databricks") + +_STATEMENTS_PATH = "/api/2.0/sql/statements" + +# `since` as an interval string is interpolated into SQL, so it is validated +# strictly rather than escaped — only a bare count plus a unit is ever accepted. +_INTERVAL_RE = re.compile(r"^\s*(\d{1,5})\s+(second|minute|hour|day|week)s?\s*$", re.I) + + +@dataclass +class DatabricksUsageRow: + """One billable row, already shaped for `emit()`. + + `usd_cost` is set only for BYOK rows, where Databricks meters the provider cost + itself in `external_model_spend`. Hosted rows leave it None: Databricks bills + those in DBUs against a rate card that exists in no system table, so there is no + per-request dollar figure to pass through and they bill as token counts. + """ + + usage: CanonicalUsage + subscription: str | None + row_id: str + kind: str + usd_cost: float | None = None + prefix: str = "dbx" + raw: dict[str, Any] = field(default_factory=dict, repr=False) + + @property + def is_byok(self) -> bool: + return self.usd_cost is not None + + @property + def event_id(self) -> str: + """Idempotency key for billing this row to the subscription its tags name.""" + return self.event_id_for(self.subscription) + + @property + def reconcile_dimensions(self) -> dict[str, str]: + """The Databricks-side grouping key for this row, to be emitted as dimensions. + + This is what makes the connector checkable: the customer opens the Databricks + page, groups Lago the same way, and reads the two side by side. Without it the + comparison fails on naming alone — our `model` is normalized + (`qwen35-122b-a10b`) where the gateway page shows `system.ai.qwen35-122b-a10b` + or even a display label (`GPT OSS 20B`). + + Each kind gets the key that its OWN Databricks surface aggregates by, and only + keys that are true of the whole row: + + * hosted — `endpoint_name`, how the AI Gateway usage page groups. + * BYOK — `bucket`, the hour, which is `external_model_spend`'s own + aggregation key. Deliberately NOT `endpoint_name` here: a spend row covers + an hour of requests, so any per-request field would be one sampled value + dressed up as a property of the bucket. + + `invocation_id` / `request_id` / `status_code` are excluded for the same reason + plus cardinality — one Lago group per request is not a comparison, it's a list. + """ + if self.kind == "spend": + bucket = _stamp(self.raw.get("bucket")) + return {"bucket": bucket} if bucket else {} + endpoint = _safe_str(self.usage.extras.get("endpoint_name")) + return {"endpoint_name": endpoint} if endpoint else {} + + def event_id_for(self, subscription: str | None) -> str: + """The same key, scoped to whichever subscription is actually billed. + + Scoping is not cosmetic: Lago's `transaction_id` is unique account-wide, so an + id built from the source row alone silently blocks that row from ever reaching + a second subscription. And the subscription billed is not always the one on the + row — an untagged row falls back to the caller's default — so the key has to be + built from the resolved value, not from `self.subscription`. + """ + return f"{self.prefix}_{self.kind}_{subscription or 'none'}_{self.row_id}" + + +def _interval_sql(since: str | datetime) -> str: + """Render a window as a SQL predicate value. Rejects anything unrecognized.""" + if isinstance(since, datetime): + # Databricks stores `event_time`/`usage_start_time` in UTC, so an aware + # datetime must be CONVERTED, not formatted as-is: `strftime` would emit local + # wall time and a Europe/Paris caller would read a window two hours in the + # future, bill nothing, and report success. A naive datetime is taken as UTC, + # which is also what the JS port's `toISOString()` does with a Date. + moment = since.astimezone(timezone.utc) if since.tzinfo is not None else since + return f"TIMESTAMP '{moment.strftime('%Y-%m-%d %H:%M:%S')}'" + m = _INTERVAL_RE.match(str(since)) + if not m: + raise ValueError( + f"since={since!r} not understood — pass a datetime, or a string like " + "'7 days' / '24 hours' / '30 minutes'" + ) + count, unit = m.group(1), m.group(2).upper() + return f"current_timestamp() - INTERVAL {count} {unit}" + + +class DatabricksSource: + """Reads Databricks AI Gateway usage over the SQL Statement Execution API. + + Needs a PAT carrying the **`sql`** scope plus a SQL warehouse — the live + `wrap()` path needs neither. Without them every warehouse route returns + `403 "does not have required scopes: sql"`. + + A SQL warehouse is a real cost centre: measured on a test workspace, warehouse + queries cost roughly 1,500x the model-serving usage they were reporting on. Read + one wide window per run; never poll in a tight loop. + """ + + def __init__( + self, + host: str, + token: str, + warehouse_id: str, + *, + timeout: float = 180.0, + wait_timeout: str = "50s", + ) -> None: + self.host = host.rstrip("/") + self.token = token + self.warehouse_id = warehouse_id + self.timeout = timeout + # Databricks rejects anything outside 0s or 5-50s. + self.wait_timeout = wait_timeout + + @classmethod + def from_env(cls, **kwargs: Any) -> DatabricksSource: + """Build from `DATABRICKS_HOST` / `DATABRICKS_TOKEN` / `DATABRICKS_WAREHOUSE_ID`.""" + import os + + missing = [ + k + for k in ("DATABRICKS_HOST", "DATABRICKS_TOKEN", "DATABRICKS_WAREHOUSE_ID") + if not os.environ.get(k) + ] + if missing: + raise ValueError(f"missing environment variable(s): {', '.join(missing)}") + return cls( + host=os.environ["DATABRICKS_HOST"], + token=os.environ["DATABRICKS_TOKEN"], + warehouse_id=os.environ["DATABRICKS_WAREHOUSE_ID"], + **kwargs, + ) + + # ------------------------------------------------------------------ + # SQL + # ------------------------------------------------------------------ + def query(self, sql: str) -> list[dict[str, Any]]: + """Run one statement and return every row as a dict. + + Handles the three things a naive reader gets wrong: the response is COLUMNAR + (`manifest.schema.columns` plus a positional `data_array`); only chunk 0 + arrives inline — the rest must be fetched, or a wide window truncates + silently; and a statement still running when `wait_timeout` elapses comes back + as HTTP 200 with `state: PENDING`, which has to be polled rather than treated + as a failure. + """ + import requests + + headers = {"Authorization": f"Bearer {self.token}"} + resp = requests.post( + f"{self.host}{_STATEMENTS_PATH}", + headers=headers, + json={ + "statement": sql, + "warehouse_id": self.warehouse_id, + "wait_timeout": self.wait_timeout, + }, + timeout=self.timeout, + ) + body = resp.json() + body = self._await_statement(body, headers) + + manifest = body.get("manifest") or {} + columns = [c["name"] for c in (manifest.get("schema") or {}).get("columns", [])] + result = body.get("result") or {} + arrays: list[list[Any]] = list(result.get("data_array") or []) + + total_chunks = int(manifest.get("total_chunk_count") or 1) + statement_id = body.get("statement_id") + for index in range(1, total_chunks): + chunk = requests.get( + f"{self.host}{_STATEMENTS_PATH}/{statement_id}/result/chunks/{index}", + headers=headers, + timeout=self.timeout, + ).json() + arrays.extend(chunk.get("data_array") or []) + if total_chunks > 1: + logger.info("lago: databricks result spanned %d chunks (%d rows)", total_chunks, len(arrays)) + + return [dict(zip(columns, row, strict=False)) for row in arrays] + + def _await_statement(self, body: dict[str, Any], headers: dict[str, str]) -> dict[str, Any]: + """Poll a statement to a terminal state, returning the body that carries results. + + A statement still executing when the request's `wait_timeout` elapses returns + **HTTP 200** with `state: PENDING`/`RUNNING` and a `statement_id` — not an error. + Treating that as fatal breaks exactly the case this class tells operators to use: + one wide window per run, which on a cold warehouse routinely takes longer than + the 50s ceiling Databricks allows for `wait_timeout`. + """ + import time + + import requests + + deadline = time.monotonic() + self.timeout + while True: + state = (body.get("status") or {}).get("state") + if state == "SUCCEEDED": + return body + if state not in ("PENDING", "RUNNING"): + raise RuntimeError(f"Databricks statement {state}: {(body.get('status') or body)}") + statement_id = body.get("statement_id") + if not statement_id or time.monotonic() >= deadline: + raise RuntimeError( + f"Databricks statement still {state} after {self.timeout}s " + f"(statement_id={statement_id}); raise `timeout` or narrow the window" + ) + time.sleep(2.0) + body = requests.get( + f"{self.host}{_STATEMENTS_PATH}/{statement_id}", + headers=headers, + timeout=self.timeout, + ).json() + + # ------------------------------------------------------------------ + # Reading + # ------------------------------------------------------------------ + def read_usage( + self, since: str | datetime = "1 day", *, event_id_prefix: str = "dbx" + ) -> Iterator[DatabricksUsageRow]: + """Yield every billable row in the window, shaped for `emit()`. + + BYOK and hosted are read from DIFFERENT tables and must not overlap, or a + call gets billed twice: + + * BYOK — `external_model_spend`, which carries Databricks' own metered + USD *and* your `request_tags`, so cost arrives already attributed per + subscription. Token counts are joined on from `ai_gateway.usage` for + reporting; they are not used to compute the price. + * hosted — `ai_gateway.usage` only, billed as token counts. + + Rows whose usage is entirely zero (failed calls are recorded with NULL token + counts) are skipped, so nothing emits an empty event. + """ + window = _interval_sql(since) + + spend = self.query(f""" + SELECT record_id, + date_trunc('HOUR', usage_start_time) AS bucket, + usage_metadata.provider AS provider, + usage_metadata.model AS model, + to_json(custom_tags.request_tags) AS request_tags, + usage_quantity + FROM system.ai_gateway.external_model_spend + WHERE usage_start_time >= {window} + """) + + usage = self.query(f""" + SELECT * FROM system.ai_gateway.usage + WHERE event_time >= {window} + ORDER BY event_time + """) + + # Extract once per row and reuse: this loop and the hosted loop below both need + # the CanonicalUsage, and extraction parses several JSON-string columns. + extracted = [(row, extract_databricks_log(row)) for row in usage] + + # Index token counts by the spend table's own grouping key, so a BYOK event + # can carry real counts alongside Databricks' dollar figure. + tokens: dict[tuple[Any, ...], CanonicalUsage] = {} + for row, u in extracted: + if u.provider == "databricks": + continue + key = ( + _bucket_of(row.get("event_time")), + u.provider, + str(row.get("destination_model") or ""), + _canonical_tags(row.get("request_tags")), + ) + prior = tokens.get(key) + tokens[key] = _merge_usage(prior, u) if prior else _as_bucket(u) + billed_keys: set[tuple[Any, ...]] = set() + + for row in spend: + usd = _safe_float(row.get("usage_quantity")) + if not usd: + continue + key = ( + _truncate_hour(_stamp(row.get("bucket"))), + str(row.get("provider") or ""), + str(row.get("model") or ""), + _canonical_tags(row.get("request_tags")), + ) + billed_keys.add(key) + joined = tokens.get(key) + usage_obj = joined or CanonicalUsage( + model=str(row.get("model") or ""), + provider=str(row.get("provider") or ""), + api="databricks_gateway", + ) + sub = resolve_databricks_subscription({"request_tags": row.get("request_tags")}) + yield DatabricksUsageRow( + usage=usage_obj, + subscription=sub, + # record_id is unique per aggregated spend row — a natural + # idempotency key. See `event_id_for` for why it is still scoped. + row_id=_row_id(row, "record_id"), + kind="spend", + usd_cost=usd, + prefix=event_id_prefix, + raw=row, + ) + + # A BYOK bucket with no spend row is billed by NEITHER loop, so say so rather + # than lose it. `external_model_spend` is an hourly aggregate that lags + # `ai_gateway.usage`, so the window's most recent hour routinely has token rows + # whose dollar row does not exist yet; a $0 metered row does the same. Re-running + # the window once Databricks has aggregated picks them up — but only if the + # operator knows to, which is what this warning is for. + unbilled = sorted(set(tokens) - billed_keys) + if unbilled: + logger.warning( + "lago: %d BYOK token bucket(s) in this window have no external_model_spend " + "row yet and were NOT billed (e.g. hour=%s provider=%s model=%s). The spend " + "table lags; re-run this window later to bill them.", + len(unbilled), + unbilled[0][0], + unbilled[0][1], + unbilled[0][2], + ) + + for row, u in extracted: + if u.provider != "databricks": + continue # BYOK already billed from spend above — never twice + if not u.nonzero_numeric(): + continue # failed calls carry NULL tokens + yield DatabricksUsageRow( + usage=u, + subscription=resolve_databricks_subscription(row), + # One request with a fallback yields several invocations, so + # invocation_id is the per-row key; request_id is the fallback for a + # row that somehow carries no invocation. + row_id=_row_id(row, "invocation_id", "request_id"), + kind="usage", + usd_cost=None, + prefix=event_id_prefix, + raw=row, + ) + + +def _row_id(row: dict[str, Any], *columns: str) -> str: + """First usable id among `columns`, falling back to a hash of the whole row. + + Two ways the obvious `_safe_str(a or b)` goes wrong, both silent and both losing + money. A row with NULL ids yields ""; so does a row whose id a driver hands back as + a UUID or int object rather than a str, because `or` selects it and `_safe_str` + rejects the type without ever trying the next column. Either way `event_id_for` + still produces a well-formed key (`dbx_usage_sub_x_`), so EVERY such row in the + window shares one `transaction_id` — Lago accepts the first and rejects the rest as + duplicates, and those calls are never billed at all. + + The content hash keeps the key deterministic, so re-running the same window is still + idempotent, which a random UUID would break. + """ + for column in columns: + value = row.get(column) + if value is None: + continue + text = str(value).strip() + if text: + return text + digest = hashlib.sha256( + json.dumps(row, sort_keys=True, default=str).encode("utf-8", "replace") + ).hexdigest() + return f"sha{digest[:32]}" + + +def _safe_float(v: Any) -> float: + """Coerce a decimal(38,18) column to float. Returns 0.0 on anything unparseable. + + `float()` raises on a non-numeric string, and this runs inside a generator whose + docstring promises one malformed row cannot take down the batch — an exception here + would abort the window mid-emit with no record of where it stopped. 0.0 means the + row is skipped like any other zero-dollar row. Mirrors the JS port, where + `Number()` yields NaN and the same `if (!usd)` skips it. + """ + try: + return float(v or 0) + except (TypeError, ValueError): + return 0.0 + + +def _truncate_hour(value: str) -> str: + """Normalize a timestamp string to its hour, for joining across the two tables.""" + return value[:13] if len(value) >= 13 else value + + +def _stamp(value: Any) -> str: + """Stringify a timestamp column, whatever the access path produced. + + The Statement Execution API returns TIMESTAMPs as strings, but + `databricks-sql-connector` returns real `datetime` objects — a documented, supported + input path. `_safe_str` would map those to "", collapsing every hour of the window + into one join bucket and dropping the `bucket` reconcile dimension entirely. + """ + if value is None: + return "" + # ISO-8601 for a real datetime, matching what the REST API returns as a string and + # what the JS port's `toISOString()` produces — so the hour prefix `_truncate_hour` + # keys on is the same across both access paths and both repos. + if isinstance(value, datetime): + return value.isoformat() + return str(value) + + +def _bucket_of(value: Any) -> str: + return _truncate_hour(_stamp(value)) + + +def _canonical_tags(value: Any) -> str: + """Stable string form of a request_tags map, for use as a join key.""" + import json + + if isinstance(value, str): + try: + value = json.loads(value or "{}") + except ValueError: + return str(value) + if isinstance(value, dict): + return json.dumps(value, sort_keys=True) + return "{}" + + +# Extras that describe the endpoint a spend bucket's requests went to, rather than any +# one of those requests. Everything else the adapter captures — `invocation_id`, +# `request_id`, `status_code` — is per-request, and carrying it on an hourly aggregate +# states one sampled request's value as if it described the whole hour. +_BUCKET_INVARIANT_EXTRAS = ( + "endpoint_name", + "endpoint_id", + "destination_type", + "destination_name", + "api_type", +) + + +def _as_bucket(u: CanonicalUsage) -> CanonicalUsage: + """One usage row restated as a spend-bucket representative. + + Applied to the FIRST row of a bucket as well as to merges, so a bucket holding one + request is described the same way as a bucket holding ten — otherwise `status_code` + would survive on single-request hours and vanish on busy ones. + """ + out = CanonicalUsage( + model=u.model, + provider=u.provider, + api=u.api, + extras={k: v for k, v in u.extras.items() if k in _BUCKET_INVARIANT_EXTRAS}, + ) + for name in CanonicalUsage.NUMERIC_FIELDS: + setattr(out, name, getattr(u, name)) + return out + + +def _merge_usage(a: CanonicalUsage, b: CanonicalUsage) -> CanonicalUsage: + """Sum the numeric fields of two rows in the same spend bucket.""" + merged = _as_bucket(a) + for name in CanonicalUsage.NUMERIC_FIELDS: + setattr(merged, name, getattr(a, name) + getattr(b, name)) + return merged diff --git a/src/lago_agent_sdk/pricing.py b/src/lago_agent_sdk/pricing.py index 1504676..88b2c6b 100644 --- a/src/lago_agent_sdk/pricing.py +++ b/src/lago_agent_sdk/pricing.py @@ -78,7 +78,16 @@ # 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"}) +# +# "mistral" belongs here for the same reason: the API is OpenAI-shaped and reports +# `prompt_tokens_details.cached_tokens` as a SUBSET of `prompt_tokens`. Mistral's own +# documented example is unambiguous — prompt_tokens=1013, cached_tokens=1008, and +# total_tokens=1043 = prompt + completion, which only reconciles if the cached tokens +# sit inside the prompt count. Omitting it double-billed the cached portion by 6.15x +# on that payload. 13 of 18 Mistral models on OpenRouter publish a cache-read rate, +# so the wrong path was reachable for most of them, including Mistral traffic routed +# through a Cloudflare gateway (the gateway adapter leaves provider="mistral" as-is). +_INPUT_INCLUDES_CACHE_READ = frozenset({"openai", "gemini", "workers-ai", "mistral"}) # Providers whose reported `output` token count ALREADY includes the reasoning # tokens (reasoning is a subset of output). For these, reasoning is billed as @@ -86,6 +95,26 @@ # are additive to output, so it's absent here.) _OUTPUT_INCLUDES_REASONING = frozenset({"openai"}) +# Providers this SDK bills as TOKEN COUNTS by design, even in price mode — because +# no per-token rate for them exists anywhere the SDK could read it. +# +# "databricks" means a Databricks-HOSTED foundation model (`system.ai.*`). Databricks +# bills those in DBUs at a per-model rate published only as an HTML page — verified +# absent from every column of all 88 system tables — so there is nothing to look up +# now and nothing a later refresh could supply. Token counts are the honest, complete +# answer for them, not a degraded one. +# +# This is a deliberate, NARROW exception to "a price miss is reported via on_error". +# It applies only where the miss is *structural and permanent*. A cold table, an +# unmatched model name, a mistyped provider — all still report, because those are +# genuine misses a customer can act on. Reporting this one on every call would be a +# permanent false alarm, and an alarm that always fires is one nobody reads. +# +# Note this keys on the PROVIDER, so it only ever covers Databricks-hosted models: +# BYOK traffic through the same gateway is stamped "openai"/"anthropic" and prices +# normally (verified exact against Databricks' own metered spend, 38 of 38 buckets). +TOKEN_BILLED_PROVIDERS = frozenset({"databricks"}) + # Canonical field -> OpenRouter pricing key. _OPENROUTER_FIELD_MAP = { "input": "prompt", @@ -140,7 +169,19 @@ _SCALE = 12 _Q = Decimal(1).scaleb(-_SCALE) # Decimal("1E-12") -_VERSION_DATE_SUFFIX = re.compile(r"-(?:\d{8}|v\d+)$") +# Vendors stamp resolved model names with a date in one of two shapes, and both +# must be strippable or the price lookup misses. Anthropic uses a COMPACT date +# ("claude-sonnet-4-5-20250929"); OpenAI uses a HYPHENATED one +# ("gpt-5-2025-08-07", "o3-2025-04-16"). OpenRouter lists the BARE id +# ("openai/gpt-5"), so a name we can't strip back to bare never matches. +# +# Handling only the compact form silently broke price mode for every current +# OpenAI model: `create(model="gpt-5")` returns model="gpt-5-2025-08-07", and +# `resolve_model` prefers the response's own name over the requested one, so +# gpt-4.1 / gpt-4.1-mini / gpt-5 / gpt-5-mini / o3 / o4-mini all fell through to +# token events. gpt-4o looked fine only by luck — OpenRouter happens to list +# "openai/gpt-4o-2024-08-06" verbatim. +_VERSION_DATE_SUFFIX = re.compile(r"-(?:\d{8}|\d{4}-\d{2}-\d{2}|v\d+)$") # ---------------------------------------------------------------------- @@ -204,7 +245,7 @@ def _alnum(s: str) -> str: def _strip_version(model: str) -> str: - """Drop a trailing -YYYYMMDD date or -vN version tag.""" + """Drop a trailing -YYYYMMDD / -YYYY-MM-DD date or -vN version tag.""" return _VERSION_DATE_SUFFIX.sub("", model) diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index 168fbe5..324b78e 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -15,6 +15,7 @@ from .exceptions import PricingUnavailableError, UnknownClientError from .lago_client import LagoClient from .pricing import ( + TOKEN_BILLED_PROVIDERS, CostBreakdown, PricingProvider, apply_markup, @@ -71,6 +72,9 @@ def __init__( ) if self.config.pricing_mode == "price": self._pricing.prime() # eager warm when price mode is the global default + # (provider, model) pairs already noted as token-billed, so the explanation is + # logged once rather than on every call. See `_note_token_billed`. + self._token_billed_noted: set[tuple[str, str]] = set() self._queue = EventQueue( sender=self._lago_client.send_batch, flush_interval=self.config.flush_interval_seconds, @@ -260,6 +264,14 @@ def emit( if usd_cost is not None: breakdown = compute_precomputed_cost(usd_cost, markup_value) + elif usage.provider in TOKEN_BILLED_PROVIDERS: + # NOT a failure, so deliberately not routed through on_error: this + # provider publishes no per-token rate at all, so token counts are the + # complete answer rather than a fallback. Said once per model instead + # of once per call. See TOKEN_BILLED_PROVIDERS for the reasoning. + self._note_token_billed(usage) + self._emit_token_events(usage, sub, dimensions, event_id) + return else: price = self._pricing.lookup(usage.provider, usage.model, usage.api) if price is None: @@ -275,6 +287,24 @@ def emit( except Exception as exc: # noqa: BLE001 — never raise from emit self._report_error(exc, "emit") + def _note_token_billed(self, usage: CanonicalUsage) -> None: + """Say it once per model, at info level. + + It is a standing fact about the provider, not an event about this call, so + repeating it per request would bury the log in something the reader can neither + fix nor act on. + """ + key = (usage.provider, usage.model) + if key in self._token_billed_noted: + return + self._token_billed_noted.add(key) + logger.info( + "lago: %s bills %r in its own units, not per token — emitting token counts " + "for it instead of a dollar cost", + usage.provider, + usage.model, + ) + def _emit_token_events( self, usage: CanonicalUsage, sub: str, dimensions: dict[str, Any] | None, event_id: str | None = None ) -> None: @@ -423,6 +453,89 @@ def warm_pricing(self, providers: Iterable[str] = ()) -> None: self._pricing.prime(providers) self._pricing.maybe_refresh() + def backfill_databricks( + self, + source: Any, + since: Any = "1 day", + *, + default_subscription: str | None = None, + unified: bool = False, + dimensions: dict[str, Any] | None = None, + event_id_prefix: str = "dbx", + ) -> dict[str, int]: + """Read a window of Databricks AI Gateway usage and bill all of it. + + The one-call entrypoint: give it a window, it does the rest. Returns counts + of what it emitted, e.g. ``{"cost": 56, "tokens": 45, "skipped": 0}``. + + ``source`` is normally a :class:`DatabricksSource`, and ``since`` the window. + It also accepts an already-read iterable of ``DatabricksUsageRow`` — pass one + when you have inspected the rows first, so the window is read ONCE. Reading + twice is not just slow: a SQL warehouse costs roughly 1,500x the model-serving + usage it reports on, and rows landing between the two reads make the summary + you printed disagree with what was billed. + + Billing follows the rule the connector establishes rather than re-deriving + it: a BYOK row carries Databricks' own metered USD and bills as a dollar + cost; a Databricks-hosted row has no per-request dollar figure anywhere in + Databricks' system tables and bills as token counts. + + ``unified=True`` bills everything to ``default_subscription``, ignoring + per-call ``request_tags`` — right when one gateway serves one customer. + Left False, each row goes to the subscription its own tags name, falling + back to ``default_subscription`` only when a row is untagged. + + Every event also carries the Databricks-side grouping key for its row — + ``endpoint_name`` for hosted, ``bucket`` for BYOK — so grouping Lago the + same way the Databricks page groups puts the two side by side. See + ``DatabricksUsageRow.reconcile_dimensions``. Anything in ``dimensions`` + is added on top and wins on a key collision. + + Idempotent: every event id is derived from the source row's own id and + scoped by subscription, so re-running the same window has Lago reject the + duplicates rather than double-bill. Does not flush — call ``flush()`` when + you want to block on delivery. + """ + counts = {"cost": 0, "tokens": 0, "skipped": 0} + rows = ( + source.read_usage(since, event_id_prefix=event_id_prefix) + if hasattr(source, "read_usage") + else source + ) + for row in rows: + sub = default_subscription if unified else (row.subscription or default_subscription) + if not sub: + # No attribution and no fallback — emit() would drop it anyway, but + # counting it here makes the gap visible instead of silent. + counts["skipped"] += 1 + continue + # Row's own reconciliation key first, so an explicit caller dimension of + # the same name wins rather than being silently overwritten. + dims = {**row.reconcile_dimensions, **(dimensions or {})} + if row.usd_cost is not None: + self.emit( + row.usage, + subscription=sub, + dimensions=dims, + mode="price", + usd_cost=row.usd_cost, + # Keyed off the subscription actually billed, not the row's own + # tag — an untagged row billed to the default must not carry an + # id that blocks it from a different default on a later run. + event_id=row.event_id_for(sub), + ) + counts["cost"] += 1 + else: + self.emit( + row.usage, + subscription=sub, + dimensions=dims, + mode="tokens", + event_id=row.event_id_for(sub), + ) + counts["tokens"] += 1 + return counts + def flush(self, timeout: float = 5.0) -> bool: return self._queue.flush(timeout=timeout) diff --git a/src/lago_agent_sdk/wrappers/openai.py b/src/lago_agent_sdk/wrappers/openai.py index 106d8f9..a47c863 100644 --- a/src/lago_agent_sdk/wrappers/openai.py +++ b/src/lago_agent_sdk/wrappers/openai.py @@ -94,6 +94,40 @@ def _is_cache_hit(raw_response: Any) -> bool: return False +# A Databricks-HOSTED foundation model answers on the unified mlflow surface. It +# has to be told apart from an OpenAI-BYOK call, which uses the SAME +# `openai.OpenAI` class against `/ai-gateway/openai/v1` — and the response gives +# no clue: a hosted call echoes a served-entity name ("meta-llama-4-maverick-040225") +# with no distinguishing marker, so `_infer_provider`'s model-string rule cannot +# see it. `base_url` is the only signal, and only the wrapper has it. +# +# Matching `/ai-gateway/mlflow/` specifically, NOT `/ai-gateway/`, is the whole +# point: the openai and anthropic surfaces live under the same prefix and must +# keep their real vendor provider so they price against OpenRouter. +_DATABRICKS_HOSTED_PATH = "/ai-gateway/mlflow/" + + +def _provider_hint_for(client: Any) -> str: + """Return a provider override implied by the client's base_url, or "". + + "databricks" matches no vendor in pricing's _VENDOR_MAP, so a hosted call + CANNOT hit a price table. `emit()` then emits token counts via + TOKEN_BILLED_PROVIDERS with no error reported — that is the complete answer for + these models, not a fallback. Deliberate: Databricks bills them in DBUs + against its own rate card, which is published only as HTML and exists in no + system table, while OpenRouter DOES list bare `openai/gpt-oss-20b` and + `meta-llama/llama-4-maverick` at 0.2-0.4x of Databricks' real rate. Left as + "openai", a rename of the served entity to an 8-digit date suffix would let + `_strip_version` strip it into a match and silently under-bill 2.5-5x. + Stamping "databricks" turns that accident into a guaranteed honest miss. + """ + try: + base_url = str(getattr(client, "base_url", "") or "") + except Exception: # noqa: BLE001 — some client variants don't expose it + return "" + return "databricks" if _DATABRICKS_HOSTED_PATH in base_url else "" + + def wrap_openai_client( sdk: Any, client: Any, @@ -108,6 +142,7 @@ def wrap_openai_client( base_dims = dict(dimensions or {}) base_sub = subscription is_async = type(client).__name__.startswith("Async") + provider_hint = _provider_hint_for(client) def _resolve_opts(lago_opts: dict[str, Any]) -> dict[str, Any]: return { @@ -119,7 +154,7 @@ def _resolve_opts(lago_opts: dict[str, Any]) -> dict[str, Any]: def _emit_from(payload: Any, model_id: str, opts: dict[str, Any]) -> None: try: - usage = extract_openai_native(payload, model_id=model_id) + usage = extract_openai_native(payload, model_id=model_id, provider_hint=provider_hint) sdk.emit(usage, **opts) except Exception as exc: # noqa: BLE001 logger.warning("lago: openai emit failed: %s", exc) diff --git a/tests/unit/adapters/test_openai_native.py b/tests/unit/adapters/test_openai_native.py index f710a8a..fc776e5 100644 --- a/tests/unit/adapters/test_openai_native.py +++ b/tests/unit/adapters/test_openai_native.py @@ -274,3 +274,110 @@ def test_real_openai_model_still_gets_openai_provider() -> None: 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" + + +# ---------------------------------------------------------------------- +# Nested drift sweep + total_tokens consistency guard +# ---------------------------------------------------------------------- + + +def test_cache_write_tokens_surfaces_in_extras_and_is_not_mapped() -> None: + """Real captured `gpt-5.6-sol` shape: `prompt_tokens_details.cache_write_tokens`. + + Two assertions, and the second is the important one. The field must be + SURFACED (it used to vanish entirely: `extras` swept only top-level keys and + `prompt_tokens_details` is itself a known top-level key, so nothing nested + was ever inspected). But it must NOT be mapped to `cache_write` — for OpenAI + these tokens sit INSIDE `prompt_tokens` and bill at the plain input rate, + while OpenRouter publishes a separate cache_write rate, so mapping them would + charge the same 3022 tokens twice ($0.0341 against a true $0.0152, 2.24x). + Anthropic is the opposite case, which is why mapping is right there. + """ + resp = { + "model": "gpt-5.6-sol", + "usage": { + "prompt_tokens": 3025, + "completion_tokens": 4, + "total_tokens": 3029, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 3022, "audio_tokens": 0}, + "completion_tokens_details": {"reasoning_tokens": 0, "audio_tokens": 0}, + }, + } + u = extract_openai_native(resp) + assert u.extras["prompt_tokens_details.cache_write_tokens"] == 3022 + assert u.cache_write == 0, "cache_write_tokens must not be billed as cache_write for OpenAI" + assert u.input == 3025 + + +def test_predicted_output_details_surface_in_extras() -> None: + """The module docstring promised customers could read the Predicted Outputs + counts from extras. They never arrived, for the same nested-sweep reason. + Now they do.""" + resp = { + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "completion_tokens_details": { + "reasoning_tokens": 0, + "accepted_prediction_tokens": 7, + "rejected_prediction_tokens": 3, + }, + } + } + u = extract_openai_native(resp) + assert u.extras["completion_tokens_details.accepted_prediction_tokens"] == 7 + assert u.extras["completion_tokens_details.rejected_prediction_tokens"] == 3 + + +def test_total_tokens_guard_recovers_unaccounted_output() -> None: + """Measured against Gemini through Google's own OpenAI-compatible layer: + prompt=57, completion=47, total=1253. The 1149 thinking tokens are reported + in NEITHER named bucket and there is no completion_tokens_details to recover + them from — only `total_tokens` proves they exist. Billing prompt+completion + drops 92% of the call, at the output rate. + + The remainder folds into `output`, deliberately NOT into `reasoning`: + compute_cost zeroes reasoning whenever provider is in + _OUTPUT_INCLUDES_REASONING, and an OpenAI-shaped payload is stamped + provider="openai" by definition, so that would recover nothing.""" + resp = { + "model": "gemini-2.5-flash", + "usage": {"prompt_tokens": 57, "completion_tokens": 47, "total_tokens": 1253}, + } + u = extract_openai_native(resp) + assert u.input == 57 + assert u.output == 1196, "47 reported + 1149 unaccounted" + assert u.extras["unaccounted_output_tokens"] == 1149 + + +def test_total_tokens_guard_is_a_noop_for_genuine_openai() -> None: + """For real OpenAI total_tokens == prompt + completion always holds, because + reasoning is a SUBSET of completion rather than additive. Verified across + every captured real response — zero deltas. The guard must therefore never + fire here, including for a reasoning model that spent its whole budget + thinking.""" + for usage in ( + { + "prompt_tokens": 31, + "completion_tokens": 220, + "total_tokens": 251, + "completion_tokens_details": {"reasoning_tokens": 220}, + }, + { + "prompt_tokens": 3026, + "completion_tokens": 2, + "total_tokens": 3028, + "prompt_tokens_details": {"cached_tokens": 2816}, + }, + {"prompt_tokens": 16, "total_tokens": 16}, # embeddings: no completion_tokens at all + ): + u = extract_openai_native({"usage": usage}) + assert u.output == (usage.get("completion_tokens") or 0) + assert "unaccounted_output_tokens" not in u.extras + + +def test_total_tokens_guard_ignores_a_negative_delta() -> None: + """A total SMALLER than the parts is nonsense, not drift — never subtract.""" + u = extract_openai_native({"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 10}}) + assert u.output == 50 + assert "unaccounted_output_tokens" not in u.extras diff --git a/tests/unit/fixtures/pricing/money_golden.json b/tests/unit/fixtures/pricing/money_golden.json index 3865382..e6f6f66 100644 --- a/tests/unit/fixtures/pricing/money_golden.json +++ b/tests/unit/fixtures/pricing/money_golden.json @@ -76,6 +76,17 @@ "base": "0.02577823", "total": "0.02577823", "total_cents": "2.577823" + }, + { + "name": "mistral: cache_read is a SUBSET of input, billed once", + "_note": "Counts from Mistral's own documented prompt-caching example (prompt_tokens=1013, cached_tokens=1008, completion_tokens=30) at live mistral-large-2512 OpenRouter rates. total_tokens=1043=prompt+completion in that payload, which only reconciles if the cached tokens sit INSIDE prompt_tokens. Only 1013-1008=5 tokens may be billed at the input rate; billing all 1013 double-charges the cached portion by 6.15x.", + "provider": "mistral", + "prices": { "input": "0.0000005", "output": "0.0000015", "cache_read": "0.00000005" }, + "counts": { "input": 1013, "output": 30, "cache_read": 1008 }, + "markup": "1", + "base": "0.0000979", + "total": "0.0000979", + "total_cents": "0.00979" } ], "precomputed_cases": [ diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_read.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_read.json new file mode 100644 index 0000000..44749f4 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_read.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "38efe7cb-c6d7-41c0-843b-4ef7cb9e016d", + "schema_version": "1", + "endpoint_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "endpoint_name": "workspace.default.anthropickey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.anthropickey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T13:31:47.000Z", + "latency_ms": "2111", + "time_to_first_byte_ms": "2111", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", + "destination_model": "claude-sonnet-4-5", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "user_agent": "Python-urllib/3.11", + "api_type": "anthropic/v1/messages", + "request_tags": "{\"lago_subscription\":\"sub_cachetest\",\"scenario\":\"fresh_1h_read\"}", + "input_tokens": "1651", + "output_tokens": "4", + "total_tokens": "1655", + "token_details": "{\"cache_read_input_tokens\":\"1642\",\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T13:31:47.741Z\",\"latency_ms\":\"2054\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T13:31:49.795Z\",\"destination_id\":\"629a34db-77bb-3c25-a003-5a947ee6af36\",\"error_code\":null,\"destination\":\"workspace.default.anthropickey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "18ef0618-a770-4cdd-ba74-f3c93e555c6b", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "service_name": "workspace.default.anthropickey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_read_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_read_1.json new file mode 100644 index 0000000..94c78f6 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_read_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "34a878d6-751a-4059-9e9f-7fd2bc45c4aa", + "schema_version": "1", + "endpoint_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "endpoint_name": "workspace.default.anthropickey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.anthropickey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T13:30:54.000Z", + "latency_ms": "2342", + "time_to_first_byte_ms": "2342", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", + "destination_model": "claude-sonnet-4-5", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "user_agent": "Python-urllib/3.11", + "api_type": "anthropic/v1/messages", + "request_tags": "{\"lago_subscription\":\"sub_cachetest\",\"scenario\":\"sonnet_nottl_read\"}", + "input_tokens": "1822", + "output_tokens": "4", + "total_tokens": "1826", + "token_details": "{\"cache_read_input_tokens\":\"1812\",\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T13:30:54.105Z\",\"latency_ms\":\"2341\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T13:30:56.446Z\",\"destination_id\":\"629a34db-77bb-3c25-a003-5a947ee6af36\",\"error_code\":null,\"destination\":\"workspace.default.anthropickey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "ffeee421-2d8d-4c2a-982e-544d1708086a", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "service_name": "workspace.default.anthropickey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_write.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_write.json new file mode 100644 index 0000000..4692ee3 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_write.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "e0f2afb9-abf9-481c-99da-a3a9a3247968", + "schema_version": "1", + "endpoint_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "endpoint_name": "workspace.default.anthropickey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.anthropickey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T13:31:40.000Z", + "latency_ms": "2478", + "time_to_first_byte_ms": "2478", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", + "destination_model": "claude-sonnet-4-5", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "user_agent": "Python-urllib/3.11", + "api_type": "anthropic/v1/messages", + "request_tags": "{\"lago_subscription\":\"sub_cachetest\",\"scenario\":\"fresh_1h_write\"}", + "input_tokens": "1651", + "output_tokens": "4", + "total_tokens": "1655", + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":\"1642\",\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T13:31:41.031Z\",\"latency_ms\":\"2215\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T13:31:43.247Z\",\"destination_id\":\"629a34db-77bb-3c25-a003-5a947ee6af36\",\"error_code\":null,\"destination\":\"workspace.default.anthropickey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "4e6397d9-97fd-41d7-a7a6-5dde9d0f6280", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "service_name": "workspace.default.anthropickey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_write_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_write_1.json new file mode 100644 index 0000000..e43fbfb --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_write_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "40e6e7d8-5f2b-46b1-8350-3e877d9f2709", + "schema_version": "1", + "endpoint_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "endpoint_name": "workspace.default.anthropickey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.anthropickey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T13:29:03.000Z", + "latency_ms": "2429", + "time_to_first_byte_ms": "2429", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", + "destination_model": "claude-sonnet-4-5", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "user_agent": "Python-urllib/3.11", + "api_type": "anthropic/v1/messages", + "request_tags": "{\"lago_subscription\":\"sub_acme\",\"team\":\"lago-sdk\",\"scenario\":\"5m_write\"}", + "input_tokens": "1825", + "output_tokens": "4", + "total_tokens": "1829", + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":\"1812\",\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T13:29:03.697Z\",\"latency_ms\":\"2367\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T13:29:06.065Z\",\"destination_id\":\"629a34db-77bb-3c25-a003-5a947ee6af36\",\"error_code\":null,\"destination\":\"workspace.default.anthropickey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "b3e38bb8-6ede-45be-a3cc-bb5e2bbacaea", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "service_name": "workspace.default.anthropickey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_plain.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_plain.json new file mode 100644 index 0000000..ad1d463 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_plain.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "a9099d14-8bba-4854-a39f-e824098c6a6a", + "schema_version": "1", + "endpoint_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "endpoint_name": "workspace.default.anthropickey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.anthropickey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-11T10:09:36.000Z", + "latency_ms": "10924", + "time_to_first_byte_ms": "10924", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", + "destination_model": "claude-sonnet-4-5", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "user_agent": "Anthropic/Python 0.103.1", + "api_type": "anthropic/v1/messages", + "request_tags": "{\"lago_subscription\":\"6cc703e3-d5e5-4258-826a-4d586a94f27a\",\"team\":\"lago-demo\"}", + "input_tokens": "25", + "output_tokens": "400", + "total_tokens": "425", + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-11T10:09:36.992Z\",\"latency_ms\":\"10308\",\"status_code\":\"200\",\"end_time\":\"2026-08-11T10:09:47.301Z\",\"destination_id\":\"629a34db-77bb-3c25-a003-5a947ee6af36\",\"error_code\":null,\"destination\":\"workspace.default.anthropickey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "e02db3d2-f6b1-4385-90a8-a940ccd7ff8e", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "service_name": "workspace.default.anthropickey", + "service_tags": "{}", + "mcp_metadata": null +} diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_cache_read.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_cache_read.json new file mode 100644 index 0000000..95ed527 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_cache_read.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "a3bad0ed-a382-4fc1-88fb-8c2048d85344", + "schema_version": "1", + "endpoint_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "endpoint_name": "workspace.default.openaikey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.openaikey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T13:41:34.000Z", + "latency_ms": "1753", + "time_to_first_byte_ms": "1753", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.openaikey", + "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", + "destination_model": "gpt-5.6", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "openai/v1/chat/completions", + "request_tags": "{\"lago_subscription\":\"sub_openai\",\"scenario\":\"cache_read\"}", + "input_tokens": "3025", + "output_tokens": "4", + "total_tokens": "3029", + "token_details": "{\"cache_read_input_tokens\":\"3022\",\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":\"0\"}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T13:41:34.377Z\",\"latency_ms\":\"1672\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T13:41:36.050Z\",\"destination_id\":\"64d6b098-69c7-30b2-82e6-28ec097c3f85\",\"error_code\":null,\"destination\":\"workspace.default.openaikey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "915ef459-5167-48b3-a5d0-1428173b2738", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "service_name": "workspace.default.openaikey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_cache_read_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_cache_read_1.json new file mode 100644 index 0000000..b2b25ce --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_cache_read_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "b70cac9c-5ee1-425a-9cac-baa401f59b81", + "schema_version": "1", + "endpoint_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "endpoint_name": "workspace.default.openaikey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.openaikey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T13:41:26.000Z", + "latency_ms": "581", + "time_to_first_byte_ms": "580", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.openaikey", + "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", + "destination_model": "gpt-4o", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "openai/v1/chat/completions", + "request_tags": "{\"lago_subscription\":\"sub_openai\",\"scenario\":\"cache_read\"}", + "input_tokens": "3026", + "output_tokens": "2", + "total_tokens": "3028", + "token_details": "{\"cache_read_input_tokens\":\"2816\",\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":\"0\"}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T13:41:26.805Z\",\"latency_ms\":\"498\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T13:41:27.304Z\",\"destination_id\":\"64d6b098-69c7-30b2-82e6-28ec097c3f85\",\"error_code\":null,\"destination\":\"workspace.default.openaikey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "81d7c512-8d68-45e0-bea7-c87053e07142", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "service_name": "workspace.default.openaikey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain.json new file mode 100644 index 0000000..0f5a30a --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "2891d65c-fa8f-4b91-aa84-0bf46c686383", + "schema_version": "1", + "endpoint_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "endpoint_name": "workspace.default.openaikey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.openaikey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:28:37.000Z", + "latency_ms": "2476", + "time_to_first_byte_ms": "2476", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.openaikey", + "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", + "destination_model": "gpt-3.5-turbo", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "openai/v1/chat/completions", + "request_tags": "{}", + "input_tokens": "8", + "output_tokens": "4", + "total_tokens": "12", + "token_details": "{\"cache_read_input_tokens\":\"0\",\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":\"0\"}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:28:37.706Z\",\"latency_ms\":\"2475\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T16:28:40.181Z\",\"destination_id\":\"64d6b098-69c7-30b2-82e6-28ec097c3f85\",\"error_code\":null,\"destination\":\"workspace.default.openaikey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "1d5cf514-88f6-44de-a9db-51e68d5b2879", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "service_name": "workspace.default.openaikey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain_1.json new file mode 100644 index 0000000..af2a880 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "4be003ee-0f7c-458b-897b-f11c947343c6", + "schema_version": "1", + "endpoint_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "endpoint_name": "workspace.default.openaikey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.openaikey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:28:36.000Z", + "latency_ms": "982", + "time_to_first_byte_ms": "982", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.openaikey", + "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", + "destination_model": "gpt-4o", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "openai/v1/chat/completions", + "request_tags": "{}", + "input_tokens": "8", + "output_tokens": "4", + "total_tokens": "12", + "token_details": "{\"cache_read_input_tokens\":\"0\",\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":\"0\"}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:28:36.302Z\",\"latency_ms\":\"981\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T16:28:37.283Z\",\"destination_id\":\"64d6b098-69c7-30b2-82e6-28ec097c3f85\",\"error_code\":null,\"destination\":\"workspace.default.openaikey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "de8b49cf-b67b-491c-af0e-ae45d52f3e6e", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "service_name": "workspace.default.openaikey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning.json new file mode 100644 index 0000000..e84c8d5 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "4d4bab09-8c9a-4f2d-bf70-4727b79f6061", + "schema_version": "1", + "endpoint_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "endpoint_name": "workspace.default.openaikey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.openaikey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T13:42:13.000Z", + "latency_ms": "3353", + "time_to_first_byte_ms": "3353", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.openaikey", + "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", + "destination_model": "o4-mini", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "openai/v1/chat/completions", + "request_tags": "{\"lago_subscription\":\"sub_initech\",\"scenario\":\"content\"}", + "input_tokens": "31", + "output_tokens": "220", + "total_tokens": "251", + "token_details": "{\"cache_read_input_tokens\":\"0\",\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":\"220\"}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T13:42:13.840Z\",\"latency_ms\":\"3351\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T13:42:17.192Z\",\"destination_id\":\"64d6b098-69c7-30b2-82e6-28ec097c3f85\",\"error_code\":null,\"destination\":\"workspace.default.openaikey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "606035f9-b735-438e-ab4b-4d139adaf67a", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "service_name": "workspace.default.openaikey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning_1.json new file mode 100644 index 0000000..9f9e701 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "4a6de506-7a48-4fcb-9298-5096e38e32ea", + "schema_version": "1", + "endpoint_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "endpoint_name": "workspace.default.openaikey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.openaikey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T13:42:10.000Z", + "latency_ms": "2801", + "time_to_first_byte_ms": "2801", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.openaikey", + "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", + "destination_model": "o3", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "openai/v1/chat/completions", + "request_tags": "{\"lago_subscription\":\"sub_globex\",\"scenario\":\"content\"}", + "input_tokens": "31", + "output_tokens": "220", + "total_tokens": "251", + "token_details": "{\"cache_read_input_tokens\":\"0\",\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":\"220\"}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T13:42:10.677Z\",\"latency_ms\":\"2735\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T13:42:13.413Z\",\"destination_id\":\"64d6b098-69c7-30b2-82e6-28ec097c3f85\",\"error_code\":null,\"destination\":\"workspace.default.openaikey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "2ca5210a-4e1c-4c36-a2b6-5b1b1fe50a13", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "2e30db2e-6a9d-46f9-8857-71572ca76870", + "service_name": "workspace.default.openaikey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens.json new file mode 100644 index 0000000..69be38e --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "153b0a43-b400-4295-820b-e0e5655a8a79", + "schema_version": "1", + "endpoint_id": "893ce06c-3539-4ba0-8be3-bde5dbcd8f90", + "endpoint_name": "system.ai.gpt-5-3-codex", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":\"2026-08-07T12:27:50.933Z\",\"fallbacks\":[],\"creator\":\"Databricks\",\"last_updated_time\":\"2026-08-07T12:27:50.933Z\",\"destinations\":[{\"name\":\"system.ai.databricks-gpt-5-3-codex\",\"traffic_percent\":\"100.0\",\"type\":\"PAY_PER_TOKEN_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:41:14.000Z", + "latency_ms": "508", + "time_to_first_byte_ms": "508", + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.databricks-gpt-5-3-codex", + "destination_id": "54e874af-9a08-3055-9e85-02d726b8c023", + "destination_model": "gpt-5-3-codex", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "mlflow/v1/chat/completions", + "request_tags": "{}", + "input_tokens": null, + "output_tokens": null, + "total_tokens": null, + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json", + "status_code": "400", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:41:14.556Z\",\"latency_ms\":\"27\",\"status_code\":\"400\",\"end_time\":\"2026-08-07T16:41:14.583Z\",\"destination_id\":\"54e874af-9a08-3055-9e85-02d726b8c023\",\"error_code\":null,\"destination\":\"system.ai.databricks-gpt-5-3-codex\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "ce9409e5-5187-4174-b51e-fc1984ae4c14", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_SERVICE", + "service_id": "893ce06c-3539-4ba0-8be3-bde5dbcd8f90", + "service_name": "system.ai.gpt-5-3-codex", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens_1.json new file mode 100644 index 0000000..921fc1e --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "02d0a7eb-a84b-46ff-b4d5-57be3faecb83", + "schema_version": "1", + "endpoint_id": "a177330a-68ff-49f3-8fcb-93fb4ca7f7ed", + "endpoint_name": "system.ai.gpt-5-5-pro", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":\"2026-08-07T12:27:50.933Z\",\"fallbacks\":[],\"creator\":\"Databricks\",\"last_updated_time\":\"2026-08-07T12:27:50.933Z\",\"destinations\":[{\"name\":\"system.ai.databricks-gpt-5-5-pro\",\"traffic_percent\":\"100.0\",\"type\":\"PAY_PER_TOKEN_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:41:13.000Z", + "latency_ms": "6", + "time_to_first_byte_ms": "6", + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.databricks-gpt-5-5-pro", + "destination_id": "3db57a18-9cab-32cf-a717-c03995f30770", + "destination_model": "gpt-5-5-pro", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "mlflow/v1/chat/completions", + "request_tags": "{}", + "input_tokens": null, + "output_tokens": null, + "total_tokens": null, + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json", + "status_code": "400", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:41:13.650Z\",\"latency_ms\":\"5\",\"status_code\":\"400\",\"end_time\":\"2026-08-07T16:41:13.656Z\",\"destination_id\":\"3db57a18-9cab-32cf-a717-c03995f30770\",\"error_code\":null,\"destination\":\"system.ai.databricks-gpt-5-5-pro\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "6617fe46-d3f9-4a3b-89a8-402099b299a0", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_SERVICE", + "service_id": "a177330a-68ff-49f3-8fcb-93fb4ca7f7ed", + "service_name": "system.ai.gpt-5-5-pro", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken.json new file mode 100644 index 0000000..d61176d --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "f5efabd9-fcac-4cce-899f-5a3df03cd642", + "schema_version": "1", + "endpoint_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "endpoint_name": "workspace.default.anthropickey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.anthropickey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:36:51.000Z", + "latency_ms": "417", + "time_to_first_byte_ms": "417", + "destination_type": null, + "destination_name": null, + "destination_id": null, + "destination_model": null, + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/gemini/v1beta/models/x:generateContent", + "user_agent": "Python-urllib/3.11", + "api_type": "gemini/v1/generateContent", + "request_tags": "{}", + "input_tokens": null, + "output_tokens": null, + "total_tokens": null, + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json", + "status_code": "403", + "routing_information": "{\"attempts\":null}", + "invocation_id": "a5d79b76-3d93-4d63-b9cb-001c7ca7be75", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "service_name": "workspace.default.anthropickey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken_1.json new file mode 100644 index 0000000..a7d844a --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "966d0588-0cf9-4a65-a94d-c9f1ea83d303", + "schema_version": "1", + "endpoint_id": "f5165653-fc22-468f-ba74-bcadd08e2089", + "endpoint_name": "workspace.default.geminikey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.geminikey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:35:12.000Z", + "latency_ms": "302", + "time_to_first_byte_ms": "302", + "destination_type": null, + "destination_name": null, + "destination_id": null, + "destination_model": null, + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/gemini/v1beta/models/gemini-2.5-flash:generateContent", + "user_agent": "google-genai-sdk/2.7.0 gl-python/3.11.15", + "api_type": "gemini/v1/generateContent", + "request_tags": "{}", + "input_tokens": null, + "output_tokens": null, + "total_tokens": null, + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json", + "status_code": "500", + "routing_information": "{\"attempts\":null}", + "invocation_id": "fc6dbb16-8ab4-4d38-8093-a5942a1b7fad", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "f5165653-fc22-468f-ba74-bcadd08e2089", + "service_name": "workspace.default.geminikey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat.json new file mode 100644 index 0000000..50e4a59 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "a387139c-4a75-43a0-bc47-b0e9f6cd0407", + "schema_version": "1", + "endpoint_id": "6dfebdae-181c-4f2d-a0e7-8dbe913a11af", + "endpoint_name": "system.ai.llama-4-maverick", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":\"2026-08-07T12:32:59.883Z\",\"fallbacks\":[],\"creator\":\"Databricks\",\"last_updated_time\":\"2026-08-07T12:32:59.883Z\",\"destinations\":[{\"name\":\"system.ai.llama-4-maverick\",\"traffic_percent\":\"100.0\",\"type\":\"PAY_PER_TOKEN_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:41:15.000Z", + "latency_ms": "190", + "time_to_first_byte_ms": "186", + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.llama-4-maverick", + "destination_id": "f0753807-2a5d-3e12-9a70-1b895d651fa5", + "destination_model": "llama-4-maverick", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "mlflow/v1/chat/completions", + "request_tags": "{}", + "input_tokens": "11", + "output_tokens": "4", + "total_tokens": "15", + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:41:15.011Z\",\"latency_ms\":\"189\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T16:41:15.201Z\",\"destination_id\":\"f0753807-2a5d-3e12-9a70-1b895d651fa5\",\"error_code\":null,\"destination\":\"system.ai.llama-4-maverick\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "a2479f77-0cc7-4f79-92d3-5dae844ac2e2", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_SERVICE", + "service_id": "6dfebdae-181c-4f2d-a0e7-8dbe913a11af", + "service_name": "system.ai.llama-4-maverick", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_1.json new file mode 100644 index 0000000..e21842b --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "6c259bd7-0f26-434c-a621-0de1be5f36eb", + "schema_version": "1", + "endpoint_id": "adb8b8a0-26c9-3e8b-9e5d-830d15809dd6", + "endpoint_name": "databricks-llama-4-maverick", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":\"2023-11-10T09:53:20.000Z\",\"fallbacks\":[],\"creator\":\"Databricks\",\"last_updated_time\":\"2023-11-10T09:53:20.000Z\",\"destinations\":[{\"name\":\"system.ai.llama-4-maverick\",\"traffic_percent\":\"100.0\",\"type\":\"PAY_PER_TOKEN_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:41:15.000Z", + "latency_ms": "176", + "time_to_first_byte_ms": "173", + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.llama-4-maverick", + "destination_id": "ae1efffe34f03464b267ca56d5f6b6dc", + "destination_model": "Llama 4 Maverick", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "mlflow/v1/chat/completions", + "request_tags": "{}", + "input_tokens": "11", + "output_tokens": "4", + "total_tokens": "15", + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:41:15.620Z\",\"latency_ms\":\"175\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T16:41:15.795Z\",\"destination_id\":\"ae1efffe34f03464b267ca56d5f6b6dc\",\"error_code\":null,\"destination\":\"system.ai.llama-4-maverick\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "665f28c2-ef6b-4b77-b279-43aeebb62ffb", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": null, + "service_id": null, + "service_name": null, + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_endpoint_prefixed_name.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_endpoint_prefixed_name.json new file mode 100644 index 0000000..853fd52 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_endpoint_prefixed_name.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "15fd1c59-1c7d-489b-b4b3-bd9d5fb5df15", + "schema_version": "1", + "endpoint_id": "7de0014f-20a9-41c5-8613-a66d3b62cb77", + "endpoint_name": "system.ai.qwen35-122b-a10b", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":\"2026-08-07T12:34:52.915Z\",\"fallbacks\":[],\"creator\":\"Databricks\",\"last_updated_time\":\"2026-08-07T12:34:52.915Z\",\"destinations\":[{\"name\":\"system.ai.databricks-qwen35-122b-a10b\",\"traffic_percent\":\"100.0\",\"type\":\"PAY_PER_TOKEN_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:41:06.000Z", + "latency_ms": "1571", + "time_to_first_byte_ms": "1569", + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.databricks-qwen35-122b-a10b", + "destination_id": "1802e050-b85b-3de0-bdfe-11728b51cf85", + "destination_model": "qwen35-122b-a10b", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "user_agent": "Python-urllib/3.11", + "api_type": "mlflow/v1/chat/completions", + "request_tags": "{\"lago_subscription\":\"sub_acme\",\"scenario\":\"content\"}", + "input_tokens": "37", + "output_tokens": "200", + "total_tokens": "237", + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:41:06.564Z\",\"latency_ms\":\"1510\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T16:41:08.075Z\",\"destination_id\":\"1802e050-b85b-3de0-bdfe-11728b51cf85\",\"error_code\":null,\"destination\":\"system.ai.databricks-qwen35-122b-a10b\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "4848f02f-5d29-4dfb-b858-fdb4b042b197", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_SERVICE", + "service_id": "7de0014f-20a9-41c5-8613-a66d3b62cb77", + "service_name": "system.ai.qwen35-122b-a10b", + "service_tags": "{}", + "mcp_metadata": null +} diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings.json new file mode 100644 index 0000000..59ff684 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "fc20021d-a6df-4510-9148-05cf29739fed", + "schema_version": "1", + "endpoint_id": "2bac1f0f-85cd-4879-b387-158ad026af1b", + "endpoint_name": "system.ai.qwen3-embedding-0-6b", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":\"2026-08-07T12:34:52.915Z\",\"fallbacks\":[],\"creator\":\"Databricks\",\"last_updated_time\":\"2026-08-07T12:34:52.915Z\",\"destinations\":[{\"name\":\"system.ai.qwen3-embedding-0-6b\",\"traffic_percent\":\"100.0\",\"type\":\"PAY_PER_TOKEN_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:41:12.000Z", + "latency_ms": "298", + "time_to_first_byte_ms": "245", + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.qwen3-embedding-0-6b", + "destination_id": "9712d5a7-3608-397e-9f55-9aa47b526f23", + "destination_model": "qwen3-embedding-0-6b", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/embeddings", + "user_agent": "Python-urllib/3.11", + "api_type": "mlflow/v1/embeddings", + "request_tags": "{\"lago_subscription\":\"sub_embed\",\"scenario\":\"embeddings\"}", + "input_tokens": "13", + "output_tokens": null, + "total_tokens": "13", + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:41:12.314Z\",\"latency_ms\":\"297\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T16:41:12.611Z\",\"destination_id\":\"9712d5a7-3608-397e-9f55-9aa47b526f23\",\"error_code\":null,\"destination\":\"system.ai.qwen3-embedding-0-6b\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "59f5ab52-db78-49ea-8892-9f5a56aa5b0b", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_SERVICE", + "service_id": "2bac1f0f-85cd-4879-b387-158ad026af1b", + "service_name": "system.ai.qwen3-embedding-0-6b", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings_1.json new file mode 100644 index 0000000..2360655 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "ba6f923a-215c-407e-8877-3206656fa2bc", + "schema_version": "1", + "endpoint_id": "4847df68-2d7e-4dec-954c-e595124cb115", + "endpoint_name": "system.ai.bge-large-en", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":\"2026-08-06T14:01:18.828Z\",\"fallbacks\":[],\"creator\":\"Databricks\",\"last_updated_time\":\"2026-08-06T14:01:18.828Z\",\"destinations\":[{\"name\":\"system.ai.bge_large_en_v1_5\",\"traffic_percent\":\"100.0\",\"type\":\"PAY_PER_TOKEN_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:41:10.000Z", + "latency_ms": "527", + "time_to_first_byte_ms": "472", + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.bge_large_en_v1_5", + "destination_id": "10b8a8ba-6702-3498-84b8-ce1077c8a898", + "destination_model": "bge_large_en_v1_5", + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/embeddings", + "user_agent": "Python-urllib/3.11", + "api_type": "mlflow/v1/embeddings", + "request_tags": "{\"lago_subscription\":\"sub_embed\",\"scenario\":\"embeddings\"}", + "input_tokens": "16", + "output_tokens": null, + "total_tokens": "16", + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": "application/json; charset=utf-8", + "status_code": "200", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:41:11.196Z\",\"latency_ms\":\"252\",\"status_code\":\"200\",\"end_time\":\"2026-08-07T16:41:11.449Z\",\"destination_id\":\"10b8a8ba-6702-3498-84b8-ce1077c8a898\",\"error_code\":null,\"destination\":\"system.ai.bge_large_en_v1_5\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "ba0aba51-3878-4e2c-9ba1-0e99d39f664c", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_SERVICE", + "service_id": "4847df68-2d7e-4dec-954c-e595124cb115", + "service_name": "system.ai.bge-large-en", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path.json new file mode 100644 index 0000000..9036c12 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "b40d1814-2329-492b-9c68-7909c3af836a", + "schema_version": "1", + "endpoint_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "endpoint_name": "workspace.default.anthropickey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.anthropickey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:37:30.000Z", + "latency_ms": "46", + "time_to_first_byte_ms": "46", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", + "destination_model": null, + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/models", + "user_agent": "Python-urllib/3.11", + "api_type": "unmanaged", + "request_tags": "{}", + "input_tokens": null, + "output_tokens": null, + "total_tokens": null, + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": null, + "status_code": "404", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:37:30.810Z\",\"latency_ms\":\"44\",\"status_code\":\"404\",\"end_time\":\"2026-08-07T16:37:30.854Z\",\"destination_id\":\"629a34db-77bb-3c25-a003-5a947ee6af36\",\"error_code\":null,\"destination\":\"workspace.default.anthropickey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "ef89365f-8454-4064-8f0d-1ac9b6a4ec57", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "service_name": "workspace.default.anthropickey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path_1.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path_1.json new file mode 100644 index 0000000..070cf58 --- /dev/null +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path_1.json @@ -0,0 +1,38 @@ +{ + "account_id": "73357fbe-b495-49ad-9eec-ad226318df4f", + "workspace_id": "7474648573314045", + "request_id": "597ef94a-4dec-43c9-a010-562b75a086de", + "schema_version": "1", + "endpoint_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "endpoint_name": "workspace.default.anthropickey", + "endpoint_tags": "{}", + "endpoint_metadata": "{\"inference_table\":null,\"creation_time\":null,\"fallbacks\":[],\"creator\":null,\"last_updated_time\":null,\"destinations\":[{\"name\":\"workspace.default.anthropickey\",\"traffic_percent\":null,\"type\":\"EXTERNAL_FOUNDATION_MODEL\"}]}", + "event_time": "2026-08-07T16:37:30.000Z", + "latency_ms": "96", + "time_to_first_byte_ms": "96", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", + "destination_model": null, + "requester": "beready1994@gmail.com", + "requester_type": "USER", + "ip_address": "88.168.101.75", + "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/nonsense", + "user_agent": "Python-urllib/3.11", + "api_type": "unmanaged", + "request_tags": "{}", + "input_tokens": null, + "output_tokens": null, + "total_tokens": null, + "token_details": "{\"cache_read_input_tokens\":null,\"cache_creation_input_tokens\":null,\"output_reasoning_tokens\":null}", + "response_content_type": null, + "status_code": "404", + "routing_information": "{\"attempts\":[{\"priority\":\"1\",\"start_time\":\"2026-08-07T16:37:30.289Z\",\"latency_ms\":\"95\",\"status_code\":\"404\",\"end_time\":\"2026-08-07T16:37:30.385Z\",\"destination_id\":\"629a34db-77bb-3c25-a003-5a947ee6af36\",\"error_code\":null,\"destination\":\"workspace.default.anthropickey\",\"action\":\"INITIAL_ATTEMPT\"}]}", + "invocation_id": "e2085681-367e-47f0-a7dd-1fee75fb3723", + "invocation_metadata": "{\"source\":\"EXTERNAL_CLIENT\",\"service_tier\":null}", + "service_type": "MODEL_PROVIDER_SERVICE", + "service_id": "c99d0710-d0c0-486e-a2e6-1fd7d1d45972", + "service_name": "workspace.default.anthropickey", + "service_tags": "{}", + "mcp_metadata": null +} \ No newline at end of file diff --git a/tests/unit/gateway/adapters/test_databricks_gateway.py b/tests/unit/gateway/adapters/test_databricks_gateway.py new file mode 100644 index 0000000..10ceebb --- /dev/null +++ b/tests/unit/gateway/adapters/test_databricks_gateway.py @@ -0,0 +1,324 @@ +"""Databricks AI Gateway usage adapter — verified against real captured table rows. + +Fixtures were read from a live workspace's `system.ai_gateway.usage` over the SQL +Statement Execution API, one file per scenario, exactly as the adapter receives them. +""" + +from __future__ import annotations + +import json +import pathlib + +from lago_agent_sdk.gateway.adapters import extract_databricks_log, resolve_databricks_subscription + +FIX = pathlib.Path(__file__).parent / "fixtures" / "databricks_gateway" + + +def _load(name: str) -> dict: + return json.loads((FIX / name).read_text()) + + +# -------------------------------------------------------------------------- +# Real fixtures — the two destination types +# -------------------------------------------------------------------------- +def test_real_hosted_chat_row() -> None: + """A Databricks-hosted (pay-per-token) foundation model via the mlflow surface.""" + u = extract_databricks_log(_load("hosted_chat.json")) + assert u.input == 11 + assert u.output == 4 + assert u.model == "llama-4-maverick" + assert u.provider == "databricks" + assert u.api == "databricks_gateway" + + +def test_real_hosted_embeddings_row() -> None: + """Embeddings report input only — `output_tokens` is NULL, not 0, and must not + become a phantom output event.""" + u = extract_databricks_log(_load("hosted_embeddings.json")) + assert u.input == 13 + assert u.output == 0 + assert u.provider == "databricks" + assert u.extras["api_type"] == "mlflow/v1/embeddings" + + +def test_real_byok_anthropic_row() -> None: + """BYOK: the model comes from `destination_model` and the provider from the + leading segment of `api_type`.""" + u = extract_databricks_log(_load("byok_anthropic_cache_read.json")) + assert u.model == "claude-sonnet-4-5" + assert u.provider == "anthropic" + assert u.extras["destination_type"] == "EXTERNAL_FOUNDATION_MODEL" + + +def test_real_byok_openai_reasoning_row() -> None: + """`token_details.output_reasoning_tokens` IS broken out in the table, even + though the mlflow response body reports no reasoning at all — the live and + backfill paths genuinely disagree on this field.""" + u = extract_databricks_log(_load("byok_openai_reasoning.json")) + assert u.provider == "openai" + assert u.reasoning == 220 + assert u.output == 220 + + +# -------------------------------------------------------------------------- +# The two naming quirks that a docs-only reading gets wrong +# -------------------------------------------------------------------------- +def test_hosted_model_comes_from_destination_name_not_destination_model() -> None: + """For hosted rows `destination_model` is unstable — the same + `destination_name` was observed reporting both `llama-4-maverick` and the + display label `Llama 4 Maverick`. `destination_name` is the stable id, so it + wins, with the `system.ai.` prefix stripped.""" + row = { + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.gpt-oss-20b", + "destination_model": "GPT OSS 20B", # display label, spaces and capitals + "api_type": "mlflow/v1/chat/completions", + "input_tokens": "102", + "output_tokens": "4", + } + u = extract_databricks_log(row) + assert u.model == "gpt-oss-20b" + assert u.provider == "databricks" + + +def test_hosted_destination_name_sheds_its_endpoint_prefix() -> None: + """Most hosted entities are named `system.ai.databricks-`, not + `system.ai.` — measured on a live workspace, 38 of 48 distinct hosted + `destination_name`s carry that inner `databricks-`. It is a serving-endpoint + artefact, not part of the model id: leaving it in emits + `databricks-qwen35-122b-a10b`, which both reads as a vendor prefix and splits + one model into two rows in Lago against the live path's own name. + + Real captured row, not hand-written.""" + u = extract_databricks_log(_load("hosted_chat_endpoint_prefixed_name.json")) + assert u.model == "qwen35-122b-a10b" + assert u.provider == "databricks" + assert u.input == 37 and u.output == 200 + # The raw name stays visible for reconciliation against Databricks' own console. + assert u.extras["destination_name"] == "system.ai.databricks-qwen35-122b-a10b" + + +def test_hosted_prefix_stripping_does_not_rename_a_genuinely_databricks_model() -> None: + """Databricks publishes models whose own names start with `databricks-` + (`databricks-dbrx-instruct`, `databricks-dolly-v2`), so an unconditional strip + would rename them. `destination_model` is the tie-breaker: it agrees with the + shed form when the prefix is an endpoint artefact, and with the full name when + the model is really called that.""" + artefact = { + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.databricks-claude-sonnet-4-5", + "destination_model": "claude-sonnet-4-5", + "api_type": "mlflow/v1/chat/completions", + "input_tokens": "5", + "output_tokens": "5", + } + assert extract_databricks_log(artefact).model == "claude-sonnet-4-5" + + real_name = {**artefact, "destination_name": "system.ai.databricks-dbrx-instruct"} + real_name["destination_model"] = "databricks-dbrx-instruct" + assert extract_databricks_log(real_name).model == "databricks-dbrx-instruct" + + # Disagreement (the unstable display-label case) keeps the raw name rather than + # guessing — an ugly id beats a wrong one. + ambiguous = {**artefact, "destination_model": "Claude Sonnet 4.5"} + assert extract_databricks_log(ambiguous).model == "databricks-claude-sonnet-4-5" + + +def test_byok_never_uses_destination_name_as_the_model() -> None: + """For BYOK rows `destination_name` is the PROVIDER SERVICE — a Unity Catalog + credential name, not a model. Falling back to it would bill + `workspace.default.anthropickey` as the model on every BYOK row.""" + row = { + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_model": "claude-opus-4-5", + "api_type": "anthropic/v1/messages", + "input_tokens": "16", + "output_tokens": "47", + } + u = extract_databricks_log(row) + assert u.model == "claude-opus-4-5" + assert "workspace.default" not in u.model + assert u.extras["destination_name"] == "workspace.default.anthropickey" + + +def test_provider_is_derived_from_api_type_leading_segment() -> None: + """`api_type` is the full ingress path, and its leading segment already IS this + SDK's provider vocabulary — so no alias table is needed.""" + for api_type, expected in ( + ("anthropic/v1/messages", "anthropic"), + ("openai/v1/chat/completions", "openai"), + ("gemini/v1/generateContent", "gemini"), + ("unmanaged", "unmanaged"), + ): + u = extract_databricks_log({"destination_type": "EXTERNAL_FOUNDATION_MODEL", "api_type": api_type}) + assert u.provider == expected + + +def test_hosted_provider_cannot_match_a_vendor_price_table() -> None: + """`provider="databricks"` is deliberate: it matches no vendor in pricing's + _VENDOR_MAP, so the lookup CANNOT hit and emit() falls back to token events. + OpenRouter does list bare `openai/gpt-oss-20b` at ~0.4x of Databricks' own DBU + rate, so an accidental match would under-bill 2.5-5x.""" + from lago_agent_sdk.pricing import lookup_openrouter, parse_openrouter + + table = parse_openrouter({"data": [{"id": "openai/gpt-oss-20b", "pricing": {"prompt": "0.00000003"}}]}) + u = extract_databricks_log(_load("hosted_chat.json")) + assert lookup_openrouter(table, u.provider, u.model) is None + + +# -------------------------------------------------------------------------- +# STRUCT / MAP columns arrive as JSON strings over the REST API +# -------------------------------------------------------------------------- +def test_token_details_parses_from_a_json_string() -> None: + """The SQL drivers hand back real dicts, but the Statement Execution API + serializes STRUCT columns as JSON strings. Both must work, or the adapter + silently reads zeros from a string it never parsed.""" + as_string = { + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "api_type": "anthropic/v1/messages", + "destination_model": "claude-sonnet-4-5", + "input_tokens": "1825", + "output_tokens": "4", + "token_details": '{"cache_read_input_tokens":"1812","cache_creation_input_tokens":null}', + } + as_dict = {**as_string, "token_details": {"cache_read_input_tokens": 1812}} + for row in (as_string, as_dict): + u = extract_databricks_log(row) + assert u.cache_read == 1812 + assert u.cache_write == 0 + + +def test_input_tokens_includes_cache_so_the_difference_is_recoverable() -> None: + """Measured, and the inverse of every provider's own response body: this table's + `input_tokens` INCLUDES cache_read and cache_write. The fixture pair below came + from calls whose response bodies reported `input_tokens: 9`. + + The adapter extracts faithfully rather than subtracting — billing takes + Databricks' own metered USD, which never touches these counts. This test pins + that the arithmetic stays recoverable: only one of read/write is ever non-zero, + so input - read - write is the true non-cached input. + """ + for name in ("byok_anthropic_cache_read.json", "byok_anthropic_cache_write.json"): + u = extract_databricks_log(_load(name)) + assert not (u.cache_read and u.cache_write), "only one direction per row" + assert u.input - u.cache_read - u.cache_write == 9 + + +def test_request_tags_parses_from_a_json_string_too() -> None: + for tags in ('{"lago_subscription":"sub_acme","team":"x"}', {"lago_subscription": "sub_acme"}): + assert resolve_databricks_subscription({"request_tags": tags}) == "sub_acme" + + +# -------------------------------------------------------------------------- +# Attribution +# -------------------------------------------------------------------------- +def test_real_row_resolves_its_subscription() -> None: + assert resolve_databricks_subscription(_load("byok_openai_cache_read.json")) == "sub_openai" + + +def test_untagged_row_has_no_subscription() -> None: + """Untagged calls do produce rows, with `request_tags` empty. Attribution is + absent, and what to do about that is the caller's decision.""" + # `hosted_chat.json` IS the untagged capture — its `request_tags` is `{}`. A separate + # `untagged.json` existed and was byte-identical, so it is gone rather than kept as a + # second name for the same bytes. + assert resolve_databricks_subscription(_load("hosted_chat.json")) is None + + +def test_missing_or_malformed_request_tags_resolve_to_none() -> None: + for tags in (None, "{}", {}, "not json", [], 7, {"lago_subscription": ""}): + assert resolve_databricks_subscription({"request_tags": tags}) is None + assert resolve_databricks_subscription({}) is None + + +# -------------------------------------------------------------------------- +# Failure rows must bill nothing +# -------------------------------------------------------------------------- +def test_failed_rows_extract_to_zero_so_nothing_is_billed() -> None: + """Failed calls are recorded with NULL token counts. They must extract to + all-zero, leaving `nonzero_numeric()` empty so the caller emits nothing — the + same way a Cloudflare cache hit extracts to zero.""" + for name in ("failed_null_tokens.json", "gemini_broken.json", "unmanaged_path.json"): + u = extract_databricks_log(_load(name)) + assert u.nonzero_numeric() == {} + + +# -------------------------------------------------------------------------- +# Robustness — one malformed row must not take down a batch +# -------------------------------------------------------------------------- +def test_empty_row_is_all_zero() -> None: + u = extract_databricks_log({}) + assert u.nonzero_numeric() == {} + assert u.model == "" + assert u.provider == "" + assert u.api == "databricks_gateway" + + +def test_negative_and_non_numeric_counts_clamp_to_zero() -> None: + u = extract_databricks_log({"input_tokens": -5, "output_tokens": "bogus", "total_tokens": "9"}) + assert u.input == 0 + assert u.output == 0 + + +def test_non_string_model_and_destination_fields_do_not_crash() -> None: + u = extract_databricks_log( + {"destination_type": 7, "destination_name": [], "destination_model": {}, "api_type": None} + ) + assert u.model == "" + assert u.provider == "" + + +def test_total_tokens_is_not_mapped() -> None: + """It is derived from input+output; mapping it would double-count. Same reason + the Cloudflare adapter skips `usage_metadata.total_tokens`.""" + u = extract_databricks_log({"input_tokens": "10", "output_tokens": "5", "total_tokens": "15"}) + assert u.nonzero_numeric() == {"input": 10, "output": 5} + + +# -------------------------------------------------------------------------- +# Sweep — every captured fixture must extract cleanly +# -------------------------------------------------------------------------- +def test_all_captured_fixtures_extract() -> None: + """Iterate the whole fixture directory, mirroring `test_all_models_sweep`. + + Without this, a capture that no named test mentions asserts nothing — 12 of the + files here were in exactly that state, shipped and inert. A sweep also means the + next capture is covered the moment it lands, rather than when someone remembers to + write a test for it. Skips cleanly if the directory is absent, so a missing capture + reads as "not covered" rather than as a pass. + """ + fixtures = sorted(FIX.glob("*.json")) + if not fixtures: + import pytest + + pytest.skip("no databricks_gateway fixtures captured") + + for path in fixtures: + row = json.loads(path.read_text()) + u = extract_databricks_log(row) + assert u.api == "databricks_gateway", path.name + # Every numeric field is a count: never negative, never a float. + for field_name in u.NUMERIC_FIELDS: + value = getattr(u, field_name) + assert isinstance(value, int) and value >= 0, f"{path.name}:{field_name}={value!r}" + # A row with tokens must name a model; a row without is a failure/rejected row. + if u.nonzero_numeric(): + assert u.model, f"{path.name} has tokens but no model" + assert u.provider, f"{path.name} has tokens but no provider" + # The subscription resolver must never raise on a real row, whatever its tags. + resolve_databricks_subscription(row) + + +def test_no_two_fixtures_are_byte_identical() -> None: + """A duplicate file is a second name for the same evidence, and it lies about + coverage: three pairs existed here, one of which ("plain" Anthropic BYOK) was + actually the cache-write capture, so the scenario it claimed to hold had never + been captured at all.""" + import hashlib + + seen: dict[str, str] = {} + for path in sorted(FIX.glob("*.json")): + digest = hashlib.sha256(path.read_bytes()).hexdigest() + assert digest not in seen, f"{path.name} is byte-identical to {seen[digest]}" + seen[digest] = path.name diff --git a/tests/unit/gateway/test_databricks_source.py b/tests/unit/gateway/test_databricks_source.py new file mode 100644 index 0000000..1e3d2c5 --- /dev/null +++ b/tests/unit/gateway/test_databricks_source.py @@ -0,0 +1,633 @@ +"""Databricks usage reader — the I/O half, exercised without touching a warehouse. + +`DatabricksSource.query` is faked here; the SQL it would run is asserted, and the +COLUMNAR response shape is reproduced exactly as the Statement Execution API returns +it (`manifest.schema.columns` plus a positional `data_array`, one chunk inline). +""" + +from __future__ import annotations + +import json +from datetime import datetime +from typing import Any + +import pytest + +from lago_agent_sdk import LagoSDK +from lago_agent_sdk.gateway.databricks import DatabricksSource, DatabricksUsageRow, _interval_sql + +# -------------------------------------------------------------------------- +# Fake rows, in the exact shapes the two tables return +# -------------------------------------------------------------------------- +_HOSTED = { + "invocation_id": "inv-hosted-1", + "request_id": "req-hosted-1", + "event_time": "2026-08-07 14:22:03.123", + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.llama-4-maverick", + "destination_model": "llama-4-maverick", + "api_type": "mlflow/v1/chat/completions", + "endpoint_name": "system.ai.llama-4-maverick", + "input_tokens": "11", + "output_tokens": "4", + "request_tags": '{"lago_subscription":"sub_hosted"}', +} + +_BYOK_USAGE = { + "invocation_id": "inv-byok-1", + "request_id": "req-byok-1", + "event_time": "2026-08-07 14:22:59.900", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_model": "claude-sonnet-4-5", + "api_type": "anthropic/v1/messages", + "endpoint_name": "workspace.default.anthropickey", + "status_code": "200", + "input_tokens": "1825", + "output_tokens": "47", + "token_details": '{"cache_read_input_tokens":1812}', + "request_tags": '{"lago_subscription":"sub_byok"}', +} + +_BYOK_SPEND = { + "record_id": "rec-1", + "bucket": "2026-08-07 14:00:00", + "provider": "anthropic", + "model": "claude-sonnet-4-5", + "request_tags": '{"lago_subscription":"sub_byok"}', + "usage_quantity": "0.0011187", +} + +_FAILED = { + "invocation_id": "inv-failed", + "event_time": "2026-08-07 14:30:00", + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.gpt-oss-20b", + "api_type": "mlflow/v1/chat/completions", + "input_tokens": None, + "output_tokens": None, + "status_code": "403", +} + + +def _source(spend: list[dict], usage: list[dict]) -> DatabricksSource: + """A source whose `query` answers from canned rows, keyed on which table.""" + src = DatabricksSource(host="https://x", token="t", warehouse_id="w") + seen: list[str] = [] + + def fake_query(sql: str) -> list[dict[str, Any]]: + seen.append(sql) + return spend if "external_model_spend" in sql else usage + + src.query = fake_query # type: ignore[method-assign] + src.queries = seen # type: ignore[attr-defined] + return src + + +# -------------------------------------------------------------------------- +# The window +# -------------------------------------------------------------------------- +def test_interval_strings_render_to_sql() -> None: + assert _interval_sql("1 day") == "current_timestamp() - INTERVAL 1 DAY" + assert _interval_sql("36 hours") == "current_timestamp() - INTERVAL 36 HOUR" + assert _interval_sql("30 minutes") == "current_timestamp() - INTERVAL 30 MINUTE" + + +def test_datetime_window_renders_as_a_literal() -> None: + assert _interval_sql(datetime(2026, 8, 7, 14, 0, 0)) == "TIMESTAMP '2026-08-07 14:00:00'" + + +@pytest.mark.parametrize( + "bad", + [ + "1 day; DROP TABLE system.ai_gateway.usage", + "1 day OR 1=1", + "yesterday", + "-1 day", + "", + ], +) +def test_unrecognized_window_is_refused_not_interpolated(bad: str) -> None: + """The window reaches SQL by interpolation, so validation is the only thing + standing between a caller's string and the warehouse. Anything but a bare + count-plus-unit is refused outright.""" + with pytest.raises(ValueError, match="not understood"): + _interval_sql(bad) + + +def test_read_usage_scopes_both_queries_to_the_window() -> None: + src = _source([], []) + list(src.read_usage("3 days")) + assert len(src.queries) == 2 # type: ignore[attr-defined] + for sql in src.queries: # type: ignore[attr-defined] + assert "current_timestamp() - INTERVAL 3 DAY" in sql + + +# -------------------------------------------------------------------------- +# The BYOK / hosted split — the double-billing guard +# -------------------------------------------------------------------------- +def test_byok_bills_once_from_spend_and_hosted_once_from_usage() -> None: + """A BYOK call appears in BOTH tables. It must yield exactly one row, carrying + Databricks' own metered USD; the token row it also has must not become a second + billable row.""" + rows = list(_source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE]).read_usage("1 day")) + assert len(rows) == 2 + + byok = [r for r in rows if r.is_byok] + hosted = [r for r in rows if not r.is_byok] + assert len(byok) == 1 and len(hosted) == 1 + assert byok[0].usd_cost == pytest.approx(0.0011187) + assert byok[0].usage.model == "claude-sonnet-4-5" + assert hosted[0].usage.model == "llama-4-maverick" + assert hosted[0].usd_cost is None + + +def test_byok_row_carries_the_token_counts_joined_from_the_usage_table() -> None: + """The dollar figure is authoritative, but the event should still report real + tokens — they are joined on (hour, provider, model, tags), the spend table's own + aggregation key.""" + (byok,) = [r for r in _source([_BYOK_SPEND], [_BYOK_USAGE]).read_usage("1 day") if r.is_byok] + assert byok.usage.input == 1825 + assert byok.usage.output == 47 + assert byok.usage.cache_read == 1812 + + +def test_several_calls_in_one_spend_bucket_have_their_tokens_summed() -> None: + """The spend table aggregates per (hour, model, provider, tags), so N calls in the + same hour collapse to ONE dollar row while `ai_gateway.usage` still holds N token + rows. Reporting only the first would understate the tokens behind a cost the + customer can see — so they sum.""" + second = {**_BYOK_USAGE, "invocation_id": "inv-byok-2", "input_tokens": "100", "output_tokens": "3"} + (byok,) = list(_source([_BYOK_SPEND], [_BYOK_USAGE, second]).read_usage("1 day")) + assert byok.usage.input == 1925 + assert byok.usage.output == 50 + assert byok.usage.cache_read == 3624 + # Still ONE event: the dollar figure already covers both calls. + assert byok.usd_cost == pytest.approx(0.0011187) + + +def test_tokens_only_merge_within_the_same_hour() -> None: + """The bucket is part of the join key, so a call in the next hour belongs to a + different spend row and must not inflate this one.""" + next_hour = {**_BYOK_USAGE, "invocation_id": "inv-byok-3", "event_time": "2026-08-07 15:04:00"} + (byok,) = list(_source([_BYOK_SPEND], [_BYOK_USAGE, next_hour]).read_usage("1 day")) + assert byok.usage.input == 1825 + + +def test_unparseable_request_tags_do_not_crash_the_join() -> None: + """A tag column that isn't JSON still has to produce a stable key rather than + raising — one malformed row must not take down the batch.""" + rows = list(_source([{**_BYOK_SPEND, "request_tags": "not json"}], [_BYOK_USAGE]).read_usage("1 day")) + assert len(rows) == 1 + assert rows[0].usd_cost == pytest.approx(0.0011187) + + +def test_byok_spend_with_no_matching_usage_still_bills_its_dollars() -> None: + """A join miss (a row aggregated across an hour boundary, say) must not drop + revenue — the cost is what Databricks charged either way, just with no tokens.""" + (byok,) = list(_source([_BYOK_SPEND], []).read_usage("1 day")) + assert byok.usd_cost == pytest.approx(0.0011187) + assert byok.usage.model == "claude-sonnet-4-5" + assert byok.usage.provider == "anthropic" + assert byok.usage.nonzero_numeric() == {} + + +def test_zero_dollar_spend_rows_are_skipped() -> None: + assert list(_source([{**_BYOK_SPEND, "usage_quantity": "0"}], []).read_usage("1 day")) == [] + + +def test_failed_calls_yield_nothing() -> None: + """403/404s are recorded with NULL token counts. Emitting them would bill an + empty event for a call that never reached a provider.""" + assert list(_source([], [_FAILED]).read_usage("1 day")) == [] + + +def test_hosted_rows_keep_the_databricks_provider() -> None: + """Which is what makes the price lookup miss deliberately rather than matching + some other vendor's rate for a DBU-billed model.""" + (hosted,) = list(_source([], [_HOSTED]).read_usage("1 day")) + assert hosted.usage.provider == "databricks" + assert hosted.usage.api == "databricks_gateway" + + +# -------------------------------------------------------------------------- +# Chunked results — the silent-truncation guard +# -------------------------------------------------------------------------- +class _FakeResponse: + def __init__(self, payload: dict) -> None: + self._payload = payload + + def json(self) -> dict: + return self._payload + + +def test_query_zips_columns_and_follows_every_chunk(monkeypatch: pytest.MonkeyPatch) -> None: + """Only chunk 0 arrives inline. A reader that stops there works on a small window + and silently bills a fraction of a large one — so all `total_chunk_count` chunks + are fetched and the columnar rows zipped back into dicts.""" + import requests + + first = { + "statement_id": "stmt-1", + "status": {"state": "SUCCEEDED"}, + "manifest": { + "schema": {"columns": [{"name": "invocation_id"}, {"name": "input_tokens"}]}, + "total_chunk_count": 3, + }, + "result": {"data_array": [["a", "1"]]}, + } + chunks = {1: {"data_array": [["b", "2"]]}, 2: {"data_array": [["c", "3"]]}} + fetched: list[str] = [] + + def fake_post(url: str, **_kw: Any) -> _FakeResponse: + return _FakeResponse(first) + + def fake_get(url: str, **_kw: Any) -> _FakeResponse: + fetched.append(url) + return _FakeResponse(chunks[int(url.rsplit("/", 1)[-1])]) + + monkeypatch.setattr(requests, "post", fake_post) + monkeypatch.setattr(requests, "get", fake_get) + + rows = DatabricksSource(host="https://x/", token="t", warehouse_id="w").query("SELECT 1") + assert rows == [ + {"invocation_id": "a", "input_tokens": "1"}, + {"invocation_id": "b", "input_tokens": "2"}, + {"invocation_id": "c", "input_tokens": "3"}, + ] + assert [u.rsplit("/", 1)[-1] for u in fetched] == ["1", "2"] + assert all(u.startswith("https://x/api/2.0/sql/statements/stmt-1/result/chunks/") for u in fetched) + + +def test_query_raises_on_a_failed_statement(monkeypatch: pytest.MonkeyPatch) -> None: + """A FAILED statement returns 200 with the failure in the body. Reading rows from + it would report an empty window as "no usage" and bill nothing.""" + import requests + + monkeypatch.setattr( + requests, + "post", + lambda *_a, **_kw: _FakeResponse({"status": {"state": "FAILED", "error": {"message": "boom"}}}), + ) + with pytest.raises(RuntimeError, match="FAILED"): + DatabricksSource(host="https://x", token="t", warehouse_id="w").query("SELECT 1") + + +# -------------------------------------------------------------------------- +# Idempotency keys +# -------------------------------------------------------------------------- +def test_event_ids_are_unique_per_row_and_scoped_by_subscription() -> None: + rows = list(_source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE]).read_usage("1 day")) + ids = [r.event_id for r in rows] + assert len(set(ids)) == len(ids) + assert "sub_byok" in [i for i in ids if "spend" in i][0] + assert "sub_hosted" in [i for i in ids if "usage" in i][0] + + +def test_event_id_prefix_namespaces_the_whole_read() -> None: + rows = list(_source([_BYOK_SPEND], [_HOSTED]).read_usage("1 day", event_id_prefix="tenant7")) + assert all(r.event_id.startswith("tenant7_") for r in rows) + + +def test_event_id_for_rescopes_without_changing_the_row_key() -> None: + """`transaction_id` is unique account-wide, so the same source row billed to two + subscriptions needs two ids — and the id must follow the subscription actually + billed, which for an untagged row is the caller's default, not the row's tag.""" + row = DatabricksUsageRow( + usage=None, # type: ignore[arg-type] + subscription=None, + row_id="rec-9", + kind="spend", + usd_cost=1.0, + ) + assert row.event_id == "dbx_spend_none_rec-9" + assert row.event_id_for("sub_a") == "dbx_spend_sub_a_rec-9" + assert row.event_id_for("sub_b") == "dbx_spend_sub_b_rec-9" + assert row.event_id_for("sub_a") != row.event_id_for("sub_b") + + +# -------------------------------------------------------------------------- +# The one-liner +# -------------------------------------------------------------------------- +class _Recorder: + """Collects delivered events, so assertions read the real emitted shape.""" + + def __init__(self) -> None: + self.batches: list[list[dict]] = [] + + @property + def events(self) -> list[dict]: + return [e for b in self.batches for e in b] + + +def _sdk() -> tuple[LagoSDK, _Recorder]: + rec = _Recorder() + sdk = LagoSDK(api_key="dummy") + sdk._queue._sender = lambda b: rec.batches.append(list(b)) # type: ignore[attr-defined] + return sdk, rec + + +def _drain(sdk: LagoSDK) -> None: + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + + +def test_backfill_counts_cost_tokens_and_skips() -> None: + sdk, q = _sdk() + src = _source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE, {**_HOSTED, "request_tags": "{}"}]) + counts = sdk.backfill_databricks(src, "1 day") + _drain(sdk) + # The untagged row has no subscription and no default to fall back on. + assert counts == {"cost": 1, "tokens": 1, "skipped": 1} + assert {e["external_subscription_id"] for e in q.events} == {"sub_byok", "sub_hosted"} + + +def test_backfill_falls_back_to_the_default_subscription() -> None: + sdk, q = _sdk() + src = _source([], [{**_HOSTED, "request_tags": "{}"}]) + assert sdk.backfill_databricks(src, "1 day", default_subscription="sub_fb")["skipped"] == 0 + _drain(sdk) + assert {e["external_subscription_id"] for e in q.events} == {"sub_fb"} + # ...and the id follows the subscription billed, not the row's absent tag. + assert all("sub_fb" in e["transaction_id"] for e in q.events) + + +def test_backfill_unified_ignores_per_row_tags() -> None: + """One gateway serving one customer: everything lands on one subscription even + though the rows carry their own tags.""" + sdk, q = _sdk() + src = _source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE]) + sdk.backfill_databricks(src, "1 day", default_subscription="sub_one", unified=True) + _drain(sdk) + assert {e["external_subscription_id"] for e in q.events} == {"sub_one"} + assert all("sub_one" in e["transaction_id"] for e in q.events) + + +def test_backfill_bills_byok_as_cost_and_hosted_as_tokens() -> None: + sdk, q = _sdk() + src = _source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE]) + sdk.backfill_databricks(src, "1 day") + _drain(sdk) + + cost = [e for e in q.events if e["code"] == "llm_cost"] + tokens = [e for e in q.events if e["code"] != "llm_cost"] + assert len(cost) == 1 + # Databricks' own $0.0011187 -> 0.11187 cents, passed through, not recomputed. + assert cost[0]["precise_total_amount_cents"].startswith("0.11187") + assert cost[0]["properties"]["price_source"] == "precomputed" + # Hosted has no dollar figure anywhere in Databricks' tables, so: token events. + assert {e["code"] for e in tokens} == {"llm_input_tokens", "llm_output_tokens"} + assert all("precise_total_amount_cents" not in e for e in tokens) + + +def test_backfill_is_idempotent_across_a_re_run() -> None: + """Re-reading the same window must produce byte-identical transaction ids, so + Lago rejects the duplicates instead of double-billing.""" + ids = [] + for _ in range(2): + sdk, q = _sdk() + sdk.backfill_databricks(_source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE]), "1 day") + _drain(sdk) + ids.append([e["transaction_id"] for e in q.events]) + assert ids[0] == ids[1] + + +def test_backfill_survives_one_malformed_row() -> None: + """Instrumentation never breaks the caller: a row that extracts to nothing usable + is skipped, and the rows around it still bill.""" + sdk, q = _sdk() + src = _source([_BYOK_SPEND], [{"nonsense": True}, _HOSTED, _BYOK_USAGE]) + counts = sdk.backfill_databricks(src, "1 day", default_subscription="sub_fb") + _drain(sdk) + assert counts["cost"] == 1 and counts["tokens"] == 1 + assert len(q.events) >= 3 + + +# -------------------------------------------------------------------------- +# Reconciliation dimensions — the whole point of the connector being checkable +# -------------------------------------------------------------------------- +def test_hosted_events_carry_the_endpoint_the_gateway_page_groups_by() -> None: + """Our `model` is normalized (`llama-4-maverick`) where the AI Gateway usage page + shows `system.ai.llama-4-maverick`. Without the endpoint on the event, grouping + Lago one way and Databricks the other fails on naming alone.""" + sdk, q = _sdk() + sdk.backfill_databricks(_source([], [_HOSTED]), "1 day") + _drain(sdk) + assert q.events + for e in q.events: + assert e["properties"]["endpoint_name"] == "system.ai.llama-4-maverick" + + +def test_byok_events_carry_the_hour_bucket_not_a_sampled_endpoint() -> None: + """A spend row covers an hour of requests, so its authoritative key is the hour — + `external_model_spend`'s own aggregation key. A per-request field here would be one + sampled value presented as a property of the whole bucket.""" + sdk, q = _sdk() + sdk.backfill_databricks(_source([_BYOK_SPEND], [_BYOK_USAGE]), "1 day") + _drain(sdk) + (event,) = q.events + assert event["properties"]["bucket"] == "2026-08-07 14:00:00" + assert "endpoint_name" not in event["properties"] + + +def test_caller_dimensions_are_added_and_win_on_a_collision() -> None: + sdk, q = _sdk() + sdk.backfill_databricks( + _source([], [_HOSTED]), + "1 day", + dimensions={"team": "platform", "endpoint_name": "mine"}, + ) + _drain(sdk) + for e in q.events: + assert e["properties"]["team"] == "platform" + # An explicit dimension is the caller's decision, so it overrides the auto key + # rather than being silently discarded. + assert e["properties"]["endpoint_name"] == "mine" + + +def test_a_row_with_no_endpoint_adds_no_empty_dimension() -> None: + """An empty string would create a phantom Lago group rather than saying nothing.""" + sdk, q = _sdk() + sdk.backfill_databricks(_source([], [{**_HOSTED, "endpoint_name": None}]), "1 day") + _drain(sdk) + for e in q.events: + assert "endpoint_name" not in e["properties"] + + +def test_merged_bucket_drops_per_request_extras_but_keeps_the_endpoint() -> None: + """`invocation_id` and `status_code` describe one request. Carrying them on an + hourly aggregate states one sampled request's value as if it covered the hour — + and once dimensions are emitted from extras, that becomes a live mis-statement.""" + second = {**_BYOK_USAGE, "invocation_id": "inv-byok-2", "input_tokens": "100"} + (byok,) = list(_source([_BYOK_SPEND], [_BYOK_USAGE, second]).read_usage("1 day")) + extras = byok.usage.extras + assert extras["endpoint_name"] == _BYOK_USAGE["endpoint_name"] + assert extras["api_type"] == "anthropic/v1/messages" + for per_request in ("invocation_id", "request_id", "status_code"): + assert per_request not in extras + + +def test_a_single_request_bucket_is_described_the_same_way() -> None: + """Otherwise `status_code` survives on quiet hours and vanishes on busy ones — + the same bucket shape reporting different fields depending on traffic.""" + (byok,) = list(_source([_BYOK_SPEND], [_BYOK_USAGE]).read_usage("1 day")) + assert "invocation_id" not in byok.usage.extras + assert byok.usage.extras["endpoint_name"] == _BYOK_USAGE["endpoint_name"] + + +def test_from_env_names_every_missing_variable(monkeypatch: pytest.MonkeyPatch) -> None: + for k in ("DATABRICKS_HOST", "DATABRICKS_TOKEN", "DATABRICKS_WAREHOUSE_ID"): + monkeypatch.delenv(k, raising=False) + with pytest.raises(ValueError) as exc: + DatabricksSource.from_env() + assert "DATABRICKS_HOST" in str(exc.value) + assert "DATABRICKS_WAREHOUSE_ID" in str(exc.value) + + +def test_from_env_trims_a_trailing_slash_off_the_host(monkeypatch: pytest.MonkeyPatch) -> None: + """Or every URL doubles its separator — Databricks 404s on `//api/2.0/...`.""" + monkeypatch.setenv("DATABRICKS_HOST", "https://dbc-x.cloud.databricks.com/") + monkeypatch.setenv("DATABRICKS_TOKEN", "dapi-x") + monkeypatch.setenv("DATABRICKS_WAREHOUSE_ID", "wh-1") + assert DatabricksSource.from_env().host == "https://dbc-x.cloud.databricks.com" + + +def test_json_string_columns_survive_the_round_trip() -> None: + """STRUCT/MAP columns arrive as JSON strings over the Statement Execution API. + The reader joins on the tag map, so it has to parse the same way the adapter + does or every BYOK row misses its token counts.""" + src = _source( + [{**_BYOK_SPEND, "request_tags": json.dumps({"lago_subscription": "sub_byok"})}], + [_BYOK_USAGE], + ) + (byok,) = list(src.read_usage("1 day")) + assert byok.usage.input == 1825 + assert byok.subscription == "sub_byok" + + +# -------------------------------------------------------------------------- +# Post-review hardening — each of these pins a bug found by code review +# -------------------------------------------------------------------------- +def test_rows_with_no_usable_id_do_not_collide() -> None: + """`_safe_str(a or b)` returned "" for a row whose ids were NULL, and also for one + whose id a driver handed back as a non-str (the `or` picks it, `_safe_str` rejects the + type, `request_id` is never tried). Every such row then shared one `transaction_id`, + so Lago billed the first and rejected the rest as duplicates — silently.""" + import uuid as _uuid + + a = {**_HOSTED, "invocation_id": None, "request_id": None, "input_tokens": "7"} + b = {**_HOSTED, "invocation_id": None, "request_id": None, "input_tokens": "9"} + rows = list(_source([], [a, b]).read_usage("1 day")) + assert len(rows) == 2 + assert rows[0].row_id and rows[1].row_id + assert rows[0].event_id != rows[1].event_id + + # A non-str id must be used, not skipped into the fallback. + ident = _uuid.uuid4() + (row,) = list(_source([], [{**_HOSTED, "invocation_id": ident}]).read_usage("1 day")) + assert row.row_id == str(ident) + + +def test_the_id_fallback_is_deterministic_so_re_runs_stay_idempotent() -> None: + """A random UUID would bill an id-less row again on every run.""" + row = {**_HOSTED, "invocation_id": None, "request_id": None} + first = list(_source([], [row]).read_usage("1 day"))[0].event_id + second = list(_source([], [row]).read_usage("1 day"))[0].event_id + assert first == second + + +def test_byok_tokens_with_no_spend_row_are_reported_not_lost(caplog) -> None: + """`external_model_spend` lags `ai_gateway.usage`, so the newest hour has token rows + whose dollar row does not exist yet. The spend loop skips them (no dollars) and the + hosted loop skips them (not databricks), so they were billed by neither and counted + by nothing. Losing them quietly is the failure "never silently under-bill" forbids.""" + import logging as _logging + + with caplog.at_level(_logging.WARNING, logger="lago_agent_sdk.gateway.databricks"): + rows = list(_source([], [_BYOK_USAGE]).read_usage("1 day")) + assert rows == [] + assert any("no external_model_spend row yet" in r.getMessage() for r in caplog.records) + + +def test_an_aware_datetime_window_is_converted_to_utc() -> None: + """`strftime` ignores tzinfo, so a Europe/Paris caller rendered local wall time + against Databricks' UTC columns — a window two hours in the future that reads + nothing and reports success. Also the JS port converts, so this kept the two repos + reading different windows from the same input.""" + from datetime import timedelta, timezone + + paris = timezone(timedelta(hours=2)) + assert _interval_sql(datetime(2026, 8, 11, 14, 0, 0, tzinfo=paris)) == "TIMESTAMP '2026-08-11 12:00:00'" + # Naive is taken as UTC, matching the JS port's Date handling. + assert _interval_sql(datetime(2026, 8, 11, 14, 0, 0)) == "TIMESTAMP '2026-08-11 14:00:00'" + + +def test_datetime_timestamp_columns_still_bucket_and_reconcile() -> None: + """`databricks-sql-connector` returns TIMESTAMPs as `datetime`, not str. `_safe_str` + mapped those to "", collapsing every hour into one join bucket and dropping the + `bucket` reconcile dimension.""" + stamp = datetime(2026, 8, 7, 14, 22, 3) + spend = {**_BYOK_SPEND, "bucket": datetime(2026, 8, 7, 14, 0, 0)} + (byok,) = list(_source([spend], [{**_BYOK_USAGE, "event_time": stamp}]).read_usage("1 day")) + assert byok.usage.input == 1825, "hour key must survive a datetime" + # ISO-8601, matching the string form the REST API returns and the JS port's output. + assert byok.reconcile_dimensions["bucket"] == "2026-08-07T14:00:00" + + +def test_a_malformed_usage_quantity_skips_its_row_instead_of_aborting() -> None: + """`float("NULL")` raised out of the generator, through `backfill_databricks`, and + into the caller — half a window emitted with no record of where it stopped, against + a docstring promising one bad row cannot take down the batch.""" + rows = list(_source([{**_BYOK_SPEND, "usage_quantity": "NULL"}], [_BYOK_USAGE]).read_usage("1 day")) + assert rows == [] + + +def test_query_polls_a_statement_that_is_still_running(monkeypatch: pytest.MonkeyPatch) -> None: + """A statement still executing when `wait_timeout` elapses returns HTTP 200 with + `state: PENDING` — not an error. Raising on it broke the exact usage this class + recommends: one wide window per run, which on a cold warehouse exceeds the 50s + ceiling Databricks allows.""" + import requests + + pending = {"statement_id": "s1", "status": {"state": "PENDING"}} + running = {"statement_id": "s1", "status": {"state": "RUNNING"}} + done = { + "statement_id": "s1", + "status": {"state": "SUCCEEDED"}, + "manifest": {"schema": {"columns": [{"name": "a"}]}, "total_chunk_count": 1}, + "result": {"data_array": [["1"]]}, + } + replies = iter([running, done]) + monkeypatch.setattr(requests, "post", lambda *_a, **_k: _FakeResponse(pending)) + monkeypatch.setattr(requests, "get", lambda *_a, **_k: _FakeResponse(next(replies))) + monkeypatch.setattr("time.sleep", lambda _s: None) + src = DatabricksSource(host="https://x", token="t", warehouse_id="w") + assert src.query("SELECT 1") == [{"a": "1"}] + + +def test_query_gives_up_on_a_statement_that_never_finishes(monkeypatch: pytest.MonkeyPatch) -> None: + import requests + + pending = {"statement_id": "s1", "status": {"state": "PENDING"}} + monkeypatch.setattr(requests, "post", lambda *_a, **_k: _FakeResponse(pending)) + monkeypatch.setattr(requests, "get", lambda *_a, **_k: _FakeResponse(pending)) + monkeypatch.setattr("time.sleep", lambda _s: None) + src = DatabricksSource(host="https://x", token="t", warehouse_id="w", timeout=0.0) + with pytest.raises(RuntimeError, match="still PENDING"): + src.query("SELECT 1") + + +def test_backfill_accepts_already_read_rows_without_querying_again() -> None: + """The demo read the window to print a summary and then handed the SOURCE to + `backfill_databricks`, re-running both warehouse queries — doubling the cost of the + expensive half, and letting the printed summary disagree with what was billed.""" + src = _source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE]) + rows = list(src.read_usage("1 day")) + queries_after_read = len(src.queries) # type: ignore[attr-defined] + + sdk, q = _sdk() + counts = sdk.backfill_databricks(rows, default_subscription="sub_x") + _drain(sdk) + assert counts == {"cost": 1, "tokens": 1, "skipped": 0} + assert len(src.queries) == queries_after_read, "must not re-read" # type: ignore[attr-defined] + assert len(q.events) >= 3 diff --git a/tests/unit/test_buffer_overflow.py b/tests/unit/test_buffer_overflow.py index d1c9907..c44be0d 100644 --- a/tests/unit/test_buffer_overflow.py +++ b/tests/unit/test_buffer_overflow.py @@ -14,10 +14,23 @@ def test_overflow_drops_oldest_at_exact_boundary(): def slow_sender(batch): paused.wait(timeout=30.0) + # max_batch_size must stay ABOVE max_buffer_size, or this test races the worker + # and fails intermittently in CI. `push` sets `_wake` whenever + # `len(buffer) >= max_batch_size`, so with the two equal the overflowing push below + # both drops i=0 AND wakes the worker — which then drains all 10,000 via + # `_take_batch`. If that lands before the next line reads the buffer, `buf` is empty + # and the assertion reads `assert 0 == 10000`. Reproduced deterministically by + # sleeping 50ms in that window; CI's scheduler does it for free under load. + # + # With the cap below the batch size the buffer can never reach it, so the worker + # only ever runs when shutdown() releases `paused` in the finally block. Nothing + # here depends on batch size — every assertion is about buffer CONTENTS. Same + # technique as test_repeated_overflow_keeps_window_sliding below, which was fixed + # for this exact reason. q = EventQueue( sender=slow_sender, flush_interval=10.0, # never timer-flush during the test - max_batch_size=10_000, # match buffer so worker takes everything once unpaused + max_batch_size=20_000, max_buffer_size=10_000, ) try: diff --git a/tests/unit/test_drift.py b/tests/unit/test_drift.py index 35df495..b8cdb78 100644 --- a/tests/unit/test_drift.py +++ b/tests/unit/test_drift.py @@ -5,6 +5,7 @@ from lago_agent_sdk.adapters import ( extract_bedrock_converse, extract_bedrock_invoke, + extract_openai_native, ) @@ -63,3 +64,144 @@ def test_invoke_openai_compat_prompt_tokens_details_lands_in_extras(): u = extract_bedrock_invoke(resp, model_id="openai.gpt-oss-safeguard-20b-1:0") assert "prompt_tokens_details" in u.extras assert u.extras["prompt_tokens_details"] == {"cached_tokens": 48} + + +# ---------------------------------------------------------------------- +# Native OpenAI adapter — drift must be caught ONE LEVEL DOWN too +# ---------------------------------------------------------------------- + + +def test_openai_native_nested_detail_drift_reaches_extras(): + """The drift contract has to hold inside the *_tokens_details sub-objects, + not just at the top level. + + This is the hole a live `gpt-5.6-sol` response found: it reports + `prompt_tokens_details.cache_write_tokens: 3022`, and because + `prompt_tokens_details` is itself a KNOWN top-level key, the old sweep never + looked inside it. 3022 real tokens were discarded with no error and no + on_error — the exact failure this module exists to prevent. Every drift test + passed, because none of them looked one level down. + """ + resp = { + "usage": { + "prompt_tokens": 3025, + "completion_tokens": 4, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 3022}, + "completion_tokens_details": {"reasoning_tokens": 0, "future_nested_xyz": 42}, + } + } + u = extract_openai_native(resp) + assert u.extras["prompt_tokens_details.cache_write_tokens"] == 3022 + assert u.extras["completion_tokens_details.future_nested_xyz"] == 42 + + +def test_openai_native_mapped_nested_fields_do_not_pollute_extras(): + """The mirror of the above: a nested key we DO map must not also appear in + extras, or every event carries a duplicate of a value already billed.""" + resp = { + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "prompt_tokens_details": {"cached_tokens": 40, "audio_tokens": 5}, + "completion_tokens_details": {"reasoning_tokens": 20, "audio_tokens": 3}, + } + } + u = extract_openai_native(resp) + assert u.cache_read == 40 and u.reasoning == 20 + assert u.audio_input == 5 and u.audio_output == 3 + for k in u.extras: + assert not k.endswith((".cached_tokens", ".reasoning_tokens", ".audio_tokens")), k + + +def test_openai_native_responses_api_nested_drift_reaches_extras(): + """Same guarantee on the Responses-API shape, whose detail containers are + named differently (`input_tokens_details` / `output_tokens_details`).""" + resp = { + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "input_tokens_details": {"cached_tokens": 2, "novel_input_detail": "x"}, + "output_tokens_details": {"reasoning_tokens": 1, "novel_output_detail": "y"}, + } + } + u = extract_openai_native(resp) + assert u.api == "responses" + assert u.extras["input_tokens_details.novel_input_detail"] == "x" + assert u.extras["output_tokens_details.novel_output_detail"] == "y" + + +def test_anthropic_service_tier_and_inference_geo_reach_extras(): + """Two fields that appeared on live Anthropic responses through the Databricks + gateway and are in no fixture predating it: `service_tier` ("standard") and + `inference_geo` ("global" for sonnet-4-6, "not_available" for the others). + + Neither is a token count, so both must land in extras — never be miscounted as + a metric, and never silently dropped.""" + from lago_agent_sdk.adapters import extract_anthropic_native + + resp = { + "model": "claude-sonnet-4-6", + "usage": { + "input_tokens": 8, + "output_tokens": 4, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "service_tier": "standard", + "inference_geo": "global", + }, + } + u = extract_anthropic_native(resp) + assert u.input == 8 and u.output == 4 + assert u.extras["service_tier"] == "standard" + assert u.extras["inference_geo"] == "global" + # and they must not have leaked into any numeric field + assert u.nonzero_numeric() == {"input": 8, "output": 4} + + +def test_responses_audio_tokens_reach_extras_because_nothing_maps_them() -> None: + """`output_tokens_details.audio_tokens` was listed as a MAPPED nested key, so it was + excluded from extras — while the Responses branch hardcodes `audio_output = 0` + because the API doesn't expose it. Both true at once means the count is neither + billed nor surfaced: 500 real tokens gone with no error, which is the precise hole + this module exists to close.""" + resp = { + "usage": { + "input_tokens": 10, + "output_tokens": 500, + "output_tokens_details": {"reasoning_tokens": 0, "audio_tokens": 500}, + } + } + u = extract_openai_native(resp) + assert u.api == "responses" + assert u.audio_output == 0, "Responses API does not expose it, so it must not be invented" + assert u.extras["output_tokens_details.audio_tokens"] == 500 + + +def test_unaccounted_total_does_not_double_bill_additive_reasoning() -> None: + """The `total_tokens` guard folds an unexplained delta into `output`. For a provider + whose reasoning is ADDITIVE (this adapter now stamps `databricks` and `workers-ai`, + not only `openai`), a payload reporting BOTH `reasoning_tokens` and an inflated total + would be charged for them twice — inside the grown output and again as a reasoning + line. Subtracting reasoning from the accounted total prevents that.""" + resp = { + "usage": { + "prompt_tokens": 57, + "completion_tokens": 47, + "total_tokens": 1253, + "completion_tokens_details": {"reasoning_tokens": 1149}, + } + } + u = extract_openai_native(resp, provider_hint="databricks") + assert u.reasoning == 1149 + assert u.output == 47, "reasoning already accounts for the delta; output must not grow" + assert "unaccounted_output_tokens" not in u.extras + + +def test_unaccounted_total_still_recovers_tokens_nobody_broke_out() -> None: + """The case the guard was written for is unchanged: a thinking model behind a proxy + that reports no breakdown at all. Measured live — prompt 57, completion 47, total + 1253, and no `completion_tokens_details` to recover the 1,149 from.""" + resp = {"usage": {"prompt_tokens": 57, "completion_tokens": 47, "total_tokens": 1253}} + u = extract_openai_native(resp) + assert u.output == 47 + 1149 + assert u.extras["unaccounted_output_tokens"] == 1149 diff --git a/tests/unit/test_pricing.py b/tests/unit/test_pricing.py index 6787e3d..e6442c2 100644 --- a/tests/unit/test_pricing.py +++ b/tests/unit/test_pricing.py @@ -15,6 +15,7 @@ HttpPricingFetcher, PricingProvider, _parse_price, + _strip_version, bedrock_model_key, coerce_markup, compute_cost, @@ -1294,6 +1295,72 @@ def test_price_unavailable_falls_back_to_token_events_and_reports() -> None: assert any(name == "PricingUnavailableError" and where == "pricing" for name, where in errors) +def test_token_billed_provider_emits_tokens_without_reporting_an_error() -> None: + """A Databricks-hosted model has no per-token rate anywhere — not a cold table, not + an unmatched name, none exists. So token counts are the complete answer, and calling + that a failure on every request trains the reader to ignore on_error entirely.""" + errors: list = [] + sdk, received = _price_sdk( + _warm_provider(), on_error=lambda exc, where: errors.append((type(exc).__name__, where)) + ) + u = CanonicalUsage( + input=11, + output=4, + model="meta-llama-4-maverick-040225", + provider="databricks", + api="chat_completions", + ) + 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 {e["code"] for e in flat} == {"llm_input_tokens", "llm_output_tokens"} + assert [e["properties"]["value"] for e in flat if e["code"] == "llm_input_tokens"] == ["11"] + assert errors == [] + + +def test_a_real_price_miss_still_reports() -> None: + """The narrow exception above must not become a blanket silence: an unmatched model + on a provider that DOES publish rates is a genuine miss the customer can act on.""" + errors: list = [] + sdk, received = _price_sdk( + _warm_provider(), on_error=lambda exc, where: errors.append((type(exc).__name__, where)) + ) + sdk.emit(CanonicalUsage(input=5, model="no-such-model", provider="anthropic", api="native")) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert any(n == "PricingUnavailableError" and w == "pricing" for n, w in errors) + + +def test_token_billed_note_is_logged_once_per_model(caplog) -> None: + """It is a standing fact about the provider, not an event about this call.""" + import logging as _logging + + sdk, _ = _price_sdk(_warm_provider()) + with caplog.at_level(_logging.INFO, logger="lago_agent_sdk"): + for _ in range(3): + sdk.emit(CanonicalUsage(input=1, model="llama-4-maverick", provider="databricks", api="x")) + sdk.emit(CanonicalUsage(input=1, model="gpt-oss-20b", provider="databricks", api="x")) + sdk.shutdown(timeout=1.0) + notes = [r for r in caplog.records if "in its own units" in r.getMessage()] + assert len(notes) == 2 # one per distinct model, not one per call + assert any("llama-4-maverick" in n.getMessage() for n in notes) + + +def test_byok_through_the_same_gateway_still_prices() -> None: + """TOKEN_BILLED_PROVIDERS keys on provider, so it covers Databricks-HOSTED models + only — BYOK traffic through the same gateway is stamped with the real vendor and + must keep pricing normally.""" + sdk, received = _price_sdk(_warm_provider()) + sdk.emit( + CanonicalUsage(input=100, output=50, model="claude-opus-4.8", provider="anthropic", api="native") + ) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert {e["code"] for e in flat} == {"llm_cost"} + + def test_per_call_price_mode_overrides_global_tokens() -> None: # global mode is tokens (default); per-call asks for price provider = _warm_provider() @@ -1321,3 +1388,76 @@ def test_default_mode_is_tokens_unchanged() -> None: sdk.shutdown(timeout=1.0) flat = [e for batch in received for e in batch] assert {e["code"] for e in flat} == {"llm_input_tokens", "llm_output_tokens"} + + +# ---------------------------------------------------------------------- +# Date-suffix shapes — both vendors' conventions must strip +# ---------------------------------------------------------------------- + +# OpenRouter lists BARE ids for the current OpenAI lineup; the API returns dated +# ones. `resolve_model` prefers the response's own name, so the dated form is what +# reaches lookup. +_BARE_OPENAI_TABLE = parse_openrouter( + { + "data": [ + {"id": f"openai/{m}", "pricing": {"prompt": "0.000001", "completion": "0.000002"}} + for m in ("gpt-4.1", "gpt-4.1-mini", "gpt-5", "gpt-5-mini", "o3", "o4-mini") + ] + } +) + + +@pytest.mark.parametrize( + "dated", + [ + "gpt-4.1-2025-04-14", + "gpt-4.1-mini-2025-04-14", + "gpt-5-2025-08-07", + "gpt-5-mini-2025-08-07", + "o3-2025-04-16", + "o4-mini-2025-04-16", + ], +) +def test_openai_hyphenated_date_suffix_strips_to_a_hit(dated: str) -> None: + """OpenAI stamps HYPHENATED dates ("gpt-5-2025-08-07"), Anthropic COMPACT ones + ("claude-sonnet-4-5-20250929"). Handling only the compact shape silently broke + price mode for every current OpenAI model — all six of these missed and fell + back to token events. Verified against the live 400-model OpenRouter table + before and after. + """ + assert lookup_openrouter(_BARE_OPENAI_TABLE, "openai", dated) is not None + + +@pytest.mark.parametrize( + "dated,bare", + [ + ("claude-sonnet-4-5-20250929", "anthropic/claude-sonnet-4.5"), + ("claude-haiku-4-5-20251001", "anthropic/claude-haiku-4.5"), + ("claude-opus-4-5-20251101", "anthropic/claude-opus-4.5"), + ], +) +def test_anthropic_compact_date_suffix_still_strips(dated: str, bare: str) -> None: + """Regression guard: widening the pattern must not break the compact form.""" + table = parse_openrouter({"data": [{"id": bare, "pricing": {"prompt": "0.000003"}}]}) + assert lookup_openrouter(table, "anthropic", dated) is not None + + +def test_non_date_suffix_is_not_stripped() -> None: + """`gpt-5.6-sol` resolves with a `-sol` suffix that is neither a date nor a + version tag. It must be left intact — OpenRouter lists it verbatim as + "openai/gpt-5.6-sol", so stripping would turn a hit into a miss.""" + assert _strip_version("gpt-5.6-sol") == "gpt-5.6-sol" + table = parse_openrouter({"data": [{"id": "openai/gpt-5.6-sol", "pricing": {"prompt": "0.000005"}}]}) + assert lookup_openrouter(table, "openai", "gpt-5.6-sol") is not None + + +def test_workers_ai_model_names_are_never_date_stripped() -> None: + """Workers AI ids carry dotted versions and fp8 suffixes, not dates. The + widened pattern must leave them untouched or the Cloudflare catalog lookup + breaks.""" + for m in ( + "@cf/meta/llama-3.2-1b-instruct", + "@cf/meta/llama-3.3-70b-instruct-fp8-fast", + "@cf/moonshotai/kimi-k2.7-code", + ): + assert _strip_version(m) == m diff --git a/tests/unit/test_wrapper_openai.py b/tests/unit/test_wrapper_openai.py index cfa88af..e57d0b4 100644 --- a/tests/unit/test_wrapper_openai.py +++ b/tests/unit/test_wrapper_openai.py @@ -619,3 +619,165 @@ async def test_async_responses_create_stream_extracts_usage_from_completed_event "looks only at event.usage (top-level), but Responses uses event.response.usage." ) assert by_code.get("llm_output_tokens") == 6 + + +# ---------------------------------------------------------------------- +# Databricks: base_url decides the provider, and streaming quirks +# ---------------------------------------------------------------------- +from lago_agent_sdk.wrappers.openai import _provider_hint_for # noqa: E402 + +_DBX = "https://dbc-0223ef70-2638.cloud.databricks.com" + + +class _FakeClient: + def __init__(self, base_url: str) -> None: + self.base_url = base_url + + +@pytest.mark.parametrize( + "base_url,expected", + [ + # Hosted foundation models — DBU-billed, must NOT reach a vendor price table. + (f"{_DBX}/ai-gateway/mlflow/v1", "databricks"), + (f"{_DBX}/ai-gateway/mlflow/v1/", "databricks"), + # BYOK surfaces keep their real vendor so they price against OpenRouter. + (f"{_DBX}/ai-gateway/openai/v1", ""), + (f"{_DBX}/ai-gateway/anthropic", ""), + # Unrelated clients are untouched. + ("https://api.openai.com/v1", ""), + ("https://gateway.ai.cloudflare.com/v1/acct/gw/compat", ""), + ("", ""), + ], +) +def test_provider_hint_keys_on_the_mlflow_path_only(base_url: str, expected: str) -> None: + """Two of Databricks' four surfaces use the SAME openai.OpenAI class, and the + response body cannot tell them apart — a hosted call echoes a served-entity + name with no marker. base_url is the only signal. + + Matching `/ai-gateway/mlflow/` and not `/ai-gateway/` is load-bearing: the + openai/anthropic BYOK surfaces share that prefix and must keep their vendor + provider, or they would stop being priceable.""" + assert _provider_hint_for(_FakeClient(base_url)) == expected + + +def test_provider_hint_survives_a_client_without_base_url() -> None: + """Some client variants don't expose it; instrumentation must never break the + customer's call over that.""" + + class NoBaseUrl: + pass + + class Raises: + @property + def base_url(self) -> str: + raise RuntimeError("boom") + + assert _provider_hint_for(NoBaseUrl()) == "" + assert _provider_hint_for(Raises()) == "" + + +def test_databricks_hosted_call_is_stamped_databricks_end_to_end() -> None: + """Through the real wrapper: a hosted model must come out as + provider="databricks" so the price lookup cannot hit. OpenRouter lists bare + `openai/gpt-oss-20b` at ~0.4x of Databricks' own DBU rate, so being stamped + "openai" would silently under-bill 2.5-5x the moment a served-entity rename + let _strip_version match it.""" + sdk, received = _new_sdk() + fake = FakeOpenAI() + fake.base_url = f"{_DBX}/ai-gateway/mlflow/v1" + client = sdk.wrap(fake) + client.chat.completions.create(model="system.ai.llama-4-maverick", messages=[]) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert received, "nothing emitted" + assert all(e["properties"]["provider"] == "databricks" for e in received) + + +def test_databricks_byok_call_keeps_its_vendor_provider() -> None: + """The mirror: the same client class against the OpenAI BYOK surface must stay + "openai", because that path IS priceable and was verified exact against + Databricks' own metered spend on 38 of 38 buckets.""" + sdk, received = _new_sdk() + fake = FakeOpenAI() + fake.base_url = f"{_DBX}/ai-gateway/openai/v1" + client = sdk.wrap(fake) + client.chat.completions.create(model="gpt-4o", messages=[]) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert all(e["properties"]["provider"] == "openai" for e in received) + + +class _DbxStreamCompletions: + """Minimal fake reproducing Databricks' streaming convention, which differs from + OpenAI's in two measured ways: usage is on EVERY frame and is CUMULATIVE, and + there is no final usage-only frame — the last frame is an ordinary delta.""" + + def __init__(self, cumulative: list[int]) -> None: + self._cumulative = cumulative + self.with_raw_response = None # force the plain .create() path + + def create(self, **kwargs: Any) -> Any: + assert kwargs.get("stream") is True + return iter( + [ + FakeStreamChunk( + { + "model": "meta-llama-4-maverick-040225", + "choices": [ + { + "index": 0, + "delta": {"content": "a"}, + "finish_reason": "stop" if n == self._cumulative[-1] else None, + } + ], + "usage": {"prompt_tokens": 14, "completion_tokens": n, "total_tokens": 14 + n}, + } + ) + for n in self._cumulative + ] + ) + + +class _DbxStreamClient: + def __init__(self, cumulative: list[int]) -> None: + self.chat = type("C", (), {"completions": _DbxStreamCompletions(cumulative)})() + self.base_url = f"{_DBX}/ai-gateway/mlflow/v1" + + +_DbxStreamClient.__module__ = "openai.fake" + + +def test_databricks_streaming_cumulative_usage_takes_the_last_frame() -> None: + """last-usage-wins lands on the correct total by construction. This pins it, + because a "sum the frames" implementation would bill 1+7+15=23 instead of 15, + and a "first frame wins" one would bill 1.""" + sdk, received = _new_sdk() + client = sdk.wrap(_DbxStreamClient([1, 7, 15])) + list(client.chat.completions.create(model="system.ai.llama-4-maverick", messages=[], stream=True)) + 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"] == 14, "cumulative input must not be summed" + assert by_code["llm_output_tokens"] == 15, "final cumulative value, not 1+7+15" + assert all(e["properties"]["provider"] == "databricks" for e in received) + + +def test_databricks_abandoned_stream_bills_the_partial_total() -> None: + """A behavioral divergence worth pinning rather than discovering later. + + Against real OpenAI, abandoning a stream yields no usage at all — it only + arrives on a final usage-only chunk — so nothing is billed. Databricks puts a + cumulative usage on every frame, so the `finally`-block emit bills whatever had + been generated when the consumer walked away. Arguably better (it bills real + work), but NOT what the OpenAI path does.""" + sdk, received = _new_sdk() + client = sdk.wrap(_DbxStreamClient([1, 7, 15])) + stream = client.chat.completions.create(model="system.ai.llama-4-maverick", messages=[], stream=True) + for i, _ in enumerate(stream): + if i == 1: # abandon after the second frame + break + stream.close() # trigger the generator's finally-block emit deterministically + 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.get("llm_output_tokens") == 7, "partial cumulative count at abandonment"