Feature/databricks gateway connector - #14
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.
ancorcruz
left a comment
ancorcruz
left a comment
There was a problem hiding this comment.
Review notes
Ten items. The first five I'd fix before merge — four of them are silent-failure paths in the money-reading half, which is the one place a mistake doesn't announce itself. The rest is hardening, one scope question, and a cleanup.
Four were reproduced rather than reasoned about: the chunk-fetch truncation, the empty-columns collapse, the false "NOT billed" warning (the repo's own 22 fixtures produce 4 phantom rows), and the prefix-strip non-determinism (fixtures hosted_chat.json / hosted_chat_1.json carry the same destination_name with destination_model of llama-4-maverick vs Llama 4 Maverick).
Two notes that aren't inline because the lines aren't in this diff:
sdk.py's precomputed branch reportsunit = usage.input + usage.output. I flagged that on the Cloudflare PR already, but this connector aggravates it: perdatabricks_gateway.py's own billing-hazard note, this table'sinputincludes cache_read/cache_write, so onbyok_anthropic_cache_write_1.json(input=1825,cache_read=1812) theunitpublished alongside Databricks' dollar figure overstates by ~3x. The notebook's reconciliation block sumsr.usage.inputthe same way.- Considered and set aside: mixing a row ordinal into
_row_id's hash fallback to break ties between two id-less identical rows.ORDER BY event_timeisn't a unique sort, so an ordinal isn't stable across runs and would trade a rare collision for a broken idempotency guarantee. Leaving it as-is looks like the better trade.
|
|
||
| total_chunks = int(manifest.get("total_chunk_count") or 1) | ||
| statement_id = body.get("statement_id") | ||
| for index in range(1, total_chunks): |
There was a problem hiding this comment.
Blocker: the chunk loop reintroduces the silent truncation this module's docstring says it exists to prevent.
Neither the POST at line 209, the poll GET at line 268, nor this chunk GET checks HTTP status — every one goes straight to .json(). So when a chunk fetch fails (503, expired statement, revoked token mid-read), the error body has no data_array, or [] swallows it, and query() returns a partial row set with no exception. The total_chunks > 1 log at line 236 then cheerfully reports "spanned 3 chunks" for a read that got 2.
That is verbatim the first bullet of the module docstring: "A naive reader works on a small window and quietly bills a fraction of a large one, with no error."
Worth noting the codebase already has the convention — pricing.py calls raise_for_status() at all five of its fetch sites. This is the one HTTP path where a silent failure costs billing rows, and it's the one without the check. test_query_zips_columns_and_follows_every_chunk covers only the success path.
A raise_for_status() on all three calls is enough; letting it raise is right here, since backfill_databricks aborting loudly is far better than under-billing quietly.
| 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] |
There was a problem hiding this comment.
Blocker: a missing/short manifest.schema makes the whole window read as zero usage and report success.
columns is built from manifest.schema.columns with .get("columns", []) on line 223, so a SUCCEEDED body without a schema yields columns == []. With strict=False, dict(zip([], row)) is {} — so query() returns [{}, {}, ...], one empty dict per real row, no exception.
Downstream every one of those degrades cleanly and wrongly: extract_databricks_log({}) gives all-zero usage, the hosted loop skips them on nonzero_numeric(), the BYOK index builds all-zero buckets keyed ('', '', '', '{}'), and backfill_databricks returns {"cost": 0, "tokens": 0, "skipped": 0} for a window that had real traffic. Every defensive layer does its job and the result is a confident zero.
strict=True would turn a column/row length mismatch into an error, but it won't catch the empty case (zip([], row) is simply empty) — so an explicit if not columns: raise RuntimeError(...) is the part that matters.
| 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")): |
There was a problem hiding this comment.
Blocker: exact-equality against a column this module documents as unstable, so one hosted model can bill under two different ids.
The comment above (lines 88-101) is right that the prefix can't be stripped unconditionally and that destination_model is the signal for telling an artefact from a real name. The problem is the exactness of the comparison, given what the module docstring says two paragraphs earlier: destination_model for hosted models flips between a slug and a human display label — measured, gpt-oss-20b / GPT OSS 20B.
The fixtures show the flip directly: hosted_chat.json and hosted_chat_1.json carry the same destination_name (system.ai.llama-4-maverick) with destination_model of llama-4-maverick and 'Llama 4 Maverick'.
llama-4-maverick is unaffected because it carries no inner prefix. But for one of the 38-of-48 prefixed names, the flip decides the outcome: a slug row emits qwen35-122b-a10b while a display-label row emits databricks-qwen35-122b-a10b. Same model, two Lago rows — the exact split the comment says the rule prevents, just triggered by row-level label variance instead of by a bad strip.
Comparing normalized forms on both sides (_alnum, or lowercase + spaces→hyphens) makes the decision stable whichever label the row carried, and still refuses to strip when the two columns genuinely disagree — Qwen35 122B A10B normalizes onto qwen35-122b-a10b, while databricks-dbrx-instruct vs dbrx-instruct still does not.
| # 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: |
There was a problem hiding this comment.
Blocker (severity is trust, not money): failed rows enter the join index, so the "NOT billed" warning is permanently wrong.
This loop filters only on u.provider == "databricks". The hosted loop at line 384 additionally guards if not u.nonzero_numeric(): continue; this one doesn't, so failed calls — which the adapter docstring says arrive with NULL tokens — become entries in tokens, and then show up in set(tokens) - billed_keys at line 369.
Ran the repo's own 22 fixtures through extract_databricks_log: 4 rows are non-hosted with entirely zero usage — gemini_broken.json, gemini_broken_1.json, unmanaged_path.json, unmanaged_path_1.json. Those are calls that never reached a provider, so external_model_spend will never have a row for them, and the warning's advice ("re-run this window later to bill them") is false forever, not just this window. Any workspace with a broken Gemini connection gets a permanent warning on every run, plus an inflated bucket count.
Which matters more than a cosmetic log: this warning is the only signal that real BYOK rows went unbilled (see the backfill_databricks comment). An alarm that always fires is one nobody reads — which is the argument TOKEN_BILLED_PROVIDERS makes in this same PR.
Same if not u.nonzero_numeric(): continue guard here fixes it.
| duplicates rather than double-bill. Does not flush — call ``flush()`` when | ||
| you want to block on delivery. | ||
| """ | ||
| counts = {"cost": 0, "tokens": 0, "skipped": 0} |
There was a problem hiding this comment.
Blocker: the one-call entrypoint returns success while knowingly under-billing.
counts reports cost / tokens / skipped, but skipped only counts rows this loop received and couldn't attribute. The BYOK token buckets that read_usage drops — the documented external_model_spend aggregation lag, so routinely the window's most recent hour — are never yielded, so they land in no counter. They're reported only to a module logger, never through config.on_error.
Net effect: backfill_databricks returns {"cost": 60, "tokens": 47, "skipped": 0} for a window where real usage went unbilled, and a caller reconciling on the return value or an on_error hook sees nothing. That's in tension with the invariant emit() is otherwise careful about — the CHANGELOG's "don't silently under-bill" and the PricingUnavailableError report on a price miss.
Two options, either works: have read_usage surface the count (return it, or yield the unbilled buckets as a sentinel) so counts can carry a deferred key; or route the warning through on_error so it reaches the same hook every other billing gap uses. The docstring's promise that the return value is "counts of what it emitted" stays true either way — the gap is that what it didn't emit is invisible.
| billed_keys: set[tuple[Any, ...]] = set() | ||
|
|
||
| for row in spend: | ||
| usd = _safe_float(row.get("usage_quantity")) |
There was a problem hiding this comment.
A negative usage_quantity passes this filter, then floors to $0 and burns the row's idempotency key.
if not usd catches 0.0 but not -0.0042. A credit or correction row therefore yields a DatabricksUsageRow(usd_cost=-0.0042), and compute_precomputed_cost does _parse_price(usd_cost) or Decimal(0) — _parse_price returns None for negatives (pricing.py:196), so base is 0.
The floor-to-zero itself is deliberate and documented at pricing.py:346. The problem is what happens here: a real llm_cost event is pushed with precise_total_amount_cents="0" under a real transaction_id. The credit is lost, nothing is logged, and the id is now consumed — so if the row is later corrected, the re-run is rejected as a duplicate and can't fix it.
if usd <= 0: continue (with a warning when it's negative) keeps the id unused, which is the recoverable state. Whether Databricks ever emits negative usage_quantity on this table I can't confirm — but the cost of guarding is one comparison, and the cost of not guarding is unrecoverable.
| WHERE usage_start_time >= {window} | ||
| """) | ||
|
|
||
| usage = self.query(f""" |
There was a problem hiding this comment.
SELECT * plus an unused ORDER BY, on the one resource this module says costs ~1500x the usage it prices.
Two separate costs:
Warehouse time. SELECT * pulls all 36 columns including endpoint_metadata, routing_information, invocation_metadata and the service_* set, while the adapter reads about twelve. The ORDER BY event_time forces a sort over the whole window, and the rows go straight into extracted and a dict index where order is never used — the spend loop iterates spend, the hosted loop iterates extracted, and neither depends on ordering.
A correctness edge. _row_id's fallback hashes the entire row (line 421), so for a row with NULL invocation_id and request_id, an unused column changes the transaction_id. invocation_metadata or a latency field differing between two reads of the same window would produce a different id for the same row — breaking the idempotency the docstring promises. Naming the needed columns removes that coupling as a side effect.
Given the module explicitly tells operators "read one wide window per run; never poll in a tight loop", narrowing the projection seems worth it here.
| # 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) |
There was a problem hiding this comment.
Scope question on the total_tokens guard: it subtracts reasoning, but nothing else that can be additive.
The reasoning subtraction is well argued, and the Gemini-behind-a-compat-layer case it's built for is real. What I'm less sure of is that reasoning is the only field that can inflate total_tokens relative to input + output.
This same diff documents (lines 52-53) that Anthropic's cache_creation_input_tokens sits OUTSIDE input_tokens, and both gateways in this SDK front Anthropic models behind OpenAI-shaped surfaces. On a payload like {prompt_tokens: 13, completion_tokens: 4, total_tokens: 1829, prompt_tokens_details: {cache_write_tokens: 1812}}, unaccounted is 1812 and those cache-write tokens get folded into output — billed at the output rate, and simultaneously surfaced in extras["prompt_tokens_details.cache_write_tokens"].
I can't point to a live payload with that shape, so this may be unreachable in practice — but the guard is unconditional on provider, and the argument for it ("no completion_tokens_details to recover them from") is specifically about a payload with no breakdown. Gating on that — only fold in the remainder when the details sub-objects are absent or empty — would keep the measured case working while making the assumption explicit. If instead you've confirmed no fronted provider reports cache-creation additively in total_tokens, a line saying so would be enough.
| _DATABRICKS_HOSTED_PATH = "/ai-gateway/mlflow/" | ||
|
|
||
|
|
||
| def _provider_hint_for(client: Any) -> str: |
There was a problem hiding this comment.
The base_url→provider rule now exists in three places, and this one only covers the OpenAI wrapper.
The reasoning is right and the hazard is real — my concern is placement. base_url is currently inspected to decide a provider in three spots that must stay in agreement: here, sdk.py::_auto_prime_pricing_for (its own read, for Cloudflare), and _infer_provider's model-string variant. Only this one produces a provider_hint, and only wrappers/openai.py passes it.
Concretely: Databricks hosts Anthropic models (system.ai.databricks-claude-sonnet-4-5). Whether a customer can reach those through anthropic.Anthropic(base_url=...) depends on whether Databricks exposes a native Anthropic-shaped surface for hosted entities — if it does, that client gets provider="anthropic", strips to claude-sonnet-4-5, hits OpenRouter, and prices a DBU-billed model at Anthropic's rate: the 2.5-5x mispricing this docstring says it makes impossible. Worth confirming one way or the other, since the answer decides whether this is a gap or just an asymmetry.
Either way, there's a bonus in consolidating: a single base_url→provider resolver consumed by wrap() for every client kind would also close the Cloudflare /compat streaming hole I flagged on the base PR — where a workers-ai/@cf/... streaming call is stamped openai because the response carries no model and the string check wants a bare @cf/. The wrapper knows the base_url there too; it just has no way to say so today.
| from typing import Any | ||
|
|
||
| from ..canonical import CanonicalUsage | ||
| from .adapters.databricks_gateway import ( |
There was a problem hiding this comment.
Cleanup (low priority): third copy of the _safe_* helpers, and a private symbol crossing a module boundary.
_safe_dict / _safe_int / _safe_str now exist in adapters/openai_native.py, gateway/adapters/cloudflare_gateway.py and gateway/adapters/databricks_gateway.py. The bodies have already diverged — only the Databricks copy tolerates a JSON-string dict, which is a genuinely useful improvement the other two don't get.
This import also makes _safe_str part of a cross-module contract despite the leading underscore. adapters/_common.py is the established home for shared adapter helpers; a gateway/_common.py mirroring it would fit the existing layout.
Also minor, same file: json is imported at line 35 and again inside _canonical_tags at line 471.
2 participants
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.