Feature/databricks gateway connector - #14
Open
anassg-lago wants to merge 8 commits into
Open
Conversation
`_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 OpenRouter table with this repo's own `lookup_openrouter`: gpt-4.1, gpt-4.1-mini, gpt-5, gpt-5-mini, o3 and o4-mini all fell through to token events, so 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.
…ovider hint Three changes to the native OpenAI adapter, all found while validating real provider responses. The drift contract did not hold one level down. `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 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 the four `*_tokens_details` containers. Deliberately NOT mapped to `CanonicalUsage.cache_write`: for OpenAI these sit INSIDE `prompt_tokens` and bill at the plain input rate — cross-checked against Databricks' own metered spend, which charged exactly what billing all 3,025 as input produces. OpenRouter publishes a separate cache-write rate, so mapping the field would charge those 3,022 tokens twice, a 2.24x over-bill. `extras` keeps it visible without touching the money. Tokens in neither named bucket were silently dropped. For genuine OpenAI, `total_tokens` always equals prompt + completion — verified across every captured response, zero deltas. Behind an OpenAI-COMPATIBLE proxy fronting a thinking model it breaks: measured against Gemini through Google's compat layer, prompt 57 / completion 47 / total 1253, with 1,149 thinking tokens reported nowhere. A positive delta now folds into `output` as `extras["unaccounted_output_tokens"]`, minus any reasoning already broken out so an additive-reasoning provider is not billed for them twice. `extract_openai_native` also gains `provider_hint`, because two of Databricks' gateway surfaces use the same `openai.OpenAI` class but need different price tables, and the response body cannot tell them apart. Only the wrapper knows the `base_url`; the adapter stays the single place `provider` is decided.
Second entry in the `gateway/` namespace, alongside Cloudflare. `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. 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), with 22 captured fixtures covering both destination types, cache read/write, reasoning, embeddings and all three failure shapes. Three mapping quirks a docs-only reading gets wrong, all caught by real rows: `destination_name` means DIFFERENT things per destination type — the model for a hosted row, the PROVIDER SERVICE (a Unity Catalog credential name) for BYOK. A single "model, falling back to name" rule bills a credential as the model on every BYOK row. `destination_model` is unstable for hosted models: the same `destination_name` reports both `gpt-oss-20b` and the display label `GPT OSS 20B`, depending on which of Databricks' two request aliases the caller used. Most hosted entities carry a second, INNER prefix — `system.ai.databricks-<model>`, on 38 of 48 distinct names. It is a serving-endpoint artefact, not part of the model id, but it cannot be stripped unconditionally because Databricks also publishes models genuinely named that way (`databricks-dbrx-instruct`). `destination_model` is the tie-breaker; disagreement keeps the raw name, since an ugly id is recoverable and a silently renamed model is not. `provider="databricks"` for hosted models is deliberately unmatchable in `_VENDOR_MAP`. Databricks bills them in DBUs against a rate card published only as HTML and present in no system table, while OpenRouter does list bare `openai/gpt-oss-20b` at 0.2-0.4x of Databricks' real rate — so being stamped "openai" would silently under-bill 2.5-5x. Same trap as Workers AI. This table's `input_tokens` INCLUDES cache_read and cache_write, the inverse of the providers' own response bodies. The adapter extracts faithfully and does not subtract; the module docstring records why, and why a computed fallback would need to correct per provider rather than uniformly. The barrel now exports gateway-scoped names, so neither gateway is the implicit default.
`gateway/databricks.py` is the one piece of gateway code that does I/O, and the
adapter beside it stays pure. Cloudflare's read is a single paginated GET and
rightly lives in its example notebook; Databricks needs a SQL warehouse, the
Statement Execution API, columnar-to-dict zipping, chunked result fetching, a
statement poll, and two tables reconciled against each other. Hand-rolled that is
~100 lines in which four money-losing mistakes are easy, and the first version of
the demo notebook made three of them:
Silent truncation — only chunk 0 arrives inline, so a window wide enough to span
`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 key
built from the source row alone blocks that row from ever reaching a second
subscription. And the subscription billed is not always the one on the row, since
an untagged row falls back to the caller's default, so `event_id_for()` builds
the key from the resolved value.
Lost rows — a row with NULL ids, or an id a driver hands back as a non-string,
collapsed to an empty key, so every such row in the window shared one
transaction_id and only the first was ever billed. Falls back to a content hash,
which stays deterministic so re-runs remain idempotent.
`LagoSDK.backfill_databricks(source, "7 days")` bills a whole window and returns
`{"cost": n, "tokens": n, "skipped": n}`. It also accepts an already-read
iterable of rows, because a SQL warehouse costs roughly 1,500x the model-serving
usage it reports on — reading the window twice to print a summary first doubles
the expensive half and lets the summary disagree with what was billed.
A BYOK bucket with no spend row is billed by neither path, which the spend
table's ~19h lag makes routine for the newest hour, so it warns rather than
vanishing. The window is validated rather than escaped, since it reaches SQL by
interpolation. Deliberately absent: scheduler, cursor store, credential store.
In price mode a hosted call logged `lago pricing failed: no price for provider='databricks' model='meta-llama-4-maverick-040225'` and routed it to `on_error` on EVERY request. 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` 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 "never silently under-bill". That invariant exists so a price miss cannot 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`. This 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. It keys on the PROVIDER, so it covers Databricks-hosted traffic only. BYOK through the same gateway is stamped openai/anthropic and prices normally — verified exact against Databricks' own metered spend on 38 of 38 buckets. The OpenAI wrapper supplies the `provider_hint` that makes this reachable, reading `base_url` once at wrap time. It must key on `/ai-gateway/mlflow/`, not `/ai-gateway/`, or the OpenAI BYOK path gets mis-stamped and priced against the wrong table.
README gains a `## Databricks AI Gateway` section covering all live ingress paths and the backfill, plus the gotchas customers would otherwise report as SDK bugs: `gpt-oss` inflates input by ~100 tokens from a server-injected preamble, `claude-opus-4-5` does not cache through this gateway at all, hosted models report three different name strings, and running the live path and the backfill over the same traffic emits token events twice. Corrects a claim that had nothing behind it: the "What gets billed" table said hosted backfill produced a dollar cost from `system.billing.usage` × `list_prices`, while the paragraph two lines below said the opposite, and neither table appears anywhere in the source tree. Hosted bills token counts on both paths. Those dollars do exist — `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. They are not billed from because they come from a different Databricks screen than the gateway view: no `request_tags`, so per-subscription splits would be ours rather than Databricks', and ~19h of lag. Every number this connector sends is one you can find on a Databricks *gateway* page, which is the property that makes it checkable. Each backfilled event carries the grouping key of the surface it came from — `endpoint_name` for hosted, `bucket` for BYOK — so grouping Lago the way the Databricks page groups puts the two side by side. Without it the comparison fails on naming alone, since our `model` is normalized and the page's is not. CONTRIBUTING gains an "Adding a gateway" recipe, the bar a read must clear to belong in the SDK rather than a notebook, and the two rules that keep a connector comparable against the gateway's own dashboard. `examples/databricks_gateway_demo.ipynb` demonstrates both halves and was re-run against a live workspace: 107 billable rows over 7 days, 60 dollar-cost events plus 88 token events, all transaction_ids unique.
…rker The test set `max_batch_size` equal to `max_buffer_size`, and `push` sets `_wake` whenever `len(buffer) >= max_batch_size`. So the overflowing push both dropped i=0 AND woke the background worker, which then drained all 10,000 events through `_take_batch`. When that landed before the next line read the buffer, `buf` came back empty and the assertion read `assert 0 == 10000`. Failed in CI on a loaded runner; reproduced deterministically by sleeping 50ms in that window, which is all the scheduler needs to do for free. Fixed the same way `test_repeated_overflow_keeps_window_sliding` already was: keep the batch size ABOVE the buffer cap so the buffer can never reach it, and the worker only runs once shutdown() releases the sender. Nothing in the test depends on batch size — every assertion is about buffer CONTENTS. Verified over 150 consecutive runs. Pre-existing; unrelated to the Databricks connector, but it is what turns this branch's CI red.
`mistral` was missing from `_INPUT_INCLUDES_CACHE_READ`, so in price mode 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 is unambiguous — prompt_tokens=1013, cached_tokens=1008, total_tokens=1043=prompt+completion, which only reconciles if the cached tokens sit inside the prompt count. Mistral bills them at 10% of the input rate. Measured 6.15x over-bill 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 routed through a Cloudflare gateway (the gateway adapter leaves provider="mistral" unmapped). Token mode was unaffected — only the price computation was wrong. money_golden.json gains a `mistral` case built from Mistral's documented payload; removing the fix makes it produce 0.0006019 against the expected 0.0000979 and fail. This is the second provider missing from that set after `workers-ai`. The set is still hand-maintained; a completeness check over every provider slug the SDK can emit remains the real fix.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add Databricks AI Gateway connector
Databricks is the second gateway connector, after Cloudflare. Everything below was
established against a live workspace, not inferred from docs: 25 models exercised,
~100 calls, 226 real usage rows read over the SQL Statement Execution API.
What it covers
Databricks differs from Cloudflare in two ways that shape the whole design:
/compatendpointfronting every provider; Databricks makes each provider reachable only through its
own native surface — and two of those use the same
openai.OpenAIclass while needingdifferent price tables, so
base_urldiscrimination is load-bearing rather thancosmetic.
base_urlwrap()/ai-gateway/openai/v1external_model_spend/ai-gateway/anthropicexternal_model_spend/ai-gateway/mlflow/v1/ai-gateway/gemini/v1betaThe decision that shapes everything: gateway parity
What makes the Cloudflare connector trustworthy is that you can put Cloudflare's own
dashboard beside Lago and see the same numbers. Databricks makes that harder, because
hosted traffic appears on two surfaces in two units:
request_tags?custom_tagsis{}This connector mirrors the gateway's own surfaces. Hosted bills token counts
(matching the AI Gateway usage page); BYOK bills Databricks' own metered USD (matching
the external-model-spend view).
Hosted money is deliberately not billed from, and this is the part most worth
challenging in review. It is obtainable —
system.billing.usage×list_prices, oraccount_pricesfor an account's contract rate — so "hosted USD is impossible" would bewrong; that applies only to the tokens→DBU rate, which is published on an HTML page and
exists in no system table (verified by searching every column of all 88 of them). The
reason not to use it is a product one: it comes from a different Databricks screen than
the gateway view, carries no attribution (so per-subscription splits would be ours rather
than Databricks'), and lags ~19h — measured,
max(usage_start_time)inbilling.usagewas
2026-08-10T17:00whilemax(event_time)inai_gateway.usagewas2026-08-11T10:09. Every number this connector sends is one you can find on a Databricksgateway page.
Two earlier approaches to hosted pricing were built and reverted, both recorded in
CHANGELOG.md: vendoring the 18-model DBU rate card (hand-maintained price data in acodebase whose every other source is live HTTP), and solving the rate from the customer's
own tables (worked — all 6 solvable endpoints recovered the published rate to three
decimals — but the SQL warehouse needed to run the solve costs ~1,500× the usage it
prices: $6.54 of SQL against $0.0043 of MODEL_SERVING in this account).
Evidence
captured response body was run through the real pipeline (
extract_*_native→compute_costat live OpenRouter rates) and compared against Databricks' ownusage_quantity, joined on Databricks' own grouping key. 13 models, both cacheconventions, four reasoning models, costs from $0.0000036 to $0.015245.
byte-identical
transaction_ids across a re-run.(3,640 in / 4,344 out), confirmed against a real Lago instance.
qwen35-122b-a10bat 48 in /204 out, which is exactly what falls out of dividing that endpoint's DBU rows by the
published rate card — the tokens billed agree with what Databricks charged DBUs for,
via a completely separate table.
Commits — the first two are not Databricks-specific
Commits 1–2 fix bugs that affect every price-mode user, found while validating this
connector, and each passes the full suite on its own (462 → 474 → 559 tests). If you
would rather land them separately, say so and I will split them into their own PR
underneath this one — no re-work needed.
_strip_versiononly matched Anthropic's compact date (-20250929), not OpenAI'shyphenated one (
gpt-5-2025-08-07). Sinceresolve_modelprefers the response's ownname,
create(model="gpt-5")resolved to a name that matched nothing and fell throughto token events. Verified against the live OpenRouter table:
gpt-4.1,gpt-4.1-mini,gpt-5,gpt-5-mini,o3,o4-miniall missed.gpt-4olooked fine only by luck.extrasswept onlytop-level usage keys, and
prompt_tokens_detailsis itself a known key, so nothingnested was ever inspected: a live
gpt-5.6-solresponse'sprompt_tokens_details.cache_write_tokens: 3022was discarded with no error. Everydrift test passed, because none looked inside a details object.
Where to look closely
gateway/databricks.py— the money paths. Four ways this read loses money silently(chunk-0 truncation, double billing across the two tables, unscoped idempotency keys,
and rows whose ids collapse to an empty string) are each guarded and each has a
regression test naming the failure.
gateway/adapters/databricks_gateway.py— three naming quirks that a docs-only readinggets wrong, all caught by real rows. In particular
destination_namemeans the modelfor hosted rows but a credential name for BYOK, so a single fallback rule bills
workspace.default.anthropickeyas the model.sdk.pyTOKEN_BILLED_PROVIDERS— a deliberate, narrow exception to "never silentlyunder-bill". Reasoning is in commit 5's message; the invariant still holds for every
miss a customer could act on.
input_tokensnote in the adapter docstring — this table's input count includescache tokens, the inverse of the providers' own response bodies. Nothing computes from
it today, and the docstring records why a computed fallback would need a per-provider
correction rather than a uniform one.
Gates
ruff check+ruff format --check+mypy --strictclean; 559 tests, 90% coverage(
gateway/databricks.pyat 100%). 22 fixtures are real captured rows — neverhand-written — and a sweep test iterates the whole directory so a capture no named test
mentions still asserts something.
Deliberately out of scope
this does not remember where it got to.
but every request past it returns
500with an empty body — including Databricks'own documented code sample, and including
:countTokens, which involves no inference.A genuine upstream failure on this gateway is richly wrapped; these carry none of that,
so the gateway throws while constructing the call. Needs a Databricks support ticket,
not SDK work. Reproduction ids and five saved 500-responses are recorded.
Known gaps
system.ai_gateway.usagestores the requested alias (gpt-5.6) while OpenRouter listsonly the resolved name (
gpt-5.6-sol), so that one model prices live and misses onbackfill — falling back to token events, never mispriced. Needs an alias step.
api="chat_completions"; the shape detector only knows ChatCompletions vs Responses. Numerically correct, mislabelled.