From a4c4ad3b8572f77156b90bb4b7a7cf93fd53e897 Mon Sep 17 00:00:00 2001 From: Anass Date: Tue, 11 Aug 2026 17:33:03 +0200 Subject: [PATCH 01/22] Strip hyphenated version dates so current OpenAI models can be priced MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_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. --- src/lago_agent_sdk/pricing.py | 17 +++++++- tests/unit/test_pricing.py | 74 +++++++++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 2 deletions(-) diff --git a/src/lago_agent_sdk/pricing.py b/src/lago_agent_sdk/pricing.py index 1504676..f132b36 100644 --- a/src/lago_agent_sdk/pricing.py +++ b/src/lago_agent_sdk/pricing.py @@ -86,6 +86,7 @@ # are additive to output, so it's absent here.) _OUTPUT_INCLUDES_REASONING = frozenset({"openai"}) + # Canonical field -> OpenRouter pricing key. _OPENROUTER_FIELD_MAP = { "input": "prompt", @@ -140,7 +141,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 +217,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/tests/unit/test_pricing.py b/tests/unit/test_pricing.py index 6787e3d..5544927 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, @@ -1321,3 +1322,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 From 2265195560f2fcb93c01f671a4859f16e6ebde85 Mon Sep 17 00:00:00 2001 From: Anass Date: Tue, 11 Aug 2026 17:33:03 +0200 Subject: [PATCH 02/22] Catch nested usage drift, account for unexplained totals, accept a provider hint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- src/lago_agent_sdk/adapters/openai_native.py | 89 +++++++++++- tests/unit/adapters/test_openai_native.py | 107 ++++++++++++++ tests/unit/test_drift.py | 142 +++++++++++++++++++ 3 files changed, 336 insertions(+), 2 deletions(-) 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/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/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 From fe1ce43aab1e0e29adcb801192b6d79f86f044c2 Mon Sep 17 00:00:00 2001 From: Anass Date: Tue, 11 Aug 2026 17:34:00 +0200 Subject: [PATCH 03/22] Add Databricks AI Gateway usage adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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-`, 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/adapters/__init__.py | 9 + .../gateway/adapters/databricks_gateway.py | 230 +++++++++++++ .../byok_anthropic_cache_read.json | 38 ++ .../byok_anthropic_cache_read_1.json | 38 ++ .../byok_anthropic_cache_write.json | 38 ++ .../byok_anthropic_cache_write_1.json | 38 ++ .../byok_anthropic_plain.json | 38 ++ .../byok_openai_cache_read.json | 38 ++ .../byok_openai_cache_read_1.json | 38 ++ .../databricks_gateway/byok_openai_plain.json | 38 ++ .../byok_openai_plain_1.json | 38 ++ .../byok_openai_reasoning.json | 38 ++ .../byok_openai_reasoning_1.json | 38 ++ .../failed_null_tokens.json | 38 ++ .../failed_null_tokens_1.json | 38 ++ .../databricks_gateway/gemini_broken.json | 38 ++ .../databricks_gateway/gemini_broken_1.json | 38 ++ .../databricks_gateway/hosted_chat.json | 38 ++ .../databricks_gateway/hosted_chat_1.json | 38 ++ .../hosted_chat_endpoint_prefixed_name.json | 38 ++ .../databricks_gateway/hosted_embeddings.json | 38 ++ .../hosted_embeddings_1.json | 38 ++ .../databricks_gateway/unmanaged_path.json | 38 ++ .../databricks_gateway/unmanaged_path_1.json | 38 ++ .../adapters/test_databricks_gateway.py | 324 ++++++++++++++++++ 25 files changed, 1399 insertions(+) create mode 100644 src/lago_agent_sdk/gateway/adapters/databricks_gateway.py create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_read.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_read_1.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_write.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_cache_write_1.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_plain.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_cache_read.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_cache_read_1.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain_1.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning_1.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens_1.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken_1.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_1.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_endpoint_prefixed_name.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings_1.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path.json create mode 100644 tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path_1.json create mode 100644 tests/unit/gateway/adapters/test_databricks_gateway.py 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/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 From b7737cf1014ecc5a821a7a3b46781d9cbc268785 Mon Sep 17 00:00:00 2001 From: Anass Date: Tue, 11 Aug 2026 17:34:00 +0200 Subject: [PATCH 04/22] Add Databricks usage reader and one-call backfill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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. --- src/lago_agent_sdk/gateway/databricks.py | 519 +++++++++++++++ src/lago_agent_sdk/sdk.py | 83 +++ tests/unit/gateway/test_databricks_source.py | 633 +++++++++++++++++++ 3 files changed, 1235 insertions(+) create mode 100644 src/lago_agent_sdk/gateway/databricks.py create mode 100644 tests/unit/gateway/test_databricks_source.py 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/sdk.py b/src/lago_agent_sdk/sdk.py index 168fbe5..8912052 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -423,6 +423,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/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 From d85ad7185ef824304dd82b6bf6aafd126b1cb264 Mon Sep 17 00:00:00 2001 From: Anass Date: Tue, 11 Aug 2026 17:34:34 +0200 Subject: [PATCH 05/22] Bill Databricks-hosted models as token counts, not as a price failure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- src/lago_agent_sdk/__init__.py | 9 +- src/lago_agent_sdk/pricing.py | 19 +++ src/lago_agent_sdk/sdk.py | 30 +++++ src/lago_agent_sdk/wrappers/openai.py | 37 +++++- tests/unit/test_pricing.py | 66 +++++++++++ tests/unit/test_wrapper_openai.py | 162 ++++++++++++++++++++++++++ 6 files changed, 321 insertions(+), 2 deletions(-) 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/pricing.py b/src/lago_agent_sdk/pricing.py index f132b36..95efffb 100644 --- a/src/lago_agent_sdk/pricing.py +++ b/src/lago_agent_sdk/pricing.py @@ -86,6 +86,25 @@ # 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 = { diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index 8912052..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: 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/test_pricing.py b/tests/unit/test_pricing.py index 5544927..e6442c2 100644 --- a/tests/unit/test_pricing.py +++ b/tests/unit/test_pricing.py @@ -1295,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() 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" From 27dbe818c6dd50e9918d1a25927d304760b54cd5 Mon Sep 17 00:00:00 2001 From: Anass Date: Tue, 11 Aug 2026 17:34:34 +0200 Subject: [PATCH 06/22] Document the Databricks AI Gateway connector MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- CHANGELOG.md | 47 ++- CONTRIBUTING.md | 78 +++++ README.md | 89 ++++++ examples/.env.example | 22 ++ examples/databricks_gateway_demo.ipynb | 382 +++++++++++++++++++++++++ 5 files changed, 609 insertions(+), 9 deletions(-) create mode 100644 examples/databricks_gateway_demo.ipynb diff --git a/CHANGELOG.md b/CHANGELOG.md index 450db6a..7d2274b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,16 +4,21 @@ 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 + +- **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 +33,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 +} From acc7d3ce3de5d84808c1b6e9333e113f01bb0044 Mon Sep 17 00:00:00 2001 From: Anass Date: Tue, 11 Aug 2026 17:50:51 +0200 Subject: [PATCH 07/22] Stop test_overflow_drops_oldest_at_exact_boundary racing the queue worker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- tests/unit/test_buffer_overflow.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) 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: From 8c4ffa9815fa590aae506d84e86d103f5f95aff5 Mon Sep 17 00:00:00 2001 From: Anass Date: Wed, 12 Aug 2026 21:46:16 +0200 Subject: [PATCH 08/22] Bill Mistral prompt-cached calls once, not twice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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. --- CHANGELOG.md | 2 ++ src/lago_agent_sdk/pricing.py | 11 ++++++++++- tests/unit/fixtures/pricing/money_golden.json | 11 +++++++++++ 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7d2274b..a1cc959 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,8 @@ All notable changes to this project will be documented here. Format follows [Kee ### 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. diff --git a/src/lago_agent_sdk/pricing.py b/src/lago_agent_sdk/pricing.py index 95efffb..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 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": [ From 03e81d06a0f96d1b4ebc05cf0301f6f43fdd6120 Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 10:04:16 +0200 Subject: [PATCH 09/22] Stop the queue respinning after isolation, and drain before refreshing pricing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two JS-port behaviours the Python queue never got. Both measured against the real `EventQueue` on a driver rather than reasoned about. **Unbounded respin after isolating a batch.** A permanent batch failure (422) whose isolated sends then fail transiently (429) re-queued the survivors and continued straight into re-taking them with the backoff reset to 0 — no delay anywhere in the cycle. Measured: **280,388 HTTP requests in 1.2s**, aimed at the server that had just asked us to slow down. `_send_individually` now returns how many it re-queued; a non-zero count arms the normal 1->2->4->...->60s backoff, and the immediate-continue is kept only for the case it was written for — isolation fully resolved the batch, so the buffer really did shrink. After: 8 calls, backoff 2.0s. The exit drain gets `requeue_transient=False` for the same reason: there is no later retry to re-queue TO, so an event failing there is reported lost rather than handed back to a buffer this same loop immediately re-takes. **Refresh ran ahead of the drain.** `maybe_refresh()` does HTTP — up to 10s per source — and ran before the buffer was drained, on every tick. Measured: a 600ms refresh pushed first delivery to 629ms; now 37ms. Drain, then refresh, then drain again so anything pushed during the fetch does not wait out another flush interval. Both the refresh and the trailing drain are skipped once `_stopping` is set, so a shutdown landing mid-tick no longer spends the caller's budget on a table nothing will read. Four tests mirroring the JS names; each confirmed to fail with its own fix reverted. ruff, mypy strict, 645 unit tests green. --- CHANGELOG.md | 4 + src/lago_agent_sdk/queue.py | 151 ++++++++++++++++++++++++++---------- tests/unit/test_queue.py | 111 ++++++++++++++++++++++++++ 3 files changed, 224 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d5365f..ab0d035 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,10 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **An isolated batch could respin without limit, hammering the server that had just throttled us.** When a batch failed permanently (422) and every isolated send then failed transiently (429), `_send_individually` re-queued the survivors and `_run` continued straight into re-taking them with `_backoff_seconds` reset to 0 — no delay anywhere in the cycle. Measured against a server returning that exact pair: **280,388 HTTP requests in 1.2 seconds**. `_send_individually` now returns the number of events it re-queued, and the drain arms the normal 1→2→4→…→60s backoff whenever that count is non-zero, keeping the immediate-continue only for the case it was meant for (isolation fully resolved the batch, so the buffer genuinely shrank). The exit drain passes `requeue_transient=False` for the same reason: there is no later retry to re-queue *to*, so an event failing there is reported as lost rather than put back on a buffer the same loop immediately re-takes. Closes a divergence with the JS port, which has guarded this since its own queue fix. + +- **A slow pricing refresh sat in front of every queued billable event.** `_run` called `maybe_refresh()` *before* draining the buffer, so up to a 10s-per-source HTTP fetch delayed delivery on every tick, and a source that kept failing repeated it indefinitely. Measured: a 600ms refresh pushed first delivery from 37ms to 629ms. The drain now runs first, the refresh second, and a second drain follows it so anything pushed while the fetch was in flight does not wait out another whole flush interval. Nothing in the drain depends on the refresh — an event's price is resolved at `emit()` time, so a fresh table only ever matters to the NEXT call. The refresh and the trailing drain are both skipped once `_stopping` is set, so a `shutdown()` landing mid-tick no longer spends the caller's shutdown budget fetching a table nothing will read. Mirrors the JS port's `drainBuffer()` shape. + - **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. diff --git a/src/lago_agent_sdk/queue.py b/src/lago_agent_sdk/queue.py index 18c59d0..0a324b4 100644 --- a/src/lago_agent_sdk/queue.py +++ b/src/lago_agent_sdk/queue.py @@ -220,7 +220,12 @@ def _report_error(self, exc: Exception, where: str = "send_batch") -> None: except Exception: # noqa: BLE001 pass - def _send_individually(self, batch: list[dict[str, Any]], batch_exc: Exception) -> None: + def _send_individually( + self, + batch: list[dict[str, Any]], + batch_exc: Exception, + requeue_transient: bool = True, + ) -> int: """Recovery path for a batch that failed with a permanent (4xx) error. Each event is sent alone: one that individually 4xxs (e.g. its own @@ -230,6 +235,13 @@ def _send_individually(self, batch: list[dict[str, Any]], batch_exc: Exception) any other event. Reports once via on_error for the batch as a whole (the original exception) so a caller isn't flooded with N callbacks for what's really one root cause. + + Returns the number of events re-queued, which the caller needs in order to + decide whether it may keep draining immediately or must back off first — + see `_drain_buffer`. `requeue_transient=False` is for the exit drain, where + there is no later retry to re-queue TO: an event that fails there is lost + and must be reported as such rather than put back on a buffer nobody will + read again. """ self._report_error(batch_exc) # Collected and re-queued ONCE at the end, not per event. `_replay_failed` @@ -250,57 +262,108 @@ def _send_individually(self, batch: list[dict[str, Any]], batch_exc: Exception) event.get("transaction_id"), exc, ) - else: + elif requeue_transient: logger.warning("lago send failed for isolated event, will retry: %s", exc) retry.append(event) + else: + self._report_error(exc) + logger.warning( + "lago: event LOST on shutdown — no retry left: transaction_id=%s: %s", + event.get("transaction_id"), + exc, + ) if retry: self._replay_failed(retry) + return len(retry) + + def _next_backoff(self) -> float: + """1s -> 2s -> 4s -> ... -> `max_retry_seconds`.""" + if self._backoff_seconds == 0: + return 1.0 + return min(self._backoff_seconds * 2, self._max_retry_seconds) + + def _drain_buffer(self) -> None: + """Send everything buffered, one batch at a time, until the buffer is empty or + a failure hands the batch to the retry backoff. + + Returns rather than looping on a failure: the caller waits out + `flush_interval` and comes back. + """ + while True: + batch = self._take_batch() + if not batch: + return + # Re-checked every iteration, not only around the backoff wait. Always + # `return` after re-queuing, never a path that abandons the batch: the exit + # drain is what reports whatever it cannot send, so the events have to be + # back on the buffer before this leaves. + if self._stopping.is_set(): + self._replay_failed(batch) + return + if self._backoff_seconds: + if self._stopping.wait(timeout=self._backoff_seconds): + self._replay_failed(batch) + return + try: + self._http_calls += 1 + self._sender(batch) + self._backoff_seconds = 0.0 + except Exception as exc: # noqa: BLE001 + if _is_permanent_failure(exc): + # Lago's batch endpoint is all-or-nothing: a single bad + # transaction_id fails the WHOLE batch, even if the rest + # are perfectly valid — re-queuing the batch as-is would + # retry (and re-fail) forever, but dropping it outright + # would silently lose those valid events too. Isolate by + # falling back to one-by-one for this batch only; only + # the events that individually 4xx get dropped. + requeued = self._send_individually(batch, exc) + if requeued == 0: + # Batch fully resolved — the buffer shrank, so keep draining. + self._backoff_seconds = 0.0 + continue + # Some isolated sends failed transiently and went back on the + # buffer. Continuing here would re-take them with no delay and + # re-fail at the speed of the failure. Measured on this exact pair + # (422 on the batch, 429 on every isolated send): 280,388 HTTP + # requests in 1.2s, aimed at the server that had just asked us to + # slow down. They must go through the normal backoff path. + self._backoff_seconds = self._next_backoff() + return + self._replay_failed(batch) + self._report_error(exc) + logger.warning("lago send_batch failed: %s", exc) + self._backoff_seconds = self._next_backoff() + return def _run(self) -> None: while not self._stopping.is_set(): self._wake.wait(timeout=self._flush_interval) self._wake.clear() + # Drain BEFORE refreshing pricing, not after. `maybe_refresh()` does HTTP — + # up to a 10s timeout per source — and refreshing first put that latency in + # front of every queued billable event, on every tick: measured, a 600ms + # refresh delayed the first delivery to 629ms. Nothing in the drain depends + # on it — an event's price was already resolved at emit() time, so a fresh + # table only ever matters to the NEXT call. + self._drain_buffer() + # Refresh pricing tables on this background thread (off the hot path). - if self._pricing is not None: + # Skipped once shutting down: a fetch here can take the full HTTP timeout, + # and it would spend the caller's shutdown budget on a table nothing will + # ever read. + if self._pricing is not None and not self._stopping.is_set(): try: self._pricing.maybe_refresh() except Exception: # noqa: BLE001 — pricing must never break the queue pass - while True: - batch = self._take_batch() - if not batch: - break - if self._backoff_seconds: - if self._stopping.wait(timeout=self._backoff_seconds): - self._replay_failed(batch) - return - try: - self._http_calls += 1 - self._sender(batch) - self._backoff_seconds = 0.0 - except Exception as exc: # noqa: BLE001 - if _is_permanent_failure(exc): - # Lago's batch endpoint is all-or-nothing: a single bad - # transaction_id fails the WHOLE batch, even if the rest - # are perfectly valid — re-queuing the batch as-is would - # retry (and re-fail) forever, but dropping it outright - # would silently lose those valid events too. Isolate by - # falling back to one-by-one for this batch only; only - # the events that individually 4xx get dropped. - self._send_individually(batch, exc) - self._backoff_seconds = 0.0 - continue - self._replay_failed(batch) - self._report_error(exc) - logger.warning("lago send_batch failed: %s", exc) - self._backoff_seconds = ( - 1.0 - if self._backoff_seconds == 0 - else min(self._backoff_seconds * 2, self._max_retry_seconds) - ) - break + # Anything pushed while the refresh was in flight would otherwise wait out + # a whole flush interval on top of it. + if not self._stopping.is_set(): + self._drain_buffer() + # Drain on exit — keep sending until the buffer is truly empty, not # just one batch's worth (a buffer holding more than max_batch_size # events at shutdown previously left the rest never even attempted). @@ -310,11 +373,11 @@ def _run(self) -> None: # previously did — that's what actually lost events, not the network # blip itself, which by itself is recoverable if it's just reported. # `_send_individually` re-queues transient sub-failures for retry — - # appropriate for the main loop, which lives on, but during this exit - # drain that could spin forever against a persistently-down network. - # Bound the whole drain by wall-clock time; whatever's still in the - # buffer once the budget is spent is logged as lost, not retried - # forever in an exiting daemon thread. + # appropriate for the main loop, which lives on, but here it would spin + # against a persistently-down network, so this drain passes + # `requeue_transient=False`. The drain is additionally bounded by wall-clock + # time; whatever's still in the buffer once the budget is spent is logged as + # lost, not retried forever in an exiting daemon thread. drain_deadline = time.monotonic() + min(self._max_retry_seconds, 10.0) while time.monotonic() < drain_deadline: batch = self._take_batch() @@ -324,7 +387,11 @@ def _run(self) -> None: self._sender(batch) except Exception as exc: # noqa: BLE001 if _is_permanent_failure(exc): - self._send_individually(batch, exc) + # `requeue_transient=False`: re-queuing here would put the event + # back on a buffer this loop immediately re-takes, with no later + # retry to reach — a hot loop for the whole drain budget. Report + # it as lost instead. + self._send_individually(batch, exc, requeue_transient=False) else: self._report_error(exc) logger.warning( diff --git a/tests/unit/test_queue.py b/tests/unit/test_queue.py index 0423f6b..d7ca109 100644 --- a/tests/unit/test_queue.py +++ b/tests/unit/test_queue.py @@ -480,3 +480,114 @@ def sender(batch): assert [e["id"] for e in q._buffer] == ["b", "c", "d"] finally: q.shutdown(timeout=1.0) + + +# ---------------------------------------------------------------------- +# No unbounded respin after isolating a batch. Mirrors +# `EventQueue — no unbounded respin after isolating a batch` in the JS port. +# ---------------------------------------------------------------------- +def test_isolation_requeue_is_paced_not_spun() -> None: + """A permanent batch error plus a transient error on every isolated send used to + loop with no delay: `_send_individually` put the events back and `_run` continued + straight into re-taking them. Measured before the fix: 280,388 HTTP requests in + 1.2s, aimed at the server that had just returned 429.""" + calls = {"n": 0} + + def sender(batch): + calls["n"] += 1 + # Permanent on the batch, transient on every isolated send: the exact pair. + raise LagoApiError(422 if len(batch) > 1 else 429, "x") + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=60.0) + try: + for name in ("a", "b", "c"): + q.push({"transaction_id": name}) + time.sleep(1.2) + # 1 batch + 3 isolated + at most a couple of paced retries. + assert calls["n"] < 40, f"expected paced retries, got {calls['n']} calls" + assert q._backoff_seconds > 0, "a partial re-queue must arm the backoff" + finally: + q.shutdown(timeout=0.5) + + +def test_exit_drain_does_not_respin_either(caplog) -> None: + """The exit drain has no later retry, so re-queuing a transient sub-failure there + means re-taking it immediately — a hot loop for the whole drain budget. Those + events must be reported as lost instead.""" + calls = {"n": 0} + + def sender(batch): + calls["n"] += 1 + raise LagoApiError(422 if len(batch) > 1 else 429, "x") + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=60.0) + for name in ("a", "b", "c"): + q.push({"transaction_id": name}) + time.sleep(0.4) + before = calls["n"] + with caplog.at_level("WARNING"): + q.shutdown(timeout=1.5) + # 1 batch + 3 isolated sends per pass, not thousands. + assert calls["n"] - before < 20, f"exit drain spun: {calls['n'] - before} calls" + assert any("LOST" in r.getMessage() for r in caplog.records) + + +def test_keeps_draining_when_isolation_fully_resolves_the_batch() -> None: + """The counterpart: nothing re-queued means the buffer shrank, so the loop should + keep going immediately rather than waiting out a whole flush interval.""" + delivered: list[str] = [] + + def sender(batch): + if len(batch) > 1: + raise LagoApiError(422, "batch rejected") + delivered.append(batch[0]["transaction_id"]) + + q = EventQueue( + sender=sender, flush_interval=0.05, max_batch_size=2, max_buffer_size=10_000, max_retry_seconds=0.5 + ) + try: + for name in ("a", "b", "c", "d"): + q.push({"transaction_id": name}) + assert q.flush(timeout=3.0) + assert set(delivered) == {"a", "b", "c", "d"} + finally: + q.shutdown(timeout=1.0) + + +# ---------------------------------------------------------------------- +# A slow pricing refresh must not delay event delivery. Mirrors +# `EventQueue — a slow pricing refresh does not delay event delivery` in JS. +# ---------------------------------------------------------------------- +def test_slow_pricing_refresh_does_not_delay_delivery() -> None: + """`maybe_refresh()` used to run BEFORE the drain, so its HTTP latency sat in front + of every queued billable event on every tick. Measured before the fix: a 600ms + refresh pushed first delivery to 629ms.""" + refresh_seconds = 0.6 + state = {"refresh_done": False, "delivered_before_refresh": None} + + class SlowPricing: + def maybe_refresh(self) -> None: + time.sleep(refresh_seconds) + state["refresh_done"] = True + raise RuntimeError("bad credential") # the failing case, repeated every tick + + def sender(batch): + if state["delivered_before_refresh"] is None: + state["delivered_before_refresh"] = not state["refresh_done"] + + q = EventQueue( + sender=sender, + flush_interval=0.025, + max_batch_size=100, + max_buffer_size=10_000, + max_retry_seconds=60.0, + pricing=SlowPricing(), + ) + try: + q.push({"transaction_id": "t1"}) + deadline = time.monotonic() + 2.0 + while time.monotonic() < deadline and state["delivered_before_refresh"] is None: + time.sleep(0.01) + assert state["delivered_before_refresh"] is True + finally: + q.shutdown(timeout=1.0) From 217fa0e53b1521f8224860f964ce805d85acc178 Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 10:13:16 +0200 Subject: [PATCH 10/22] Compact the swept fixtures 151 -> 38, and scrub personal data from the captures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three machine-swept directories held captures that re-asserted the same thing. The sweep tests assert DISPATCH and non-zero usage, so two captures sharing a usage shape, an adapter family and a pricing provider exercise one code path twice. Reduced to one per distinct key: bedrock/converse 39 -> 12 shape x pricing provider bedrock/invoke 39 -> 14 shape x invoke family x pricing provider mistral_native/all_models 73 -> 12 shape x model family x vision flag Every fixture a dedicated test names is kept as its group's representative, so nothing referenced by name disappeared — verified by extracting every literal `*.json` in both repos' tests first. The sweeps' own coverage assertions are unchanged and still pass: all 7 InvokeModel families, all 8 Mistral families, 4 vision captures. A COVERAGE.md per directory lists every model that was captured and which committed fixture now stands for it, so the breadth of the live verification stays reviewable as a table instead of as 113 near-identical JSON blobs. Hand-curated directories are untouched. Several Databricks pairs that look redundant are not: hosted_chat.json / hosted_chat_1.json differ only in destination_model ("llama-4-maverick" vs "Llama 4 Maverick"), which is the captured evidence that the column is unstable. Separately — personal data removed from all 22 Databricks fixtures. `system.ai_gateway.usage` logs the caller's account email and source IP on every row, and no adapter reads either, so both repos were publishing a personal Gmail address, a residential IP, a real workspace subdomain and one live Lago subscription id to a public index. Replaced with RFC 2606/5737 reserved values; no credentials were ever committed. The gateway fixtures have no capture script to hang a scrub step off, so a hygiene test is the durable fix — confirmed to fail on all three when the originals are put back. --- CHANGELOG.md | 6 ++ .../fixtures/bedrock/converse/COVERAGE.md | 17 ++++ .../converse/eu.amazon.nova-2-lite-v1_0.json | 25 ------ .../converse/eu.amazon.nova-micro-v1_0.json | 25 ------ .../converse/eu.amazon.nova-pro-v1_0.json | 25 ------ ...hropic.claude-haiku-4-5-20251001-v1_0.json | 29 ------- ...thropic.claude-opus-4-5-20251101-v1_0.json | 29 ------- .../eu.anthropic.claude-opus-4-6-v1.json | 29 ------- ...ropic.claude-sonnet-4-5-20250929-v1_0.json | 29 ------- .../converse/google.gemma-3-27b-it.json | 25 ------ .../converse/google.gemma-3-4b-it.json | 25 ------ .../converse/minimax.minimax-m2.5.json | 29 ------- .../bedrock/converse/minimax.minimax-m2.json | 29 ------- .../converse/mistral.devstral-2-123b.json | 25 ------ .../mistral.magistral-small-2509.json | 25 ------ .../mistral.ministral-3-14b-instruct.json | 25 ------ .../mistral.ministral-3-3b-instruct.json | 25 ------ .../mistral.ministral-3-8b-instruct.json | 25 ------ .../mistral.mistral-7b-instruct-v0_2.json | 25 ------ .../mistral.mixtral-8x7b-instruct-v0_1.json | 25 ------ .../mistral.voxtral-mini-3b-2507.json | 25 ------ .../mistral.voxtral-small-24b-2507.json | 25 ------ .../converse/nvidia.nemotron-nano-3-30b.json | 25 ------ .../converse/nvidia.nemotron-nano-9b-v2.json | 25 ------ .../converse/openai.gpt-oss-120b-1_0.json | 29 ------- .../openai.gpt-oss-safeguard-120b.json | 29 ------- .../qwen.qwen3-coder-30b-a3b-v1_0.json | 25 ------ .../converse/qwen.qwen3-next-80b-a3b.json | 25 ------ .../converse/qwen.qwen3-vl-235b-a22b.json | 25 ------ .../fixtures/bedrock/invoke/COVERAGE.md | 21 +++++ .../invoke/eu.amazon.nova-2-lite-v1_0.json | 23 ----- .../invoke/eu.amazon.nova-micro-v1_0.json | 23 ----- .../invoke/eu.amazon.nova-pro-v1_0.json | 23 ----- ...hropic.claude-haiku-4-5-20251001-v1_0.json | 27 ------ ...thropic.claude-opus-4-5-20251101-v1_0.json | 27 ------ .../eu.anthropic.claude-opus-4-6-v1.json | 27 ------ ...ropic.claude-sonnet-4-5-20250929-v1_0.json | 27 ------ .../bedrock/invoke/google.gemma-3-27b-it.json | 27 ------ .../bedrock/invoke/google.gemma-3-4b-it.json | 27 ------ .../bedrock/invoke/minimax.minimax-m2.5.json | 27 ------ .../bedrock/invoke/minimax.minimax-m2.json | 27 ------ .../invoke/mistral.magistral-small-2509.json | 27 ------ .../mistral.ministral-3-14b-instruct.json | 27 ------ .../mistral.ministral-3-3b-instruct.json | 27 ------ .../mistral.ministral-3-8b-instruct.json | 27 ------ .../mistral.mistral-7b-instruct-v0_2.json | 11 --- .../mistral.mixtral-8x7b-instruct-v0_1.json | 11 --- .../invoke/mistral.voxtral-mini-3b-2507.json | 27 ------ .../mistral.voxtral-small-24b-2507.json | 27 ------ .../invoke/nvidia.nemotron-nano-3-30b.json | 27 ------ .../invoke/nvidia.nemotron-nano-9b-v2.json | 27 ------ .../invoke/openai.gpt-oss-120b-1_0.json | 27 ------ .../invoke/qwen.qwen3-coder-30b-a3b-v1_0.json | 27 ------ .../invoke/qwen.qwen3-next-80b-a3b.json | 27 ------ .../invoke/qwen.qwen3-vl-235b-a22b.json | 27 ------ .../mistral_native/all_models/COVERAGE.md | 20 +++++ .../all_models/codestral-latest.json | 44 ---------- .../all_models/devstral-latest.json | 44 ---------- .../all_models/devstral-medium-2507.json | 44 ---------- .../all_models/devstral-medium-latest.json | 44 ---------- .../all_models/devstral-small-2507.json | 44 ---------- .../all_models/magistral-medium-latest.json | 55 ------------ .../magistral-medium-latest__vision.json | 56 ------------ .../all_models/magistral-small-2509.json | 55 ------------ .../magistral-small-2509__vision.json | 56 ------------ .../all_models/magistral-small-latest.json | 44 ---------- .../magistral-small-latest__vision.json | 45 ---------- .../all_models/ministral-14b-latest.json | 44 ---------- .../ministral-14b-latest__vision.json | 45 ---------- .../all_models/ministral-3b-2512.json | 44 ---------- .../all_models/ministral-3b-2512__vision.json | 45 ---------- .../all_models/ministral-3b-latest.json | 44 ---------- .../ministral-3b-latest__vision.json | 45 ---------- .../all_models/ministral-8b-2512.json | 44 ---------- .../all_models/ministral-8b-2512__vision.json | 45 ---------- .../all_models/ministral-8b-latest.json | 44 ---------- .../ministral-8b-latest__vision.json | 45 ---------- .../all_models/mistral-large-2512.json | 44 ---------- .../all_models/mistral-large-latest.json | 44 ---------- .../mistral-large-latest__vision.json | 45 ---------- .../mistral-large-pixtral-2411.json | 44 ---------- .../all_models/mistral-medium-2505.json | 44 ---------- .../mistral-medium-2505__vision.json | 45 ---------- .../all_models/mistral-medium-2508.json | 44 ---------- .../mistral-medium-2508__vision.json | 45 ---------- .../all_models/mistral-medium-2604.json | 44 ---------- .../mistral-medium-2604__vision.json | 45 ---------- .../all_models/mistral-medium-3-5.json | 44 ---------- .../mistral-medium-3-5__vision.json | 45 ---------- .../all_models/mistral-medium-3.5.json | 44 ---------- .../mistral-medium-3.5__vision.json | 45 ---------- .../all_models/mistral-medium-3.json | 44 ---------- .../all_models/mistral-medium-3__vision.json | 45 ---------- .../mistral-medium-c21211-r0-75.json | 44 ---------- .../mistral-medium-c21211-r0-75__vision.json | 45 ---------- .../all_models/mistral-medium-latest.json | 44 ---------- .../mistral-medium-latest__vision.json | 45 ---------- .../all_models/mistral-medium.json | 44 ---------- .../all_models/mistral-medium__vision.json | 45 ---------- .../all_models/mistral-small-2506.json | 44 ---------- .../mistral-small-2506__vision.json | 45 ---------- .../all_models/mistral-small-2603.json | 44 ---------- .../mistral-small-2603__vision.json | 45 ---------- .../all_models/mistral-small-latest.json | 44 ---------- .../mistral-small-latest__vision.json | 45 ---------- .../all_models/mistral-tiny-2407.json | 44 ---------- .../all_models/mistral-tiny-latest.json | 44 ---------- .../all_models/mistral-vibe-cli-fast.json | 44 ---------- .../mistral-vibe-cli-fast__vision.json | 45 ---------- .../all_models/mistral-vibe-cli-latest.json | 44 ---------- .../mistral-vibe-cli-latest__vision.json | 45 ---------- .../mistral-vibe-cli-with-tools.json | 44 ---------- .../mistral-vibe-cli-with-tools__vision.json | 45 ---------- .../all_models/open-mistral-nemo.json | 44 ---------- .../all_models/voxtral-mini-latest.json | 44 ---------- .../all_models/voxtral-small-2507.json | 44 ---------- .../all_models/voxtral-small-latest.json | 44 ---------- .../byok_anthropic_cache_read.json | 6 +- .../byok_anthropic_cache_read_1.json | 6 +- .../byok_anthropic_cache_write.json | 6 +- .../byok_anthropic_cache_write_1.json | 6 +- .../byok_anthropic_plain.json | 8 +- .../byok_openai_cache_read.json | 6 +- .../byok_openai_cache_read_1.json | 6 +- .../databricks_gateway/byok_openai_plain.json | 6 +- .../byok_openai_plain_1.json | 6 +- .../byok_openai_reasoning.json | 6 +- .../byok_openai_reasoning_1.json | 6 +- .../failed_null_tokens.json | 6 +- .../failed_null_tokens_1.json | 6 +- .../databricks_gateway/gemini_broken.json | 6 +- .../databricks_gateway/gemini_broken_1.json | 6 +- .../databricks_gateway/hosted_chat.json | 6 +- .../databricks_gateway/hosted_chat_1.json | 6 +- .../hosted_chat_endpoint_prefixed_name.json | 6 +- .../databricks_gateway/hosted_embeddings.json | 6 +- .../hosted_embeddings_1.json | 6 +- .../databricks_gateway/unmanaged_path.json | 6 +- .../databricks_gateway/unmanaged_path_1.json | 6 +- tests/unit/test_fixture_hygiene.py | 85 +++++++++++++++++++ 140 files changed, 216 insertions(+), 4157 deletions(-) create mode 100644 tests/unit/adapters/fixtures/bedrock/converse/COVERAGE.md delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-2-lite-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-micro-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-pro-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-haiku-4-5-20251001-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-opus-4-5-20251101-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-opus-4-6-v1.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-sonnet-4-5-20250929-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/google.gemma-3-27b-it.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/google.gemma-3-4b-it.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/minimax.minimax-m2.5.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/minimax.minimax-m2.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/mistral.devstral-2-123b.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/mistral.magistral-small-2509.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-14b-instruct.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-3b-instruct.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-8b-instruct.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/mistral.mistral-7b-instruct-v0_2.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/mistral.mixtral-8x7b-instruct-v0_1.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/mistral.voxtral-mini-3b-2507.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/mistral.voxtral-small-24b-2507.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/nvidia.nemotron-nano-3-30b.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/nvidia.nemotron-nano-9b-v2.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/openai.gpt-oss-120b-1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/openai.gpt-oss-safeguard-120b.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-coder-30b-a3b-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-next-80b-a3b.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-vl-235b-a22b.json create mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/COVERAGE.md delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-2-lite-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-micro-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-pro-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-haiku-4-5-20251001-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-opus-4-5-20251101-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-opus-4-6-v1.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-sonnet-4-5-20250929-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/google.gemma-3-27b-it.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/google.gemma-3-4b-it.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/minimax.minimax-m2.5.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/minimax.minimax-m2.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/mistral.magistral-small-2509.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-14b-instruct.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-3b-instruct.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-8b-instruct.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/mistral.mistral-7b-instruct-v0_2.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/mistral.mixtral-8x7b-instruct-v0_1.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/mistral.voxtral-mini-3b-2507.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/mistral.voxtral-small-24b-2507.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/nvidia.nemotron-nano-3-30b.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/nvidia.nemotron-nano-9b-v2.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/openai.gpt-oss-120b-1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-coder-30b-a3b-v1_0.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-next-80b-a3b.json delete mode 100644 tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-vl-235b-a22b.json create mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/COVERAGE.md delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/codestral-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/devstral-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/devstral-medium-2507.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/devstral-medium-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/devstral-small-2507.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/magistral-medium-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/magistral-medium-latest__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-2509.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-2509__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-latest__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-14b-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-14b-latest__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-2512.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-2512__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-latest__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-2512.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-2512__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-latest__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-2512.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-latest__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-pixtral-2411.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2505.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2505__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2508.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2508__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2604.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2604__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3-5.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3-5__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.5.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.5__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-c21211-r0-75.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-c21211-r0-75__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-latest__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2506.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2506__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2603.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2603__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-latest__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-tiny-2407.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-tiny-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-fast.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-fast__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-latest__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-with-tools.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-with-tools__vision.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/open-mistral-nemo.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-mini-latest.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-small-2507.json delete mode 100644 tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-small-latest.json create mode 100644 tests/unit/test_fixture_hygiene.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ab0d035..00d71fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,12 @@ All notable changes to this project will be documented here. Format follows [Kee - **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. +### Changed + +- **Fixture set compacted 151 -> 38, and personal data scrubbed from the Databricks captures.** The three machine-swept fixture directories (`bedrock/converse` 39, `bedrock/invoke` 39, `mistral_native/all_models` 73) held many captures that re-asserted the same thing: the sweep tests assert DISPATCH and non-zero usage, so two captures sharing a usage shape, an adapter family and a pricing provider exercise one code path twice. Reduced to one fixture per distinct key — 12 / 14 / 12 — with every fixture a dedicated test names kept as its group's representative, so nothing referenced by name disappeared. A `COVERAGE.md` in each directory lists every model that was captured and which committed fixture now stands for it, so the breadth of the live verification is still reviewable without the bulk. The sweep tests' own coverage assertions are unchanged and still satisfied: all 7 InvokeModel families, all 8 Mistral families, and 4 vision captures. + - The **hand-curated** directories are untouched: `anthropic_native`, `openai_native`, `gemini_native`, `cloudflare_gateway` and `databricks_gateway` each already carry one fixture per behaviour, and several Databricks pairs that look redundant are not — `hosted_chat.json` / `hosted_chat_1.json` differ only in `destination_model` (`llama-4-maverick` vs `Llama 4 Maverick`), which is the captured evidence that the column is unstable. + - **Personal data removed from all 22 Databricks fixtures.** `system.ai_gateway.usage` records the caller's account email and source IP on every row, and neither field is read by any adapter — so both repos were publishing a personal Gmail address, a residential IP, a real workspace subdomain and one live Lago subscription id to a public package index. Replaced with RFC 2606/5737 reserved values. No credentials were ever committed. New `test_fixture_hygiene.py` fails if any of the three reappears; there is no capture script for the gateway fixtures to hang a scrub step off, so the guard is the durable fix. + ### Fixed - **An isolated batch could respin without limit, hammering the server that had just throttled us.** When a batch failed permanently (422) and every isolated send then failed transiently (429), `_send_individually` re-queued the survivors and `_run` continued straight into re-taking them with `_backoff_seconds` reset to 0 — no delay anywhere in the cycle. Measured against a server returning that exact pair: **280,388 HTTP requests in 1.2 seconds**. `_send_individually` now returns the number of events it re-queued, and the drain arms the normal 1→2→4→…→60s backoff whenever that count is non-zero, keeping the immediate-continue only for the case it was meant for (isolation fully resolved the batch, so the buffer genuinely shrank). The exit drain passes `requeue_transient=False` for the same reason: there is no later retry to re-queue *to*, so an event failing there is reported as lost rather than put back on a buffer the same loop immediately re-takes. Closes a divergence with the JS port, which has guarded this since its own queue fix. diff --git a/tests/unit/adapters/fixtures/bedrock/converse/COVERAGE.md b/tests/unit/adapters/fixtures/bedrock/converse/COVERAGE.md new file mode 100644 index 0000000..528be66 --- /dev/null +++ b/tests/unit/adapters/fixtures/bedrock/converse/COVERAGE.md @@ -0,0 +1,17 @@ +# Bedrock Converse — captured-model coverage + +Every model below was captured live and run through the adapter. 39 captures reduced to 12 committed fixtures: one per row of this table, chosen because the sweep tests assert DISPATCH and non-zero usage, so a second capture with the same key re-asserts the same thing. + +Recapture with the `capture*.py` / `capture*.ts` script in this tree; the sweep tests skip cleanly when the directory is absent, so a missing capture reads as "not covered" rather than as a pass. + +| Committed fixture | Distinguishing key | Models it stands for | +|---|---|---| +| `eu.amazon.nova-lite-v1_0.json` | usage shape + pricing provider `amazon` | `eu.amazon.nova-2-lite-v1:0`, `eu.amazon.nova-lite-v1:0`, `eu.amazon.nova-micro-v1:0`, `eu.amazon.nova-pro-v1:0` | +| `eu.anthropic.claude-opus-4-7.json` | usage shape + pricing provider `anthropic` | `eu.anthropic.claude-haiku-4-5-20251001-v1:0`, `eu.anthropic.claude-opus-4-5-20251101-v1:0`, `eu.anthropic.claude-opus-4-6-v1`, `eu.anthropic.claude-opus-4-7`, `eu.anthropic.claude-sonnet-4-5-20250929-v1:0`, `eu.anthropic.claude-sonnet-4-6` | +| `eu.mistral.pixtral-large-2502-v1_0.json` | usage shape + pricing provider `mistral` | `eu.mistral.pixtral-large-2502-v1:0`, `mistral.devstral-2-123b`, `mistral.magistral-small-2509`, `mistral.ministral-3-14b-instruct`, `mistral.ministral-3-3b-instruct`, `mistral.ministral-3-8b-instruct`, `mistral.mistral-7b-instruct-v0:2`, `mistral.mistral-large-2402-v1:0`, `mistral.mixtral-8x7b-instruct-v0:1`, `mistral.voxtral-mini-3b-2507`, `mistral.voxtral-small-24b-2507` | +| `google.gemma-3-12b-it.json` | usage shape + pricing provider `google` | `google.gemma-3-12b-it`, `google.gemma-3-27b-it`, `google.gemma-3-4b-it` | +| `minimax.minimax-m2.1.json` | usage shape + pricing provider `minimax` | `minimax.minimax-m2.1`, `minimax.minimax-m2.5`, `minimax.minimax-m2` | +| `nvidia.nemotron-nano-12b-v2.json` | usage shape + pricing provider `nvidia` | `nvidia.nemotron-nano-12b-v2`, `nvidia.nemotron-nano-3-30b`, `nvidia.nemotron-nano-9b-v2` | +| `openai.gpt-oss-20b-1_0.json` | usage shape + pricing provider `openai` | `openai.gpt-oss-120b-1:0`, `openai.gpt-oss-20b-1:0`, `openai.gpt-oss-safeguard-120b`, `openai.gpt-oss-safeguard-20b` | +| `qwen.qwen3-32b-v1_0.json` | usage shape + pricing provider `qwen` | `qwen.qwen3-32b-v1:0`, `qwen.qwen3-coder-30b-a3b-v1:0`, `qwen.qwen3-next-80b-a3b`, `qwen.qwen3-vl-235b-a22b` | +| `zai.glm-4.7-flash.json` | usage shape + pricing provider `zai` | `zai.glm-4.7-flash` | diff --git a/tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-2-lite-v1_0.json b/tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-2-lite-v1_0.json deleted file mode 100644 index a56284c..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-2-lite-v1_0.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "eu.amazon.nova-2-lite-v1:0", - "_response": { - "metrics": { - "latencyMs": 641 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication systems, and strong social bonds within their pods." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 51, - "outputTokens": 29, - "serverToolUsage": {}, - "totalTokens": 80 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-micro-v1_0.json b/tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-micro-v1_0.json deleted file mode 100644 index 3f6b140..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-micro-v1_0.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "eu.amazon.nova-micro-v1:0", - "_response": { - "metrics": { - "latencyMs": 609 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are intelligent marine mammals known for their playful behavior and complex social structures, often found in warm ocean waters around the world." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 5, - "outputTokens": 26, - "serverToolUsage": {}, - "totalTokens": 31 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-pro-v1_0.json b/tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-pro-v1_0.json deleted file mode 100644 index 4aae9ae..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/eu.amazon.nova-pro-v1_0.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "eu.amazon.nova-pro-v1:0", - "_response": { - "metrics": { - "latencyMs": 498 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 5, - "outputTokens": 17, - "serverToolUsage": {}, - "totalTokens": 22 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-haiku-4-5-20251001-v1_0.json b/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-haiku-4-5-20251001-v1_0.json deleted file mode 100644 index 5a8ac56..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-haiku-4-5-20251001-v1_0.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "_model_id": "eu.anthropic.claude-haiku-4-5-20251001-v1:0", - "_response": { - "metrics": { - "latencyMs": 1139 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication, and remarkable ability to navigate using echolocation." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "cacheReadInputTokenCount": 0, - "cacheReadInputTokens": 0, - "cacheWriteInputTokenCount": 0, - "cacheWriteInputTokens": 0, - "inputTokens": 12, - "outputTokens": 30, - "serverToolUsage": {}, - "totalTokens": 42 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-opus-4-5-20251101-v1_0.json b/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-opus-4-5-20251101-v1_0.json deleted file mode 100644 index f1829c5..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-opus-4-5-20251101-v1_0.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "_model_id": "eu.anthropic.claude-opus-4-5-20251101-v1:0", - "_response": { - "metrics": { - "latencyMs": 1860 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and remarkable ability to communicate using clicks and whistles." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "cacheReadInputTokenCount": 0, - "cacheReadInputTokens": 0, - "cacheWriteInputTokenCount": 0, - "cacheWriteInputTokens": 0, - "inputTokens": 12, - "outputTokens": 32, - "serverToolUsage": {}, - "totalTokens": 44 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-opus-4-6-v1.json b/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-opus-4-6-v1.json deleted file mode 100644 index 971e8e4..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-opus-4-6-v1.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "_model_id": "eu.anthropic.claude-opus-4-6-v1", - "_response": { - "metrics": { - "latencyMs": 1900 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and remarkable ability to communicate using a series of clicks and whistles." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "cacheReadInputTokenCount": 0, - "cacheReadInputTokens": 0, - "cacheWriteInputTokenCount": 0, - "cacheWriteInputTokens": 0, - "inputTokens": 12, - "outputTokens": 35, - "serverToolUsage": {}, - "totalTokens": 47 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-sonnet-4-5-20250929-v1_0.json b/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-sonnet-4-5-20250929-v1_0.json deleted file mode 100644 index 7dd805d..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/eu.anthropic.claude-sonnet-4-5-20250929-v1_0.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "_model_id": "eu.anthropic.claude-sonnet-4-5-20250929-v1:0", - "_response": { - "metrics": { - "latencyMs": 1627 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and remarkable ability to communicate using clicks, whistles, and body language." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "cacheReadInputTokenCount": 0, - "cacheReadInputTokens": 0, - "cacheWriteInputTokenCount": 0, - "cacheWriteInputTokens": 0, - "inputTokens": 12, - "outputTokens": 36, - "serverToolUsage": {}, - "totalTokens": 48 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/google.gemma-3-27b-it.json b/tests/unit/adapters/fixtures/bedrock/converse/google.gemma-3-27b-it.json deleted file mode 100644 index 1caa2fd..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/google.gemma-3-27b-it.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "google.gemma-3-27b-it", - "_response": { - "metrics": { - "latencyMs": 828 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication skills, and streamlined bodies perfectly adapted for life in the ocean.\n\n\n\n" - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 14, - "outputTokens": 31, - "serverToolUsage": {}, - "totalTokens": 45 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/google.gemma-3-4b-it.json b/tests/unit/adapters/fixtures/bedrock/converse/google.gemma-3-4b-it.json deleted file mode 100644 index 444f45c..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/google.gemma-3-4b-it.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "google.gemma-3-4b-it", - "_response": { - "metrics": { - "latencyMs": 404 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 14, - "outputTokens": 19, - "serverToolUsage": {}, - "totalTokens": 33 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/minimax.minimax-m2.5.json b/tests/unit/adapters/fixtures/bedrock/converse/minimax.minimax-m2.5.json deleted file mode 100644 index ad71fdb..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/minimax.minimax-m2.5.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "_model_id": "minimax.minimax-m2.5", - "_response": { - "metrics": { - "latencyMs": 30251 - }, - "output": { - "message": { - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "The user asked: \"One sentence about dolphins.\" This is a straightforward request. There's no policy conflict. I can comply by providing one sentence describing dolphins. Use proper grammar, keep to one sentence." - } - } - } - ], - "role": "assistant" - } - }, - "stopReason": "max_tokens", - "usage": { - "inputTokens": 43, - "outputTokens": 40, - "serverToolUsage": {}, - "totalTokens": 83 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/minimax.minimax-m2.json b/tests/unit/adapters/fixtures/bedrock/converse/minimax.minimax-m2.json deleted file mode 100644 index af4dfb0..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/minimax.minimax-m2.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "_model_id": "minimax.minimax-m2", - "_response": { - "metrics": { - "latencyMs": 409 - }, - "output": { - "message": { - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "The user explicitly requests \"One sentence about dolphins.\" I need to comply and keep it concise, so I'll deliver a single sentence. Options for this sentence could include:\n\n- \"Dolphins are highly" - } - } - } - ], - "role": "assistant" - } - }, - "stopReason": "max_tokens", - "usage": { - "inputTokens": 27, - "outputTokens": 40, - "serverToolUsage": {}, - "totalTokens": 67 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/mistral.devstral-2-123b.json b/tests/unit/adapters/fixtures/bedrock/converse/mistral.devstral-2-123b.json deleted file mode 100644 index ade320b..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/mistral.devstral-2-123b.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "mistral.devstral-2-123b", - "_response": { - "metrics": { - "latencyMs": 586 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 9, - "outputTokens": 19, - "serverToolUsage": {}, - "totalTokens": 28 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/mistral.magistral-small-2509.json b/tests/unit/adapters/fixtures/bedrock/converse/mistral.magistral-small-2509.json deleted file mode 100644 index 41b4b4c..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/mistral.magistral-small-2509.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "mistral.magistral-small-2509", - "_response": { - "metrics": { - "latencyMs": 709 - }, - "output": { - "message": { - "content": [ - { - "text": "\"Dolphins are highly intelligent marine mammals known for their playful behavior, sophisticated communication, and strong social bonds.\"" - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 9, - "outputTokens": 24, - "serverToolUsage": {}, - "totalTokens": 33 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-14b-instruct.json b/tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-14b-instruct.json deleted file mode 100644 index bba1161..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-14b-instruct.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "mistral.ministral-3-14b-instruct", - "_response": { - "metrics": { - "latencyMs": 324 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication skills, and strong social bonds within their pods." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 9, - "outputTokens": 27, - "serverToolUsage": {}, - "totalTokens": 36 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-3b-instruct.json b/tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-3b-instruct.json deleted file mode 100644 index 6c6eb57..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-3b-instruct.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "mistral.ministral-3-3b-instruct", - "_response": { - "metrics": { - "latencyMs": 385 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, sophisticated communication, and strong social bonds." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 9, - "outputTokens": 23, - "serverToolUsage": {}, - "totalTokens": 32 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-8b-instruct.json b/tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-8b-instruct.json deleted file mode 100644 index 9929204..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/mistral.ministral-3-8b-instruct.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "mistral.ministral-3-8b-instruct", - "_response": { - "metrics": { - "latencyMs": 431 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent, social marine mammals known for their complex communication, playful behavior, and advanced problem-solving skills." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 9, - "outputTokens": 26, - "serverToolUsage": {}, - "totalTokens": 35 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/mistral.mistral-7b-instruct-v0_2.json b/tests/unit/adapters/fixtures/bedrock/converse/mistral.mistral-7b-instruct-v0_2.json deleted file mode 100644 index ed41399..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/mistral.mistral-7b-instruct-v0_2.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "mistral.mistral-7b-instruct-v0:2", - "_response": { - "metrics": { - "latencyMs": 469 - }, - "output": { - "message": { - "content": [ - { - "text": " Dolphins are highly intelligent, social marine mammals known for their distinctive curved dorsal fins and playful behavior in the wild." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 15, - "outputTokens": 30, - "serverToolUsage": {}, - "totalTokens": 45 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/mistral.mixtral-8x7b-instruct-v0_1.json b/tests/unit/adapters/fixtures/bedrock/converse/mistral.mixtral-8x7b-instruct-v0_1.json deleted file mode 100644 index f4225bf..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/mistral.mixtral-8x7b-instruct-v0_1.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "mistral.mixtral-8x7b-instruct-v0:1", - "_response": { - "metrics": { - "latencyMs": 584 - }, - "output": { - "message": { - "content": [ - { - "text": " Dolphins are highly intelligent and social marine mammals known for their playful behavior, sophisticated communication skills, and ability to perform acrobatic displays in the water." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 15, - "outputTokens": 35, - "serverToolUsage": {}, - "totalTokens": 50 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/mistral.voxtral-mini-3b-2507.json b/tests/unit/adapters/fixtures/bedrock/converse/mistral.voxtral-mini-3b-2507.json deleted file mode 100644 index ce4ab04..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/mistral.voxtral-mini-3b-2507.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "mistral.voxtral-mini-3b-2507", - "_response": { - "metrics": { - "latencyMs": 170 - }, - "output": { - "message": { - "content": [ - { - "text": "They are declining due to stressors such as pollution, climate change, and entanglement." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 9, - "outputTokens": 19, - "serverToolUsage": {}, - "totalTokens": 28 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/mistral.voxtral-small-24b-2507.json b/tests/unit/adapters/fixtures/bedrock/converse/mistral.voxtral-small-24b-2507.json deleted file mode 100644 index ba5a616..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/mistral.voxtral-small-24b-2507.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "mistral.voxtral-small-24b-2507", - "_response": { - "metrics": { - "latencyMs": 697 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their complex communication system, playful behavior, and exceptional swimming skills." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 9, - "outputTokens": 24, - "serverToolUsage": {}, - "totalTokens": 33 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/nvidia.nemotron-nano-3-30b.json b/tests/unit/adapters/fixtures/bedrock/converse/nvidia.nemotron-nano-3-30b.json deleted file mode 100644 index 33fa531..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/nvidia.nemotron-nano-3-30b.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "nvidia.nemotron-nano-3-30b", - "_response": { - "metrics": { - "latencyMs": 418 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals that use echolocation and complex vocalizations to communicate and hunt." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 22, - "outputTokens": 22, - "serverToolUsage": {}, - "totalTokens": 44 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/nvidia.nemotron-nano-9b-v2.json b/tests/unit/adapters/fixtures/bedrock/converse/nvidia.nemotron-nano-9b-v2.json deleted file mode 100644 index 9b65312..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/nvidia.nemotron-nano-9b-v2.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "nvidia.nemotron-nano-9b-v2", - "_response": { - "metrics": { - "latencyMs": 425 - }, - "output": { - "message": { - "content": [ - { - "text": "Okay, the user wants a single sentence about dolphins. Let me start by thinking about what defines dolphins. They're marine mammals, right? So I should mention that. Also, they're" - } - ], - "role": "assistant" - } - }, - "stopReason": "max_tokens", - "usage": { - "inputTokens": 18, - "outputTokens": 40, - "serverToolUsage": {}, - "totalTokens": 58 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/openai.gpt-oss-120b-1_0.json b/tests/unit/adapters/fixtures/bedrock/converse/openai.gpt-oss-120b-1_0.json deleted file mode 100644 index e927c72..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/openai.gpt-oss-120b-1_0.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "_model_id": "openai.gpt-oss-120b-1:0", - "_response": { - "metrics": { - "latencyMs": 524 - }, - "output": { - "message": { - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "The user asks: \"One sentence about dolphins.\" So just one sentence. Need to respond concisely with a single sentence about dolphins. Could be factual. Provide a single sentence." - } - } - } - ], - "role": "assistant" - } - }, - "stopReason": "max_tokens", - "usage": { - "inputTokens": 72, - "outputTokens": 40, - "serverToolUsage": {}, - "totalTokens": 112 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/openai.gpt-oss-safeguard-120b.json b/tests/unit/adapters/fixtures/bedrock/converse/openai.gpt-oss-safeguard-120b.json deleted file mode 100644 index 6fa0b17..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/openai.gpt-oss-safeguard-120b.json +++ /dev/null @@ -1,29 +0,0 @@ -{ - "_model_id": "openai.gpt-oss-safeguard-120b", - "_response": { - "metrics": { - "latencyMs": 305 - }, - "output": { - "message": { - "content": [ - { - "reasoningContent": { - "reasoningText": { - "text": "The task: The user says \"One sentence about dolphins.\" So produce a single sentence about dolphins. That's straightforward. Just a sentence. Could be interesting. Provide a sentence.\n\nMake sure" - } - } - } - ], - "role": "assistant" - } - }, - "stopReason": "max_tokens", - "usage": { - "inputTokens": 72, - "outputTokens": 40, - "serverToolUsage": {}, - "totalTokens": 112 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-coder-30b-a3b-v1_0.json b/tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-coder-30b-a3b-v1_0.json deleted file mode 100644 index 4741e38..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-coder-30b-a3b-v1_0.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "qwen.qwen3-coder-30b-a3b-v1:0", - "_response": { - "metrics": { - "latencyMs": 408 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their complex social structures, echolocation abilities, and playful behavior." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 13, - "outputTokens": 24, - "serverToolUsage": {}, - "totalTokens": 37 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-next-80b-a3b.json b/tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-next-80b-a3b.json deleted file mode 100644 index 122c7ca..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-next-80b-a3b.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "qwen.qwen3-next-80b-a3b", - "_response": { - "metrics": { - "latencyMs": 605 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent, social marine mammals known for their playful behavior, complex communication, and remarkable echolocation abilities." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 13, - "outputTokens": 26, - "serverToolUsage": {}, - "totalTokens": 39 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-vl-235b-a22b.json b/tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-vl-235b-a22b.json deleted file mode 100644 index fe6b08a..0000000 --- a/tests/unit/adapters/fixtures/bedrock/converse/qwen.qwen3-vl-235b-a22b.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "_model_id": "qwen.qwen3-vl-235b-a22b", - "_response": { - "metrics": { - "latencyMs": 684 - }, - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent, social marine mammals known for their playful behavior, echolocation abilities, and complex communication." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 13, - "outputTokens": 25, - "serverToolUsage": {}, - "totalTokens": 38 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/COVERAGE.md b/tests/unit/adapters/fixtures/bedrock/invoke/COVERAGE.md new file mode 100644 index 0000000..3b4fb66 --- /dev/null +++ b/tests/unit/adapters/fixtures/bedrock/invoke/COVERAGE.md @@ -0,0 +1,21 @@ +# Bedrock InvokeModel — captured-model coverage + +Every model below was captured live and run through the adapter. 39 captures reduced to 14 committed fixtures: one per row of this table, chosen because the sweep tests assert DISPATCH and non-zero usage, so a second capture with the same key re-asserts the same thing. + +Recapture with the `capture*.py` / `capture*.ts` script in this tree; the sweep tests skip cleanly when the directory is absent, so a missing capture reads as "not covered" rather than as a pass. + +| Committed fixture | Distinguishing key | Models it stands for | +|---|---|---| +| `eu.amazon.nova-lite-v1_0.json` | `nova` family + pricing provider `amazon` | `eu.amazon.nova-2-lite-v1:0`, `eu.amazon.nova-lite-v1:0`, `eu.amazon.nova-micro-v1:0`, `eu.amazon.nova-pro-v1:0` | +| `eu.anthropic.claude-sonnet-4-6.json` | `anthropic` family + pricing provider `anthropic` | `eu.anthropic.claude-haiku-4-5-20251001-v1:0`, `eu.anthropic.claude-opus-4-5-20251101-v1:0`, `eu.anthropic.claude-opus-4-6-v1`, `eu.anthropic.claude-sonnet-4-5-20250929-v1:0`, `eu.anthropic.claude-sonnet-4-6` | +| `eu.anthropic.claude-opus-4-7.json` | `opus_4_7` family + pricing provider `anthropic` | `eu.anthropic.claude-opus-4-7` | +| `eu.mistral.pixtral-large-2502-v1_0.json` | `pixtral` family + pricing provider `mistral` | `eu.mistral.pixtral-large-2502-v1:0` | +| `google.gemma-3-12b-it.json` | `openai_compat_basic` family + pricing provider `google` | `google.gemma-3-12b-it`, `google.gemma-3-27b-it`, `google.gemma-3-4b-it` | +| `minimax.minimax-m2.1.json` | `openai_compat_with_details` family + pricing provider `minimax` | `minimax.minimax-m2.1`, `minimax.minimax-m2.5`, `minimax.minimax-m2` | +| `mistral.devstral-2-123b.json` | `openai_compat_basic` family + pricing provider `mistral` | `mistral.devstral-2-123b`, `mistral.magistral-small-2509`, `mistral.ministral-3-14b-instruct`, `mistral.ministral-3-3b-instruct`, `mistral.ministral-3-8b-instruct`, `mistral.voxtral-mini-3b-2507`, `mistral.voxtral-small-24b-2507` | +| `mistral.mistral-large-2402-v1_0.json` | `mistral_legacy` family + pricing provider `mistral` | `mistral.mistral-7b-instruct-v0:2`, `mistral.mistral-large-2402-v1:0`, `mistral.mixtral-8x7b-instruct-v0:1` | +| `nvidia.nemotron-nano-12b-v2.json` | `openai_compat_basic` family + pricing provider `nvidia` | `nvidia.nemotron-nano-12b-v2`, `nvidia.nemotron-nano-3-30b`, `nvidia.nemotron-nano-9b-v2` | +| `openai.gpt-oss-20b-1_0.json` | `openai_compat_basic` family + pricing provider `openai` | `openai.gpt-oss-120b-1:0`, `openai.gpt-oss-20b-1:0` | +| `openai.gpt-oss-safeguard-120b.json` | `openai_compat_with_details` family + pricing provider `openai` | `openai.gpt-oss-safeguard-120b`, `openai.gpt-oss-safeguard-20b` | +| `qwen.qwen3-32b-v1_0.json` | `openai_compat_basic` family + pricing provider `qwen` | `qwen.qwen3-32b-v1:0`, `qwen.qwen3-coder-30b-a3b-v1:0`, `qwen.qwen3-next-80b-a3b`, `qwen.qwen3-vl-235b-a22b` | +| `zai.glm-4.7-flash.json` | `openai_compat_basic` family + pricing provider `zai` | `zai.glm-4.7-flash` | diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-2-lite-v1_0.json b/tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-2-lite-v1_0.json deleted file mode 100644 index a889d24..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-2-lite-v1_0.json +++ /dev/null @@ -1,23 +0,0 @@ -{ - "_model_id": "eu.amazon.nova-2-lite-v1:0", - "_response": { - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication systems, and strong social bonds within their pods." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 51, - "outputTokens": 29, - "totalTokens": 80, - "cacheReadInputTokenCount": 0, - "cacheWriteInputTokenCount": 0 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-micro-v1_0.json b/tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-micro-v1_0.json deleted file mode 100644 index cfea521..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-micro-v1_0.json +++ /dev/null @@ -1,23 +0,0 @@ -{ - "_model_id": "eu.amazon.nova-micro-v1:0", - "_response": { - "output": { - "message": { - "content": [ - { - "text": "Dolphins are intelligent marine mammals known for their playful behavior, complex social structures, and remarkable ability to communicate using a variety of vocalizations." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 5, - "outputTokens": 29, - "totalTokens": 34, - "cacheReadInputTokenCount": 0, - "cacheWriteInputTokenCount": 0 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-pro-v1_0.json b/tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-pro-v1_0.json deleted file mode 100644 index 70dedd1..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/eu.amazon.nova-pro-v1_0.json +++ /dev/null @@ -1,23 +0,0 @@ -{ - "_model_id": "eu.amazon.nova-pro-v1:0", - "_response": { - "output": { - "message": { - "content": [ - { - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures." - } - ], - "role": "assistant" - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 5, - "outputTokens": 17, - "totalTokens": 22, - "cacheReadInputTokenCount": 0, - "cacheWriteInputTokenCount": 0 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-haiku-4-5-20251001-v1_0.json b/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-haiku-4-5-20251001-v1_0.json deleted file mode 100644 index 9f0b53a..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-haiku-4-5-20251001-v1_0.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "eu.anthropic.claude-haiku-4-5-20251001-v1:0", - "_response": { - "model": "claude-haiku-4-5-20251001", - "id": "msg_bdrk_01JL3rpyU2r2cCoVr7PoNKvF", - "type": "message", - "role": "assistant", - "content": [ - { - "type": "text", - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication, and remarkable ability to navigate using echolocation." - } - ], - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { - "input_tokens": 12, - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "cache_creation": { - "ephemeral_5m_input_tokens": 0, - "ephemeral_1h_input_tokens": 0 - }, - "output_tokens": 30 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-opus-4-5-20251101-v1_0.json b/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-opus-4-5-20251101-v1_0.json deleted file mode 100644 index 35358c1..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-opus-4-5-20251101-v1_0.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "eu.anthropic.claude-opus-4-5-20251101-v1:0", - "_response": { - "model": "claude-opus-4-5-20251101", - "id": "msg_bdrk_01564FnWiTr7na2Ggi44vYTA", - "type": "message", - "role": "assistant", - "content": [ - { - "type": "text", - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior and sophisticated communication through clicks and whistles." - } - ], - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { - "input_tokens": 12, - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "cache_creation": { - "ephemeral_5m_input_tokens": 0, - "ephemeral_1h_input_tokens": 0 - }, - "output_tokens": 25 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-opus-4-6-v1.json b/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-opus-4-6-v1.json deleted file mode 100644 index a417b78..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-opus-4-6-v1.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "eu.anthropic.claude-opus-4-6-v1", - "_response": { - "model": "claude-opus-4-6", - "id": "msg_bdrk_01PweEaYn9JEkuexQRnJFuq6", - "type": "message", - "role": "assistant", - "content": [ - { - "type": "text", - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and remarkable ability to communicate using a series of clicks and whistles." - } - ], - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { - "input_tokens": 12, - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "cache_creation": { - "ephemeral_5m_input_tokens": 0, - "ephemeral_1h_input_tokens": 0 - }, - "output_tokens": 35 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-sonnet-4-5-20250929-v1_0.json b/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-sonnet-4-5-20250929-v1_0.json deleted file mode 100644 index 0e13c3f..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/eu.anthropic.claude-sonnet-4-5-20250929-v1_0.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "eu.anthropic.claude-sonnet-4-5-20250929-v1:0", - "_response": { - "model": "claude-sonnet-4-5-20250929", - "id": "msg_bdrk_01J5rmTpPAXG4ne7cw62JmkN", - "type": "message", - "role": "assistant", - "content": [ - { - "type": "text", - "text": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and remarkable communication abilities using clicks and whistles." - } - ], - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { - "input_tokens": 12, - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "cache_creation": { - "ephemeral_5m_input_tokens": 0, - "ephemeral_1h_input_tokens": 0 - }, - "output_tokens": 31 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/google.gemma-3-27b-it.json b/tests/unit/adapters/fixtures/bedrock/invoke/google.gemma-3-27b-it.json deleted file mode 100644 index 6fce6e9..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/google.gemma-3-27b-it.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "google.gemma-3-27b-it", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication skills, and remarkable ability to navigate and hunt using echolocation.\n\n\n\n", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395243, - "id": "chatcmpl-f574e827-4334-47cb-9843-c16785bee2fd", - "model": "google.gemma-3-27b-it", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 31, - "prompt_tokens": 14, - "total_tokens": 45 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/google.gemma-3-4b-it.json b/tests/unit/adapters/fixtures/bedrock/invoke/google.gemma-3-4b-it.json deleted file mode 100644 index 4fb0311..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/google.gemma-3-4b-it.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "google.gemma-3-4b-it", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex communication skills.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395245, - "id": "chatcmpl-4a431d4b-7c2e-4f1c-b01a-1716184c1f77", - "model": "google.gemma-3-4b-it", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 19, - "prompt_tokens": 14, - "total_tokens": 33 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/minimax.minimax-m2.5.json b/tests/unit/adapters/fixtures/bedrock/invoke/minimax.minimax-m2.5.json deleted file mode 100644 index d91a0a8..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/minimax.minimax-m2.5.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "minimax.minimax-m2.5", - "_response": { - "choices": [ - { - "finish_reason": "length", - "index": 0, - "logprobs": null, - "message": { - "content": "We need to respond to the user: \"One sentence about dolphins.\" So just one sentence describing dolphins, perhaps something like: \"Dolphins are highly intelligent marine mammals known for their playful behavior,", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395248, - "id": "chatcmpl-ed00ef62-8104-4f49-847b-429e042f1dc2", - "model": "minimax.minimax-m2.5", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 40, - "prompt_tokens": 43, - "total_tokens": 83 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/minimax.minimax-m2.json b/tests/unit/adapters/fixtures/bedrock/invoke/minimax.minimax-m2.json deleted file mode 100644 index b8d1e03..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/minimax.minimax-m2.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "minimax.minimax-m2", - "_response": { - "choices": [ - { - "finish_reason": "length", - "index": 0, - "logprobs": null, - "message": { - "content": "The user asks \"One sentence about dolphins.\" I need to produce one sentence describing dolphins. It needs to be concise. The user's request is straightforward. No disallowed content. So I'll output a sentence.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395245, - "id": "chatcmpl-aee04882-5d78-4a5f-a71b-86141ec194b8", - "model": "minimax.minimax-m2", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 40, - "prompt_tokens": 27, - "total_tokens": 67 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.magistral-small-2509.json b/tests/unit/adapters/fixtures/bedrock/invoke/mistral.magistral-small-2509.json deleted file mode 100644 index 4f27e7e..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.magistral-small-2509.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "mistral.magistral-small-2509", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "\"Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication, and strong social bonds.\"", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395267, - "id": "chatcmpl-b8db991b-fcb4-4123-8170-aacdac8a5a17", - "model": "mistral.magistral-small-2509", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 24, - "prompt_tokens": 9, - "total_tokens": 33 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-14b-instruct.json b/tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-14b-instruct.json deleted file mode 100644 index 22d8249..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-14b-instruct.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "mistral.ministral-3-14b-instruct", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are highly intelligent marine mammals known for their playful nature, complex social behaviors, and powerful ultrasonic communication.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395268, - "id": "chatcmpl-f351f8ff-f65a-47db-86bc-a5fa62334d07", - "model": "mistral.ministral-3-14b-instruct", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 24, - "prompt_tokens": 9, - "total_tokens": 33 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-3b-instruct.json b/tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-3b-instruct.json deleted file mode 100644 index 499efa4..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-3b-instruct.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "mistral.ministral-3-3b-instruct", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are intelligent, playful marine mammals known for their razor-sharp mental prowess, iconic playful behavior, and strong social bonds.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395269, - "id": "chatcmpl-72f25314-ab13-4e05-9243-20181b0f385f", - "model": "mistral.ministral-3-3b-instruct", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 29, - "prompt_tokens": 9, - "total_tokens": 38 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-8b-instruct.json b/tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-8b-instruct.json deleted file mode 100644 index 61e0595..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.ministral-3-8b-instruct.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "mistral.ministral-3-8b-instruct", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins, highly intelligent mammals, are known for their complex social behaviors, sonar-based echolocation, and cooperative hunting strategies in the ocean.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395269, - "id": "chatcmpl-55ba26f7-e719-4069-8fd4-91d510b7fc81", - "model": "mistral.ministral-3-8b-instruct", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 30, - "prompt_tokens": 9, - "total_tokens": 39 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.mistral-7b-instruct-v0_2.json b/tests/unit/adapters/fixtures/bedrock/invoke/mistral.mistral-7b-instruct-v0_2.json deleted file mode 100644 index 99ca7cb..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.mistral-7b-instruct-v0_2.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "_model_id": "mistral.mistral-7b-instruct-v0:2", - "_response": { - "outputs": [ - { - "text": " Dolphins are highly intelligent, social marine mammals known for their acrobatic displays and echolocation abilities, which they use for communication and hunting.", - "stop_reason": "stop" - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.mixtral-8x7b-instruct-v0_1.json b/tests/unit/adapters/fixtures/bedrock/invoke/mistral.mixtral-8x7b-instruct-v0_1.json deleted file mode 100644 index 4afbc8c..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.mixtral-8x7b-instruct-v0_1.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "_model_id": "mistral.mixtral-8x7b-instruct-v0:1", - "_response": { - "outputs": [ - { - "text": " Dolphins are highly intelligent and social marine mammals known for their playful behavior, complex communication skills, and ability to perform acrobatic displays in the water.", - "stop_reason": "stop" - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.voxtral-mini-3b-2507.json b/tests/unit/adapters/fixtures/bedrock/invoke/mistral.voxtral-mini-3b-2507.json deleted file mode 100644 index 7d70c70..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.voxtral-mini-3b-2507.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "mistral.voxtral-mini-3b-2507", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are intelligent marine mammals known for their playful behavior and sophisticated communication skills.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395273, - "id": "chatcmpl-8887a18a-59b8-4496-98a8-6eb6074af24e", - "model": "mistral.voxtral-mini-3b-2507", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 18, - "prompt_tokens": 9, - "total_tokens": 27 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.voxtral-small-24b-2507.json b/tests/unit/adapters/fixtures/bedrock/invoke/mistral.voxtral-small-24b-2507.json deleted file mode 100644 index 6d37875..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/mistral.voxtral-small-24b-2507.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "mistral.voxtral-small-24b-2507", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are highly intelligent marine mammals known for their complex behaviors and communication abilities.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395274, - "id": "chatcmpl-1e203e42-f066-435a-a628-fa74aad7904a", - "model": "mistral.voxtral-small-24b-2507", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 18, - "prompt_tokens": 9, - "total_tokens": 27 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/nvidia.nemotron-nano-3-30b.json b/tests/unit/adapters/fixtures/bedrock/invoke/nvidia.nemotron-nano-3-30b.json deleted file mode 100644 index eb8f682..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/nvidia.nemotron-nano-3-30b.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "nvidia.nemotron-nano-3-30b", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are highly social, intelligent marine mammals that communicate through clicks, whistles, and body language.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395276, - "id": "chatcmpl-40b86f35-0b35-4b50-9a29-193b9057a037", - "model": "nvidia.nemotron-nano-3-30b", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 24, - "prompt_tokens": 22, - "total_tokens": 46 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/nvidia.nemotron-nano-9b-v2.json b/tests/unit/adapters/fixtures/bedrock/invoke/nvidia.nemotron-nano-9b-v2.json deleted file mode 100644 index 6b8b151..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/nvidia.nemotron-nano-9b-v2.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "nvidia.nemotron-nano-9b-v2", - "_response": { - "choices": [ - { - "finish_reason": "length", - "index": 0, - "logprobs": null, - "message": { - "content": "Okay, the user wants one sentence about dolphins. Let me start by thinking about the key points related to dolphins. They are marine mammals, right? So maybe mention their habitat. Also,", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395276, - "id": "chatcmpl-e143d94c-ea90-47f3-a26e-e85af62d5fee", - "model": "nvidia.nemotron-nano-9b-v2", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 40, - "prompt_tokens": 18, - "total_tokens": 58 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/openai.gpt-oss-120b-1_0.json b/tests/unit/adapters/fixtures/bedrock/invoke/openai.gpt-oss-120b-1_0.json deleted file mode 100644 index 9b226f2..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/openai.gpt-oss-120b-1_0.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "openai.gpt-oss-120b-1:0", - "_response": { - "choices": [ - { - "finish_reason": "length", - "index": 0, - "logprobs": null, - "message": { - "content": "The user asks: \"One sentence about dolphins.\" Provide a single sentence about dolphins. Should be concise. Provide a fact or description. Ensure it's one sentence.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395337, - "id": "chatcmpl-45d0b931-b7d6-4be4-8464-a9f40214e1b2", - "model": "openai.gpt-oss-120b-1:0", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 40, - "prompt_tokens": 72, - "total_tokens": 112 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-coder-30b-a3b-v1_0.json b/tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-coder-30b-a3b-v1_0.json deleted file mode 100644 index 61860c3..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-coder-30b-a3b-v1_0.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "qwen.qwen3-coder-30b-a3b-v1:0", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication systems, and sophisticated echolocation abilities that help them navigate and hunt in the ocean.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395341, - "id": "chatcmpl-ec888cb2-eadd-40a1-ab2c-88d6cce2726a", - "model": "qwen.qwen3-coder-30b-a3b-v1:0", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 34, - "prompt_tokens": 13, - "total_tokens": 47 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-next-80b-a3b.json b/tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-next-80b-a3b.json deleted file mode 100644 index c983257..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-next-80b-a3b.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "qwen.qwen3-next-80b-a3b", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are highly intelligent, social marine mammals known for their playful behavior, complex communication, and echolocation abilities.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395342, - "id": "chatcmpl-0b9097ca-ca35-406d-9ae2-0fc0f4c2965c", - "model": "qwen.qwen3-next-80b-a3b", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 25, - "prompt_tokens": 13, - "total_tokens": 38 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-vl-235b-a22b.json b/tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-vl-235b-a22b.json deleted file mode 100644 index fdd7f97..0000000 --- a/tests/unit/adapters/fixtures/bedrock/invoke/qwen.qwen3-vl-235b-a22b.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "_model_id": "qwen.qwen3-vl-235b-a22b", - "_response": { - "choices": [ - { - "finish_reason": "stop", - "index": 0, - "logprobs": null, - "message": { - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and ability to communicate using a variety of clicks, whistles, and body movements.", - "refusal": null, - "role": "assistant" - } - } - ], - "created": 1777395343, - "id": "chatcmpl-615e4fba-b8b3-4e78-a2fb-40cb3de8e3d6", - "model": "qwen.qwen3-vl-235b-a22b", - "object": "chat.completion", - "service_tier": "default", - "usage": { - "completion_tokens": 36, - "prompt_tokens": 13, - "total_tokens": 49 - } - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/COVERAGE.md b/tests/unit/adapters/fixtures/mistral_native/all_models/COVERAGE.md new file mode 100644 index 0000000..9dfc85e --- /dev/null +++ b/tests/unit/adapters/fixtures/mistral_native/all_models/COVERAGE.md @@ -0,0 +1,20 @@ +# Mistral native — captured-model coverage + +Every model below was captured live and run through the adapter. 73 captures reduced to 12 committed fixtures: one per row of this table, chosen because the sweep tests assert DISPATCH and non-zero usage, so a second capture with the same key re-asserts the same thing. + +Recapture with the `capture*.py` / `capture*.ts` script in this tree; the sweep tests skip cleanly when the directory is absent, so a missing capture reads as "not covered" rather than as a pass. + +| Committed fixture | Distinguishing key | Models it stands for | +|---|---|---| +| `codestral-2508.json` | family `codestral` | `codestral-2508`, `codestral-latest` | +| `devstral-2512.json` | family `devstral` | `devstral-2512`, `devstral-latest`, `devstral-medium-2507`, `devstral-medium-latest`, `devstral-small-2507` | +| `magistral-medium-2509.json` | family `magistral` | `magistral-medium-2509`, `magistral-medium-latest`, `magistral-small-2509`, `magistral-small-latest` | +| `magistral-medium-2509__vision.json` | family `magistral` + vision call | `magistral-medium-2509`, `magistral-medium-latest`, `magistral-small-2509`, `magistral-small-latest` | +| `ministral-14b-2512.json` | family `ministral` | `ministral-14b-2512`, `ministral-14b-latest`, `ministral-3b-2512`, `ministral-3b-latest`, `ministral-8b-2512`, `ministral-8b-latest` | +| `ministral-14b-2512__vision.json` | family `ministral` + vision call | `ministral-14b-2512`, `ministral-14b-latest`, `ministral-3b-2512`, `ministral-3b-latest`, `ministral-8b-2512`, `ministral-8b-latest` | +| `mistral-large-2411.json` | family `mistral` | `mistral-large-2411`, `mistral-large-2512`, `mistral-large-latest`, `mistral-large-pixtral-2411`, `mistral-medium-2505`, `mistral-medium-2508`, `mistral-medium-2604`, `mistral-medium-3-5`, `mistral-medium-3.5`, `mistral-medium-3`, `mistral-medium-c21211-r0-75`, `mistral-medium-latest`, `mistral-medium`, `mistral-small-2506`, `mistral-small-2603`, `mistral-small-latest`, `mistral-tiny-2407`, `mistral-tiny-latest`, `mistral-vibe-cli-fast`, `mistral-vibe-cli-latest`, `mistral-vibe-cli-with-tools` | +| `mistral-large-2512__vision.json` | family `mistral` + vision call | `mistral-large-2512`, `mistral-large-latest`, `mistral-medium-2505`, `mistral-medium-2508`, `mistral-medium-2604`, `mistral-medium-3-5`, `mistral-medium-3.5`, `mistral-medium-3`, `mistral-medium-c21211-r0-75`, `mistral-medium-latest`, `mistral-medium`, `mistral-small-2506`, `mistral-small-2603`, `mistral-small-latest`, `mistral-vibe-cli-fast`, `mistral-vibe-cli-latest`, `mistral-vibe-cli-with-tools` | +| `open-mistral-nemo-2407.json` | family `open` | `open-mistral-nemo-2407`, `open-mistral-nemo` | +| `pixtral-large-2411.json` | family `pixtral` | `pixtral-large-2411` | +| `pixtral-large-latest__vision.json` | family `pixtral` + vision call | `pixtral-large-latest` | +| `voxtral-mini-2507.json` | family `voxtral` | `voxtral-mini-2507`, `voxtral-mini-latest`, `voxtral-small-2507`, `voxtral-small-latest` | diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/codestral-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/codestral-latest.json deleted file mode 100644 index d20dd6f..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/codestral-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "codestral-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": true, - "fine_tuning": false, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "0472bd6f552d4bd1b7adb62713bd2230", - "object": "chat.completion", - "model": "codestral-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 24, - "total_tokens": 34, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985803, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their complex social structures, playful behavior, and advanced communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-latest.json deleted file mode 100644 index cdb83b5..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "devstral-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "69764576a7f64714bb6dbed590782d7a", - "object": "chat.completion", - "model": "devstral-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 24, - "total_tokens": 34, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985805, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and remarkable communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-medium-2507.json b/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-medium-2507.json deleted file mode 100644 index 8346f49..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-medium-2507.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "devstral-medium-2507", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "dbaaaf4b379441cebb9360f7e60b50a2", - "object": "chat.completion", - "model": "devstral-medium-2507", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 19, - "total_tokens": 29, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985863, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and exceptional communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-medium-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-medium-latest.json deleted file mode 100644 index 8ac40de..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-medium-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "devstral-medium-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "7c1ff459703547d0ab485d874decc182", - "object": "chat.completion", - "model": "devstral-medium-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 24, - "total_tokens": 34, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985805, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and remarkable communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-small-2507.json b/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-small-2507.json deleted file mode 100644 index 1519a32..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/devstral-small-2507.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "devstral-small-2507", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "49b56cf6e77e40b28f531f3861a24c2d", - "object": "chat.completion", - "model": "devstral-small-2507", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 19, - "total_tokens": 29, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985861, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-medium-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-medium-latest.json deleted file mode 100644 index 5d566cf..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-medium-latest.json +++ /dev/null @@ -1,55 +0,0 @@ -{ - "_model_id": "magistral-medium-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "d3564afe820340acba2a282db61ed58c", - "object": "chat.completion", - "model": "magistral-medium-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 40, - "total_tokens": 50, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985814, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": [ - { - "thinking": [ - { - "text": "Okay, I know that dolphins are marine mammals known for their intelligence and playful behavior. Let me draft a sentence that captures this. \"Dolphins are highly intelligent marine mammals known for their", - "type": "text" - } - ], - "type": "thinking", - "closed": true - } - ], - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-medium-latest__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-medium-latest__vision.json deleted file mode 100644 index 9aa2b31..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-medium-latest__vision.json +++ /dev/null @@ -1,56 +0,0 @@ -{ - "_model_id": "magistral-medium-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "dd66a1940d4e41d9aa8c756ee225bedf", - "object": "chat.completion", - "model": "magistral-medium-latest", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 40, - "total_tokens": 51, - "prompt_tokens_details": { - "cached_tokens": 5 - } - }, - "created": 1777985815, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": [ - { - "thinking": [ - { - "text": "Okay, the user is asking about the contents of an image. Since I don't have the ability to view or analyze images, I need to let the user know that I can't assist with", - "type": "text" - } - ], - "type": "thinking", - "closed": true - } - ], - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-2509.json b/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-2509.json deleted file mode 100644 index a04e91a..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-2509.json +++ /dev/null @@ -1,55 +0,0 @@ -{ - "_model_id": "magistral-small-2509", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "12f41d9f385241f0b34d3749eabf2a5b", - "object": "chat.completion", - "model": "magistral-small-2509", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 40, - "total_tokens": 50, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985866, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": [ - { - "thinking": [ - { - "text": "Alright, the user has asked to write one sentence about dolphins. Let me think about what interesting fact or description I can provide in a concise manner.\n\nDolphins are marine mammals known", - "type": "text" - } - ], - "type": "thinking", - "closed": true - } - ], - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-2509__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-2509__vision.json deleted file mode 100644 index 7eb162d..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-2509__vision.json +++ /dev/null @@ -1,56 +0,0 @@ -{ - "_model_id": "magistral-small-2509", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "4748127459d24f28b9cf9a3b19c5fde9", - "object": "chat.completion", - "model": "magistral-small-2509", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 40, - "total_tokens": 51, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985866, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": [ - { - "thinking": [ - { - "text": "Okay, the user has provided an image, but since I can't actually see images, I need to figure out how to handle this. Maybe they described something in the image or perhaps they expect", - "type": "text" - } - ], - "type": "thinking", - "closed": true - } - ], - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-latest.json deleted file mode 100644 index 9f44851..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "magistral-small-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "58946a0efb6146dd85faee1836d9392c", - "object": "chat.completion", - "model": "magistral-small-latest", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985810, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-latest__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-latest__vision.json deleted file mode 100644 index 2701331..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/magistral-small-latest__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "magistral-small-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "32e86e4313514a4f8241f6624bf86547", - "object": "chat.completion", - "model": "magistral-small-latest", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 40, - "total_tokens": 63, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985811, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": "The image appears to be a close-up of a textured surface, possibly fabric or a woven material. The texture is characterized by a series of diagonal lines forming a diamond pattern, with a central point", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-14b-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-14b-latest.json deleted file mode 100644 index bd7d339..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-14b-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "ministral-14b-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "4221a87d43564760a8af099276b30986", - "object": "chat.completion", - "model": "ministral-14b-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 33, - "total_tokens": 43, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985826, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and remarkable ability to communicate using a range of sounds and body movements.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-14b-latest__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-14b-latest__vision.json deleted file mode 100644 index fdf2b5b..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-14b-latest__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "ministral-14b-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "b8358a64ac524e10a2d470ef89c64aad", - "object": "chat.completion", - "model": "ministral-14b-latest", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 40, - "total_tokens": 51, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985826, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": "The image you provided appears to be completely blank, showing only a solid white or light-colored background with no visible objects, text, or details. There is nothing identifiable or discernible in it.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-2512.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-2512.json deleted file mode 100644 index 1a68701..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-2512.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "ministral-3b-2512", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "6a4c4f50397e47ffbd6d8e6fc3813dd7", - "object": "chat.completion", - "model": "ministral-3b-2512", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 25, - "total_tokens": 35, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985822, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, strong social bonds, and exceptional echolocation abilities.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-2512__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-2512__vision.json deleted file mode 100644 index 5487c25..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-2512__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "ministral-3b-2512", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "6dc816666a9842789e7f7d1468f61308", - "object": "chat.completion", - "model": "ministral-3b-2512", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 40, - "total_tokens": 51, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985822, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": "The image you provided appears to be a blank or empty white square with no discernible content or objects.\n\nIf you are referring to a typical icon or symbol, it might be a placeholder or a generic", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-latest.json deleted file mode 100644 index a9f2fb4..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "ministral-3b-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "b0c0b6c89df342bea40d0fb70bafbe57", - "object": "chat.completion", - "model": "ministral-3b-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 25, - "total_tokens": 35, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985822, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, strong social bonds, and exceptional echolocation abilities.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-latest__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-latest__vision.json deleted file mode 100644 index 1c1f2bd..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-3b-latest__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "ministral-3b-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "8492b3f5c33747a88266b37daa74fe6c", - "object": "chat.completion", - "model": "ministral-3b-latest", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 40, - "total_tokens": 51, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985823, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": "The image you provided appears to be a blank, plain white background with no discernible objects, text, or details.\n\nIf you're looking for something specific or want to analyze it for a particular purpose", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-2512.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-2512.json deleted file mode 100644 index 7a6c3b6..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-2512.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "ministral-8b-2512", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "e14cb01a846148d990bf3d1ab668a1ea", - "object": "chat.completion", - "model": "ministral-8b-2512", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 24, - "total_tokens": 34, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985823, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their complex social behaviors, playful nature, and remarkable communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-2512__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-2512__vision.json deleted file mode 100644 index 588d4d2..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-2512__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "ministral-8b-2512", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "3ffa06ef59a940e69a57c88ce5108bfc", - "object": "chat.completion", - "model": "ministral-8b-2512", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 40, - "total_tokens": 51, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985823, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": "The image you provided is completely blank\u2014it is a solid, uniform color with no visible objects, text, or patterns.\n\nIf you intended to upload something else, please double-check your upload. If", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-latest.json deleted file mode 100644 index 82f9923..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "ministral-8b-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "d3e2d2dd8f3e472ab1155c2ca38a7f33", - "object": "chat.completion", - "model": "ministral-8b-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 24, - "total_tokens": 34, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985824, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex social structures, and remarkable communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-latest__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-latest__vision.json deleted file mode 100644 index 047303b..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/ministral-8b-latest__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "ministral-8b-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "058bee50b17f40ffb0998f86615c1a9e", - "object": "chat.completion", - "model": "ministral-8b-latest", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 34, - "total_tokens": 45, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985824, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "This image is completely blank\u2014it is a solid, uniform color with no visible objects, text, or patterns. There is nothing to describe in terms of content.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-2512.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-2512.json deleted file mode 100644 index 74d3b91..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-2512.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-large-2512", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "ad1b7b3ab0af4f08a33e7ec27fe4ebcc", - "object": "chat.completion", - "model": "mistral-large-2512", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 19, - "total_tokens": 29, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985817, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and remarkable communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-latest.json deleted file mode 100644 index 77b10af..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-large-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "67604d63da2c42d4bb1afb2973c56e33", - "object": "chat.completion", - "model": "mistral-large-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 19, - "total_tokens": 29, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985819, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and remarkable communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-latest__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-latest__vision.json deleted file mode 100644 index 23719ac..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-latest__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-large-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "6690b9da0c19479196004a6b6a0204a8", - "object": "chat.completion", - "model": "mistral-large-latest", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 38, - "total_tokens": 49, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985820, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "The image is a plain white square with no visible content or features. There are no objects, text, or discernible patterns present. It appears to be a blank or empty image.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-pixtral-2411.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-pixtral-2411.json deleted file mode 100644 index 0a1a406..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-large-pixtral-2411.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-large-pixtral-2411", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "85af241cf1c94309b338d621d43ae420", - "object": "chat.completion", - "model": "mistral-large-pixtral-2411", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 21, - "total_tokens": 32, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985852, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2505.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2505.json deleted file mode 100644 index b5e90b4..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2505.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-medium-2505", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "4df7edf6297b43c8960c36fc8e42ed56", - "object": "chat.completion", - "model": "mistral-medium-2505", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 19, - "total_tokens": 29, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985788, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2505__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2505__vision.json deleted file mode 100644 index 893cd1c..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2505__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-medium-2505", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "5aed118415fb4d0a935497e7223d9561", - "object": "chat.completion", - "model": "mistral-medium-2505", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 32, - "total_tokens": 43, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985789, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "The image you provided appears to be a plain white background with no discernible objects, text, or features. It is essentially an empty or blank image.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2508.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2508.json deleted file mode 100644 index 11262b3..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2508.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-medium-2508", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "d5b80197561e4f388e100439ae69d110", - "object": "chat.completion", - "model": "mistral-medium-2508", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 27, - "total_tokens": 37, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985790, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent, social marine mammals known for their playful behavior, complex communication, and strong bonds within their pods.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2508__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2508__vision.json deleted file mode 100644 index c13b794..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2508__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-medium-2508", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "560cd1540c734452b106db10104fa1ce", - "object": "chat.completion", - "model": "mistral-medium-2508", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 28, - "total_tokens": 39, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985790, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "This image appears to be completely blank or empty, showing only a solid light blue color with no other objects, text, or details.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2604.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2604.json deleted file mode 100644 index 3bcd92c..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2604.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-medium-2604", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "5aae28cf344749448b9259b60824a7c2", - "object": "chat.completion", - "model": "mistral-medium-2604", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985831, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2604__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2604__vision.json deleted file mode 100644 index 474e719..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-2604__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-medium-2604", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "476bdb936ea1427f95c7355fe4024294", - "object": "chat.completion", - "model": "mistral-medium-2604", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 32, - "total_tokens": 55, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985832, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "The image you've provided is a solid green color with no discernible objects, text, or other features. It appears to be a plain green background.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3-5.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3-5.json deleted file mode 100644 index 01cedfc..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3-5.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-medium-3-5", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "5cf5231c56e04089994452a30e451d5b", - "object": "chat.completion", - "model": "mistral-medium-3-5", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985827, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3-5__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3-5__vision.json deleted file mode 100644 index 8ff7cb8..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3-5__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-medium-3-5", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "057355081e55407888136a84a81f403b", - "object": "chat.completion", - "model": "mistral-medium-3-5", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 40, - "total_tokens": 63, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985827, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": "The image you've provided is a solid green square. There are no other objects, text, or discernible features within the image. It is a simple, uniform color with no additional details.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.5.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.5.json deleted file mode 100644 index e70a940..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.5.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-medium-3.5", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "218a733253f041928227771067083dee", - "object": "chat.completion", - "model": "mistral-medium-3.5", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985828, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.5__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.5__vision.json deleted file mode 100644 index 4e2ef11..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.5__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-medium-3.5", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "c837469f5b364db8a748aa5d9bd94d5b", - "object": "chat.completion", - "model": "mistral-medium-3.5", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 40, - "total_tokens": 63, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985830, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": "The image you've provided is a solid green color with no other discernible features or objects. It appears to be a plain, uniform green background. There are no texts, shapes, or other elements", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.json deleted file mode 100644 index 372f636..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-medium-3", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "e7779f84ea544ffe8ff4d4374bfbb7ed", - "object": "chat.completion", - "model": "mistral-medium-3", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985830, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3__vision.json deleted file mode 100644 index 760d1a7..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-3__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-medium-3", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "65c35fec4d0441fc9845acccdbcbd7ba", - "object": "chat.completion", - "model": "mistral-medium-3", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 31, - "total_tokens": 54, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985831, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "The image appears to be a solid green color with no discernible objects, text, or other features. It is a plain, uniform green background.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-c21211-r0-75.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-c21211-r0-75.json deleted file mode 100644 index a48e409..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-c21211-r0-75.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-medium-c21211-r0-75", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "d21ae4e5c54a4d74a6c0c3befe768d5d", - "object": "chat.completion", - "model": "mistral-medium-c21211-r0-75", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985832, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-c21211-r0-75__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-c21211-r0-75__vision.json deleted file mode 100644 index 6d87723..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-c21211-r0-75__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-medium-c21211-r0-75", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "5e65295662cd4257bc51b99808a2946d", - "object": "chat.completion", - "model": "mistral-medium-c21211-r0-75", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 31, - "total_tokens": 54, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985833, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "The image appears to be a solid green color with no discernible objects, text, or other features. It is a plain, uniform green background.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-latest.json deleted file mode 100644 index 7d5959c..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-medium-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "79eda1a396c74270a7f7aadcdb1ee57d", - "object": "chat.completion", - "model": "mistral-medium-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 27, - "total_tokens": 37, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985791, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent, social marine mammals known for their playful behavior, complex communication, and strong bonds within their pods.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-latest__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-latest__vision.json deleted file mode 100644 index afc0563..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium-latest__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-medium-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "866d45b54c454d7e82dd568f117c0c59", - "object": "chat.completion", - "model": "mistral-medium-latest", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 27, - "total_tokens": 38, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985792, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "This image appears to be completely blank, showing only a solid light blue color with no other visible details, objects, or text.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium.json deleted file mode 100644 index c38dff5..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-medium", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "7f8b88a6436b4058a43aac973137720c", - "object": "chat.completion", - "model": "mistral-medium", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 27, - "total_tokens": 37, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985793, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent, social marine mammals known for their playful behavior, complex communication, and strong bonds within their pods.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium__vision.json deleted file mode 100644 index d6acd02..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-medium__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-medium", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "5965428a57ba4691947f0f2bc94db706", - "object": "chat.completion", - "model": "mistral-medium", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 25, - "total_tokens": 36, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985796, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "This image appears to be completely blank and white\u2014there are no visible objects, text, or other elements in it.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2506.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2506.json deleted file mode 100644 index 497ae83..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2506.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-small-2506", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "02eb48d3bc884e55bf6f317a6560fac3", - "object": "chat.completion", - "model": "mistral-small-2506", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 23, - "total_tokens": 33, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985867, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior, complex communication, and strong social bonds.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2506__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2506__vision.json deleted file mode 100644 index bd66f08..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2506__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-small-2506", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "a23dfa9f3fda440fad11a59182288c4d", - "object": "chat.completion", - "model": "mistral-small-2506", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 40, - "total_tokens": 51, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985868, - "choices": [ - { - "index": 0, - "finish_reason": "length", - "message": { - "role": "assistant", - "content": "The image you provided is a solid, light blue color with no distinct objects, text, or other features. It appears to be a plain, uniform background. If you were expecting something different or have", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2603.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2603.json deleted file mode 100644 index 1c70bd8..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2603.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-small-2603", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "be829b4a24f641f19c8694e91632f2c1", - "object": "chat.completion", - "model": "mistral-small-2603", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985806, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2603__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2603__vision.json deleted file mode 100644 index 19710fc..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-2603__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-small-2603", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "aea707319e5048c6a706edf3d84d32fe", - "object": "chat.completion", - "model": "mistral-small-2603", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 29, - "total_tokens": 52, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985807, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "The image appears to be a blank or empty white square. There are no discernible objects, text, or features present in the image.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-latest.json deleted file mode 100644 index 2a32171..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-small-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "f69d0fa5c193483b8abb5d9f16aa60b7", - "object": "chat.completion", - "model": "mistral-small-latest", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985808, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-latest__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-latest__vision.json deleted file mode 100644 index 8a9e3e0..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-small-latest__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-small-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "3738d838d2594cfb81af2743ba8feb5f", - "object": "chat.completion", - "model": "mistral-small-latest", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 28, - "total_tokens": 51, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985808, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "The image appears to be a blank or empty image. There are no discernible objects, text, or figures present in the frame.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-tiny-2407.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-tiny-2407.json deleted file mode 100644 index 4b5d466..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-tiny-2407.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-tiny-2407", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "2e467124d8ff4409abec61d95eff4f58", - "object": "chat.completion", - "model": "mistral-tiny-2407", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 25, - "total_tokens": 35, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985802, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins, known for their intelligence and playful nature, are the only mammals besides humans that have sex for pleasure.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-tiny-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-tiny-latest.json deleted file mode 100644 index 023ed5f..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-tiny-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-tiny-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "92e40f91b5e7438286b5505477e1e5c5", - "object": "chat.completion", - "model": "mistral-tiny-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 19, - "total_tokens": 29, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985803, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex communication skills.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-fast.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-fast.json deleted file mode 100644 index a4358ad..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-fast.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-vibe-cli-fast", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "cd6c2b1ab3444d70b7e5fd05121fc2a2", - "object": "chat.completion", - "model": "mistral-vibe-cli-fast", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985809, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-fast__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-fast__vision.json deleted file mode 100644 index 76d2524..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-fast__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-vibe-cli-fast", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "891c55614b884b2d8bd27a83aae33dcd", - "object": "chat.completion", - "model": "mistral-vibe-cli-fast", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 34, - "total_tokens": 57, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985810, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "The image appears to be a plain, solid white background with no discernible objects, text, or features. It's a simple and unadorned visual.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-latest.json deleted file mode 100644 index 07a8513..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-vibe-cli-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "9ee43efdc5284e869ab7ec11ad23620d", - "object": "chat.completion", - "model": "mistral-vibe-cli-latest", - "usage": { - "prompt_tokens": 22, - "completion_tokens": 19, - "total_tokens": 41, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985834, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-latest__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-latest__vision.json deleted file mode 100644 index 7c613ad..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-latest__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-vibe-cli-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": true, - "completion_fim": false, - "fine_tuning": false, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "0611699d584d4f7cb3bd8a069ff71d14", - "object": "chat.completion", - "model": "mistral-vibe-cli-latest", - "usage": { - "prompt_tokens": 23, - "completion_tokens": 26, - "total_tokens": 49, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985834, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "The image you provided is a solid green color with no other elements or details. It appears to be a plain green background.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-with-tools.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-with-tools.json deleted file mode 100644 index eeb958a..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-with-tools.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "mistral-vibe-cli-with-tools", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "27dd67add3e54c4e88c59e5177411271", - "object": "chat.completion", - "model": "mistral-vibe-cli-with-tools", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 27, - "total_tokens": 37, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985797, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent, social marine mammals known for their playful behavior, complex communication, and strong bonds within their pods.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-with-tools__vision.json b/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-with-tools__vision.json deleted file mode 100644 index 6e52c2f..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/mistral-vibe-cli-with-tools__vision.json +++ /dev/null @@ -1,45 +0,0 @@ -{ - "_model_id": "mistral-vibe-cli-with-tools", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": true, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_kind": "vision", - "_response": { - "id": "4f9eeff578394201b56e9d6b6297c9ec", - "object": "chat.completion", - "model": "mistral-vibe-cli-with-tools", - "usage": { - "prompt_tokens": 11, - "completion_tokens": 27, - "total_tokens": 38, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985800, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "This image appears to be completely blank, showing only a solid light blue color with no other visible elements, objects, or details.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/open-mistral-nemo.json b/tests/unit/adapters/fixtures/mistral_native/all_models/open-mistral-nemo.json deleted file mode 100644 index 2706ace..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/open-mistral-nemo.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "open-mistral-nemo", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": true, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": false, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "128efcfe9ea2497bab12496c519eaff5", - "object": "chat.completion", - "model": "open-mistral-nemo", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 23, - "total_tokens": 33, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985801, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins, known for their intelligence and playful nature, are the most socially complex marine mammals after humans.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-mini-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-mini-latest.json deleted file mode 100644 index ec75a95..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-mini-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "voxtral-mini-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": false, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": true, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "087d6af5e6274876992de81b70e1316b", - "object": "chat.completion", - "model": "voxtral-mini-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 17, - "total_tokens": 27, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985865, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful and social nature.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-small-2507.json b/tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-small-2507.json deleted file mode 100644 index 90182bb..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-small-2507.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "voxtral-small-2507", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": true, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "a03ee13f92964236a4b98c4e77b8e275", - "object": "chat.completion", - "model": "voxtral-small-2507", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 19, - "total_tokens": 29, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985816, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file diff --git a/tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-small-latest.json b/tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-small-latest.json deleted file mode 100644 index 10d58a0..0000000 --- a/tests/unit/adapters/fixtures/mistral_native/all_models/voxtral-small-latest.json +++ /dev/null @@ -1,44 +0,0 @@ -{ - "_model_id": "voxtral-small-latest", - "_capabilities": { - "completion_chat": true, - "function_calling": true, - "reasoning": false, - "completion_fim": false, - "fine_tuning": false, - "vision": false, - "ocr": false, - "classification": false, - "moderation": false, - "audio": true, - "audio_transcription": false, - "audio_transcription_realtime": false, - "audio_speech": false - }, - "_response": { - "id": "a92369ae583c45888f237384d246f742", - "object": "chat.completion", - "model": "voxtral-small-latest", - "usage": { - "prompt_tokens": 10, - "completion_tokens": 19, - "total_tokens": 29, - "prompt_tokens_details": { - "cached_tokens": 0 - } - }, - "created": 1777985817, - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": "Dolphins are highly intelligent marine mammals known for their playful behavior and complex social structures.", - "tool_calls": null, - "prefix": false - } - } - ] - } -} \ No newline at end of file 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 index 44749f4..1ca3181 100644 --- 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 @@ -14,10 +14,10 @@ "destination_name": "workspace.default.anthropickey", "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", "destination_model": "claude-sonnet-4-5", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", 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 index 94c78f6..3806311 100644 --- 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 @@ -14,10 +14,10 @@ "destination_name": "workspace.default.anthropickey", "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", "destination_model": "claude-sonnet-4-5", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", 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 index 4692ee3..968f40f 100644 --- 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 @@ -14,10 +14,10 @@ "destination_name": "workspace.default.anthropickey", "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", "destination_model": "claude-sonnet-4-5", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", 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 index e43fbfb..85fd692 100644 --- 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 @@ -14,10 +14,10 @@ "destination_name": "workspace.default.anthropickey", "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", "destination_model": "claude-sonnet-4-5", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", 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 index ad1d463..d1c894b 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_plain.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_anthropic_plain.json @@ -14,13 +14,13 @@ "destination_name": "workspace.default.anthropickey", "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", "destination_model": "claude-sonnet-4-5", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/messages", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", + "request_tags": "{\"lago_subscription\":\"sub_demo\",\"team\":\"lago-demo\"}", "input_tokens": "25", "output_tokens": "400", "total_tokens": "425", 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 index 95ed527..a07abe7 100644 --- 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 @@ -14,10 +14,10 @@ "destination_name": "workspace.default.openaikey", "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", "destination_model": "gpt-5.6", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", 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 index b2b25ce..142b9e4 100644 --- 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 @@ -14,10 +14,10 @@ "destination_name": "workspace.default.openaikey", "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", "destination_model": "gpt-4o", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", 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 index 0f5a30a..a43b584 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_plain.json @@ -14,10 +14,10 @@ "destination_name": "workspace.default.openaikey", "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", "destination_model": "gpt-3.5-turbo", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", "user_agent": "Python-urllib/3.11", "api_type": "openai/v1/chat/completions", "request_tags": "{}", 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 index af2a880..9d4c396 100644 --- 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 @@ -14,10 +14,10 @@ "destination_name": "workspace.default.openaikey", "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", "destination_model": "gpt-4o", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", "user_agent": "Python-urllib/3.11", "api_type": "openai/v1/chat/completions", "request_tags": "{}", 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 index e84c8d5..5245469 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/byok_openai_reasoning.json @@ -14,10 +14,10 @@ "destination_name": "workspace.default.openaikey", "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", "destination_model": "o4-mini", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", 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 index 9f9e701..1b18d70 100644 --- 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 @@ -14,10 +14,10 @@ "destination_name": "workspace.default.openaikey", "destination_id": "64d6b098-69c7-30b2-82e6-28ec097c3f85", "destination_model": "o3", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/openai/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", 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 index 69be38e..1381e95 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/failed_null_tokens.json @@ -14,10 +14,10 @@ "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": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", "user_agent": "Python-urllib/3.11", "api_type": "mlflow/v1/chat/completions", "request_tags": "{}", 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 index 921fc1e..c66926f 100644 --- 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 @@ -14,10 +14,10 @@ "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": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", "user_agent": "Python-urllib/3.11", "api_type": "mlflow/v1/chat/completions", "request_tags": "{}", diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken.json index d61176d..5ac9ea6 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken.json @@ -14,10 +14,10 @@ "destination_name": null, "destination_id": null, "destination_model": null, - "requester": "beready1994@gmail.com", + "requester": "analyst@example.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", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.cloud.databricks.com/ai-gateway/gemini/v1beta/models/x:generateContent", "user_agent": "Python-urllib/3.11", "api_type": "gemini/v1/generateContent", "request_tags": "{}", 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 index a7d844a..c2832f6 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken_1.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/gemini_broken_1.json @@ -14,10 +14,10 @@ "destination_name": null, "destination_id": null, "destination_model": null, - "requester": "beready1994@gmail.com", + "requester": "analyst@example.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", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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": "{}", diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat.json index 50e4a59..42d0037 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat.json @@ -14,10 +14,10 @@ "destination_name": "system.ai.llama-4-maverick", "destination_id": "f0753807-2a5d-3e12-9a70-1b895d651fa5", "destination_model": "llama-4-maverick", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", "user_agent": "Python-urllib/3.11", "api_type": "mlflow/v1/chat/completions", "request_tags": "{}", 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 index e21842b..dbe846f 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_1.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_chat_1.json @@ -14,10 +14,10 @@ "destination_name": "system.ai.llama-4-maverick", "destination_id": "ae1efffe34f03464b267ca56d5f6b6dc", "destination_model": "Llama 4 Maverick", - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", "user_agent": "Python-urllib/3.11", "api_type": "mlflow/v1/chat/completions", "request_tags": "{}", 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 index 853fd52..4a2ffe8 100644 --- 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 @@ -14,10 +14,10 @@ "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": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/chat/completions", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings.json index 59ff684..05c4188 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings.json @@ -14,10 +14,10 @@ "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": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/embeddings", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", 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 index 2360655..df0a66d 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings_1.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/hosted_embeddings_1.json @@ -14,10 +14,10 @@ "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": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/embeddings", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.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\"}", diff --git a/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path.json b/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path.json index 9036c12..4ecf968 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path.json @@ -14,10 +14,10 @@ "destination_name": "workspace.default.anthropickey", "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", "destination_model": null, - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/mlflow/v1/models", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.cloud.databricks.com/ai-gateway/mlflow/v1/models", "user_agent": "Python-urllib/3.11", "api_type": "unmanaged", "request_tags": "{}", 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 index 070cf58..c8b617c 100644 --- a/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path_1.json +++ b/tests/unit/gateway/adapters/fixtures/databricks_gateway/unmanaged_path_1.json @@ -14,10 +14,10 @@ "destination_name": "workspace.default.anthropickey", "destination_id": "629a34db-77bb-3c25-a003-5a947ee6af36", "destination_model": null, - "requester": "beready1994@gmail.com", + "requester": "analyst@example.com", "requester_type": "USER", - "ip_address": "88.168.101.75", - "url": "https://dbc-0223ef70-2638.cloud.databricks.com/ai-gateway/anthropic/v1/nonsense", + "ip_address": "203.0.113.10", + "url": "https://dbc-00000000-0000.cloud.databricks.com/ai-gateway/anthropic/v1/nonsense", "user_agent": "Python-urllib/3.11", "api_type": "unmanaged", "request_tags": "{}", diff --git a/tests/unit/test_fixture_hygiene.py b/tests/unit/test_fixture_hygiene.py new file mode 100644 index 0000000..71c8559 --- /dev/null +++ b/tests/unit/test_fixture_hygiene.py @@ -0,0 +1,85 @@ +"""Committed fixtures must carry no personal or real-account data. + +These fixtures are captured from live provider and gateway calls, and both repos +publish to a PUBLIC package index. A capture therefore arrives carrying whatever the +provider chose to log about the operator who made it — `system.ai_gateway.usage` +records the caller's account email and source IP on every row, neither of which any +adapter reads. Twenty-two Databricks fixtures shipped with a personal Gmail address, a +residential IP, a real workspace subdomain and one live Lago subscription id before +this test existed. + +There is no capture script for the gateway fixtures (they come out of a SQL warehouse +query), so there is nowhere to put a scrub step that a future recapture would run. +This test is the guard instead: it fails on the way back in. + +Kept in step with `fixture_hygiene.test.ts` in the JS port. +""" + +from __future__ import annotations + +import ipaddress +import pathlib +import re + +FIXTURE_ROOT = pathlib.Path(__file__).parent + +# Only `example.com` / `example.org` — the RFC 2606 reserved names — are acceptable. +_EMAIL = re.compile(r"[A-Za-z0-9._%+-]+@([A-Za-z0-9.-]+\.[A-Za-z]{2,})") +_IPV4 = re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b") +# A Databricks workspace subdomain is a real, addressable host. +_DBX_HOST = re.compile(r"\bdbc-[0-9a-f]{4,}-[0-9a-f]{4,}\b") + +_ALLOWED_EMAIL_DOMAINS = {"example.com", "example.org", "example.net"} +_PLACEHOLDER_DBX_HOST = "dbc-00000000-0000" + + +def _fixtures() -> list[pathlib.Path]: + return sorted(FIXTURE_ROOT.rglob("*.json")) + + +def _is_public_ip(text: str) -> bool: + """True only for a globally routable address — a real host somewhere.""" + try: + addr = ipaddress.ip_address(text) + except ValueError: + return False + # RFC 5737 documentation ranges are the intended replacements and are not global. + return addr.is_global + + +def test_no_real_email_addresses() -> None: + offenders = [ + f"{p.relative_to(FIXTURE_ROOT)}: {domain}" + for p in _fixtures() + for domain in _EMAIL.findall(p.read_text(encoding="utf-8")) + if domain.lower() not in _ALLOWED_EMAIL_DOMAINS + ] + assert not offenders, ( + "fixtures carry email addresses outside the RFC 2606 reserved domains — " + f"replace with an example.com address: {offenders}" + ) + + +def test_no_publicly_routable_ip_addresses() -> None: + offenders = [ + f"{p.relative_to(FIXTURE_ROOT)}: {ip}" + for p in _fixtures() + for ip in set(_IPV4.findall(p.read_text(encoding="utf-8"))) + if _is_public_ip(ip) + ] + assert not offenders, ( + "fixtures carry globally routable IP addresses — replace with an RFC 5737 " + f"documentation address such as 203.0.113.10: {offenders}" + ) + + +def test_no_real_databricks_workspace_hosts() -> None: + offenders = [ + f"{p.relative_to(FIXTURE_ROOT)}: {host}" + for p in _fixtures() + for host in set(_DBX_HOST.findall(p.read_text(encoding="utf-8"))) + if host != _PLACEHOLDER_DBX_HOST + ] + assert not offenders, ( + f"fixtures name a real Databricks workspace — use {_PLACEHOLDER_DBX_HOST}: {offenders}" + ) From 9e22961473831d6e53bbe7da046d703c4160fdb0 Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 11:19:14 +0200 Subject: [PATCH 11/22] Raise on a failed result chunk instead of billing a short window MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Statement Execution API returns only chunk 0 inline; query() fetched the rest but never checked their HTTP status. A failed chunk fetch returns a JSON error body with no `data_array`, so `or []` appended zero rows, the loop moved to the next index, and query() returned a PARTIAL result reporting success — verbatim the "silent truncation" this module's docstring leads with. Measured against a live warehouse on a genuinely chunked read (9,000 rows over 2 chunks): a 403, 404 and 503 on chunk 1 each returned 6,750 rows with no exception. 25% of the window billed as if it were all of it, three times out of three. Chunk fetches now raise through a new _raise_for_api_error, which reads the API's body rather than calling raise_for_status(): Databricks puts the cause there and requests shows only the status line. The assembled row count is also asserted against manifest.total_row_count, which catches truncation no per-request status check can see — a chunk returning HTTP 200 with fewer rows than promised, or a manifest/chunk disagreement. The other two calls were already loud, contrary to how this was reported: a non-OK submission or poll has no status.state, so _await_statement's own guard already raised. What changes for them is only legibility — the real cause instead of "Databricks statement None: {...}". The 403 "does not have required scopes: sql" this class warns operators about is the error most likely to hit a first-time setup, so it is the one that had to read clearly. Rows arriving with no manifest.schema.columns now raise rather than zipping to {} each, which every layer downstream degrades cleanly and wrongly on, ending in a confident {"cost": 0, "tokens": 0} for a window that had real traffic. Not observed on this API — every SELECT returns a full schema, zero-row reads included — so this guards the decode, not a known bug. strict=False is kept on the zip deliberately: a length mismatch is an API-contract violation the row-count check already catches, and strict=True would newly reject reads that work today. _FakeResponse gained status_code, which a real requests response always carries. Its absence is why this gap survived: no test could express a non-OK response. --- CHANGELOG.md | 9 ++ src/lago_agent_sdk/gateway/databricks.py | 70 ++++++++++++- tests/unit/gateway/test_databricks_source.py | 104 ++++++++++++++++++- 3 files changed, 177 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 00d71fa..351998c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,15 @@ All notable changes to this project will be documented here. Format follows [Kee ## [Unreleased] +### Fixed + +- **A failed result-chunk fetch no longer truncates a Databricks window silently.** The Statement Execution API returns only chunk 0 inline; `query()` fetched the rest but never checked their HTTP status. A failure returns a JSON error body with no `data_array`, so `or []` appended zero rows, the loop moved to the next index, and `query()` returned a **partial** result reporting success — the exact "silent truncation" this module's docstring leads with. Measured against a live warehouse: a 403, 404 and 503 on chunk 1 of 2 each returned 6,750 of 9,000 rows with no exception, i.e. 25% of the window billed as if it were all of it. + - Chunk fetches now raise via a new `_raise_for_api_error`, which reads the API's **body** rather than calling `raise_for_status()`, because that is where Databricks puts the cause (`{"error_code": "PERMISSION_DENIED", "message": "... does not have required scopes: sql"}`) and `requests` shows only the status line. + - The assembled row count is also asserted against `manifest.total_row_count` before anything is billed. That catches truncation no per-request status check can see — a chunk returning HTTP 200 with fewer rows than promised, or a manifest/chunk disagreement. + - **The other two calls were already loud, contrary to how this was reported.** A non-OK submission or poll has no `status.state`, so `_await_statement`'s own guard already raised. What changes for them is only legibility: `Databricks statement submission failed: HTTP 404: {"error_code": "NOT_FOUND", "message": "The warehouse … was not found."}` instead of `Databricks statement None: {…}`. The `403 does not have required scopes: sql` that this class warns operators about is the error most likely to hit a first-time setup, so it is the one that had to read clearly. + - Rows arriving with no `manifest.schema.columns` now raise instead of zipping to `{}` each — which every layer downstream degrades cleanly and wrongly on, ending in a confident `{"cost": 0, "tokens": 0}` for a window that had real traffic. Not observed on this API (every SELECT returns a full schema, zero-row reads included, 36 columns); this guards the decode, not a known bug. `strict=False` is kept on the zip deliberately: a length mismatch is an API-contract violation the row-count check already catches, and `strict=True` would newly reject reads that work today. + - The `_FakeResponse` double in `test_databricks_source.py` gained `status_code`, which a real `requests` response always carries. Its absence is why this gap survived: no test could express a non-OK response. + ### 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. diff --git a/src/lago_agent_sdk/gateway/databricks.py b/src/lago_agent_sdk/gateway/databricks.py index 2522816..b06d544 100644 --- a/src/lago_agent_sdk/gateway/databricks.py +++ b/src/lago_agent_sdk/gateway/databricks.py @@ -56,6 +56,27 @@ _INTERVAL_RE = re.compile(r"^\s*(\d{1,5})\s+(second|minute|hour|day|week)s?\s*$", re.I) +def _raise_for_api_error(resp: Any, what: str) -> None: + """Raise with the API's own error text when a Statement Execution call is not OK. + + Deliberately NOT `raise_for_status()`: Databricks puts the useful part in the BODY + (`{"error_code": "PERMISSION_DENIED", "message": "... does not have required scopes: + sql"}`) and `requests` shows only the status line. The `403 does not have required + scopes: sql` that this class's docstring warns operators about is the error most + likely to hit a first-time setup, so it is the one that must read clearly. + + Truncated because these bodies can carry a multi-KB `details` array. + """ + status = getattr(resp, "status_code", 200) + if 200 <= int(status) < 300: + return + try: + detail = json.dumps(resp.json()) + except ValueError: + detail = getattr(resp, "text", "") or "" + raise RuntimeError(f"Databricks {what} failed: HTTP {status}: {detail[:500]}") + + @dataclass class DatabricksUsageRow: """One billable row, already shaped for `emit()`. @@ -216,6 +237,7 @@ def query(self, sql: str) -> list[dict[str, Any]]: }, timeout=self.timeout, ) + _raise_for_api_error(resp, "statement submission") body = resp.json() body = self._await_statement(body, headers) @@ -227,15 +249,46 @@ def query(self, sql: str) -> list[dict[str, Any]]: 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( + chunk_resp = 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 []) + ) + # THE check this whole method exists for. A failed chunk fetch returns a JSON + # error body with no `data_array`, so `or []` would append zero rows, the loop + # would continue, and `query()` would return a PARTIAL result reporting + # success — measured live: a 403/404/503 on chunk 1 of 2 silently dropped 25% + # of the window. Billing a fraction of a window with no error is the single + # worst outcome this reader can produce, so it must raise. + _raise_for_api_error(chunk_resp, f"result chunk {index} of {total_chunks}") + arrays.extend((chunk_resp.json()).get("data_array") or []) if total_chunks > 1: logger.info("lago: databricks result spanned %d chunks (%d rows)", total_chunks, len(arrays)) + # End-to-end truncation check, independent of cause: catches a short read that no + # per-request status could reveal (a chunk that returns HTTP 200 with fewer rows + # than promised, a manifest/chunk disagreement). `total_row_count` is absent on + # some statement kinds, so only assert when Databricks actually stated a count. + promised = manifest.get("total_row_count") + if promised is not None and int(promised) != len(arrays): + raise RuntimeError( + f"Databricks returned {len(arrays)} row(s) but the manifest promised " + f"{int(promised)} across {total_chunks} chunk(s) — refusing to bill a " + f"partial window (statement_id={statement_id})" + ) + # A row set with no column names decodes to `{}` per row, which every layer + # downstream degrades cleanly and wrongly on: all-zero usage, and a confident + # `{"cost": 0, "tokens": 0}` for a window that had real traffic. Not observed on + # this API (every SELECT returns a full schema, zero-row reads included) — this + # guards the decode, not a known bug. Deliberately keeps `strict=False` on the + # zip below: a length mismatch is an API-contract violation the row-count check + # above already catches, and strict=True would newly reject reads that work today. + if arrays and not columns: + raise RuntimeError( + "Databricks returned rows with no `manifest.schema.columns` — cannot " + f"decode {len(arrays)} row(s) (statement_id={statement_id})" + ) + 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]: @@ -265,11 +318,18 @@ def _await_statement(self, body: dict[str, Any], headers: dict[str, str]) -> dic f"(statement_id={statement_id}); raise `timeout` or narrow the window" ) time.sleep(2.0) - body = requests.get( + poll = requests.get( f"{self.host}{_STATEMENTS_PATH}/{statement_id}", headers=headers, timeout=self.timeout, - ).json() + ) + # This path already failed loudly without a status check — a non-OK body has + # no `status.state`, so the loop's own `state not in (PENDING, RUNNING)` branch + # raised. Checking here only changes WHAT the operator reads: the real cause + # ("Invalid access token", "statement expired") instead of + # `Databricks statement None: {...}`. + _raise_for_api_error(poll, "statement poll") + body = poll.json() # ------------------------------------------------------------------ # Reading diff --git a/tests/unit/gateway/test_databricks_source.py b/tests/unit/gateway/test_databricks_source.py index 1e3d2c5..77613b3 100644 --- a/tests/unit/gateway/test_databricks_source.py +++ b/tests/unit/gateway/test_databricks_source.py @@ -214,8 +214,9 @@ def test_hosted_rows_keep_the_databricks_provider() -> None: # Chunked results — the silent-truncation guard # -------------------------------------------------------------------------- class _FakeResponse: - def __init__(self, payload: dict) -> None: + def __init__(self, payload: dict, status_code: int = 200) -> None: self._payload = payload + self.status_code = status_code def json(self) -> dict: return self._payload @@ -259,6 +260,107 @@ def fake_get(url: str, **_kw: Any) -> _FakeResponse: assert all(u.startswith("https://x/api/2.0/sql/statements/stmt-1/result/chunks/") for u in fetched) +def test_query_raises_when_a_chunk_fetch_fails(monkeypatch: pytest.MonkeyPatch) -> None: + """A failed chunk fetch must NOT be swallowed into a short row set. + + The error body carries no `data_array`, so `or []` would append nothing, the loop + would move to the next index, and `query()` would return a partial window reporting + success. Measured against a live warehouse: a 403 on chunk 1 of 2 returned 6,750 of + 9,000 rows with no exception — 25% of the window billed as if it were all of it. + """ + import requests + + first = { + "statement_id": "stmt-1", + "status": {"state": "SUCCEEDED"}, + "manifest": { + "schema": {"columns": [{"name": "invocation_id"}, {"name": "input_tokens"}]}, + "total_chunk_count": 2, + }, + "result": {"data_array": [["a", "1"]]}, + } + # exactly what the API returns for an expired statement / revoked token mid-read + denied = {"error_code": "PERMISSION_DENIED", "message": "does not have required scopes: sql"} + + monkeypatch.setattr(requests, "post", lambda url, **_kw: _FakeResponse(first)) + monkeypatch.setattr(requests, "get", lambda url, **_kw: _FakeResponse(denied, status_code=403)) + + src = DatabricksSource(host="https://x/", token="t", warehouse_id="w") + with pytest.raises(RuntimeError) as excinfo: + src.query("SELECT 1") + message = str(excinfo.value) + assert "result chunk 1 of 2" in message + # the operator must see the API's own cause, not just a status line + assert "does not have required scopes: sql" in message + + +def test_query_raises_when_the_row_count_misses_the_manifest(monkeypatch: pytest.MonkeyPatch) -> None: + """A chunk that returns HTTP 200 with fewer rows than promised is still truncation. + No per-request status check can catch that, so the assembled count is compared with + `manifest.total_row_count` before any of it is billed.""" + import requests + + first = { + "statement_id": "stmt-1", + "status": {"state": "SUCCEEDED"}, + "manifest": { + "schema": {"columns": [{"name": "invocation_id"}]}, + "total_chunk_count": 2, + "total_row_count": 3, + }, + "result": {"data_array": [["a"]]}, + } + monkeypatch.setattr(requests, "post", lambda url, **_kw: _FakeResponse(first)) + # HTTP 200, but one row short of the promised three + monkeypatch.setattr(requests, "get", lambda url, **_kw: _FakeResponse({"data_array": [["b"]]})) + + src = DatabricksSource(host="https://x/", token="t", warehouse_id="w") + with pytest.raises(RuntimeError, match=r"returned 2 row\(s\) but the manifest promised 3"): + src.query("SELECT 1") + + +def test_query_raises_when_rows_arrive_with_no_columns(monkeypatch: pytest.MonkeyPatch) -> None: + """Rows with no column names zip to `{}` each, which every layer downstream degrades + cleanly and wrongly on — all-zero usage and a confident `{"cost": 0, "tokens": 0}` + for a window that had real traffic. Not observed on this API; guards the decode.""" + import requests + + first = { + "statement_id": "stmt-1", + "status": {"state": "SUCCEEDED"}, + "manifest": {"total_chunk_count": 1}, + "result": {"data_array": [["a"], ["b"]]}, + } + monkeypatch.setattr(requests, "post", lambda url, **_kw: _FakeResponse(first)) + + src = DatabricksSource(host="https://x/", token="t", warehouse_id="w") + with pytest.raises(RuntimeError, match="no `manifest.schema.columns`"): + src.query("SELECT 1") + + +def test_query_error_carries_the_api_error_code(monkeypatch: pytest.MonkeyPatch) -> None: + """A non-OK submission used to surface as `Databricks statement None: {...}` — the + state was absent, so the poll loop's own guard raised with a misleading prefix. The + cause was in the body all along; name it.""" + import requests + + monkeypatch.setattr( + requests, + "post", + lambda url, **_kw: _FakeResponse( + {"error_code": "NOT_FOUND", "message": "The warehouse w was not found."}, + status_code=404, + ), + ) + src = DatabricksSource(host="https://x/", token="t", warehouse_id="w") + with pytest.raises(RuntimeError) as excinfo: + src.query("SELECT 1") + message = str(excinfo.value) + assert "statement submission failed: HTTP 404" in message + assert "The warehouse w was not found." in message + assert "statement None" not in message + + 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.""" From 5d9754983c9aa289d01155f17395b6207edba37b Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 11:58:00 +0200 Subject: [PATCH 12/22] Key gateway cache de-overlap on the surface, not the vendor name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit system.ai_gateway.usage re-reports every vendor in the OpenAI shape. Measured across 246 rows and 6 vendors: total_tokens == input + output for EVERY group, with cache_read AND cache_write inside input and reasoning inside output. The two billing paths decided that from the provider name, which is right for a native call and wrong for a table row. Anthropic's own API reports cache additively — measured live, cache_read=3962 against input=9 — so an Anthropic row read from this table had its cached tokens counted twice. On a real backfill over 2026-08-06: 48,798 tokens billed against 31,091 consumed, 1.570x. The same window now reports 31,018, the same ~0.2% lag-only shortfall openai already had. The api is the honest key. A gateway row reuses the live vendor names, so nothing in the name separates a table row from a direct call. workers-ai stays a provider entry by contrast: it names a vendor reachable through exactly one surface, so there the name is sufficient. This also adds the first cache_write de-overlap the SDK has had. It is surface-only by design — Anthropic is the one vendor whose native API bills cache writes at all, and it reports them additively, so no native response needs the correction. Databricks-hosted models were a latent 1.991x rather than a live one: they bill as token counts and 0 of 96 hosted rows carry cache today. Fixed before it fires. Cloudflare AI Gateway is deliberately excluded and now pinned by a test. Its logs preserve each vendor's native shape — a real Anthropic entry reads input=10, output=4, total=14 with input_cached_tokens=3429 outside that total — so adding it would under-bill the cached portion. Five cases added to money_golden.json, byte-identical in both repos, including controls that an openai row is subtracted exactly once (correcting it twice was a measured 13% under-bill) and that the native Anthropic path is untouched. All five new assertions confirmed failing with the fix reverted. --- CHANGELOG.md | 8 + src/lago_agent_sdk/pricing.py | 86 ++++++-- tests/unit/fixtures/pricing/money_golden.json | 190 ++++++++++++++++-- tests/unit/test_pricing.py | 83 +++++++- 4 files changed, 333 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 351998c..5922d75 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,14 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **Databricks gateway rows were de-overlapped by the vendor's token semantics instead of the surface's.** `system.ai_gateway.usage` re-reports every vendor in the OpenAI shape: measured across 246 rows and 6 vendors, `total_tokens == input + output` for **every** group, with `cache_read` *and* `cache_write` inside `input` and reasoning inside `output`. `compute_cost`/`deoverlapped_token_total` keyed that decision on the provider name — correct for a native call, wrong for a table row. Anthropic's own API reports cache additively (measured live: `cache_read=3962` against `input=9`), so an Anthropic row read from this table had its cached tokens counted a second time. On a real backfill that billed **48,798 tokens against 31,091 consumed, 1.570x**; the same window now reports 31,018, i.e. the same ~0.2% lag-only shortfall `openai` already had. + - New `_OPENAI_SHAPED_APIS` keys on `CanonicalUsage.api` rather than the provider, because a gateway row reuses the live vendor names and nothing in the name separates the two. `workers-ai` stays a *provider* entry by contrast: it names a vendor reachable through exactly one surface, so the name is sufficient there. + - Adds this SDK's first `cache_write ⊆ input` correction. It is surface-only by design — Anthropic is the one vendor whose native API bills cache writes at all, and it reports them additively, so no native response needs it. + - Both billing paths now read the same `_token_semantics` helper, so the token basis and the priced basis cannot drift apart. + - **Databricks-hosted models were a latent 1.991x over-bill, not a live one.** They bill as token counts (`TOKEN_BILLED_PROVIDERS`) and 0 of 96 hosted rows carry cache today, so this fixes it before it can fire rather than after. + - **Cloudflare AI Gateway is deliberately excluded, and that is now pinned by a test.** Its logs preserve each vendor's native shape — a real Anthropic entry reads `input=10, output=4, total=14` with `input_cached_tokens=3429` sitting outside that total — so adding it would UNDER-bill the cached portion. + - Five cases added to `money_golden.json` (both repos, byte-identical), including controls that an `openai` row is subtracted exactly **once** — correcting it twice was a measured 13% under-bill — and that the native Anthropic path is untouched. + - **A failed result-chunk fetch no longer truncates a Databricks window silently.** The Statement Execution API returns only chunk 0 inline; `query()` fetched the rest but never checked their HTTP status. A failure returns a JSON error body with no `data_array`, so `or []` appended zero rows, the loop moved to the next index, and `query()` returned a **partial** result reporting success — the exact "silent truncation" this module's docstring leads with. Measured against a live warehouse: a 403, 404 and 503 on chunk 1 of 2 each returned 6,750 of 9,000 rows with no exception, i.e. 25% of the window billed as if it were all of it. - Chunk fetches now raise via a new `_raise_for_api_error`, which reads the API's **body** rather than calling `raise_for_status()`, because that is where Databricks puts the cause (`{"error_code": "PERMISSION_DENIED", "message": "... does not have required scopes: sql"}`) and `requests` shows only the status line. - The assembled row count is also asserted against `manifest.total_row_count` before anything is billed. That catches truncation no per-request status check can see — a chunk returning HTTP 200 with fewer rows than promised, or a manifest/chunk disagreement. diff --git a/src/lago_agent_sdk/pricing.py b/src/lago_agent_sdk/pricing.py index 9072a0b..bf0c76f 100644 --- a/src/lago_agent_sdk/pricing.py +++ b/src/lago_agent_sdk/pricing.py @@ -105,6 +105,32 @@ # guard, and Cloudflare hosts reasoning models (deepseek-r1, qwen, glm). _OUTPUT_INCLUDES_REASONING = frozenset({"openai", "workers-ai"}) +# Gateway SURFACES that re-report every vendor's usage in the OpenAI shape: `input` +# already contains cache_read AND cache_write, and `output` already contains +# reasoning, no matter which vendor actually served the call. +# +# This keys on `CanonicalUsage.api` rather than the provider because on a gateway +# it is the SURFACE that decides the shape, and a surface row reuses the live +# vendor names. A `provider="anthropic"` row read from Databricks' system table +# needs the correction; a `provider="anthropic"` response from Anthropic's own API +# must NOT get it. The vendor name cannot tell those two apart, so it is the wrong +# key — unlike "workers-ai" above, which names a vendor reachable through exactly +# one surface and so works as a provider entry. +# +# Measured on `system.ai_gateway.usage`, 246 rows across 6 vendors: `total_tokens +# == input + output` for EVERY vendor group, with cache_read and cache_write inside +# input and reasoning inside output. Anthropic's own API reports the exact opposite +# (cache_read=3962 against input=9, additive), which is why keying on the vendor +# over-billed a real backfill 1.570x — 48,798 tokens reported against 31,091 +# consumed, the excess being exactly cache_read + cache_write. +# +# Cloudflare AI Gateway is deliberately ABSENT: its logs preserve each vendor's +# native shape instead of normalising them. A real Anthropic entry there reads +# input=10, output=4, total=14 with input_cached_tokens=3429 sitting OUTSIDE that +# total — additive, exactly like the native API — so the provider-keyed sets are +# already right for it and adding it here would UNDER-bill the cached portion. +_OPENAI_SHAPED_APIS = frozenset({"databricks_gateway"}) + # 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. # @@ -344,6 +370,28 @@ class CostBreakdown: fields: dict[str, dict[str, str]] # field -> {tokens, unit_price, cost} +def _token_semantics(usage: Any) -> tuple[bool, bool, bool]: + """Which of a record's subsets are ALREADY inside their parent count. + + Returns ``(input_includes_cache_read, input_includes_cache_write, + output_includes_reasoning)``, the three overlaps the billing paths have to + remove. The SURFACE wins over the vendor: a gateway that re-reports usage in + its own shape has already decided the convention, so `api` is checked first + and the provider-keyed sets only answer for a native call. + + `cache_write` is surface-only by design and has no provider set to consult: + Anthropic is the one vendor whose native API bills cache writes at all, and it + reports them additively, so no native response needs the correction. + """ + provider = (getattr(usage, "provider", "") or "").lower() + shaped = (getattr(usage, "api", "") or "").lower() in _OPENAI_SHAPED_APIS + return ( + shaped or provider in _INPUT_INCLUDES_CACHE_READ, + shaped, + shaped or provider in _OUTPUT_INCLUDES_REASONING, + ) + + def compute_cost(usage: CanonicalUsage, price: ModelPrice, markup: Decimal) -> CostBreakdown: """Compute ``Σ(unit_price × count) × markup`` for the priced fields present. @@ -351,18 +399,24 @@ def compute_cost(usage: CanonicalUsage, price: ModelPrice, markup: Decimal) -> C call whose only counts are unpriced yields total "0" so it stays accounted for. """ - provider = (usage.provider or "").lower() counts = {f: (getattr(usage, f, 0) or 0) for f in PRICED_FIELDS} - # Remove double-counting where a provider's `input`/`output` already include - # a separately-listed subset (see the _INCLUDES_ sets above): + # Remove double-counting where the reported `input`/`output` already include a + # separately-listed subset (see `_token_semantics` and the sets above): # • reasoning ⊆ output → bill it as output only (drop the separate line). # • cache_read ⊆ input → bill the cached portion at the cache-read rate, # so subtract it from input (only when a cache_read price exists; with no # cache price the cached tokens stay in input at the prompt rate). - if provider in _OUTPUT_INCLUDES_REASONING: + # • cache_write ⊆ input → same treatment, on the surfaces that report it + # that way. Only one of cache_read/cache_write is non-zero on a given + # Databricks row, but both are subtracted unconditionally so a surface + # that does report both at once still reconciles. + inc_cache_read, inc_cache_write, inc_reasoning = _token_semantics(usage) + if inc_reasoning: counts["reasoning"] = 0 - if provider in _INPUT_INCLUDES_CACHE_READ and price.get("cache_read") is not None: + if inc_cache_read and price.get("cache_read") is not None: counts["input"] = max(0, counts["input"] - counts["cache_read"]) + if inc_cache_write and price.get("cache_write") is not None: + counts["input"] = max(0, counts["input"] - counts["cache_write"]) base = Decimal(0) fields: dict[str, dict[str, str]] = {} @@ -401,16 +455,18 @@ def _finalize_breakdown( def deoverlapped_token_total(usage: Any) -> int: - """Total tokens a call actually consumed, with per-provider overlaps removed. + """Total tokens a call actually consumed, with the reported overlaps removed. Sums the same PRICED_FIELDS the split cost path emits one event each for, so the single-event `unit` equals the sum of the split path's `unit`s instead of - reporting a different basis. Both `_INCLUDES_` sets are applied, because a - subset counted twice inflates the reported quantity exactly as it would inflate - a price: + reporting a different basis. Every overlap `_token_semantics` reports is + applied, because a subset counted twice inflates the reported quantity exactly + as it would inflate a price: - * reasoning ⊆ output for providers in _OUTPUT_INCLUDES_REASONING - * cache_read ⊆ input for providers in _INPUT_INCLUDES_CACHE_READ + * reasoning ⊆ output — providers in _OUTPUT_INCLUDES_REASONING, or any + row from a surface in _OPENAI_SHAPED_APIS + * cache_read ⊆ input — providers in _INPUT_INCLUDES_CACHE_READ, likewise + * cache_write ⊆ input — surfaces in _OPENAI_SHAPED_APIS only Deliberately NOT gated on a unit price existing, unlike `compute_cost`'s subtraction — this is a token count, so whether a rate happens to be published @@ -424,12 +480,14 @@ def deoverlapped_token_total(usage: Any) -> int: a breakdown OF `cache_write`, so including any of them would not be a token total. This mirrors price mode's documented five-field scope. """ - provider = (getattr(usage, "provider", "") or "").lower() counts = {f: (getattr(usage, f, 0) or 0) for f in PRICED_FIELDS} - if provider in _OUTPUT_INCLUDES_REASONING: + inc_cache_read, inc_cache_write, inc_reasoning = _token_semantics(usage) + if inc_reasoning: counts["reasoning"] = 0 - if provider in _INPUT_INCLUDES_CACHE_READ: + if inc_cache_read: counts["cache_read"] = 0 + if inc_cache_write: + counts["cache_write"] = 0 return sum(int(v or 0) for v in counts.values()) diff --git a/tests/unit/fixtures/pricing/money_golden.json b/tests/unit/fixtures/pricing/money_golden.json index e6f6f66..3bf652a 100644 --- a/tests/unit/fixtures/pricing/money_golden.json +++ b/tests/unit/fixtures/pricing/money_golden.json @@ -3,8 +3,14 @@ "cases": [ { "name": "input+output, no markup", - "prices": { "input": "0.000003", "output": "0.000015" }, - "counts": { "input": 1000, "output": 500 }, + "prices": { + "input": "0.000003", + "output": "0.000015" + }, + "counts": { + "input": 1000, + "output": 500 + }, "markup": "1", "base": "0.0105", "total": "0.0105", @@ -12,8 +18,14 @@ }, { "name": "input+output, 1.2x markup", - "prices": { "input": "0.000003", "output": "0.000015" }, - "counts": { "input": 1000, "output": 500 }, + "prices": { + "input": "0.000003", + "output": "0.000015" + }, + "counts": { + "input": 1000, + "output": 500 + }, "markup": "1.2", "base": "0.0105", "total": "0.0126", @@ -21,8 +33,14 @@ }, { "name": "cache_read + reasoning, 1.5x markup", - "prices": { "cache_read": "0.0000005", "reasoning": "0.00001" }, - "counts": { "cache_read": 2000, "reasoning": 100 }, + "prices": { + "cache_read": "0.0000005", + "reasoning": "0.00001" + }, + "counts": { + "cache_read": 2000, + "reasoning": 100 + }, "markup": "1.5", "base": "0.002", "total": "0.003", @@ -30,8 +48,14 @@ }, { "name": "free model -> zero", - "prices": { "input": "0", "output": "0" }, - "counts": { "input": 1000, "output": 200 }, + "prices": { + "input": "0", + "output": "0" + }, + "counts": { + "input": 1000, + "output": 200 + }, "markup": "1", "base": "0", "total": "0", @@ -39,8 +63,12 @@ }, { "name": "12dp exact precision", - "prices": { "input": "0.000000333333" }, - "counts": { "input": 3 }, + "prices": { + "input": "0.000000333333" + }, + "counts": { + "input": 3 + }, "markup": "1", "base": "0.000000999999", "total": "0.000000999999", @@ -48,8 +76,12 @@ }, { "name": "floor-12dp truncation on markup", - "prices": { "input": "0.000001" }, - "counts": { "input": 1000 }, + "prices": { + "input": "0.000001" + }, + "counts": { + "input": 1000 + }, "markup": "1.333333333333", "base": "0.001", "total": "0.001333333333", @@ -59,8 +91,14 @@ "name": "workers-ai: cache_read is a SUBSET of input, billed once", "_note": "Real counts from a live OpenAI-shaped cached call (prompt 23233, cached 23168) at live Cloudflare @cf/moonshotai/kimi-k2.6 rates. Only 23233-23168=65 tokens may be billed at the input rate; billing all 23233 double-charges the cached portion.", "provider": "workers-ai", - "prices": { "input": "0.00000095", "cache_read": "0.00000016" }, - "counts": { "input": 23233, "cache_read": 23168 }, + "prices": { + "input": "0.00000095", + "cache_read": "0.00000016" + }, + "counts": { + "input": 23233, + "cache_read": 23168 + }, "markup": "1", "base": "0.00376863", "total": "0.00376863", @@ -70,8 +108,14 @@ "name": "anthropic: cache_read is ADDITIVE, input not reduced", "_note": "Same counts and rates as the workers-ai case above; the only difference is the provider's token semantics. Anthropic reports input exclusive of cache, so all 23233 input tokens are billed.", "provider": "anthropic", - "prices": { "input": "0.00000095", "cache_read": "0.00000016" }, - "counts": { "input": 23233, "cache_read": 23168 }, + "prices": { + "input": "0.00000095", + "cache_read": "0.00000016" + }, + "counts": { + "input": 23233, + "cache_read": 23168 + }, "markup": "1", "base": "0.02577823", "total": "0.02577823", @@ -81,12 +125,122 @@ "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 }, + "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" + }, + { + "name": "databricks_gateway: anthropic cache_read is INSIDE input on this surface", + "_note": "Real shape from system.ai_gateway.usage (input=1822 containing cache_read=1812, output=4, total_tokens=1826). Anthropic's OWN api reports cache additively, so the provider name is the wrong key here: without the api override this bills all 1822 at the input rate and charges the cached portion twice.", + "provider": "anthropic", + "api": "databricks_gateway", + "prices": { + "input": "0.000003", + "output": "0.000015", + "cache_read": "0.0000003" + }, + "counts": { + "input": 1822, + "output": 4, + "cache_read": 1812 + }, + "markup": "1", + "base": "0.0006336", + "total": "0.0006336", + "total_cents": "0.06336" + }, + { + "name": "databricks_gateway: cache_write is INSIDE input too", + "_note": "The mechanism no provider-keyed set has: on this surface cache_write also sits inside input. Only one of cache_read/cache_write is non-zero per row, so this is the write half of the same real shape.", + "provider": "anthropic", + "api": "databricks_gateway", + "prices": { + "input": "0.000003", + "output": "0.000015", + "cache_write": "0.00000375" + }, + "counts": { + "input": 1825, + "output": 4, + "cache_write": 1812 + }, + "markup": "1", + "base": "0.006894", + "total": "0.006894", + "total_cents": "0.6894" + }, + { + "name": "native anthropic: cache_read stays ADDITIVE (the surface set must not leak)", + "_note": "Same counts as the databricks_gateway case above but api='native'. Measured live: Anthropic's own API returns cache_read=3962 against input=9, i.e. outside input. All 1822 input tokens are billable here. Guards the override against widening into the native path.", + "provider": "anthropic", + "api": "native", + "prices": { + "input": "0.000003", + "output": "0.000015", + "cache_read": "0.0000003" + }, + "counts": { + "input": 1822, + "output": 4, + "cache_read": 1812 + }, + "markup": "1", + "base": "0.0060696", + "total": "0.0060696", + "total_cents": "0.60696" + }, + { + "name": "databricks_gateway: an openai row is subtracted exactly ONCE", + "_note": "openai is already in _INPUT_INCLUDES_CACHE_READ, so the surface override must not subtract a second time. Correcting an openai row twice was measured at a 13% UNDER-bill.", + "provider": "openai", + "api": "databricks_gateway", + "prices": { + "input": "0.000003", + "output": "0.000015", + "cache_read": "0.0000003" + }, + "counts": { + "input": 1000, + "output": 100, + "cache_read": 400 + }, + "markup": "1", + "base": "0.00342", + "total": "0.00342", + "total_cents": "0.342" + }, + { + "name": "databricks_gateway: reasoning is INSIDE output for a vendor not in the provider set", + "_note": "gemini reports thoughts ADDITIVELY on its own API, but this surface folds reasoning into output for every vendor (total_tokens == input + output on all 246 measured rows). Billing the reasoning line as well double-charges it.", + "provider": "gemini", + "api": "databricks_gateway", + "prices": { + "input": "0.000003", + "output": "0.000015", + "cache_read": "0.0000003", + "reasoning": "0.00003" + }, + "counts": { + "input": 500, + "output": 200, + "cache_read": 100, + "reasoning": 50 + }, + "markup": "1", + "base": "0.00423", + "total": "0.00423", + "total_cents": "0.423" } ], "precomputed_cases": [ diff --git a/tests/unit/test_pricing.py b/tests/unit/test_pricing.py index c632e4d..b8d68e3 100644 --- a/tests/unit/test_pricing.py +++ b/tests/unit/test_pricing.py @@ -499,6 +499,81 @@ def test_workers_ai_provider_inferred_from_both_spellings(requested: str) -> Non 30, "tool_calls excluded", ), + # --- gateway SURFACES that re-shape every vendor (_OPENAI_SHAPED_APIS) --- + # Real shape from system.ai_gateway.usage: input CONTAINS cache_read even for + # Anthropic, whose own API reports it additively. Keying on the vendor billed + # 48,798 tokens against 31,091 consumed on a real backfill (1.570x). The honest + # total is the table's own total_tokens, i.e. input + output. + ( + CanonicalUsage( + input=1822, + output=4, + cache_read=1812, + provider="anthropic", + api="databricks_gateway", + model="m", + ), + 1826, + "databricks_gateway folds cache_read into input for every vendor", + ), + # The write half of the same shape — the overlap no provider-keyed set covers, + # because no vendor's native API reports cache_write inside input. + ( + CanonicalUsage( + input=1825, + output=4, + cache_write=1812, + provider="anthropic", + api="databricks_gateway", + model="m", + ), + 1829, + "databricks_gateway folds cache_write into input too", + ), + # Hosted Databricks models bill as TOKEN COUNTS (TOKEN_BILLED_PROVIDERS), so + # this path IS the bill. Latent today (0 of 96 hosted rows carry cache) and a + # direct 1.991x over-bill the day one does. + ( + CanonicalUsage( + input=1825, + output=4, + cache_read=1812, + provider="databricks", + api="databricks_gateway", + model="m", + ), + 1829, + "hosted databricks rows carry the surface's shape, not a vendor's", + ), + # A vendor the surface set must NOT reach: reasoning is inside output here even + # though gemini reports thoughts additively on its own API. + ( + CanonicalUsage( + input=500, + output=200, + reasoning=50, + provider="gemini", + api="databricks_gateway", + model="m", + ), + 700, + "databricks_gateway folds reasoning into output for every vendor", + ), + # Cloudflare is deliberately NOT in _OPENAI_SHAPED_APIS: measured on real logs, + # an anthropic entry reads input=10, output=4, total=14 with cache OUTSIDE that + # total. Adding it to the set would UNDER-bill by the cached portion. + ( + CanonicalUsage( + input=10, + output=4, + cache_read=3429, + provider="anthropic", + api="cloudflare_gateway", + model="m", + ), + 3443, + "cloudflare preserves each vendor's native shape", + ), ], ) def test_deoverlapped_token_total(usage: CanonicalUsage, expected: int, why: str) -> None: @@ -995,8 +1070,12 @@ def test_money_golden_cases() -> None: price = ModelPrice(source="openrouter", **prices) # `provider` is optional and defaults to a name in no _INCLUDES_ set, so # the pre-existing cases keep their original semantics; cases that pin - # per-provider token semantics set it explicitly. - usage = CanonicalUsage(model="m", provider=c.get("provider", "p"), api="native", **c["counts"]) + # per-provider token semantics set it explicitly. `api` defaults to + # "native" for the same reason — only the cases pinning a gateway + # surface's own token shape (_OPENAI_SHAPED_APIS) set it. + usage = CanonicalUsage( + model="m", provider=c.get("provider", "p"), api=c.get("api", "native"), **c["counts"] + ) b = compute_cost(usage, price, Decimal(c["markup"])) assert b.base == c["base"], f"{c['name']}: base {b.base} != {c['base']}" assert b.total == c["total"], f"{c['name']}: total {b.total} != {c['total']}" From 2254a48a754eace6218bffc636a58ce6c24e024a Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 13:05:45 +0200 Subject: [PATCH 13/22] Stamp backfilled events with the source row's time, not the run's MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit emit() read time.time() at each of its three push sites, so backfill_databricks billed a whole window into whatever period the script happened to run in. Measured on a live backfill over 2026-08-06: 128 events read off rows spanning 2026-08-06 to 2026-08-11 all carried one timestamp, the run's — up to 13.9 days of drift. Once the event is in Lago nothing can tell which period the usage actually belonged to. The same window now emits 25 distinct timestamps, one per distinct source time, 128 of 128 matching their own row. emit() takes a new timestamp= (a datetime, a naive one read as UTC, or epoch seconds) and backfill_databricks passes each row's own time through a new DatabricksUsageRow.occurred_at: event_time for a usage row, and for a spend row the bucket hour START. The start is the only instant certain to sit inside the hour that row aggregates; the hour's end would push a bucket closing exactly on a period boundary into the following period. Resolved once per call, ahead of every branch, rather than at each push site. A price-lookup miss falls through to the token path, so one usage row can reach two of those sites, and two separate clock reads there let a call straddling a period boundary land half in each period. A value that cannot be read is reported through on_error (where="timestamp") and the call still bills, at now. Stamping the wrong period is a reconciliation problem the operator can see and fix; dropping the event is revenue that never appears at all. An ISO-8601 string is deliberately not accepted, and neither is a numeric one. Python 3.10 is still supported and its fromisoformat rejects the trailing "Z" that gateway APIs emit, while the JS port's new Date() accepts it — a string would parse in one repo and fail in the other. int("1786112523") would likewise coerce where the JS port's typeof check refuses. Connectors parse their own source column instead, where the shapes it really returns are known and tested: _epoch takes both the API's "…Z" strings and databricks-sql-connector's datetime objects, and reads an offset-less stamp as UTC so the JS port's Date cannot treat it as local time. All five real column shapes verified to produce byte-identical epochs in both repos. The live wrap() path passes no timestamp and stamps now exactly as before, pinned by its own test. 11 tests added; the 10 behavioural ones confirmed failing with the fix reverted. 555 pass, ruff + format + mypy clean. --- CHANGELOG.md | 6 + src/lago_agent_sdk/gateway/databricks.py | 60 ++++++++++ src/lago_agent_sdk/sdk.py | 83 +++++++++++-- tests/unit/gateway/test_databricks_source.py | 65 ++++++++++- tests/unit/test_sdk.py | 115 +++++++++++++++++++ 5 files changed, 320 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5922d75..069c69f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **Backfilled events were stamped with the run's clock instead of the usage's own time.** `emit()` read `time.time()` at each of its three push sites, so `backfill_databricks` billed a whole window into whatever period the script happened to run in. Measured on a live backfill: 128 events read off rows spanning `2026-08-06` to `2026-08-11` all carried one timestamp, the run's — up to **13.9 days** of drift, and once the event is in Lago nothing can tell which period the usage actually belonged to. `emit()` now takes `timestamp=` (a `datetime`, a naive one read as UTC, or epoch seconds), and `backfill_databricks` passes each row's own time via a new `DatabricksUsageRow.occurred_at`: `event_time` for a usage row, and for a spend row the `bucket` hour **start**, which is the only instant certain to sit inside the hour that row aggregates — the hour's end would push a bucket closing exactly on a period boundary into the following period. + - Resolved **once per call**, ahead of every branch, rather than at each push site. A price-lookup miss falls through to the token path, so one usage row can reach two of those sites; two separate clock reads there let a call straddling a period boundary land half in each period. + - A value that cannot be read is reported through `on_error` (`where="timestamp"`) and the call still bills, at `now`. Never silently under-bill: stamping the wrong period is a reconciliation problem the operator can see and fix, while dropping the event is revenue that never appears at all. + - **An ISO-8601 string is deliberately not accepted.** Python 3.10 is still supported here and its `fromisoformat` rejects the trailing `Z` that gateway APIs emit, while the JS port's `new Date()` accepts it — so a string would parse in one repo and fail in the other. Connectors parse their own source column instead, where the shapes that column really returns are known and tested: the Databricks reader takes both the REST API's `"…Z"` strings and `databricks-sql-connector`'s `datetime` objects, and reads an offset-less stamp as UTC in both repos rather than letting JS's `Date` treat it as local time. + - The live `wrap()` path passes no timestamp and stamps `now` exactly as before, pinned by its own test. + - **Databricks gateway rows were de-overlapped by the vendor's token semantics instead of the surface's.** `system.ai_gateway.usage` re-reports every vendor in the OpenAI shape: measured across 246 rows and 6 vendors, `total_tokens == input + output` for **every** group, with `cache_read` *and* `cache_write` inside `input` and reasoning inside `output`. `compute_cost`/`deoverlapped_token_total` keyed that decision on the provider name — correct for a native call, wrong for a table row. Anthropic's own API reports cache additively (measured live: `cache_read=3962` against `input=9`), so an Anthropic row read from this table had its cached tokens counted a second time. On a real backfill that billed **48,798 tokens against 31,091 consumed, 1.570x**; the same window now reports 31,018, i.e. the same ~0.2% lag-only shortfall `openai` already had. - New `_OPENAI_SHAPED_APIS` keys on `CanonicalUsage.api` rather than the provider, because a gateway row reuses the live vendor names and nothing in the name separates the two. `workers-ai` stays a *provider* entry by contrast: it names a vendor reachable through exactly one surface, so the name is sufficient there. - Adds this SDK's first `cache_write ⊆ input` correction. It is surface-only by design — Anthropic is the one vendor whose native API bills cache writes at all, and it reports them additively, so no native response needs it. diff --git a/src/lago_agent_sdk/gateway/databricks.py b/src/lago_agent_sdk/gateway/databricks.py index b06d544..ded252a 100644 --- a/src/lago_agent_sdk/gateway/databricks.py +++ b/src/lago_agent_sdk/gateway/databricks.py @@ -132,6 +132,28 @@ def reconcile_dimensions(self) -> dict[str, str]: endpoint = _safe_str(self.usage.extras.get("endpoint_name")) return {"endpoint_name": endpoint} if endpoint else {} + @property + def occurred_at(self) -> int | None: + """When this row's usage actually happened, as unix seconds for `emit()`. + + The whole point of a backfill is that it runs long after the usage it bills, + so the run's own clock is never the right answer: a window reaching back a + week must bill into the periods those calls fell in, not into the period the + script happens to run in. + + Each kind reports the time its OWN surface is keyed by: + + * usage — `event_time`, the request's own instant. + * spend — `bucket`, the START of the hour it aggregates. An hourly total + covers [bucket, bucket + 1h), so the start is the only instant certain to + sit inside the row's own coverage; the hour's end would push a bucket + closing exactly on a period boundary into the following period. + + None when the column is absent or unreadable, which leaves `emit()` to stamp + `now` — the pre-existing behaviour, and better than dropping the event. + """ + return _epoch(self.raw.get("bucket") if self.kind == "spend" else self.raw.get("event_time")) + def event_id_for(self, subscription: str | None) -> str: """The same key, scoped to whichever subscription is actually billed. @@ -522,6 +544,44 @@ def _stamp(value: Any) -> str: return str(value) +def _epoch(value: Any) -> int | None: + """A timestamp column as unix seconds, from either access path. + + The Statement Execution API returns TIMESTAMPs as ISO-8601 strings ending in "Z"; + `databricks-sql-connector` returns real `datetime` objects. Both are supported + input paths and both have to yield the same instant. + + The trailing "Z" is rewritten by hand because Python 3.10 — still supported here — + rejects it in `fromisoformat`, and every string this table returns carries one. A + stamp with no offset at all is read as UTC, which is both the warehouse's own + session timezone and what the JS port does with the same string. + + Unparseable returns None rather than raising: a bad timestamp column must not cost + the caller the whole row. + """ + if value is None: + return None + if isinstance(value, datetime): + moment = value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) + return int(moment.timestamp()) + text = str(value).strip() + if not text: + return None + if text[-1] in "Zz": + text = f"{text[:-1]}+00:00" + try: + moment = datetime.fromisoformat(text) + except ValueError: + return None + if moment.tzinfo is None: + moment = moment.replace(tzinfo=timezone.utc) + try: + return int(moment.timestamp()) + # A year outside the C time range; platform dependent, hence caught here too. + except (OverflowError, OSError): + return None + + def _bucket_of(value: Any) -> str: return _truncate_hour(_stamp(value)) diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index a21091e..59a0a8f 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -7,6 +7,7 @@ import time import uuid from collections.abc import Iterable +from datetime import datetime, timezone from typing import Any from .canonical import CanonicalUsage @@ -34,6 +35,24 @@ ) +def _to_epoch_seconds(value: int | float | datetime) -> int: + """A caller-supplied event time as the unix seconds Lago's `timestamp` wants.""" + if isinstance(value, datetime): + # A naive datetime is taken as UTC — the same rule `_interval_sql` documents + # for the window bound, so a caller who reads a window and bills it cannot + # have the two disagree by their machine's UTC offset. + moment = value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) + return int(moment.timestamp()) + # `not bool`: it is an `int` subclass, so `True` would otherwise bill at epoch 1. + # The JS port's `typeof value === "number"` rejects it, and the two must agree. + if isinstance(value, (int, float)) and not isinstance(value, bool): + return int(value) + raise TypeError( + f"timestamp={value!r} not understood — pass a datetime or epoch seconds " + "(an ISO-8601 string is deliberately not accepted; see emit())" + ) + + class LagoSDK: def __init__( self, @@ -240,6 +259,7 @@ def emit( markup: float | None = None, usd_cost: float | None = None, event_id: str | None = None, + timestamp: int | float | datetime | None = None, ) -> None: """Emit usage to Lago. @@ -262,6 +282,19 @@ def emit( same window never double-bills. A live, one-shot call has no natural id to reuse and should leave this as None. + ``timestamp``: bill the events at this instant instead of at now — pass + the source row's own time when replaying/backfilling from a gateway's + logs, or a window reaching back a week bills every one of its calls into + the period the script happens to run in. Accepts a ``datetime`` (a naive + one is read as UTC) or epoch seconds. Deliberately NOT an ISO-8601 + string: Python 3.10 is still supported here and its ``fromisoformat`` + rejects the trailing "Z" that gateway APIs emit, while the JS port's + ``new Date()`` accepts it — so a string would parse in one repo and fail + in the other. Connectors parse their own source column instead, where the + shapes that column really returns are known and tested (see + ``DatabricksUsageRow.occurred_at``). A live call has no source time and + should leave this as None. + Both multi-event paths suffix per field so they don't collide with each other, and they use DIFFERENT namespaces so they can't collide across modes either: @@ -280,6 +313,11 @@ def emit( were never billed, only the raw token counts, and nothing surfaced it. """ try: + # Resolved ONCE, ahead of every branch: a price-lookup miss falls through + # to the token path, so one usage row can reach two of the push paths + # below. Two separate `time.time()` reads there let a call that straddles + # a billing-period boundary land half in each period. + at = self._event_time(timestamp) sub = self._resolve_subscription(subscription) if not sub: # `_report_error` is the single channel: it invokes on_error AND @@ -317,7 +355,7 @@ def emit( ), "pricing", ) - self._emit_token_events(usage, sub, dimensions, event_id) + self._emit_token_events(usage, sub, dimensions, event_id, at) return markup_value, ok = coerce_markup(markup if markup is not None else self.config.markup) @@ -337,7 +375,7 @@ def emit( # 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) + self._emit_token_events(usage, sub, dimensions, event_id, at) return else: price = self._pricing.lookup(usage.provider, usage.model, usage.api) @@ -346,14 +384,31 @@ def emit( self._report_error( PricingUnavailableError(usage.provider, usage.model, usage.api), "pricing" ) - self._emit_token_events(usage, sub, dimensions, event_id) + self._emit_token_events(usage, sub, dimensions, event_id, at) return breakdown = compute_cost(usage, price, markup_value) - self._push_cost_event(usage, breakdown, sub, dimensions, event_id) + self._push_cost_event(usage, breakdown, sub, dimensions, event_id, at) except Exception as exc: # noqa: BLE001 — never raise from emit self._report_error(exc, "emit") + def _event_time(self, timestamp: int | float | datetime | None) -> int: + """The instant to stamp this call's events with — the caller's, or now. + + A value we cannot read is reported and falls back to `now` rather than + dropping the call. Stamping the wrong period is a reconciliation problem the + operator can see and fix; losing the event is revenue that never appears at + all. Same trade-off as a missed price lookup. + """ + if timestamp is not None: + try: + return _to_epoch_seconds(timestamp) + # OverflowError/OSError: `datetime.timestamp()` raises them, platform + # dependently, for a year outside the C time range. + except (TypeError, ValueError, OverflowError, OSError) as exc: + self._report_error(exc, "timestamp") + return int(time.time()) + def _note_token_billed(self, usage: CanonicalUsage) -> None: """Say it once per model, at info level. @@ -373,7 +428,12 @@ def _note_token_billed(self, usage: CanonicalUsage) -> None: ) def _emit_token_events( - self, usage: CanonicalUsage, sub: str, dimensions: dict[str, Any] | None, event_id: str | None = None + self, + usage: CanonicalUsage, + sub: str, + dimensions: dict[str, Any] | None, + event_id: str | None = None, + at: int | None = None, ) -> None: nonzero = usage.nonzero_numeric() # A negative count is silently unbillable — Lago would otherwise sum it into @@ -391,7 +451,9 @@ def _emit_token_events( if not nonzero: # Mistral legacy / empty — nothing to bill return - now = int(time.time()) + # `emit` already resolved the instant; the fallback covers nothing today and + # is kept only so this stays callable on its own without stamping the epoch. + now = at if at is not None else int(time.time()) for field_name, value in nonzero.items(): code = self.config.metric_codes.get(field_name) if not code: @@ -422,6 +484,7 @@ def _push_cost_event( sub: str, dimensions: dict[str, Any] | None, event_id: str | None = None, + at: int | None = None, ) -> None: """Push one llm_cost event — or several, one per token_type, when a real per-field breakdown exists. @@ -441,7 +504,8 @@ def _push_cost_event( grouped by model only; no `token_type` at all rather than a fabricated one. """ - now = int(time.time()) + # See the note in `_emit_token_events` — `emit` is the one authority on this. + now = at if at is not None else int(time.time()) # Caller dimensions are spread LAST in each `properties` below, not here — # they must win over every SDK-computed key, exactly as they already do in # `_emit_token_events`. Spreading them into `base_properties` put them @@ -624,6 +688,10 @@ def backfill_databricks( # 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), + # The row's own time, not the run's — see `occurred_at`. A + # backfill that stamps `now` bills last week's usage into this + # week's period, and nothing in Lago can tell afterwards. + timestamp=row.occurred_at, ) counts["cost"] += 1 else: @@ -633,6 +701,7 @@ def backfill_databricks( dimensions=dims, mode="tokens", event_id=row.event_id_for(sub), + timestamp=row.occurred_at, ) counts["tokens"] += 1 return counts diff --git a/tests/unit/gateway/test_databricks_source.py b/tests/unit/gateway/test_databricks_source.py index 77613b3..2e03f71 100644 --- a/tests/unit/gateway/test_databricks_source.py +++ b/tests/unit/gateway/test_databricks_source.py @@ -8,12 +8,13 @@ from __future__ import annotations import json -from datetime import datetime +import time +from datetime import datetime, timezone from typing import Any import pytest -from lago_agent_sdk import LagoSDK +from lago_agent_sdk import CanonicalUsage, LagoSDK from lago_agent_sdk.gateway.databricks import DatabricksSource, DatabricksUsageRow, _interval_sql # -------------------------------------------------------------------------- @@ -733,3 +734,63 @@ def test_backfill_accepts_already_read_rows_without_querying_again() -> None: 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 + + +# -------------------------------------------------------------------------- +# Event time — a backfill runs long after the usage it bills +# -------------------------------------------------------------------------- +def _utc(*args: int) -> int: + return int(datetime(*args, tzinfo=timezone.utc).timestamp()) # type: ignore[arg-type] + + +def test_occurred_at_reads_each_kinds_own_time_column() -> None: + """A usage row's own instant; a spend row's hour START — the only instant certain + to sit inside the hour that row aggregates.""" + hosted, byok = list(_source([_BYOK_SPEND], [_HOSTED]).read_usage("1 day")) + assert hosted.kind == "spend" and byok.kind == "usage" + assert hosted.occurred_at == _utc(2026, 8, 7, 14, 0, 0) + assert byok.occurred_at == _utc(2026, 8, 7, 14, 22, 3) + + +def test_occurred_at_reads_a_datetime_column_the_same_way() -> None: + """`databricks-sql-connector` returns TIMESTAMPs as `datetime`, the REST API as + ISO-8601 strings ending in "Z". Both are supported access paths, so both must + resolve to the same instant.""" + rows = list( + _source( + [{**_BYOK_SPEND, "bucket": datetime(2026, 8, 7, 14, 0, 0)}], + [{**_HOSTED, "event_time": "2026-08-07T14:22:03.123Z"}], + ).read_usage("1 day") + ) + assert {r.occurred_at for r in rows} == {_utc(2026, 8, 7, 14, 0, 0), _utc(2026, 8, 7, 14, 22, 3)} + + +def test_a_row_with_no_readable_time_has_no_occurred_at() -> None: + """None leaves `emit()` stamping `now`, which is wrong but billed — better than + losing the row over a bad column.""" + for bad in (None, "", "not a timestamp"): + row = DatabricksUsageRow( + usage=CanonicalUsage(model="m", provider="databricks", api="databricks_gateway"), + subscription="sub", + row_id="r", + kind="usage", + raw={"event_time": bad}, + ) + assert row.occurred_at is None + + +def test_backfill_events_carry_the_source_rows_time_not_the_run_time() -> None: + """Live-proven before the fix: 128 events off one window spanning 2026-08-06 to + 2026-08-11 all carried the run's own clock, billing historical usage into the + current period.""" + sdk, q = _sdk() + sdk.backfill_databricks(_source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE]), "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 cost and tokens + # The spend row's hour, and the hosted request's own second. + assert {e["timestamp"] for e in cost} == {_utc(2026, 8, 7, 14, 0, 0)} + assert {e["timestamp"] for e in tokens} == {_utc(2026, 8, 7, 14, 22, 3)} + assert all(e["timestamp"] < int(time.time()) - 86400 for e in q.events), "not the run time" diff --git a/tests/unit/test_sdk.py b/tests/unit/test_sdk.py index 64711c1..c0b619e 100644 --- a/tests/unit/test_sdk.py +++ b/tests/unit/test_sdk.py @@ -3,6 +3,8 @@ from __future__ import annotations import logging +import time +from datetime import datetime, timezone import pytest @@ -353,3 +355,116 @@ def boom(name, *args, **kwargs): assert sdk.config.verify_ssl is False finally: sdk.shutdown(timeout=1.0) + + +# -------------------------------------------------------------------------- +# Event time — a backfill must bill into the period the usage happened in +# -------------------------------------------------------------------------- +def test_emit_stamps_the_given_instant_on_every_event_not_now() -> None: + """Without this, a replay of last week's logs billed every call into the period + the script happened to run in, and nothing in Lago could tell afterwards.""" + sdk, received = _new_sdk(default_sub="sub") + when = datetime(2026, 8, 7, 14, 22, 3, tzinfo=timezone.utc) + u = CanonicalUsage(input=10, output=20, model="m", provider="p", api="bedrock_invoke") + sdk.emit(u, timestamp=when) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert len(flat) == 2 + # One instant for the whole call: Lago sums these into a period, so a call must + # never straddle two of them because two `time.time()` reads disagreed. + assert {e["timestamp"] for e in flat} == {int(when.timestamp())} + + +def test_emit_stamps_a_cost_event_too() -> None: + """The cost path reads its own clock, so it needed threading separately from the + token path — and a backfill of BYOK spend goes down this one.""" + sdk, received = _new_sdk(default_sub="sub") + when = datetime(2026, 8, 7, 14, 0, 0, tzinfo=timezone.utc) + u = CanonicalUsage(input=10, output=20, model="m", provider="anthropic", api="native") + sdk.emit(u, mode="price", usd_cost=0.0011187, timestamp=when) + 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"] + assert flat[0]["timestamp"] == int(when.timestamp()) + + +def test_emit_accepts_epoch_seconds() -> None: + sdk, received = _new_sdk(default_sub="sub") + u = CanonicalUsage(input=1, model="m", provider="p", api="bedrock_invoke") + sdk.emit(u, timestamp=1786112523) + # A float is what `datetime.timestamp()` hands back, so it must not be refused. + sdk.emit(u, timestamp=1786112523.987) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert {e["timestamp"] for e in flat} == {1786112523} + + +def test_a_naive_timestamp_is_read_as_utc() -> None: + """Same rule as `_interval_sql`'s window bound, and the same rule the JS port + applies to a `Date` — otherwise a caller who reads a window and bills it has the + two disagree by their machine's UTC offset.""" + sdk, received = _new_sdk(default_sub="sub") + u = CanonicalUsage(input=1, model="m", provider="p", api="bedrock_invoke") + sdk.emit(u, timestamp=datetime(2026, 8, 7, 14, 22, 3)) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert flat[0]["timestamp"] == int(datetime(2026, 8, 7, 14, 22, 3, tzinfo=timezone.utc).timestamp()) + + +def test_an_unreadable_timestamp_is_reported_and_still_bills() -> None: + """Never silently under-bill: a bad timestamp is a reconciliation problem the + operator can see and fix, while dropping the event is revenue that never appears. + An ISO string is the likely mistake, and is deliberately not accepted.""" + errors: list = [] + received: list = [] + cfg = LagoConfig( + api_key="dummy", + default_subscription_id="sub", + on_error=lambda exc, where: errors.append((str(exc), where)), + ) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: received.extend(b) # type: ignore[attr-defined] + before = int(time.time()) + sdk.emit(CanonicalUsage(input=1, model="m", provider="p", api="bedrock_invoke"), timestamp="2026-08-07Z") + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + + assert errors, "an unreadable timestamp must reach on_error" + msg, where = errors[0] + assert "2026-08-07Z" in msg and where == "timestamp" + # ...and the call is still billed, at now. + assert len(received) == 1 + assert before <= received[0]["timestamp"] <= int(time.time()) + + +def test_a_numeric_string_is_refused_not_coerced() -> None: + """`int("1786112523")` would sail through where the isinstance check rejects it — + the same input must not bill in one repo and report an error in the other. The JS + port's `Number()` is the one that would coerce, so this is pinned on both sides.""" + errors: list = [] + cfg = LagoConfig( + api_key="dummy", + default_subscription_id="sub", + on_error=lambda exc, where: errors.append((str(exc), where)), + ) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: None # type: ignore[attr-defined] + sdk.emit(CanonicalUsage(input=1, model="m", provider="p", api="bedrock_invoke"), timestamp="1786112523") + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + assert [where for _, where in errors] == ["timestamp"] + + +def test_no_timestamp_still_stamps_now() -> None: + """The live `wrap()` path passes nothing and must be unchanged by all of this.""" + sdk, received = _new_sdk(default_sub="sub") + before = int(time.time()) + sdk.emit(CanonicalUsage(input=1, model="m", provider="p", api="bedrock_invoke")) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert before <= flat[0]["timestamp"] <= int(time.time()) From f9179d91a237b19fa913c7d4b070dc0f73c5ff96 Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 13:49:27 +0200 Subject: [PATCH 14/22] Read one Databricks window, floored, excluding the open hour MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _interval_sql returned a SQL string, so current_timestamp() - INTERVAL 1 DAY was re-evaluated per statement — 5.1s of drift measured between the spend read and the usage read on a warm warehouse. Spend runs first, so the usage window was the narrower one, and a Databricks-hosted row landing in that gap was read by neither statement: hosted bills from ai_gateway.usage alone, so the call was simply never billed. Both statements now carry one pair of literals, resolved once here rather than twice in SQL. The window is floored to the hour. external_model_spend is an hourly aggregate whose usage_start_time is always the hour START — 65 of 65 live rows, none unaligned — so a mid-hour bound failed the predicate for the hour containing it while ai_gateway.usage happily returned that same hour's rows. Measured live: a since of 13:30 read 11 of 65 spend rows, dropping $0.1256 of $0.1723 (73%) of the window's metered dollars, while still reading 35 BYOK usage rows (31,815 tokens) from inside the hour it had dropped — tokens that then tripped the "no spend row" warning as if the table were lagging. The same read now returns all 65, both statements carrying the identical floored pair. Flooring the lower bound can read rows slightly older than asked for. That is safe in the only direction that matters: every transaction_id is derived from the source row, so a row already billed is rejected as a duplicate, and one not billed yet should be. Under-reading is what loses money. The still-aggregating hour is excluded. A spend row cannot be complete before its hour closes: the 08:00-09:00 row appeared ~7 min after 09:00, so the wait is not a fixed lag but however long is left in the hour — ~9 min for a call at :58, ~66 min for one at :01. Billing the open hour bills a fraction of it under that hour's record_id, and the corrected re-run is then rejected by Lago as a duplicate transaction_id, so the remainder is never billed at all. The bound applies to both tables, because a window whose halves cover different hours is the first bug over again. The caller-visible consequence is that the newest hour arrives on the next run; this reader keeps no cursor, so pass a window comfortably wider than the run interval. A since resolving entirely inside the excluded hour ("30 minutes") now warns and reads nothing, rather than spending warehouse time to return zero rows that say nothing about whether there was traffic. Closes the session-timezone dependency too. The bound no longer goes through current_timestamp(), and the literal is rendered zone-explicit (TIMESTAMP '... +00:00'): a bare literal is parsed in the warehouse's own spark.sql.session.timeZone, so on a workspace set to anything but UTC the identical literal named a different instant and the whole window slid by that offset. Verified live that the suffixed form is accepted and resolves to the same epoch under this warehouse's Etc/UTC. The interval string no longer reaches SQL at all, but stays validated strictly: a window quietly read as something other than what the caller wrote under-reads. Re-verified live end-to-end: the same backfill over 2026-08-06 still bills anthropic 31,018 against 31,091 consumed and openai 15,589 against 15,637 — the byte-identical figures from the cache de-overlap fix, so the money path is untouched. --- CHANGELOG.md | 8 +- README.md | 2 + src/lago_agent_sdk/gateway/databricks.py | 130 +++++++++++++++---- src/lago_agent_sdk/sdk.py | 2 +- tests/unit/gateway/test_databricks_source.py | 105 ++++++++++++--- tests/unit/test_sdk.py | 2 +- 6 files changed, 205 insertions(+), 44 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 069c69f..599ea01 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **The Databricks reader read two different windows, misaligned with the table it was reading.** `_interval_sql` returned a SQL *string*, so `current_timestamp() - INTERVAL 1 DAY` was re-evaluated per statement — measured **5.1s of drift** between the spend read and the usage read on a warm warehouse. Spend runs first, so the usage window was the narrower one, and a Databricks-**hosted** row landing in that gap was read by neither statement: hosted bills from `ai_gateway.usage` alone, so the call was simply never billed. Both statements now carry one pair of literals resolved once in Python/JS. + - **The window is floored to the hour.** `external_model_spend` is an hourly aggregate whose `usage_start_time` is always the hour *start* — 65 of 65 live rows, none unaligned — so a mid-hour bound failed the predicate for the hour *containing* it while `ai_gateway.usage` happily returned that same hour's rows. Measured live: `since` of `13:30` read **11 of 65** spend rows, dropping **$0.1256 of $0.1723 (73%)** of the window's metered dollars, while still reading 35 BYOK usage rows (31,815 tokens) from inside the hour it had dropped — tokens that then tripped the "no spend row" warning as if the table were lagging. Flooring the lower bound is what makes the two tables agree on one window; it can read rows slightly older than asked for, which is safe in the only direction that matters, because every `transaction_id` is derived from the source row and Lago rejects a row already billed. + - **The still-aggregating hour is excluded** (`< date_trunc('HOUR', now)`, in effect). A spend row cannot be complete before its hour closes: the `08:00–09:00` row appeared ~7 min **after** 09:00, i.e. the wait is not a fixed lag but "however long is left in the hour" — from ~9 min for a call at :58 to ~66 min for one at :01. Billing the open hour bills a fraction of it under that hour's `record_id`, and the corrected re-run is then rejected by Lago as a **duplicate `transaction_id`** — so the remainder is never billed at all. The upper bound applies to both tables, because a window whose halves cover different hours is the first bug over again. The practical consequence for a caller is that the newest hour arrives on the next run; this reader keeps no cursor, so pass a window comfortably wider than your run interval. + - A `since` that resolves *entirely* inside the excluded hour (`"30 minutes"`) now warns and reads nothing, rather than spending warehouse time to return zero rows that say nothing about whether there was traffic. + - **Closes the session-timezone dependency too.** The bound no longer goes through `current_timestamp()`, and the literal is rendered zone-explicit (`TIMESTAMP '… +00:00'`) — a bare literal is parsed in the warehouse's own `spark.sql.session.timeZone`, so on a workspace set to anything but UTC the identical literal named a different instant and the whole window slid by that offset. Verified live that the suffixed form is accepted and resolves to the same epoch under this warehouse's `Etc/UTC`. + - **Backfilled events were stamped with the run's clock instead of the usage's own time.** `emit()` read `time.time()` at each of its three push sites, so `backfill_databricks` billed a whole window into whatever period the script happened to run in. Measured on a live backfill: 128 events read off rows spanning `2026-08-06` to `2026-08-11` all carried one timestamp, the run's — up to **13.9 days** of drift, and once the event is in Lago nothing can tell which period the usage actually belonged to. `emit()` now takes `timestamp=` (a `datetime`, a naive one read as UTC, or epoch seconds), and `backfill_databricks` passes each row's own time via a new `DatabricksUsageRow.occurred_at`: `event_time` for a usage row, and for a spend row the `bucket` hour **start**, which is the only instant certain to sit inside the hour that row aggregates — the hour's end would push a bucket closing exactly on a period boundary into the following period. - Resolved **once per call**, ahead of every branch, rather than at each push site. A price-lookup miss falls through to the token path, so one usage row can reach two of those sites; two separate clock reads there let a call straddling a period boundary land half in each period. - A value that cannot be read is reported through `on_error` (`where="timestamp"`) and the call still bills, at `now`. Never silently under-bill: stamping the wrong period is a reconciliation problem the operator can see and fix, while dropping the event is revenue that never appears at all. @@ -141,7 +147,7 @@ All notable changes to this project will be documented here. Format follows [Kee - **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. + - **The window is validated, and no longer reaches SQL at all.** Only a bare count plus unit, or a `datetime`, is accepted, so `read_usage("1 day; DROP TABLE …")` is refused outright; what reaches SQL is a literal rendered from the resolved instant (see the window fix above). - **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. diff --git a/README.md b/README.md index 42c12ab..6744a24 100644 --- a/README.md +++ b/README.md @@ -235,6 +235,8 @@ sdk.flush() 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. +The window reads **whole closed hours only**: it is floored to the hour at both ends and the current, still-aggregating hour is excluded, because `external_model_spend` is an hourly aggregate whose row for an hour does not exist until that hour closes. So the newest hour of traffic arrives on the next run — pass a window comfortably wider than your run interval, since this reader keeps no cursor. + 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()`: diff --git a/src/lago_agent_sdk/gateway/databricks.py b/src/lago_agent_sdk/gateway/databricks.py index ded252a..8b2ef9f 100644 --- a/src/lago_agent_sdk/gateway/databricks.py +++ b/src/lago_agent_sdk/gateway/databricks.py @@ -37,7 +37,7 @@ import re from collections.abc import Iterator from dataclasses import dataclass, field -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Any from ..canonical import CanonicalUsage @@ -51,8 +51,11 @@ _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. +# An interval string no longer reaches SQL at all — `_window_bounds` resolves it to an +# instant and `_timestamp_sql` renders the literal — but it stays validated strictly +# rather than loosely parsed. Only a bare count plus a unit is accepted: a window quietly +# read as something other than what the caller wrote under-reads, and under-reading is +# the one direction that loses money. _INTERVAL_RE = re.compile(r"^\s*(\d{1,5})\s+(second|minute|hour|day|week)s?\s*$", re.I) @@ -166,24 +169,82 @@ def event_id_for(self, subscription: str | None) -> str: 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.""" +def _as_utc(moment: datetime) -> datetime: + """A naive datetime is taken as UTC; an aware one is CONVERTED, never reformatted. + + Databricks stores `event_time`/`usage_start_time` in UTC, so formatting an aware + datetime as-is would emit local wall time and a Europe/Paris caller would read a + window two hours in the future, bill nothing, and report success. Same rule as + `_epoch` below and as `sdk.py`'s event timestamps, and the same rule the JS port's + `Date` arithmetic follows. + """ + return ( + moment.astimezone(timezone.utc) if moment.tzinfo is not None else moment.replace(tzinfo=timezone.utc) + ) + + +def _floor_hour(moment: datetime) -> datetime: + """The start of the hour containing `moment`. + + Distinct from `_truncate_hour`, which trims a timestamp STRING to build a join key. + This one moves an instant, and it decides what gets read at all. + """ + return moment.replace(minute=0, second=0, microsecond=0) + + +def _timestamp_sql(moment: datetime) -> str: + """Render an instant as a zone-explicit SQL TIMESTAMP literal. + + The `+00:00` is not decoration. A bare `TIMESTAMP '2026-08-07 13:00:00'` is parsed + in the warehouse's own `spark.sql.session.timeZone`, so on a workspace set to + anything but UTC the identical literal names a different instant and the whole + window slides by that offset. Verified live: the suffixed form is accepted and + resolves to the same epoch as the bare form under this warehouse's `Etc/UTC`. + """ + return f"TIMESTAMP '{moment.strftime('%Y-%m-%d %H:%M:%S')}+00:00'" + + +def _window_bounds(since: str | datetime, *, now: datetime | None = None) -> tuple[datetime, datetime]: + """Resolve the read window to ONE pair of instants, both floored to the hour. + + Rejects anything unrecognized. Three money bugs live in leaving any part of this + to SQL, all three measured on real gateway tables: + + * **Two statements, two windows.** `current_timestamp() - INTERVAL 1 DAY` is a + SQL *string*, so it is re-evaluated per statement — 5.1s of drift measured + between the spend read and the usage read. Spend runs first, so the usage + window is the narrower one, and a hosted row in the gap is read by neither + statement. Hosted is billed from `usage` alone, so that row is simply lost. + * **The boundary hour.** `external_model_spend` is an hourly aggregate whose + `usage_start_time` is always the hour START (65 of 65 rows). Compared against + a mid-hour bound, the hour CONTAINING that bound fails the predicate while its + usage rows pass: live, a `since` of 13:30 read 11 of 65 spend rows and dropped + $0.1256 of $0.1723, while still reading 35 BYOK usage rows from inside the + hour it dropped. Flooring is what makes the two tables agree on one window. + * **The open hour.** A spend row cannot be complete before its hour closes (the + 08:00–09:00 row appeared ~7 min AFTER 09:00). Billing it early bills a + fraction of the hour under that hour's `record_id`, and the corrected re-run + is then rejected by Lago as a duplicate `transaction_id` — so the remainder is + never billed at all. Hence the upper bound, and hence both tables get it: a + window whose halves cover different hours is the first bug again. + + Flooring the lower bound can read rows slightly older than the caller asked for. + That is deliberate: every `transaction_id` is derived from the source row, so a row + already billed is rejected as a duplicate and one not billed yet SHOULD be. + Under-reading is the only direction that loses money. + """ + moment = _as_utc(now) if now is not None else datetime.now(timezone.utc) 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}" + lower = _as_utc(since) + else: + 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'" + ) + lower = moment - timedelta(**{f"{m.group(2).lower()}s": int(m.group(1))}) + return _floor_hour(lower), _floor_hour(moment) class DatabricksSource: @@ -372,8 +433,31 @@ def read_usage( Rows whose usage is entirely zero (failed calls are recorded with NULL token counts) are skipped, so nothing emits an empty event. + + Reads **whole closed hours only**: the window is floored to the hour at both + ends and the current, still-aggregating hour is excluded, because a spend row + for it cannot be complete yet. `_window_bounds` documents why each of those + three properties is load-bearing. The practical consequence for a caller is + that the newest hour of traffic arrives on the NEXT run, so pass a window + comfortably wider than your run interval — this reader keeps no cursor. """ - window = _interval_sql(since) + lower, upper = _window_bounds(since) + if lower >= upper: + # Not an error, but it must not read as success either: the caller asked + # for a window that lies entirely inside the hour this reader excludes, so + # zero rows here says nothing about whether there was traffic. + logger.warning( + "lago: since=%r resolves to [%s, %s), which is empty — the window falls " + "inside the current, still-aggregating hour that this reader excludes. " + "Nothing was read; widen the window past the hour boundary.", + since, + lower.isoformat(), + upper.isoformat(), + ) + return + # One pair of literals, both statements — see `_window_bounds`. Resolving the + # bounds here rather than in SQL is what makes the two reads the same window. + window, ceiling = _timestamp_sql(lower), _timestamp_sql(upper) spend = self.query(f""" SELECT record_id, @@ -383,12 +467,12 @@ def read_usage( to_json(custom_tags.request_tags) AS request_tags, usage_quantity FROM system.ai_gateway.external_model_spend - WHERE usage_start_time >= {window} + WHERE usage_start_time >= {window} AND usage_start_time < {ceiling} """) usage = self.query(f""" SELECT * FROM system.ai_gateway.usage - WHERE event_time >= {window} + WHERE event_time >= {window} AND event_time < {ceiling} ORDER BY event_time """) diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index 59a0a8f..81cefe3 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -38,7 +38,7 @@ def _to_epoch_seconds(value: int | float | datetime) -> int: """A caller-supplied event time as the unix seconds Lago's `timestamp` wants.""" if isinstance(value, datetime): - # A naive datetime is taken as UTC — the same rule `_interval_sql` documents + # A naive datetime is taken as UTC — the same rule `_as_utc` documents # for the window bound, so a caller who reads a window and bills it cannot # have the two disagree by their machine's UTC offset. moment = value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) diff --git a/tests/unit/gateway/test_databricks_source.py b/tests/unit/gateway/test_databricks_source.py index 2e03f71..50381d8 100644 --- a/tests/unit/gateway/test_databricks_source.py +++ b/tests/unit/gateway/test_databricks_source.py @@ -8,14 +8,21 @@ from __future__ import annotations import json +import logging import time -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Any import pytest from lago_agent_sdk import CanonicalUsage, LagoSDK -from lago_agent_sdk.gateway.databricks import DatabricksSource, DatabricksUsageRow, _interval_sql +from lago_agent_sdk.gateway.databricks import ( + DatabricksSource, + DatabricksUsageRow, + _floor_hour, + _timestamp_sql, + _window_bounds, +) # -------------------------------------------------------------------------- # Fake rows, in the exact shapes the two tables return @@ -88,14 +95,49 @@ def fake_query(sql: str) -> list[dict[str, Any]]: # -------------------------------------------------------------------------- # 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" +_NOW = datetime(2026, 8, 21, 13, 34, 12, tzinfo=timezone.utc) + + +def test_interval_strings_resolve_to_instants_not_sql() -> None: + """Resolved in Python, so both statements can share one bound. A `current_timestamp()` + expression is re-evaluated per statement and the two reads then cover different + windows — measured 5.1s apart, and a hosted row in the gap is billed by neither.""" + assert _window_bounds("1 day", now=_NOW) == ( + datetime(2026, 8, 20, 13, 0, tzinfo=timezone.utc), + datetime(2026, 8, 21, 13, 0, tzinfo=timezone.utc), + ) + assert _window_bounds("36 hours", now=_NOW)[0] == datetime(2026, 8, 20, 1, 0, tzinfo=timezone.utc) + assert _window_bounds("2 weeks", now=_NOW)[0] == datetime(2026, 8, 7, 13, 0, tzinfo=timezone.utc) 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'" + assert _timestamp_sql(datetime(2026, 8, 7, 14, 0, 0)) == "TIMESTAMP '2026-08-07 14:00:00+00:00'" + + +def test_the_window_is_floored_to_the_hour_at_both_ends() -> None: + """`external_model_spend` is an hourly aggregate whose `usage_start_time` is always + the hour start (65 of 65 live rows), so a mid-hour bound drops the hour CONTAINING + it while the usage table still yields that hour's rows. Live: `since` of 13:30 read + 11 of 65 spend rows, dropping $0.1256 of $0.1723, and still read 35 BYOK usage rows + from inside the dropped hour.""" + lower, upper = _window_bounds(datetime(2026, 8, 7, 13, 30, 45, tzinfo=timezone.utc), now=_NOW) + assert lower == datetime(2026, 8, 7, 13, 0, tzinfo=timezone.utc) + assert upper == datetime(2026, 8, 21, 13, 0, tzinfo=timezone.utc) + + +def test_the_still_aggregating_hour_is_excluded() -> None: + """A spend row cannot be complete before its hour closes — the 08:00–09:00 row + appeared ~7 min AFTER 09:00. Billing the open hour bills a fraction of it under + that hour's `record_id`, and Lago then rejects the corrected re-run as a duplicate + `transaction_id`, so the remainder is never billed.""" + assert _window_bounds("1 day", now=_NOW)[1] == _floor_hour(_NOW) < _NOW + + +def test_timestamp_literals_name_their_zone() -> None: + """A bare literal is parsed in the warehouse's `spark.sql.session.timeZone`, so on a + non-UTC workspace the same literal names a different instant and the window slides + by the offset. Verified live that the suffixed form is accepted.""" + assert _timestamp_sql(_NOW).endswith("+00:00'") @pytest.mark.parametrize( @@ -109,19 +151,46 @@ def test_datetime_window_renders_as_a_literal() -> None: ], ) 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.""" + """Anything but a bare count-plus-unit is refused outright. The string no longer + reaches SQL — the bound is resolved to an instant first — so this is no longer an + injection guard; it is what stops a window being quietly read as something other + than what the caller wrote.""" with pytest.raises(ValueError, match="not understood"): - _interval_sql(bad) + _window_bounds(bad) -def test_read_usage_scopes_both_queries_to_the_window() -> None: +def test_read_usage_scopes_both_queries_to_one_shared_window() -> None: + """The point of resolving the bounds in Python: both statements must carry the SAME + two literals. Two `current_timestamp()` expressions drift, and the read that runs + second covers the narrower window — a hosted row in the gap is lost, since hosted + is billed from `usage` alone.""" src = _source([], []) + before = _floor_hour(datetime.now(timezone.utc)) list(src.read_usage("3 days")) + after = _floor_hour(datetime.now(timezone.utc)) 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 + spend, usage = src.queries # type: ignore[attr-defined] + assert "current_timestamp()" not in spend and "current_timestamp()" not in usage + ceilings = {_timestamp_sql(before), _timestamp_sql(after)} + for sql, column in ((spend, "usage_start_time"), (usage, "event_time")): + lower = _timestamp_sql(before - timedelta(days=3)) + assert ( + f"{column} >= {lower}" in sql or f"{column} >= {_timestamp_sql(after - timedelta(days=3))}" in sql + ) + assert any(f"{column} < {c}" in sql for c in ceilings) + + +def test_a_window_entirely_inside_the_open_hour_reads_nothing_and_says_so( + caplog: pytest.LogCaptureFixture, +) -> None: + """Excluding the open hour means a sub-hour window can resolve to nothing. Zero rows + then says nothing about whether there was traffic, so it must not pass silently — + and it must not spend warehouse time either.""" + src = _source([], []) + with caplog.at_level(logging.WARNING): + assert list(src.read_usage("30 minutes")) == [] + assert src.queries == [] # type: ignore[attr-defined] + assert "widen the window" in caplog.text # -------------------------------------------------------------------------- @@ -657,12 +726,12 @@ def test_an_aware_datetime_window_is_converted_to_utc() -> None: 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'" + aware = _window_bounds(datetime(2026, 8, 11, 14, 0, 0, tzinfo=paris), now=_NOW)[0] + assert _timestamp_sql(aware) == "TIMESTAMP '2026-08-11 12:00:00+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'" + naive = _window_bounds(datetime(2026, 8, 11, 14, 0, 0), now=_NOW)[0] + assert _timestamp_sql(naive) == "TIMESTAMP '2026-08-11 14:00:00+00:00'" def test_datetime_timestamp_columns_still_bucket_and_reconcile() -> None: diff --git a/tests/unit/test_sdk.py b/tests/unit/test_sdk.py index c0b619e..581d998 100644 --- a/tests/unit/test_sdk.py +++ b/tests/unit/test_sdk.py @@ -403,7 +403,7 @@ def test_emit_accepts_epoch_seconds() -> None: def test_a_naive_timestamp_is_read_as_utc() -> None: - """Same rule as `_interval_sql`'s window bound, and the same rule the JS port + """Same rule as `_as_utc`'s window bound, and the same rule the JS port applies to a `Date` — otherwise a caller who reads a window and bills it has the two disagree by their machine's UTC offset.""" sdk, received = _new_sdk(default_sub="sub") From 15d14336737406ec4b547ede92a228fb0b4c63f9 Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 14:12:20 +0200 Subject: [PATCH 15/22] Keep failed Databricks calls out of the BYOK join index MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A rejected external call is logged with NULL tokens and an empty `destination_model`, so its key can never match an `external_model_spend` row — the call bought nothing and Databricks meters no dollars for it. The BYOK indexing loop had no zero-usage guard, so those rows became buckets that fell through to the "no spend row yet ... re-run this window later to bill them" warning, which for them is advice that can never work. Measured live over a fully-aggregated window: 29 buckets reported, 28 of them phantoms (all 83 zero-usage rows in the window were 4xx/5xx), and the one example bucket the warning names for the operator was a phantom — so the single genuinely lagging bucket was the least visible thing in the message. The hosted loop below already applied this guard. Also fixes a wall-clock-dependent test shipped with the window work: it drove the open-hour guard through the real clock, so "30 minutes" spanned two hours at :34 and one at :04 and the assertion passed or failed on the current minute. The collapse is now pinned against the frozen `_NOW`, and the guard itself is driven by a `since` half an hour ahead of the clock. --- CHANGELOG.md | 3 + src/lago_agent_sdk/gateway/databricks.py | 9 +++ tests/unit/gateway/test_databricks_source.py | 60 +++++++++++++++++++- 3 files changed, 69 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 599ea01..794d7bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,9 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **Failed Databricks calls were reported to the operator as revenue waiting to be billed.** The BYOK indexing loop had no zero-usage guard, so a rejected external call — NULL tokens, and an empty `destination_model` because the gateway never resolved one — became a key in the token index. It can never match a spend row (the call bought nothing, so Databricks meters no dollars for it), so it fell straight through to the "no `external_model_spend` row yet … re-run this window later to bill them" warning, which for these buckets is advice that can never work. Measured live over a fully-aggregated window: **29 buckets reported, 28 of them phantoms** — all 83 zero-usage rows in the window were 4xx/5xx — and the single example bucket the warning names for the operator was one of the phantoms, leaving the one genuinely lagging bucket the least visible thing in the message. The hosted loop already applied this guard; the BYOK loop now applies the same one, for the same reason. + - Also stops a failed call from becoming a bucket's representative in `_merge_usage`, which keeps the first row's non-numeric fields. Unobserved today — a failed row's empty `destination_model` puts it in its own bucket, and `status_code` is already excluded from a bucket's extras as per-request — so this closes it before the join key or that exclusion list changes rather than after. + - **The Databricks reader read two different windows, misaligned with the table it was reading.** `_interval_sql` returned a SQL *string*, so `current_timestamp() - INTERVAL 1 DAY` was re-evaluated per statement — measured **5.1s of drift** between the spend read and the usage read on a warm warehouse. Spend runs first, so the usage window was the narrower one, and a Databricks-**hosted** row landing in that gap was read by neither statement: hosted bills from `ai_gateway.usage` alone, so the call was simply never billed. Both statements now carry one pair of literals resolved once in Python/JS. - **The window is floored to the hour.** `external_model_spend` is an hourly aggregate whose `usage_start_time` is always the hour *start* — 65 of 65 live rows, none unaligned — so a mid-hour bound failed the predicate for the hour *containing* it while `ai_gateway.usage` happily returned that same hour's rows. Measured live: `since` of `13:30` read **11 of 65** spend rows, dropping **$0.1256 of $0.1723 (73%)** of the window's metered dollars, while still reading 35 BYOK usage rows (31,815 tokens) from inside the hour it had dropped — tokens that then tripped the "no spend row" warning as if the table were lagging. Flooring the lower bound is what makes the two tables agree on one window; it can read rows slightly older than asked for, which is safe in the only direction that matters, because every `transaction_id` is derived from the source row and Lago rejects a row already billed. - **The still-aggregating hour is excluded** (`< date_trunc('HOUR', now)`, in effect). A spend row cannot be complete before its hour closes: the `08:00–09:00` row appeared ~7 min **after** 09:00, i.e. the wait is not a fixed lag but "however long is left in the hour" — from ~9 min for a call at :58 to ~66 min for one at :01. Billing the open hour bills a fraction of it under that hour's `record_id`, and the corrected re-run is then rejected by Lago as a **duplicate `transaction_id`** — so the remainder is never billed at all. The upper bound applies to both tables, because a window whose halves cover different hours is the first bug over again. The practical consequence for a caller is that the newest hour arrives on the next run; this reader keeps no cursor, so pass a window comfortably wider than your run interval. diff --git a/src/lago_agent_sdk/gateway/databricks.py b/src/lago_agent_sdk/gateway/databricks.py index 8b2ef9f..5e825ed 100644 --- a/src/lago_agent_sdk/gateway/databricks.py +++ b/src/lago_agent_sdk/gateway/databricks.py @@ -486,6 +486,15 @@ def read_usage( for row, u in extracted: if u.provider == "databricks": continue + if not u.nonzero_numeric(): + # A failed call carries NULL tokens and bought nothing, so Databricks + # meters no dollars for it — its key can never match a spend row. Indexed, + # it becomes a bucket the warning below tells the operator to re-run the + # window for, which can never bill it: measured live over 2026-08-06, + # 28 of the 29 reported buckets were these phantoms (all 83 zero-usage + # rows were 4xx/5xx), including the one example row the warning names. + # The hosted loop applies the same guard for the same reason. + continue key = ( _bucket_of(row.get("event_time")), u.provider, diff --git a/tests/unit/gateway/test_databricks_source.py b/tests/unit/gateway/test_databricks_source.py index 50381d8..2fb429b 100644 --- a/tests/unit/gateway/test_databricks_source.py +++ b/tests/unit/gateway/test_databricks_source.py @@ -77,6 +77,21 @@ "status_code": "403", } +# The BYOK half of the same thing. Shaped from the live table: a rejected external call +# is logged with NULL tokens and — because the gateway never got far enough to resolve +# one — an EMPTY `destination_model`. +_FAILED_BYOK = { + "invocation_id": "inv-failed-byok", + "event_time": "2026-08-07 14:31:00", + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "destination_name": "workspace.default.anthropickey", + "destination_model": None, + "api_type": "anthropic/v1/messages", + "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.""" @@ -180,15 +195,28 @@ def test_read_usage_scopes_both_queries_to_one_shared_window() -> None: assert any(f"{column} < {c}" in sql for c in ceilings) +def test_a_sub_hour_interval_collapses_to_an_empty_window() -> None: + """Both bounds floor to the same hour, so there is nothing left to read. Asserted + against a pinned clock: driving this through the real one makes the test pass or + fail on the current MINUTE — "30 minutes" spans two hours at 13:34 and one at + 13:04, which is a coin flip in CI, not a property of the code.""" + lower, upper = _window_bounds("30 minutes", now=_NOW) + assert lower == upper == datetime(2026, 8, 21, 13, 0, tzinfo=timezone.utc) + + def test_a_window_entirely_inside_the_open_hour_reads_nothing_and_says_so( caplog: pytest.LogCaptureFixture, ) -> None: - """Excluding the open hour means a sub-hour window can resolve to nothing. Zero rows + """Excluding the open hour means such a window can resolve to nothing. Zero rows then says nothing about whether there was traffic, so it must not pass silently — - and it must not spend warehouse time either.""" + and it must not spend warehouse time either. + + `since` is half an hour AHEAD of the real clock so the guard fires whatever minute + this runs at; the collapse itself is pinned in the test above.""" src = _source([], []) + inside_the_open_hour = datetime.now(timezone.utc) + timedelta(minutes=30) with caplog.at_level(logging.WARNING): - assert list(src.read_usage("30 minutes")) == [] + assert list(src.read_usage(inside_the_open_hour)) == [] assert src.queries == [] # type: ignore[attr-defined] assert "widen the window" in caplog.text @@ -272,6 +300,32 @@ def test_failed_calls_yield_nothing() -> None: assert list(_source([], [_FAILED]).read_usage("1 day")) == [] +def test_failed_byok_calls_do_not_enter_the_join_index( + caplog: pytest.LogCaptureFixture, +) -> None: + """A rejected external call bought nothing, so Databricks meters no dollars for it + and its key can never match a spend row. Indexing it manufactures a bucket the + unbilled warning then tells the operator to re-run the window for — advice that can + never bill it, because there is nothing to bill.""" + with caplog.at_level(logging.WARNING): + assert list(_source([], [_FAILED_BYOK]).read_usage("1 day")) == [] + assert "NOT billed" not in caplog.text + + +def test_the_unbilled_warning_counts_only_buckets_with_real_tokens( + caplog: pytest.LogCaptureFixture, +) -> None: + """The warning exists to name genuine spend-table lag. Live over 2026-08-06 it + reported 29 buckets of which 28 were failed calls, and the single example row it + showed the operator was one of them — so the one real lagging bucket was the thing + least likely to be read.""" + lagging = {**_BYOK_USAGE, "invocation_id": "inv-lagging", "destination_model": "claude-opus-4-1"} + with caplog.at_level(logging.WARNING): + list(_source([], [_FAILED_BYOK, lagging]).read_usage("1 day")) + assert "1 BYOK token bucket(s)" in caplog.text + assert "model=claude-opus-4-1" in caplog.text + + 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.""" From a1249626c1e002af140eb02c22aaa43b01ea00d7 Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 14:28:11 +0200 Subject: [PATCH 16/22] Report the backfill's unbilled buckets instead of returning success MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `backfill_databricks` returned {"cost": …, "tokens": …, "skipped": 0} while BYOK buckets whose external_model_spend row had not landed went unbilled, and `on_error` never fired. Measured live with one hour's spend withheld — the shape of real spend-table lag — 54 buckets were unbilled and the only visible difference was `cost` falling 66 -> 12, which a caller cannot tell from a quieter window. The return value now carries `deferred`, and both it and `skipped` are reported through `on_error` (where="backfill"), the hook every other billing gap uses. `DatabricksSource.deferred_buckets` exposes the same buckets to a caller who reads the window itself, since an already-read list cannot report a bucket the reader never yielded. --- CHANGELOG.md | 4 ++ src/lago_agent_sdk/gateway/databricks.py | 16 +++++ src/lago_agent_sdk/sdk.py | 58 +++++++++++++++-- tests/unit/gateway/test_databricks_source.py | 67 ++++++++++++++++++-- 4 files changed, 134 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 794d7bd..cc2f6c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **`backfill_databricks` reported success while under-billing.** It returned `{"cost": …, "tokens": …, "skipped": 0}` — a shape that reads as "the whole window billed" — with no mention of the BYOK buckets the reader could not bill, and `config.on_error` never fired for them. Measured live over a window with one hour's spend rows withheld, which is exactly the shape of real `external_model_spend` lag: **54 buckets provably unbilled**, `on_error` called **0 times**, and the only thing the caller could see was `cost` falling from 66 to 12 — which is indistinguishable from a quieter window. The return value now carries `deferred`, and both it and `skipped` are reported through `on_error` (`where="backfill"`), the hook every other billing gap in this SDK already uses. The same window after the fix returns `{"cost": 12, "tokens": 54, "skipped": 0, "deferred": 54}` with one `on_error` call naming the hour to re-run, while the untouched window reports `deferred: 0` and stays silent. + - The two gaps are counted separately because they fail differently: a `skipped` row had no resolvable subscription and stays lost until it is tagged or a default is set, while a `deferred` bucket is billable revenue that the NEXT run of the same window collects once Databricks has aggregated its spend row. A run with both at 0 is the only one that billed everything it read. + - `DatabricksSource.deferred_buckets` carries the same buckets, so a caller who reads the window itself gets them too — the documented "inspect the rows first" path hands `backfill_databricks` an already-read list, which by construction cannot report a bucket the reader never yielded. Rewritten per read rather than appended to, so a healthy window read after a lagging one cannot report the older window's buckets as if they were current. + - **Failed Databricks calls were reported to the operator as revenue waiting to be billed.** The BYOK indexing loop had no zero-usage guard, so a rejected external call — NULL tokens, and an empty `destination_model` because the gateway never resolved one — became a key in the token index. It can never match a spend row (the call bought nothing, so Databricks meters no dollars for it), so it fell straight through to the "no `external_model_spend` row yet … re-run this window later to bill them" warning, which for these buckets is advice that can never work. Measured live over a fully-aggregated window: **29 buckets reported, 28 of them phantoms** — all 83 zero-usage rows in the window were 4xx/5xx — and the single example bucket the warning names for the operator was one of the phantoms, leaving the one genuinely lagging bucket the least visible thing in the message. The hosted loop already applied this guard; the BYOK loop now applies the same one, for the same reason. - Also stops a failed call from becoming a bucket's representative in `_merge_usage`, which keeps the first row's non-numeric fields. Unobserved today — a failed row's empty `destination_model` puts it in its own bucket, and `status_code` is already excluded from a bucket's extras as per-request — so this closes it before the join key or that exclusion list changes rather than after. diff --git a/src/lago_agent_sdk/gateway/databricks.py b/src/lago_agent_sdk/gateway/databricks.py index 5e825ed..5371645 100644 --- a/src/lago_agent_sdk/gateway/databricks.py +++ b/src/lago_agent_sdk/gateway/databricks.py @@ -274,6 +274,12 @@ def __init__( self.timeout = timeout # Databricks rejects anything outside 0s or 5-50s. self.wait_timeout = wait_timeout + # Buckets the most recent `read_usage` could not bill, in the shape its + # warning names them. A log line is not something a caller can act on: + # `backfill_databricks` turns this into a count in its return value and an + # `on_error` report, and a caller reading the window itself can re-run + # exactly these hours. See `read_usage` for why they go unbilled. + self.deferred_buckets: list[dict[str, str]] = [] @classmethod def from_env(cls, **kwargs: Any) -> DatabricksSource: @@ -441,6 +447,11 @@ def read_usage( that the newest hour of traffic arrives on the NEXT run, so pass a window comfortably wider than your run interval — this reader keeps no cursor. """ + # Rewritten per read rather than appended to, so a later read of a healthy + # window cannot leave an earlier read's gap standing as if it were current. + # Cleared here, ahead of the empty-window return below, so that path clears + # it too. Complete once the caller has drained the iterator. + self.deferred_buckets = [] lower, upper = _window_bounds(since) if lower >= upper: # Not an error, but it must not read as success either: the caller asked @@ -542,6 +553,11 @@ def read_usage( # 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) + # Same facts as the warning, in a shape a caller can act on rather than grep + # for — `backfill_databricks` reports the count through `on_error`. + self.deferred_buckets = [ + {"hour": k[0], "provider": k[1], "model": k[2], "request_tags": k[3]} for k in unbilled + ] if unbilled: logger.warning( "lago: %d BYOK token bucket(s) in this window have no external_model_spend " diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index 81cefe3..e83c893 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -631,7 +631,17 @@ def backfill_databricks( """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}``. + of what it handed to ``emit()``, e.g. + ``{"cost": 56, "tokens": 45, "skipped": 0, "deferred": 0}``. + + The last two are billing GAPS, and both are also reported through + ``config.on_error`` (``where="backfill"``) — the hook every other gap in this + SDK uses — so a caller does not have to inspect the return value to notice + one. They fail differently: ``skipped`` rows had no resolvable subscription + and stay lost until they are tagged or a default is set, while ``deferred`` + buckets are billable revenue that the NEXT run of the same window collects + once Databricks has aggregated their spend row. A run with both at 0 is the + only one that billed the whole window. ``source`` is normally a :class:`DatabricksSource`, and ``since`` the window. It also accepts an already-read iterable of ``DatabricksUsageRow`` — pass one @@ -661,12 +671,9 @@ def backfill_databricks( 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 - ) + counts = {"cost": 0, "tokens": 0, "skipped": 0, "deferred": 0} + reader = source if hasattr(source, "read_usage") else None + rows = reader.read_usage(since, event_id_prefix=event_id_prefix) if reader else source for row in rows: sub = default_subscription if unified else (row.subscription or default_subscription) if not sub: @@ -704,6 +711,43 @@ def backfill_databricks( timestamp=row.occurred_at, ) counts["tokens"] += 1 + + # Both gaps below were counted but never reported: measured live over a + # window with one hour's spend rows withheld — the shape of real spend-table + # lag — this returned `{'cost': 12, 'tokens': 54, 'skipped': 0}` while 54 + # BYOK buckets went unbilled and `on_error` fired zero times. `cost` alone + # dropping from 66 to 12 is not something an automated caller can read as a + # gap, so route both through the hook that already means "billing gap". + if counts["skipped"]: + self._report_error( + ValueError( + f"{counts['skipped']} Databricks row(s) had no resolvable subscription " + f"and were NOT billed. Pass default_subscription=..., set " + f"LagoConfig.default_subscription_id, or tag the calls." + ), + "backfill", + ) + # Only the reader knows about a bucket it never yielded, so a caller who + # passed an already-read iterable gets 0 here — they hold the source and can + # read `deferred_buckets` off it directly. `getattr` because `source` is + # duck-typed: a caller's own reader need not carry the attribute. + deferred = list(getattr(reader, "deferred_buckets", ())) if reader is not None else [] + counts["deferred"] = len(deferred) + if deferred: + first = deferred[0] + # `read_usage` logs this too. That is deliberate, not a stutter: a caller who + # reads the window itself never reaches this line, and one who ran the backfill + # needs it on the channel they reconcile against. Worded from the RUN's side so + # the two read as one gap seen from two layers rather than as two gaps. + self._report_error( + ValueError( + f"this run left {len(deferred)} Databricks BYOK bucket(s) unbilled: no " + f"external_model_spend row yet (e.g. hour={first['hour']} " + f"provider={first['provider']} model={first['model']}). The spend table " + f"lags; re-run this window later to bill them." + ), + "backfill", + ) return counts def flush(self, timeout: float = 5.0) -> bool: diff --git a/tests/unit/gateway/test_databricks_source.py b/tests/unit/gateway/test_databricks_source.py index 2fb429b..4e9be37 100644 --- a/tests/unit/gateway/test_databricks_source.py +++ b/tests/unit/gateway/test_databricks_source.py @@ -16,6 +16,7 @@ import pytest from lago_agent_sdk import CanonicalUsage, LagoSDK +from lago_agent_sdk.config import LagoConfig from lago_agent_sdk.gateway.databricks import ( DatabricksSource, DatabricksUsageRow, @@ -546,9 +547,14 @@ def events(self) -> list[dict]: return [e for b in self.batches for e in b] -def _sdk() -> tuple[LagoSDK, _Recorder]: +def _sdk(errors: list[tuple[str, str]] | None = None) -> tuple[LagoSDK, _Recorder]: rec = _Recorder() - sdk = LagoSDK(api_key="dummy") + cfg = ( + LagoConfig(api_key="dummy", on_error=lambda exc, where: errors.append((where, str(exc)))) + if errors is not None + else None + ) + sdk = LagoSDK(api_key="dummy", config=cfg) sdk._queue._sender = lambda b: rec.batches.append(list(b)) # type: ignore[attr-defined] return sdk, rec @@ -564,7 +570,7 @@ def test_backfill_counts_cost_tokens_and_skips() -> None: 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 counts == {"cost": 1, "tokens": 1, "skipped": 1, "deferred": 0} assert {e["external_subscription_id"] for e in q.events} == {"sub_byok", "sub_hosted"} @@ -629,6 +635,59 @@ def test_backfill_survives_one_malformed_row() -> None: assert len(q.events) >= 3 +# -------------------------------------------------------------------------- +# Billing gaps must reach the caller, not just the log +# -------------------------------------------------------------------------- +def test_backfill_reports_a_deferred_bucket_through_on_error() -> None: + """A BYOK bucket whose spend row has not landed is billed by neither loop. Measured + live with one hour's spend withheld, the return value was + `{'cost': 12, 'tokens': 54, 'skipped': 0}` — indistinguishable from a clean run — + while 54 buckets went unbilled and `on_error` fired zero times. Only a log line said + so, which no caller can reconcile against.""" + errors: list[tuple[str, str]] = [] + sdk, q = _sdk(errors) + # BYOK usage with NO matching spend row, plus one hosted row that bills normally. + counts = sdk.backfill_databricks(_source([], [_HOSTED, _BYOK_USAGE]), "1 day") + _drain(sdk) + assert counts == {"cost": 0, "tokens": 1, "skipped": 0, "deferred": 1} + backfill = [msg for where, msg in errors if where == "backfill"] + assert len(backfill) == 1, errors + # Names the hour to re-run and the model, so the report is actionable on its own. + assert "no external_model_spend row yet" in backfill[0] + # The hour key is the source column truncated, so it keeps that column's own form. + assert "hour=2026-08-07 14" in backfill[0] + assert "model=claude-sonnet-4-5" in backfill[0] + # The gap is a deferral, not a drop: the hosted row still billed. + assert q.events + + +def test_backfill_reports_an_unattributed_row_through_on_error() -> None: + """`skipped` was counted and returned but never routed anywhere, so a run that + attributed nothing looked like a run with nothing to attribute. It never reaches + `emit()`, which is where every other dropped event is reported from.""" + errors: list[tuple[str, str]] = [] + sdk, _ = _sdk(errors) + counts = sdk.backfill_databricks(_source([], [{**_HOSTED, "request_tags": "{}"}]), "1 day") + _drain(sdk) + assert counts["skipped"] == 1 and counts["deferred"] == 0 + backfill = [msg for where, msg in errors if where == "backfill"] + assert len(backfill) == 1, errors + assert "no resolvable subscription" in backfill[0] + + +def test_a_second_read_clears_the_previous_reads_deferred_buckets() -> None: + """The gap belongs to one read, so it is rewritten per read rather than appended + to. Left accumulating, a healthy window read after a lagging one would report the + older window's buckets — the phantom-warning shape the zero-usage guard just + removed, reintroduced through a different door.""" + src = _source([], [_BYOK_USAGE]) + list(src.read_usage("1 day")) + assert len(src.deferred_buckets) == 1 + src.query = lambda sql: [_BYOK_SPEND] if "external_model_spend" in sql else [_BYOK_USAGE] # type: ignore[method-assign] + list(src.read_usage("1 day")) + assert src.deferred_buckets == [] + + # -------------------------------------------------------------------------- # Reconciliation dimensions — the whole point of the connector being checkable # -------------------------------------------------------------------------- @@ -854,7 +913,7 @@ def test_backfill_accepts_already_read_rows_without_querying_again() -> None: sdk, q = _sdk() counts = sdk.backfill_databricks(rows, default_subscription="sub_x") _drain(sdk) - assert counts == {"cost": 1, "tokens": 1, "skipped": 0} + assert counts == {"cost": 1, "tokens": 1, "skipped": 0, "deferred": 0} assert len(src.queries) == queries_after_read, "must not re-read" # type: ignore[attr-defined] assert len(q.events) >= 3 From a6dc42222f402b759ca2d27e479c631eda3aeeea Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 14:40:29 +0200 Subject: [PATCH 17/22] Read only the Databricks columns the billing extraction uses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `read_usage` ran `SELECT * FROM system.ai_gateway.usage` — 36 columns to bill off 14. The Statement Execution API's default `disposition=INLINE` FAILS a statement whose response exceeds 25 MiB rather than paginating past it, so the width of the projection sets the largest window this reader can handle, for a module whose own guidance is "read one wide window per run". Measured live over 247 real rows: 1,411 bytes/row for `SELECT *` against 435 for the columns actually read — a ceiling of ~18k rows where it should be ~60k. The dropped columns are the wide ones nothing bills off (`routing_information`, `endpoint_metadata`, `url`, `user_agent`). Re-read the same window both ways: 120 billable events, identical as a multiset, 120 unique transaction ids, same deferred buckets. The coupling runs the other way too, which is why this needs a test: a column missing from the projection reaches the adapter as absent, and every field degrades to zero/empty rather than raising — a silently under-billed event. So the test feeds the canned rows THROUGH the statement's own column list and asserts the events match an unprojected read. `ORDER BY event_time` dropped with it. Nothing downstream reads the row order: the BYOK join is keyed, the unbilled-bucket report is sorted, and each event's `transaction_id` derives from the row's own ids. --- CHANGELOG.md | 5 ++ src/lago_agent_sdk/gateway/databricks.py | 38 ++++++++++++- tests/unit/gateway/test_databricks_source.py | 57 ++++++++++++++++++++ 3 files changed, 98 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cc2f6c1..2a1dbbc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,11 @@ All notable changes to this project will be documented here. Format follows [Kee ### Changed +- **The Databricks usage read no longer selects 36 columns to bill off 14.** `read_usage` ran `SELECT * FROM system.ai_gateway.usage`, and the Statement Execution API's default `disposition=INLINE` **fails** a statement whose response exceeds 25 MiB rather than paginating past it — so the width of the projection is what sets the largest window this reader can handle, for a module whose own guidance is "read one wide window per run". Measured on the live table over 247 real rows: **1,411 bytes/row for `SELECT *` against 435** for the 14 columns the extraction actually reads, i.e. a ceiling of ~18k rows where it should be ~60k. The dropped columns are the wide ones nothing bills off (`routing_information`, `endpoint_metadata`, `url`, `user_agent`). Billing output is byte-identical. + - The projection is one named list next to the interval regex, and the coupling runs the other way too: a column missing from it reaches the adapter as *absent*, which every field degrades to zero/empty on rather than raising — a silently under-billed event. So the test that pins it feeds the canned rows **through** the statement's own column list and asserts the events match an unprojected read, which fails if the projection is trimmed too far. + - `ORDER BY event_time` dropped with it: nothing downstream reads the row order — the BYOK join is keyed, the unbilled-bucket report is sorted, and each event's `transaction_id` derives from the row's own ids — so it only bought the warehouse a sort over the widest read this module makes. + - The one thing order *could* reach is which row represents a merged spend bucket, since `_merge_usage` keeps the first row's `_BUCKET_INVARIANT_EXTRAS`. What bills is order-free either way — the numeric fields are summed — and those 5 extras were measured not to disagree inside a bucket (0 disagreements over 66 live BYOK buckets, though only 2 of them held more than one row, so read that as unobserved rather than impossible). If they ever do disagree, the earliest row's value was never more correct than any other's, only more repeatable; making a bucket refuse to state a field its rows disagree on is the real answer and is deliberately not in this change. + - **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. diff --git a/src/lago_agent_sdk/gateway/databricks.py b/src/lago_agent_sdk/gateway/databricks.py index 5371645..a7de6a1 100644 --- a/src/lago_agent_sdk/gateway/databricks.py +++ b/src/lago_agent_sdk/gateway/databricks.py @@ -58,6 +58,37 @@ # the one direction that loses money. _INTERVAL_RE = re.compile(r"^\s*(\d{1,5})\s+(second|minute|hour|day|week)s?\s*$", re.I) +# Every `system.ai_gateway.usage` column the extraction reads, and nothing else. The +# table is 36 columns wide and this module's own guidance is "one wide window per run", +# so `SELECT *` is not merely untidy: results come back with the API's default +# `disposition=INLINE`, whose 25 MiB response cap FAILS the statement rather than +# paginating past it. Measured on the live table, 247 real rows: 1,411 bytes/row for +# `SELECT *` against 435 for these 14 columns — a ceiling of ~18k rows instead of ~60k +# before a window becomes unreadable, for identical billing output. The wide columns are +# the ones nothing bills off (`routing_information`, `endpoint_metadata`, `url`, +# `user_agent`). +# +# Keep this in sync with `extract_databricks_log` / `resolve_databricks_subscription`: +# a column dropped from here reaches the adapter as absent, which every field degrades +# to zero/empty on rather than raising — an under-billed event, silently. That coupling +# is what `test_the_projection_covers_every_column_the_extraction_reads` pins. +_USAGE_COLUMNS = ( + "event_time", # the hour bucket the BYOK join keys on + "request_id", + "invocation_id", + "endpoint_name", + "endpoint_id", + "destination_type", # hosted-vs-BYOK, which decides model AND provider + "destination_name", + "destination_model", + "api_type", # its leading segment IS the BYOK provider + "status_code", + "input_tokens", + "output_tokens", + "token_details", # cache_read / cache_write / reasoning + "request_tags", # carries `lago_subscription` +) + def _raise_for_api_error(resp: Any, what: str) -> None: """Raise with the API's own error text when a Statement Execution call is not OK. @@ -481,10 +512,13 @@ def read_usage( WHERE usage_start_time >= {window} AND usage_start_time < {ceiling} """) + # No ORDER BY: nothing downstream reads the row order — the BYOK join is keyed, + # the unbilled report is `sorted()`, and each event's `transaction_id` derives + # from the row's own ids — so it only buys the warehouse a sort over the widest + # read this module makes. usage = self.query(f""" - SELECT * FROM system.ai_gateway.usage + SELECT {", ".join(_USAGE_COLUMNS)} FROM system.ai_gateway.usage WHERE event_time >= {window} AND event_time < {ceiling} - ORDER BY event_time """) # Extract once per row and reuse: this loop and the hosted loop below both need diff --git a/tests/unit/gateway/test_databricks_source.py b/tests/unit/gateway/test_databricks_source.py index 4e9be37..952deca 100644 --- a/tests/unit/gateway/test_databricks_source.py +++ b/tests/unit/gateway/test_databricks_source.py @@ -18,6 +18,7 @@ from lago_agent_sdk import CanonicalUsage, LagoSDK from lago_agent_sdk.config import LagoConfig from lago_agent_sdk.gateway.databricks import ( + _USAGE_COLUMNS, DatabricksSource, DatabricksUsageRow, _floor_hour, @@ -222,6 +223,62 @@ def test_a_window_entirely_inside_the_open_hour_reads_nothing_and_says_so( assert "widen the window" in caplog.text +# -------------------------------------------------------------------------- +# The projection +# -------------------------------------------------------------------------- +def _projection_of(sql: str) -> list[str]: + """The column names one of this module's statements actually selects.""" + body = sql.split("SELECT", 1)[1].split("FROM", 1)[0] + return [c.strip() for c in body.split(",") if c.strip()] + + +def test_the_usage_read_projects_only_the_columns_the_adapter_uses() -> None: + """`system.ai_gateway.usage` is 36 columns wide and the caller is told to read one + wide window per run, but the API's default `disposition=INLINE` FAILS a statement + whose response exceeds 25 MiB rather than paginating. Live, over 247 real rows: + 1,411 bytes/row for `SELECT *` against 435 for these — ~18k readable rows instead + of ~60k, for byte-identical billing output.""" + src = _source([], [_HOSTED]) + list(src.read_usage("3 days")) + _, usage = src.queries # type: ignore[attr-defined] + assert "*" not in usage + assert _projection_of(usage) == list(_USAGE_COLUMNS) + # Nothing reads the row order, so the sort is pure warehouse cost on the widest + # read this module makes. + assert "ORDER BY" not in usage + + +def test_the_projection_covers_every_column_the_extraction_reads() -> None: + """The other direction of the same coupling, and the reason the narrowing needs a + test at all: a column missing from the projection does NOT raise — the adapter + degrades every absent field to zero/empty — so it would land as a silently + under-billed event. Here the canned rows are projected THROUGH the statement, so + dropping a needed column fails this rather than shipping.""" + projected = _source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE]) + inner = projected.query + + def project(sql: str) -> list[dict[str, Any]]: + rows = inner(sql) + if "external_model_spend" in sql: + return rows + keep = _projection_of(sql) + return [{k: v for k, v in row.items() if k in keep} for row in rows] + + projected.query = project # type: ignore[method-assign] + through_projection = list(projected.read_usage("3 days")) + every_column = list(_source([_BYOK_SPEND], [_HOSTED, _BYOK_USAGE]).read_usage("3 days")) + + assert [(r.usage, r.subscription, r.row_id, r.kind, r.usd_cost) for r in through_projection] == [ + (r.usage, r.subscription, r.row_id, r.kind, r.usd_cost) for r in every_column + ] + # Spelled out, because the equality above would also hold if BOTH sides were empty. + hosted = next(r for r in through_projection if r.kind == "usage") + assert (hosted.usage.model, hosted.usage.input, hosted.usage.output) == ("llama-4-maverick", 11, 4) + assert hosted.subscription == "sub_hosted" + byok = next(r for r in through_projection if r.kind == "spend") + assert (byok.usage.cache_read, byok.subscription, byok.usd_cost) == (1812, "sub_byok", 0.0011187) + + # -------------------------------------------------------------------------- # The BYOK / hosted split — the double-billing guard # -------------------------------------------------------------------------- From c7c1fdcabfefdff9bb33c8ccbbbbb61cb3f6610c Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 14:53:43 +0200 Subject: [PATCH 18/22] Surface unknown `token_details` keys instead of dropping them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Databricks drift sweep stopped at the row's columns. `token_details` is a STRUCT read field-by-field, so a key the adapter does not name reached neither a CanonicalUsage metric nor `extras`. Measured against the live table with the struct evolved by two fields (`cache_read_5m_input_tokens: 77`, `output_audio_tokens: 42`): 119 real tokens vanished with no error and no on_error — the exact failure the drift contract exists to prevent, and every drift test passed because none of them looked inside the struct. Latent today: the live struct has exactly the three fields the adapter maps, verified with DESCRIBE. Fixed anyway because this is the only column on the table that breaks tokens out by kind, so a new cache tier or output modality can land nowhere else — and this table's schema does evolve (`service_type`, `mcp_metadata`, `invocation_metadata` are later additions, and old rows still read `service_type = NULL`). Dotted key, not the container swept whole under `extras["token_details"]`: three of its keys ARE mapped, so publishing the container would re-emit counts already billed. Same shape openai_native already uses for its `*_tokens_details` containers. Both directions are pinned — one test fails if the sweep goes, another if it stops excluding the mapped keys — plus one for the JSON-string path, which is the one the backfill actually uses: the Statement Execution API serializes every STRUCT column as a string, measured, never as a dict. Scope, measured over a real window: the swept key reaches every hosted event (54 of 54) and no BYOK event (0 of 66), because a BYOK event is an hourly spend aggregate whose per-request extras are dropped by design. Both tables share this struct, so a new field still surfaces wherever there is hosted traffic. --- CHANGELOG.md | 6 ++ .../gateway/adapters/databricks_gateway.py | 74 ++++++++++++++----- tests/unit/test_drift.py | 71 ++++++++++++++++++ 3 files changed, 133 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2a1dbbc..4f5c10c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **The Databricks drift sweep stopped at the row's columns, missing the one nested column that carries token counts.** `token_details` is a STRUCT read field-by-field, so a key the adapter does not name reached neither a `CanonicalUsage` metric nor `extras`. Measured against the live table with the struct evolved by two fields (`cache_read_5m_input_tokens: 77`, `output_audio_tokens: 42`): **119 real tokens vanished** with no error and no `on_error` — the exact failure the drift contract exists to prevent, and every drift test passed because none of them looked inside the struct. Unmapped keys now surface under a dotted key (`extras["token_details.output_audio_tokens"]`), the shape `openai_native` already uses for its `*_tokens_details` containers. + - **Latent, not firing.** The live struct has exactly the three fields the adapter maps, verified with `DESCRIBE`. It is fixed anyway because this is the only column on the table that breaks tokens out by kind — a new cache tier or output modality can land nowhere else — and this table's schema does evolve: `service_type`, `mcp_metadata` and `invocation_metadata` are later additions, and old rows still read `service_type = NULL`. + - **Dotted, not the container swept whole** under `extras["token_details"]`: three of its keys ARE mapped, so publishing the container would re-emit counts already billed. Both directions are pinned — one test fails if the sweep is removed, another if it stops excluding the mapped keys. + - Swept on the **JSON-string path** too, which is the one the backfill actually uses: the Statement Execution API serializes every STRUCT column as a string (measured — `token_details` arrives as a string there, never as a dict), so a sweep handling only driver-native dicts would have missed the real reader. + - Scope, measured over a real window: the swept key reaches every hosted event (**54 of 54**) and no BYOK event (**0 of 66**), because a BYOK event is an hourly spend aggregate and `_as_bucket` drops per-request extras by design. Both tables share this struct, so a new field still surfaces wherever the workspace has hosted traffic. + - **`backfill_databricks` reported success while under-billing.** It returned `{"cost": …, "tokens": …, "skipped": 0}` — a shape that reads as "the whole window billed" — with no mention of the BYOK buckets the reader could not bill, and `config.on_error` never fired for them. Measured live over a window with one hour's spend rows withheld, which is exactly the shape of real `external_model_spend` lag: **54 buckets provably unbilled**, `on_error` called **0 times**, and the only thing the caller could see was `cost` falling from 66 to 12 — which is indistinguishable from a quieter window. The return value now carries `deferred`, and both it and `skipped` are reported through `on_error` (`where="backfill"`), the hook every other billing gap in this SDK already uses. The same window after the fix returns `{"cost": 12, "tokens": 54, "skipped": 0, "deferred": 54}` with one `on_error` call naming the hour to re-run, while the untouched window reports `deferred: 0` and stays silent. - The two gaps are counted separately because they fail differently: a `skipped` row had no resolvable subscription and stays lost until it is tagged or a default is set, while a `deferred` bucket is billable revenue that the NEXT run of the same window collects once Databricks has aggregated its spend row. A run with both at 0 is the only one that billed everything it read. - `DatabricksSource.deferred_buckets` carries the same buckets, so a caller who reads the window itself gets them too — the documented "inspect the rows first" path hands `backfill_databricks` an already-read list, which by construction cannot report a bucket the reader never yielded. Rewritten per read rather than appended to, so a healthy window read after a lagging one cannot report the older window's buckets as if they were current. diff --git a/src/lago_agent_sdk/gateway/adapters/databricks_gateway.py b/src/lago_agent_sdk/gateway/adapters/databricks_gateway.py index 8488b20..6600e10 100644 --- a/src/lago_agent_sdk/gateway/adapters/databricks_gateway.py +++ b/src/lago_agent_sdk/gateway/adapters/databricks_gateway.py @@ -17,6 +17,7 @@ token_details.cache_read_input_tokens → cache_read token_details.cache_creation_input_tokens → cache_write token_details.output_reasoning_tokens → reasoning + token_details. → extras["token_details."] destination_type + destination_name/_model → model, provider (see below) api → hardcoded "databricks_gateway" extras → routing/identity columns @@ -101,6 +102,29 @@ # emitted under a slightly ugly id is recoverable, a silently renamed one is not. _HOSTED_ENDPOINT_PREFIX = "databricks-" +# Keys inside the `token_details` STRUCT that this adapter MAPS onto a CanonicalUsage +# field. Anything else nested there is drift and is surfaced in `extras` under a dotted +# key rather than dropped. +# +# The column is read by name, so nothing ever inspects a key the mapping below does not +# name — measured against the real table with the struct evolved by two fields +# (`cache_read_5m_input_tokens: 77`, `output_audio_tokens: 42`): they reached neither a +# numeric field nor `extras`, so 119 tokens vanished with no error and no `on_error`. +# Latent today — the live struct has exactly the three fields listed here — but this +# table's schema does evolve (`service_type`, `mcp_metadata`, `invocation_metadata` are +# newer additions, and old rows still read `service_type = NULL`). +# +# Dotted key, not the dict swept whole under `extras["token_details"]`, for the same +# reason as openai_native's `_MAPPED_DETAIL_FIELDS`: three of the struct's keys ARE +# mapped, so emitting the container whole would re-publish counts already billed. +_MAPPED_DETAIL_FIELDS = frozenset( + { + "cache_read_input_tokens", + "cache_creation_input_tokens", + "output_reasoning_tokens", + } +) + def _safe_dict(v: Any) -> dict[str, Any]: """Coerce a STRUCT/MAP column to a dict, accepting either shape it arrives in. @@ -180,6 +204,37 @@ def extract_databricks_log(row: dict[str, Any]) -> CanonicalUsage: details = _safe_dict(row.get("token_details")) model, provider = _model_and_provider(row) + extras: dict[str, Any] = { + # `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"), + } + + # Drift sweep one level down, into `token_details` (see _MAPPED_DETAIL_FIELDS). + # A new key there is a token count nobody has classified yet: it must not be + # miscounted as one of the metrics above, and it must not disappear either. + # + # Scope, measured over the live 2026-08-06 window: it reaches every hosted event + # (54 of 54) and no BYOK event (0 of 66), because `_as_bucket` strips per-request + # extras from an hourly spend aggregate on purpose. Both tables share this struct, + # so a new field still surfaces wherever the workspace has hosted traffic. + for k, v in details.items(): + if k not in _MAPPED_DETAIL_FIELDS: + extras[f"token_details.{k}"] = v + return CanonicalUsage( input=_safe_int(row.get("input_tokens")), output=_safe_int(row.get("output_tokens")), @@ -189,24 +244,7 @@ def extract_databricks_log(row: dict[str, Any]) -> CanonicalUsage: 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"), - }, + extras=extras, ) diff --git a/tests/unit/test_drift.py b/tests/unit/test_drift.py index b8cdb78..66f275a 100644 --- a/tests/unit/test_drift.py +++ b/tests/unit/test_drift.py @@ -7,6 +7,7 @@ extract_bedrock_invoke, extract_openai_native, ) +from lago_agent_sdk.gateway.adapters import extract_databricks_log def test_converse_unknown_top_level_usage_field_goes_to_extras(): @@ -205,3 +206,73 @@ def test_unaccounted_total_still_recovers_tokens_nobody_broke_out() -> None: u = extract_openai_native(resp) assert u.output == 47 + 1149 assert u.extras["unaccounted_output_tokens"] == 1149 + + +# ---------------------------------------------------------------------- +# Databricks gateway adapter — the same guarantee inside `token_details` +# ---------------------------------------------------------------------- + + +def test_databricks_gateway_token_details_drift_reaches_extras() -> None: + """`token_details` is a STRUCT read by name, so an added field is invisible. + + Measured against the live table with the struct evolved by two fields: 119 real + tokens reached neither a numeric field nor `extras`. The column is the one place + this table publishes per-token-kind counts, so drift there is money-relevant by + construction — a new cache tier or modality lands nowhere else. + """ + row = { + "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_read_5m_input_tokens": 77, + "output_audio_tokens": 42, + }, + } + u = extract_databricks_log(row) + assert u.extras["token_details.cache_read_5m_input_tokens"] == 77 + assert u.extras["token_details.output_audio_tokens"] == 42 + # Never MISCOUNTED as a metric we do map — that would bill an unclassified count + # at a rate nobody chose for it. + assert u.cache_read == 1812 and u.cache_write == 0 and u.reasoning == 0 + + +def test_databricks_gateway_mapped_token_details_do_not_pollute_extras() -> None: + """The mirror: a nested key we DO map must not also appear in extras, or every + event carries a duplicate of a count already billed.""" + row = { + "destination_type": "EXTERNAL_FOUNDATION_MODEL", + "api_type": "openai/v1/chat/completions", + "destination_model": "gpt-4o", + "input_tokens": 100, + "output_tokens": 50, + "token_details": { + "cache_read_input_tokens": 40, + "cache_creation_input_tokens": 0, + "output_reasoning_tokens": 20, + }, + } + u = extract_databricks_log(row) + assert u.cache_read == 40 and u.reasoning == 20 + assert not [k for k in u.extras if k.startswith("token_details.")] + + +def test_databricks_gateway_token_details_drift_survives_the_json_string_path() -> None: + """Schema evolution arrives over the REST path as a longer JSON STRING, which is + how `DatabricksSource.query` reads every STRUCT column — measured on the live + warehouse, `token_details` is a `str` there, never a dict. A sweep that only + worked on driver-native dicts would miss the exact path the backfill uses.""" + row = { + "destination_type": "PAY_PER_TOKEN_FOUNDATION_MODEL", + "destination_name": "system.ai.gpt-oss-20b", + "input_tokens": "300", + "output_tokens": "12", + "token_details": '{"output_reasoning_tokens":"9","output_audio_tokens":"42"}', + } + u = extract_databricks_log(row) + assert u.reasoning == 9 + assert u.extras["token_details.output_audio_tokens"] == "42" From 7acad140b0240d7c24802254061a99a2eab0387d Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 15:07:40 +0200 Subject: [PATCH 19/22] Record why the provider hint cannot be forgotten here MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The JS port's stream wrapper is a module-level generator taking `providerHint` as an argument, which was defaulted to `""` — a legitimate value, so an omission silently billed a Databricks-HOSTED call as `provider: "openai"`. That argument is now required there. This wrapper's equivalent is nested inside `wrap_openai_client` and reaches the adapter only through `_emit_from`, so the hint is closed over and there is nothing to forget. Comment only; the four end-to-end tests that pin the stamp already live in `test_wrapper_openai.py`. --- src/lago_agent_sdk/wrappers/openai.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/lago_agent_sdk/wrappers/openai.py b/src/lago_agent_sdk/wrappers/openai.py index 6c69020..f61730d 100644 --- a/src/lago_agent_sdk/wrappers/openai.py +++ b/src/lago_agent_sdk/wrappers/openai.py @@ -142,6 +142,14 @@ def wrap_openai_client( base_dims = dict(dimensions or {}) base_sub = subscription is_async = type(client).__name__.startswith("Async") + # Resolved once, here, and reached only through `_emit_from` below — every emit + # path in this wrapper closes over it, so no call site can forget to pass it. That + # matters because "" is a legitimate value (it is what every non-Databricks client + # resolves to), so an omission would be indistinguishable at runtime from a real + # answer: the call would bill as provider="openai" for a Databricks-HOSTED model, + # dropping it out of TOKEN_BILLED_PROVIDERS. The JS port's stream wrapper is a + # module-level generator rather than a closure, so it keeps the same guarantee by + # making its `providerHint` parameter required. provider_hint = _provider_hint_for(client) def _resolve_opts(lago_opts: dict[str, Any]) -> dict[str, Any]: From 74b93b913f168ab194291602427ab2a01ce8f7ca Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 15:38:43 +0200 Subject: [PATCH 20/22] Skip negative spend rows, and reject a config passed as the api key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two guards for failures that are cheap to prevent and unrecoverable once they happen. A negative `usage_quantity` billed a $0 event and burnt the row's idempotency key with it. The spend loop guarded on `if not usd`, which skips 0.0 but passes -0.0042 straight through. `_parse_price` rejects a negative, so `compute_precomputed_cost` floors the event to $0 — and that $0 event still consumes the `record_id`-derived `transaction_id`, which Lago enforces unique account-wide. Driven end to end against real Lago on a real spend row: a $0.015245 row restated negative was stored as `value: "0"`, and re-running the window once Databricks had corrected it came back `422 value_already_exist` with Lago still holding "0" — the same figure billed fine only under a different event-id prefix. After the fix the negative row yields nothing, Lago 404s on that id, and the corrected re-run stores `value: "0.015245"`. The row is logged with its figure, model and hour rather than dropped quietly: it means Databricks issued a credit or restatement, which this connector cannot represent as an event, so skipping it leaves the customer billed more than Databricks metered. Its bucket then surfaces in the deferred report — the same window returns `{"cost": 65, ..., "deferred": 1}` with one on_error naming the hour, where before it returned `deferred: 0` and looked complete. Unobserved on the live table: 0 of 64 spend rows are negative, minimum $0.0000036. Guarded anyway — one comparison against a failure with no recovery. `LagoSDK(cfg)` sent every event to production Lago with an unusable key. The first positional parameter is `api_key`, so the config becomes the bearer token while `config` stays None and a fresh default replaces every field the caller set. Driven live with a config naming a local Lago: the SDK posted to api.getlago.com, every event 401'd, `flush()` still returned True, and the caller's own on_error was never invoked — it was one of the discarded fields. The only trace was a WARNING per event. Now a TypeError at construction, naming the correct call. Also drops the duplicate `import json` in `_canonical_tags`; the module has imported it at the top since it grew a row-hash fallback. --- CHANGELOG.md | 7 +++++ src/lago_agent_sdk/gateway/databricks.py | 26 +++++++++++++-- src/lago_agent_sdk/sdk.py | 18 +++++++++++ tests/unit/gateway/test_databricks_source.py | 33 ++++++++++++++++++++ tests/unit/test_sdk.py | 19 +++++++++++ 5 files changed, 100 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4f5c10c..da753cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,13 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **A negative Databricks `usage_quantity` billed a $0 event and burnt the row's idempotency key with it.** The spend loop guarded on `if not usd`, which skips `0.0` but passes `-0.0042` straight through. `_parse_price` rejects a negative, so `compute_precomputed_cost` floors the event to $0 — and that $0 event still consumes the `record_id`-derived `transaction_id`, which Lago enforces unique account-wide. Driven end to end against real Lago on a real spend row: a `$0.015245` row restated negative was stored as `value: "0"`, and re-running the window once Databricks had corrected it came back `422 value_already_exist`, leaving Lago holding `"0"` permanently — the same figure billed fine only under a different event-id prefix. Rows at or below zero are now skipped, so the id stays available and the restatement bills on the next run (verified live: the same window then stores `value: "0.015245"`). + - A negative row is **logged with its figure, model and hour** rather than dropped quietly. It means Databricks issued a credit or a restatement, which this connector has no way to represent as an event, and skipping it leaves the customer billed more than Databricks metered. + - The bucket it belonged to then surfaces in the `deferred` report — measured on the same window, `{"cost": 65, …, "deferred": 1}` with one `on_error` naming the hour. That reads as "these tokens went unbilled", which is exactly true; before the fix the same window returned `deferred: 0` and looked complete. + - **Unobserved on the live table** — 0 of 64 spend rows are negative, and the minimum is `$0.0000036`. Guarded anyway: it is one comparison, and the failure is unrecoverable once it happens. + +- **`LagoSDK(cfg)` sent every event to production Lago with an unusable key, and said nothing.** The first positional parameter is `api_key`, not `config`, so a `LagoConfig` passed positionally becomes the bearer token while `config` stays `None` and a fresh default config replaces every field the caller set. Driven live with a config naming a local Lago: the SDK posted to **`https://api.getlago.com/api/v1`**, every event came back `401`, `flush()` still returned `True`, and the caller's own `on_error` was never invoked — because that hook was one of the discarded fields. The only trace was one WARNING per event. It now raises `TypeError` at construction, naming the correct call. + - **The Databricks drift sweep stopped at the row's columns, missing the one nested column that carries token counts.** `token_details` is a STRUCT read field-by-field, so a key the adapter does not name reached neither a `CanonicalUsage` metric nor `extras`. Measured against the live table with the struct evolved by two fields (`cache_read_5m_input_tokens: 77`, `output_audio_tokens: 42`): **119 real tokens vanished** with no error and no `on_error` — the exact failure the drift contract exists to prevent, and every drift test passed because none of them looked inside the struct. Unmapped keys now surface under a dotted key (`extras["token_details.output_audio_tokens"]`), the shape `openai_native` already uses for its `*_tokens_details` containers. - **Latent, not firing.** The live struct has exactly the three fields the adapter maps, verified with `DESCRIBE`. It is fixed anyway because this is the only column on the table that breaks tokens out by kind — a new cache tier or output modality can land nowhere else — and this table's schema does evolve: `service_type`, `mcp_metadata` and `invocation_metadata` are later additions, and old rows still read `service_type = NULL`. - **Dotted, not the container swept whole** under `extras["token_details"]`: three of its keys ARE mapped, so publishing the container would re-emit counts already billed. Both directions are pinned — one test fails if the sweep is removed, another if it stops excluding the mapped keys. diff --git a/src/lago_agent_sdk/gateway/databricks.py b/src/lago_agent_sdk/gateway/databricks.py index a7de6a1..6b82a51 100644 --- a/src/lago_agent_sdk/gateway/databricks.py +++ b/src/lago_agent_sdk/gateway/databricks.py @@ -552,7 +552,29 @@ def read_usage( for row in spend: usd = _safe_float(row.get("usage_quantity")) - if not usd: + if usd <= 0: + # `if not usd` skipped 0.0 but let a NEGATIVE straight through, and a + # negative is far worse than a zero. `_parse_price` rejects it, so + # `compute_precomputed_cost` floors the event to $0 — and that $0 event + # still CONSUMES the `record_id`-derived `transaction_id`, which Lago + # enforces unique account-wide. Measured against real Lago on a real + # spend row: a $0.015245 row restated negative billed as `value: "0"`, + # and re-running the window once Databricks had corrected it came back + # `422 value_already_exist` — the same figure billed fine only under a + # fresh prefix. Skipping leaves the id unburnt, so a restatement bills + # normally on the next run. The bucket then appears in the deferred + # report below, which is the honest reading: its tokens went unbilled. + if usd < 0: + logger.warning( + "lago: skipping Databricks spend row with a NEGATIVE " + "usage_quantity (%s, model=%s, hour=%s). A credit or restatement " + "cannot be billed as an event, and billing it at $0 would burn " + "the row's transaction_id so the corrected figure could never " + "land.", + row.get("usage_quantity"), + row.get("model"), + _truncate_hour(_stamp(row.get("bucket"))), + ) continue key = ( _truncate_hour(_stamp(row.get("bucket"))), @@ -731,8 +753,6 @@ def _bucket_of(value: Any) -> str: 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 "{}") diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index e83c893..3fff53b 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -77,6 +77,24 @@ def __init__( at all: a local instance behind a self-signed cert (Traefik's default) is reachable with ``LagoSDK(api_key=..., api_url=..., verify_ssl=False)``. """ + if isinstance(api_key, LagoConfig): + # `LagoSDK(cfg)` is the natural-looking call and it is silently, totally + # wrong: the config becomes the BEARER TOKEN while ``config`` stays None, so + # a fresh default ``LagoConfig`` is built and every field the caller set is + # discarded. Measured live: a config naming a local Lago produced an SDK + # posting to PRODUCTION ``api.getlago.com`` with an unusable key — every + # event 401, ``flush()`` still returning True, and the caller's own + # ``on_error`` never invoked, because that hook was one of the discarded + # fields. The only trace was a WARNING per event. Whether the queue then + # drops those events or holds them is beside the point: no key the caller + # passed is ever used, so nothing downstream can recover. It has to fail at + # construction. + raise TypeError( + "LagoSDK's first positional parameter is `api_key`, not `config`. " + "Passing a LagoConfig here makes it the bearer token and leaves the " + "rest of your config unused, so every event is sent to the default " + "api_url with a key that 401s. Use LagoSDK(config.api_key, config=config)." + ) self.config = config or LagoConfig(api_key=api_key) # explicit args win over `config` — guarded on "was it actually passed?" # rather than on truthiness, so a config value survives when it wasn't. diff --git a/tests/unit/gateway/test_databricks_source.py b/tests/unit/gateway/test_databricks_source.py index 952deca..42c4e7d 100644 --- a/tests/unit/gateway/test_databricks_source.py +++ b/tests/unit/gateway/test_databricks_source.py @@ -352,6 +352,39 @@ def test_zero_dollar_spend_rows_are_skipped() -> None: assert list(_source([{**_BYOK_SPEND, "usage_quantity": "0"}], []).read_usage("1 day")) == [] +def test_negative_spend_rows_are_skipped_rather_than_billed_at_zero( + caplog: pytest.LogCaptureFixture, +) -> None: + """A Databricks credit or restatement arrives as a negative `usage_quantity`. + + `if not usd` skipped 0 but passed a negative, which `_parse_price` then rejects, so + the event billed at $0 while still consuming the `record_id`-derived + `transaction_id`. Verified against real Lago: the corrected positive figure came + back `422 value_already_exist` and could only be billed under a different prefix. + Skipping keeps the id available for the restatement. + """ + negative = {**_BYOK_SPEND, "usage_quantity": "-0.0011187"} + with caplog.at_level(logging.WARNING, logger="lago_agent_sdk.gateway.databricks"): + rows = list(_source([negative], [_BYOK_USAGE]).read_usage("1 day")) + # Only the hosted/usage half may survive; nothing may carry the spend row's id. + assert [r.row_id for r in rows if r.kind == "spend"] == [] + assert "NEGATIVE usage_quantity" in caplog.text + assert "-0.0011187" in caplog.text, "the operator needs the figure, not just the fact" + + +def test_a_restated_spend_row_bills_under_the_id_the_negative_would_have_burnt() -> None: + """The whole point of skipping, as one scenario: read the negative, then read the + same row once Databricks has corrected it. Anything emitted for the negative — even + at $0 — consumes that `transaction_id` account-wide, and Lago then rejects the + correction as a duplicate, so the id the second read needs must still be free.""" + negative = {**_BYOK_SPEND, "usage_quantity": "-0.0011187"} + burnt = {row.event_id_for("sub_byok") for row in _source([negative], [_BYOK_USAGE]).read_usage("1 day")} + corrected = list(_source([_BYOK_SPEND], [_BYOK_USAGE]).read_usage("1 day")) + (spend_row,) = [r for r in corrected if r.kind == "spend"] + assert spend_row.event_id_for("sub_byok") not in burnt + assert spend_row.usd_cost == pytest.approx(0.0011187) + + 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.""" diff --git a/tests/unit/test_sdk.py b/tests/unit/test_sdk.py index 581d998..cebd3d3 100644 --- a/tests/unit/test_sdk.py +++ b/tests/unit/test_sdk.py @@ -157,6 +157,25 @@ def test_explicit_verify_ssl_wins_over_config(): sdk.shutdown(timeout=1.0) +def test_a_config_passed_positionally_raises_instead_of_401ing_everything(): + """`LagoSDK(cfg)` reads correctly and is total: the config becomes the bearer + token, `config` stays None, and a fresh default LagoConfig replaces every field + the caller set. Driven live it posted to PRODUCTION api.getlago.com, 401'd every + event, still returned True from `flush()`, and never called the caller's + `on_error` — because that hook was one of the discarded fields.""" + cfg = LagoConfig( + api_key="k", + api_url="https://api.lago.dev/api/v1", + on_error=lambda exc, where: None, + ) + with pytest.raises(TypeError) as excinfo: + LagoSDK(cfg) # type: ignore[arg-type] + message = str(excinfo.value) + assert "api_key" in message and "config" in message + # The message has to carry the fix, not just the diagnosis. + assert "LagoSDK(config.api_key, config=config)" in message + + def test_ignored_usd_cost_is_reported_not_silently_dropped(): """A caller who supplies a real metered cost while the effective mode isn't 'price' had it discarded with no log and no on_error — so a hand-rolled From 0172754eb2f7d6f1d8273b06f8d6bd91cc6654a5 Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 16:10:56 +0200 Subject: [PATCH 21/22] Classify a 4xx by whether the batch is what is wrong MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The two ports disagreed about which HTTP failures destroy events, and measuring both showed neither was right. `_PERMANENT_STATUSES` routes a batch to `_send_individually`, where each event that fails again is logged and dropped for good. Python listed {400,401,402,403,404,409,413,415,422}; JS listed {400,404,409,422}. Driven over a real socket at a server returning each status (`probes/t11_status_matrix`), the two conditional cases point opposite ways: a key rotated back after 3s PY destroyed all 5 events inside the first second, none ever reached Lago. JS held them and delivered all 5 when it healed. nginx-style 413 above a byte PY's split path delivered all 5. JS held the limit, 200 below it oversized batch, delivered 0, and stalled at the backoff ceiling forever. Both ports now use {400, 409, 413, 422} and the whole matrix is identical row-for-row. The rule is written down rather than enumerated: is what makes this fail a property of the BATCH? A different payload is the only fix -> permanent, and splitting saves the good events. An out-of-band fix — a key restored, an invoice paid, a URL corrected, a proxy reconfigured -> transient, because dropping is unrecoverable while holding is bounded by `max_buffer_size`, oldest-first and reported. Per status, against a real Lago instance rather than from the RFCs: 404 -> transient. Lago answers 404 `resource_not_found` for a wrong PATH, i.e. a mistyped api_url. Neither port held it, so a typo destroyed every event. Same class as the 405/410 both already held. 415 -> transient. Splitting provably cannot help: this client always sends application/json, so every isolated send fails identically (measured, all 5 dropped). Lago answers 422 to a bad content-type, so a 415 only comes from a proxy someone can fix. 413 -> stays permanent, and it was missing here on the JS side. Lago answers 422 `too_many_events` to an oversized batch (probed at 20k events / 3.5 MiB), so a 413 only comes from something like nginx's `client_max_body_size` — exactly where splitting recovers the events. 402 -> transient. Payment required is a property of the account and stops being true the moment someone pays. Measured: 5 in, 6 HTTP calls out, 0 recoverable, one on_error for the lot. 401 403 -> transient. Already correct in JS, never ported here. Unchanged and re-confirmed live: a replayed transaction_id is 422 and stays permanent, so one bad id still does not take its batch down with it. Five new held-until-it-heals cases, confirmed to fail with the old set. The 402 half of this supersedes PR #21, which is now redundant on that point only — its empty-api_url report and its Cloudflare usage_metadata drift sweep are NOT on this branch. --- CHANGELOG.md | 10 +++++ src/lago_agent_sdk/queue.py | 78 +++++++++++++++++++++++++------------ tests/unit/test_queue.py | 62 +++++++++++++++++++++++++---- 3 files changed, 119 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index da753cb..0d83f67 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,16 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **The two ports disagreed about which HTTP failures destroy events, and both were wrong.** `_PERMANENT_STATUSES` routes a batch to `_send_individually`, where each event that fails again is logged and **dropped for good**. Python listed `{400, 401, 402, 403, 404, 409, 413, 415, 422}`; JS listed `{400, 404, 409, 422}`. Measured by driving the real queue over a real socket at a server returning each status (`probes/t11_status_matrix`), the two conditional cases settle it in opposite directions, so neither port was simply "the right one": + - **A key rotated back after 3s** — Python destroyed all 5 events inside the first second and none ever reached Lago; JS held them and delivered all 5 when it healed. + - **An nginx-style server answering 413 above a byte limit and 200 below it** — Python's split path delivered all 5; JS held the oversized batch and delivered 0, stalling at the backoff ceiling forever. + - Both now use `{400, 409, 413, 422}` and the matrix is identical row-for-row across the ports. The rule is written down rather than enumerated: **is what makes this fail a property of the batch?** A different payload is the only fix → permanent, and splitting can save the good events. An out-of-band fix (a key restored, an invoice paid, a URL corrected, a proxy reconfigured) → transient, because dropping is unrecoverable while holding is bounded by `max_buffer_size`, oldest-first and reported. + - **`404` moved to transient**, which neither port had. Probed against a real instance, Lago answers `404 resource_not_found` for a wrong PATH — that is a mistyped `api_url`, fixed out-of-band exactly like a rotated key, and it was destroying every event. It is also the same class as the `405`/`410` both ports already held for that stated reason. + - **`415` moved to transient** because splitting provably cannot help: this client always sends `application/json`, so every isolated send fails identically — measured, all 5 dropped. Lago itself answers `422` to a bad content-type (probed live), so a 415 only ever comes from a proxy, which someone can fix. + - **`413` stays permanent, and JS gains it.** Lago answers `422 too_many_events` to an oversized batch (probed live at 20k events / 3.5 MiB), so a 413 only comes from something like nginx's `client_max_body_size` — and that is precisely where splitting recovers the events. + - **`402` moved to transient** in Python (JS never had it): payment required is a property of the account and stops being true the moment someone pays. Measured against a 402 server, 5 events in, 6 HTTP calls out, 0 recoverable, one `on_error` for the lot. + - Unchanged and re-confirmed live: a replayed `transaction_id` is `422`, still permanent, so one bad id in a batch still does not take the valid events down with it. + - **A negative Databricks `usage_quantity` billed a $0 event and burnt the row's idempotency key with it.** The spend loop guarded on `if not usd`, which skips `0.0` but passes `-0.0042` straight through. `_parse_price` rejects a negative, so `compute_precomputed_cost` floors the event to $0 — and that $0 event still consumes the `record_id`-derived `transaction_id`, which Lago enforces unique account-wide. Driven end to end against real Lago on a real spend row: a `$0.015245` row restated negative was stored as `value: "0"`, and re-running the window once Databricks had corrected it came back `422 value_already_exist`, leaving Lago holding `"0"` permanently — the same figure billed fine only under a different event-id prefix. Rows at or below zero are now skipped, so the id stays available and the restatement bills on the next run (verified live: the same window then stores `value: "0.015245"`). - A negative row is **logged with its figure, model and hour** rather than dropped quietly. It means Databricks issued a credit or a restatement, which this connector has no way to represent as an event, and skipping it leaves the customer billed more than Databricks metered. - The bucket it belonged to then surfaces in the `deferred` report — measured on the same window, `{"cost": 65, …, "deferred": 1}` with one `on_error` naming the hour. That reads as "these tokens went unbilled", which is exactly true; before the fix the same window returned `deferred: 0` and looked complete. diff --git a/src/lago_agent_sdk/queue.py b/src/lago_agent_sdk/queue.py index 0a324b4..02796d6 100644 --- a/src/lago_agent_sdk/queue.py +++ b/src/lago_agent_sdk/queue.py @@ -31,36 +31,66 @@ from .exceptions import LagoApiError -# Statuses where re-sending the SAME batch can never succeed: the request itself -# is the problem (malformed body, bad credentials, a transaction_id Lago has -# already accepted). Deliberately an explicit list rather than the 400-499 range, -# because two 4xx statuses mean "try again, later": 429 (rate limited) and 408 -# (request timeout). Treating those as permanent dropped billable events AND fanned -# one throttled batch out into up to `max_batch_size` extra requests aimed at the -# server that had just asked us to slow down. +# Statuses where re-sending the SAME batch can never succeed, because the BATCH is what +# is wrong. Deliberately an explicit list, not the 400-499 range. # -# 413/402/415 are in the set for the OPPOSITE reason to 429: re-sending the same batch -# provably cannot succeed (too large, payment required, wrong media type), so treating -# them as transient re-prepended the identical batch at the head of the FIFO and backed -# off to 60s forever, blocking every event behind it until the buffer overflowed. Being -# "permanent" here routes them to `_send_individually`, which SPLITS the batch and -# delivers what is deliverable — so a 413 on a 100-event batch becomes 100 single-event -# sends rather than a stalled queue. That is the behaviour we want, and the batch that -# most needs splitting was the one that never reached it. 405/410 stay transient: they -# usually indicate a misrouted or retired endpoint, which a deploy can fix. -_PERMANENT_STATUSES = frozenset({400, 401, 402, 403, 404, 409, 413, 415, 422}) +# The test is: **is what makes this fail a property of the batch?** If a DIFFERENT +# PAYLOAD is what it takes to succeed, the batch is doomed and belongs here, where +# `_send_individually` splits it and delivers whatever is deliverable. If an OUT-OF-BAND +# change fixes it — someone rotates a key back, pays an invoice, corrects a URL, fixes a +# proxy — the events are still perfectly billable and must be HELD, because dropping +# them is unrecoverable while holding them is bounded (`max_buffer_size`, oldest-first, +# reported through `on_error`). +# +# Every line below was measured by driving this queue over a real socket at a server +# returning that status, counting events actually delivered (`probes/t11_status_matrix`): +# +# 400 malformed body — a different payload is the only fix. PERMANENT. +# 413 too large — the size IS the batch. Isolating it is a real recovery, not a +# formality: against an nginx-style server answering 413 over a byte limit and +# 200 under it, the split path delivered 5 of 5. Held instead, it delivered 0 and +# stalled at the backoff ceiling forever. Not reachable from Lago itself — an +# oversized batch there answers 422 `too_many_events` (probed live, 20k events / +# 3.5 MiB) — so this exists for `client_max_body_size` in front of Lago. +# 409 a conflicting id. Lago answers 422 for a replayed `transaction_id`, not 409 +# (probed live); 409 stays as defence against an intermediary that uses it. +# 422 Lago's real answer for a duplicate id, an oversized batch and a bad +# content-type. PERMANENT — but note it reaches `_send_individually`, which is +# what lets the valid events in a batch survive one bad transaction_id. +# +# Everything else is transient, including these, which used to be here and lost money: +# +# 401/403 a rotated or revoked key. Measured with a server that healed after 3s — +# i.e. the key put back — classified permanent this destroyed all 5 events +# inside the first second, and none of them ever reached Lago. Held, all 5 +# were delivered when it healed. +# 402 payment required — a property of the ACCOUNT; it stops being true the +# moment someone pays. Measured against a 402 server: 5 events in, 6 HTTP +# calls out, 0 recoverable, one `on_error` for the lot. +# 404 the endpoint, not the events: Lago answers 404 `resource_not_found` for a +# wrong PATH (probed live), which is a mistyped `api_url` — fixed out-of-band +# like a rotated key, and the same class as the 405/410 that were already +# transient here for exactly that reason. It was destroying every event. +# 415 a wrong media type comes from a proxy, and splitting cannot help: this +# client always sends `application/json`, so every isolated send fails the +# same way — measured, all 5 dropped. Held, they survive the proxy being +# fixed. (Lago itself answers 422 to a bad content-type, probed live.) +# 429/408 throttling — fanning a batch into N isolated sends aims more traffic at a +# server that just asked us to slow down. +# +# An unrecognized 4xx is transient too: waiting on an event that would have been dropped +# costs a delay, dropping one that would have been accepted costs revenue. +_PERMANENT_STATUSES = frozenset({400, 409, 413, 422}) def _is_permanent_failure(exc: Exception) -> bool: """True when re-sending this exact batch can never succeed. - A validation 4xx (bad request, duplicate transaction_id, revoked key) will - fail identically forever, so it is isolated and dropped. Everything else — - 5xx, a network-level exception (timeout, connection error, no LagoApiError at - all), and the throttling 4xxs 429/408 — might succeed later and stays - retryable. An unrecognized 4xx is treated as transient: waiting on an event - that would have been dropped costs a delay, dropping one that would have been - accepted costs revenue. + Only a malformed or unacceptable BATCH qualifies — see `_PERMANENT_STATUSES` for + the test and for what each status cost when it was on the wrong side of it. + Everything else (5xx, a network-level exception with no LagoApiError at all, a + credential or account or endpoint 4xx, an unrecognized 4xx) might succeed later + and stays retryable. """ return isinstance(exc, LagoApiError) and exc.status in _PERMANENT_STATUSES diff --git a/tests/unit/test_queue.py b/tests/unit/test_queue.py index d7ca109..aa013a3 100644 --- a/tests/unit/test_queue.py +++ b/tests/unit/test_queue.py @@ -268,14 +268,17 @@ def test_overflow_is_reported_through_on_error(): # ---------------------------------------------------------------------- -# The throttling 4xxs. 429 and 408 sit inside the 400-499 range but mean "try -# again, later" — classifying them as permanent dropped billable events and -# aimed `max_batch_size` extra requests at a server that had just asked us to -# slow down. +# Which side of the permanent/transient line each 4xx belongs on. The test is +# whether a DIFFERENT PAYLOAD is what it would take to succeed (permanent, so +# `_send_individually` can split the batch and save what is savable) or whether +# an OUT-OF-BAND change fixes it (transient, so the events must be held — +# dropping them is unrecoverable, holding them is bounded by `max_buffer_size`). +# See `_PERMANENT_STATUSES` for what each status cost when it was on the wrong +# side, measured over a real socket. # ---------------------------------------------------------------------- -@pytest.mark.parametrize("status", [413, 402, 415]) +@pytest.mark.parametrize("status", [413]) def test_batch_only_4xx_is_split_not_head_of_line_blocked(status: int): - """For these the SAME batch can never succeed, but its events can individually. + """For this one the SAME batch can never succeed, but its events can individually. Treating them as transient re-prepended the identical batch at the head of the FIFO and backed off to 60s forever, blocking everything behind it. Routing them to @@ -351,6 +354,51 @@ def sender(batch): q.shutdown(timeout=1.0) +@pytest.mark.parametrize( + ("status", "cause"), + [ + (401, "a rotated or revoked key"), + (403, "a key that lost its scope"), + (402, "an unpaid account"), + (404, "a mistyped api_url"), + (415, "a proxy rejecting the media type"), + ], +) +def test_out_of_band_4xx_is_held_until_it_heals(status: int, cause: str) -> None: + """None of these is a property of the BATCH, so dropping the events is unrecoverable + while holding them is not. + + Each was in `_PERMANENT_STATUSES`, which routes to `_send_individually`: the batch + fails, every isolated send fails the same way, and each event is logged and dropped + for good. Measured over a real socket at a server returning 401 for 3s and then 200 — + the shape of a key being put back — all 5 events were destroyed inside the first + second and none ever reached Lago. Held, all 5 were delivered when it healed. + + So this asserts recovery, not merely "not dropped": the events must survive the + outage AND still arrive, as one batch rather than fanned out per event. + """ + attempts = {"n": 0} + delivered: list = [] + + def sender(batch): + attempts["n"] += 1 + if attempts["n"] == 1: + raise LagoApiError(status, cause) + delivered.extend(batch) # the out-of-band fix lands + + q = EventQueue(sender=sender, flush_interval=0.05, max_batch_size=10, max_retry_seconds=0.5) + try: + q.push({"id": "a"}) + q.push({"id": "b"}) + deadline = time.monotonic() + 3.0 + while time.monotonic() < deadline and not delivered: + time.sleep(0.05) + assert [e["id"] for e in delivered] == ["a", "b"], f"{cause}: events must survive it" + assert attempts["n"] == 2, "held as one batch, never fanned out into per-event sends" + finally: + q.shutdown(timeout=2.0) + + def test_unrecognized_4xx_is_treated_as_transient(): """Only the enumerated validation statuses are permanent. An unfamiliar 4xx errs toward retrying: a needless delay costs latency, a wrong drop costs @@ -375,7 +423,7 @@ def sender(batch): q.shutdown(timeout=2.0) -@pytest.mark.parametrize("status", [400, 401, 403, 404, 409, 422]) +@pytest.mark.parametrize("status", [400, 409, 413, 422]) def test_validation_4xx_still_isolates_and_drops(status: int): """The statuses that genuinely cannot succeed on a re-send keep the isolate-one-by-one behaviour, so a single bad transaction_id still doesn't From 4a144b481fdecbb1217c115604563d1a5d134389 Mon Sep 17 00:00:00 2001 From: Anass Date: Fri, 21 Aug 2026 16:38:34 +0200 Subject: [PATCH 22/22] Report a discarded api_url, sweep gateway drift, and respect the pricing TTL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four fixes from the review round that never reached this branch. All four were verified live against the real code before and after, not reasoned about. An explicitly-passed falsy `api_url` silently resolved to PRODUCTION. Preferring the config value over "" is right — `requests` raises MissingSchema, which is not a LagoApiError, so the queue classified it transient, re-prepended the batch and retried at the 60s ceiling forever, stopping all billing with a growing buffer as the only symptom. But LagoConfig's default is the production URL, so `api_url=os.environ.get("LAGO_API_URL", "")` with the var unset resolved there with no on_error and no log: resolved api_url https://api.getlago.com/api/v1 client POST target https://api.getlago.com/api/v1/events/batch on_error invocations 0 For a CI job or a developer holding a real production key that writes live billing data, and ingested events cannot be un-ingested. The fallback is unchanged, so the original config-clobber bug stays fixed; it is now reported under `config.api_url`. An unpassed api_url stays silent — None means the caller never mentioned it, and reporting the common case would train customers to ignore the channel this fix depends on. `usage_metadata` from the Cloudflare gateway got no drift sweep, and had already lost two counters. `extras` was a fixed three-key dict, so any counter the adapter does not map vanished with no error and no on_error — the one place violating the contract test_drift.py enforces for the native adapters. Replaying the 14 captured fixtures through the adapter: neurons dropped in 4 entries Cloudflare's Workers AI billing unit input_text_tokens dropped in 1 entry and a live Logs API pull also returns `units`, a cost quantity that appears in no fixture at all — the hand-maintained enumeration in the module docstring had already drifted past reality, which is exactly the failure mode a snapshot invites. Unmapped keys now sweep into extras["usage_metadata"] against an explicit _MAPPED_USAGE_KEYS set. Deliberately NESTED rather than merged flat: the poller reads extras["cached"] to decide whether to skip billing a request Cloudflare served for free, so a future usage_metadata key called `cached` or `step` must not be able to shadow it. The regression test iterates the fixture directory rather than a fixed key list, so a recapture that introduces a new counter fails it with no test edit. Closes #16. `prime()` re-downloaded the ~400-model OpenRouter catalogue regardless of the TTL. It set `_openrouter_stale` unconditionally, and it is reached from `_auto_prime_pricing_for` on a matching wrap() and from warm_pricing() — both of which a server can run per request — so pricing_ttl_seconds never applied on that path at all, on the thread the queue drains events from. With the shipped 1-hour TTL: 4 prime()+maybe_refresh() cycles produced 4 full downloads where 1 was correct; now 1. Gated on the same "no table, or past the TTL" test lookup() uses, so priming and looking up cannot disagree — and a table that genuinely ages out is still re-primed, so prices do not freeze at the first fetch. A failed pricing fetch retried every tick, forever, with no backoff, ahead of the drain. Only the success path cleared a source's stale flag, so a bad credential re-attempted on every tick, each attempt costing up to the 10s `_get_json` timeout, all of it before the drain. Measured with a failing Cloudflare fetch: 5 ticks produced 5 real requests and 5 on_error reports; now 1, and it still recovers once the window expires — a backoff, not a permanent give-up, the same reasoning that makes a 401 transient in the queue. Per-source 1->2->4->...->60s, matching the queue's own send backoff, so one bad credential cannot delay the three healthy tables; a success clears the window rather than letting it keep doubling across unrelated outages. The four fetches deliberately stay sequential here: this refresh runs on a blocking daemon thread, where a thread pool for four fetches is a larger change than the problem warrants, and the per-source backoff already removes the harm. JS additionally parallelises them under Promise.allSettled, which its event loop makes free — language-inherent, like os.register_at_fork vs AsyncLocalStorage. Every fix has a test that fails when the fix alone is reverted (verified by reverting each one in turn). 601 unit tests, coverage 91.89%, ruff and mypy strict clean. --- CHANGELOG.md | 4 + .../gateway/adapters/cloudflare_gateway.py | 72 ++++++++- src/lago_agent_sdk/pricing.py | 97 ++++++++++- src/lago_agent_sdk/sdk.py | 21 +++ .../adapters/test_cloudflare_gateway.py | 150 ++++++++++++++++++ tests/unit/test_pricing.py | 111 +++++++++++++ tests/unit/test_sdk.py | 37 +++++ 7 files changed, 481 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d83f67..2b8d529 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ All notable changes to this project will be documented here. Format follows [Kee ### Fixed +- **An explicitly-passed falsy `api_url` silently resolved to PRODUCTION Lago.** Preferring the config value over `""` is right — `requests` raises `MissingSchema`, which is not a `LagoApiError`, so the queue classified it transient, re-prepended the batch and retried at the 60s ceiling forever, stopping all billing with only a growing buffer as the symptom. But `LagoConfig`'s default is the production URL, so `api_url=os.environ.get("LAGO_API_URL", "")` with the var unset resolved to production with **no `on_error` and no log** — verified live: 0 reports, 0 log lines, and a client posting to `api.getlago.com`. For a CI job or a developer holding a real production key that writes live billing data, and ingested events cannot be un-ingested. The fallback is unchanged, so the original config-clobber bug stays fixed; it is now reported under `config.api_url` through the same log-plus-callback floor as every other drop path. An *unpassed* `api_url` stays silent — `None` means the caller never mentioned it, and reporting the common case would train customers to ignore the channel this fix depends on. +- **`usage_metadata` from the Cloudflare gateway got no drift sweep, and had already lost two counters.** `extras` was a fixed three-key dict, so any counter the adapter does not map vanished with no error and no `on_error` — the one place violating the drift contract `test_drift.py` enforces for the native adapters. Not hypothetical: replaying the **14 captured fixtures** through the adapter drops **`neurons`** (Cloudflare's Workers AI billing unit) in 4 entries and **`input_text_tokens`** in 1, and a live Logs API pull also returns **`units`**, a cost quantity that appears in no fixture at all — the hand-maintained key enumeration in the module docstring had already drifted past reality, which is exactly the failure mode a snapshot invites. Unmapped keys are now swept into `extras["usage_metadata"]` against an explicit `_MAPPED_USAGE_KEYS` set, so a ninth spelling surfaces on its own instead of needing another 14-fixture audit. Deliberately **nested** rather than merged flat into `extras`: the poller reads `extras["cached"]` to decide whether to skip billing a request Cloudflare served for free, so a future `usage_metadata` key called `cached` or `step` must not be able to shadow it. The regression test iterates the fixture directory rather than a fixed key list, so a recapture that introduces a new counter fails it with no test edit. Closes #16. +- **`prime()` re-downloaded the ~400-model OpenRouter catalogue regardless of the TTL.** It set `_openrouter_stale` unconditionally, and it is reached from `_auto_prime_pricing_for` on a matching `wrap()` and from `warm_pricing()` — both of which a server can run per request — so `pricing_ttl_seconds` never applied on that path at all, and the catalogue was refetched on essentially every flush tick, on the thread the queue drains events from. Measured with the shipped 1-hour TTL: **4 `prime()`+`maybe_refresh()` cycles produced 4 full downloads where 1 was correct; now 1.** Gated on the same "no table, or past the TTL" test `lookup()` already uses, so priming and looking up cannot disagree about what needs fetching — and a table that genuinely ages out is still re-primed, so prices do not freeze at the first fetch. +- **A failed pricing fetch retried every tick, forever, with no backoff, ahead of the drain.** Only the success path cleared a source's stale flag, so a bad credential re-attempted on every queue tick, each attempt costing up to the 10s `_get_json` timeout, all of it before the drain. Measured with a failing Cloudflare fetch: **5 ticks produced 5 real requests and 5 `on_error` reports; now 1**, and it still recovers once the window expires — this is a backoff, not a permanent give-up, the same reasoning that makes a 401 transient in the event queue. Per-source 1→2→4→…→60s, matching the queue's own send backoff, so one bad credential cannot delay the three healthy tables; a success clears the window rather than letting it keep doubling across unrelated outages. The four fetches deliberately stay **sequential** here: this refresh runs on a blocking daemon thread, where introducing a thread pool for four fetches would be a larger change than the problem warrants, and the per-source backoff already removes the harm. JS additionally parallelises them under `Promise.allSettled`, which its event loop makes free — a language-inherent divergence, like `os.register_at_fork` vs `AsyncLocalStorage`. - **The two ports disagreed about which HTTP failures destroy events, and both were wrong.** `_PERMANENT_STATUSES` routes a batch to `_send_individually`, where each event that fails again is logged and **dropped for good**. Python listed `{400, 401, 402, 403, 404, 409, 413, 415, 422}`; JS listed `{400, 404, 409, 422}`. Measured by driving the real queue over a real socket at a server returning each status (`probes/t11_status_matrix`), the two conditional cases settle it in opposite directions, so neither port was simply "the right one": - **A key rotated back after 3s** — Python destroyed all 5 events inside the first second and none ever reached Lago; JS held them and delivered all 5 when it healed. - **An nginx-style server answering 413 above a byte limit and 200 below it** — Python's split path delivered all 5; JS held the oversized batch and delivered 0, stalling at the backoff ceiling forever. diff --git a/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py b/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py index 8b4d21f..0a0236a 100644 --- a/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py +++ b/src/lago_agent_sdk/gateway/adapters/cloudflare_gateway.py @@ -14,12 +14,18 @@ Cloudflare reports its OWN counter vocabulary here, not the provider's. Across all 14 captured fixtures — Anthropic, Workers AI, Mistral and Gemini, via every ingress -method — the only keys that ever appear are `input_tokens`, `output_tokens`, -`total_tokens`, `input_cached_tokens`, `input_cache_creation_tokens`, `neurons`, -`input_text_tokens` and `reasoningTokens`. Not one provider-native key shows up: -no Anthropic `cache_read_input_tokens`, no Gemini `thoughtsTokenCount` or +method — the keys that appear are `input_tokens`, `output_tokens`, `total_tokens`, +`input_cached_tokens`, `input_cache_creation_tokens`, `neurons`, `input_text_tokens` +and `reasoningTokens`. Not one provider-native key shows up: no Anthropic +`cache_read_input_tokens`, no Gemini `thoughtsTokenCount` or `cachedContentTokenCount`. +That list is a snapshot and has already been overtaken once: a live Logs API pull +also returned `units`, which appears in none of the fixtures. Treat the enumeration +as illustrative, not exhaustive — `_MAPPED_USAGE_KEYS` plus the drift sweep into +`extras["usage_metadata"]` is what actually keeps an unrecognized counter from being +lost, and it needs no re-audit to stay correct. + That vocabulary is *mostly* snake_case, with `reasoningTokens` as a camelCase outlier — Cloudflare's own inconsistency, not a provider key leaking through (Gemini's native spelling for the same quantity is `thoughtsTokenCount`, which @@ -62,6 +68,47 @@ def _safe_str(v: Any) -> str: return v if isinstance(v, str) else "" +# Every `usage_metadata` spelling this adapter accounts for: the ones `_first_int` +# consults below, plus the three that are redundant with the top-level `tokens_in` / +# `tokens_out` the adapter reads directly. Anything NOT in here is swept into +# `extras["usage_metadata"]` rather than dropped — see the drift note on `extras`. +# +# Keep this in sync with the `_first_int` calls. It is the mechanism that makes the +# module docstring's key enumeration self-maintaining instead of a hand-audited +# snapshot: a spelling nobody has seen shows up in `extras` on its own. +_MAPPED_USAGE_KEYS = frozenset( + { + # cache_read + "input_cached_tokens", + "inputCachedTokens", + "cachedContentTokenCount", + "cache_read_input_tokens", + # cache_write + "input_cache_creation_tokens", + "inputCacheCreationTokens", + "cache_creation_input_tokens", + # reasoning + "reasoningTokens", + "reasoning_tokens", + "thoughtsTokenCount", + # Read from the top level instead, so not drift when they appear here. + "input_tokens", + "output_tokens", + "total_tokens", + } +) + + +def _unmapped_usage(usage_meta: dict[str, Any]) -> dict[str, Any]: + """Any `usage_metadata` key this adapter does not account for, wrapped for `extras`. + + Returns an empty dict when everything was recognized, so the key is absent rather + than present-and-empty in the overwhelmingly common case. + """ + unmapped = {k: v for k, v in usage_meta.items() if k not in _MAPPED_USAGE_KEYS} + return {"usage_metadata": unmapped} if unmapped else {} + + def _first_int(meta: dict[str, Any], *names: str) -> int: """First of `names` present in `meta` with a usable value, as an int. @@ -173,6 +220,23 @@ def extract_cloudflare_log(entry: dict[str, Any]) -> CanonicalUsage: "cached": entry.get("cached"), "step": entry.get("step"), "log_id": entry.get("id"), + # Drift sweep — the same contract `adapters/openai_native.py` enforces, and + # for the same reason: a counter this adapter does not map must not vanish + # without an error or an on_error. `extras` used to be exactly the three + # keys above, so `usage_metadata` got no sweep at all, and that was not + # hypothetical — a live Logs API pull found `neurons` (Cloudflare's Workers + # AI billing unit) and `units` (a cost quantity) being dropped on every row, + # and `units` appears in NO captured fixture, so the hand-maintained + # enumeration had already drifted past what this file claimed to know. A + # money-relevant counter going missing this way surfaces first as a + # reconciliation gap, not as a failure. + # + # NESTED, not merged flat into `extras`: the poller reads `extras["cached"]` + # to decide whether to skip billing a request Cloudflare served for free, so + # a future `usage_metadata` key called `cached` or `step` must not be able to + # shadow it. Omitted entirely when there is no drift, to keep the common case + # identical to what callers already see. + **_unmapped_usage(usage_meta), }, ) diff --git a/src/lago_agent_sdk/pricing.py b/src/lago_agent_sdk/pricing.py index bf0c76f..6849f31 100644 --- a/src/lago_agent_sdk/pricing.py +++ b/src/lago_agent_sdk/pricing.py @@ -55,6 +55,11 @@ logger = logging.getLogger("lago_agent_sdk.pricing") +# Ceiling for a pricing source's post-failure backoff — the same 60s cap the event +# queue uses for send retries, so a persistently-broken credential settles into one +# attempt a minute instead of one per flush tick. +_MAX_PRICING_BACKOFF_SECONDS = 60.0 + OPENROUTER_URL = "https://openrouter.ai/api/v1/models" AWS_PRICING_HOST = "https://pricing.us-east-1.amazonaws.com" AWS_BEDROCK_REGION_INDEX = f"{AWS_PRICING_HOST}/offers/v1.0/aws/AmazonBedrock/current/region_index.json" @@ -1102,6 +1107,9 @@ def __init__( # ever requiring a separate LagoConfig.mistral_api_key. self._mistral_api_key_override: str | None = None self._refreshing: set[str] = set() + # Per-source post-failure backoff — see `_in_backoff`. + self._failure_backoff_until: dict[str, float] = {} + self._failure_backoff_seconds: dict[str, float] = {} def _heal_fork(self) -> None: """Self-heal after a fork: a lock copied from the parent may be held by a @@ -1143,14 +1151,61 @@ def prime(self, providers: Iterable[str] = ()) -> None: cold-start cost. Unknown provider names are silently ignored (no source is warmed) rather than raising, since this is a hint, not a contract.""" + # Gated on "is this table actually cold?", NOT unconditional. `prime()` is called + # from `_auto_prime_pricing_for` on a matching `wrap()` and from `warm_pricing()`, + # both of which a server can run per request — and flagging an in-TTL table stale + # meant the ~400-model OpenRouter catalogue was re-downloaded on essentially every + # flush tick, so `pricing_ttl_seconds` never applied on this path at all. Measured + # against the live catalogue with the shipped 1-hour TTL: 4 prime()+maybe_refresh() + # cycles produced 4 full downloads where 1 was correct. + # + # "Cold" is the same test `lookup()` already uses — no table, or past the TTL — so + # priming and looking up cannot disagree about what needs fetching. with self._lock: - self._openrouter_stale = True + if self._is_cold(self._openrouter, self._openrouter_fetched): + self._openrouter_stale = True for p in providers: key = (p or "").lower() if key == "workers-ai": - self._cloudflare_stale = True + if self._is_cold(self._cloudflare_workers_ai, self._cloudflare_fetched): + self._cloudflare_stale = True elif key == "mistral": - self._mistral_stale = True + if self._is_cold(self._mistral_aliases, self._mistral_fetched): + self._mistral_stale = True + + def _is_cold(self, table: Any, fetched_at: float) -> bool: + """True when a table needs fetching: absent, or older than the TTL. + + Caller must hold `self._lock`. + """ + return table is None or (time.time() - fetched_at) >= self._ttl + + def _in_backoff(self, source: str) -> bool: + """True while `source` is inside its post-failure backoff window. + + A failed fetch used to leave its stale flag set and nothing else, so the next tick + retried immediately — every tick, forever, with no delay, each attempt costing up + to the 10s `_get_json` timeout, all of it on the queue thread AHEAD of the drain. + Measured with a bad Cloudflare token: 5 ticks produced 5 real requests and 5 + `on_error` reports. + + Same 1→2→4→…→60s shape as the queue's own send backoff, tracked per source so one + bad credential cannot delay the three healthy tables. + """ + with self._lock: + return time.time() < self._failure_backoff_until.get(source, 0.0) + + def _note_failure(self, source: str) -> None: + with self._lock: + prev = self._failure_backoff_seconds.get(source, 0.0) + nxt = 1.0 if prev == 0.0 else min(prev * 2, _MAX_PRICING_BACKOFF_SECONDS) + self._failure_backoff_seconds[source] = nxt + self._failure_backoff_until[source] = time.time() + nxt + + def _note_success(self, source: str) -> None: + with self._lock: + self._failure_backoff_seconds.pop(source, None) + self._failure_backoff_until.pop(source, None) def learn_mistral_api_key(self, api_key: str) -> None: """Adopt a Mistral API key discovered from a wrapped client, so @@ -1226,16 +1281,36 @@ def maybe_refresh(self) -> None: ): return with self._lock: - do_openrouter = self._openrouter_stale and "openrouter" not in self._refreshing + now = time.time() + + def _ready(source: str) -> bool: + # Inlined rather than calling `_in_backoff`, which takes the lock we hold. + return now >= self._failure_backoff_until.get(source, 0.0) + + do_openrouter = ( + self._openrouter_stale and "openrouter" not in self._refreshing and _ready("openrouter") + ) if do_openrouter: self._refreshing.add("openrouter") - do_cloudflare = self._cloudflare_stale and "cloudflare_workers_ai" not in self._refreshing + do_cloudflare = ( + self._cloudflare_stale + and "cloudflare_workers_ai" not in self._refreshing + and _ready("cloudflare_workers_ai") + ) if do_cloudflare: self._refreshing.add("cloudflare_workers_ai") - do_mistral = self._mistral_stale and "mistral_aliases" not in self._refreshing + do_mistral = ( + self._mistral_stale + and "mistral_aliases" not in self._refreshing + and _ready("mistral_aliases") + ) if do_mistral: self._refreshing.add("mistral_aliases") - regions = [r for r in self._bedrock_stale if f"bedrock:{r}" not in self._refreshing] + regions = [ + r + for r in self._bedrock_stale + if f"bedrock:{r}" not in self._refreshing and _ready(f"bedrock:{r}") + ] for r in regions: self._refreshing.add(f"bedrock:{r}") @@ -1246,7 +1321,9 @@ def maybe_refresh(self) -> None: self._openrouter = table self._openrouter_fetched = time.time() self._openrouter_stale = False + self._note_success("openrouter") except Exception as exc: # noqa: BLE001 + self._note_failure("openrouter") self._report(exc, "pricing.fetch_openrouter") finally: with self._lock: @@ -1259,7 +1336,9 @@ def maybe_refresh(self) -> None: self._cloudflare_workers_ai = table_cf self._cloudflare_fetched = time.time() self._cloudflare_stale = False + self._note_success("cloudflare_workers_ai") except Exception as exc: # noqa: BLE001 + self._note_failure("cloudflare_workers_ai") self._report(exc, "pricing.fetch_cloudflare_workers_ai") finally: with self._lock: @@ -1274,7 +1353,9 @@ def maybe_refresh(self) -> None: self._mistral_aliases = aliases self._mistral_fetched = time.time() self._mistral_stale = False + self._note_success("mistral_aliases") except Exception as exc: # noqa: BLE001 + self._note_failure("mistral_aliases") self._report(exc, "pricing.fetch_mistral_aliases") finally: with self._lock: @@ -1287,7 +1368,9 @@ def maybe_refresh(self) -> None: self._bedrock[r] = table self._bedrock_fetched[r] = time.time() self._bedrock_stale.discard(r) + self._note_success(f"bedrock:{r}") except Exception as exc: # noqa: BLE001 + self._note_failure(f"bedrock:{r}") self._report(exc, "pricing.fetch_bedrock") finally: with self._lock: diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index 3fff53b..4631c74 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -108,6 +108,8 @@ def __init__( # classifies it transient, re-prepends the batch and retries at the 60s ceiling # forever. All billing stops, nothing is dropped or escalated, and the only # symptom is a growing buffer. + # + # Falling back is right, but it must not be SILENT — see the report below. if api_url: self.config.api_url = api_url if default_subscription_id is not None: @@ -115,6 +117,25 @@ def __init__( if verify_ssl is not None: self.config.verify_ssl = verify_ssl + # A caller who passed `api_url` explicitly MEANT to point somewhere specific. + # Discarding a falsy one is the safe choice for delivery, but doing it silently + # is the one outcome that must not happen here: `LagoConfig`'s default is + # PRODUCTION, so `api_url=os.environ.get("LAGO_API_URL", "")` with the var unset + # now resolves to production Lago and every event is accepted. For a CI job or a + # developer holding a real production key that writes live billing data, and + # ingested events cannot be un-ingested. `on_error` is opt-in, so this reports + # through the same log-plus-callback floor as every other drop path rather than + # trusting a callback to exist. + if api_url is not None and not api_url: + self._report_error( + ValueError( + f"api_url was explicitly set to an empty value; falling back to " + f"{self.config.api_url}. Set LAGO_API_URL (or pass config.api_url) " + f"if you did not intend to send events there." + ), + "config.api_url", + ) + self._lago_client = LagoClient( api_key=self.config.api_key, api_url=self.config.api_url, diff --git a/tests/unit/gateway/adapters/test_cloudflare_gateway.py b/tests/unit/gateway/adapters/test_cloudflare_gateway.py index 1e6c8cc..f8ac9a4 100644 --- a/tests/unit/gateway/adapters/test_cloudflare_gateway.py +++ b/tests/unit/gateway/adapters/test_cloudflare_gateway.py @@ -10,6 +10,7 @@ from lago_agent_sdk import CanonicalUsage from lago_agent_sdk.gateway.adapters import extract_cloudflare_log, resolve_subscription +from lago_agent_sdk.gateway.adapters.cloudflare_gateway import _MAPPED_USAGE_KEYS from lago_agent_sdk.pricing import compute_cost, lookup_openrouter, parse_openrouter FIX = pathlib.Path(__file__).parent / "fixtures" / "cloudflare_gateway" @@ -462,3 +463,152 @@ def test_provider_native_cache_and_reasoning_spellings_are_accepted() -> None: } ) assert both.cache_read == 11 + + +# ---------------------------------------------------------------------- +# Drift contract for `usage_metadata`. +# +# `extras` used to be a fixed three-key dict, so any counter this adapter did not map +# was silently dropped. That was not hypothetical: a live Logs API pull found `neurons` +# and `units` vanishing on every row, and `units` appears in no captured fixture — the +# hand-maintained enumeration in the module docstring had already drifted past reality. +# Same contract `test_drift.py` pins for the native adapters. +# ---------------------------------------------------------------------- +def test_drift_keeps_an_unmapped_counter_instead_of_dropping_it() -> None: + """Exactly the shape seen live (entry 01M0FEZ2Y7QMQR1HT11GVT2HCE).""" + u = extract_cloudflare_log( + { + "id": "log_1", + "cached": False, + "step": 0, + "tokens_in": 37, + "tokens_out": 2, + "provider": "workers-ai", + "model": "@cf/meta/llama-3.3-70b-instruct-fp8-fast", + "usage_metadata": { + "input_tokens": 37, + "output_tokens": 2, + "total_tokens": 39, + "input_cached_tokens": 0, + "neurons": 1.396314412355423, + "units": 0.00001535945853590965, + }, + } + ) + # Cloudflare's Workers AI billing unit, and a cost quantity — both money-relevant. + assert u.extras["usage_metadata"] == { + "neurons": 1.396314412355423, + "units": 0.00001535945853590965, + } + + +def test_drift_sweeps_a_counter_nobody_has_ever_seen() -> None: + u = extract_cloudflare_log( + { + "tokens_in": 10, + "tokens_out": 1, + "provider": "anthropic", + "usage_metadata": {"input_tokens": 10, "audio_input_tokens": 512}, + } + ) + assert u.extras["usage_metadata"] == {"audio_input_tokens": 512} + + +def test_drift_never_shadows_the_pollers_own_billing_inputs() -> None: + """`extras["cached"]` decides whether to skip billing a request Cloudflare served + for free. A usage_metadata key of the same name must not be able to overwrite it — + which is why the sweep is nested rather than merged flat into extras.""" + u = extract_cloudflare_log( + { + "id": "log_2", + "cached": True, + "step": 3, + "tokens_in": 5, + "tokens_out": 1, + "provider": "anthropic", + "usage_metadata": {"cached": False, "step": 99, "log_id": "spoofed"}, + } + ) + assert u.extras["cached"] is True + assert u.extras["step"] == 3 + assert u.extras["log_id"] == "log_2" + assert u.extras["usage_metadata"] == {"cached": False, "step": 99, "log_id": "spoofed"} + + +def test_drift_omits_the_key_entirely_when_there_is_none() -> None: + """The common case must look exactly as it did before the sweep existed.""" + u = extract_cloudflare_log( + { + "id": "log_3", + "cached": False, + "step": 0, + "tokens_in": 9, + "tokens_out": 21, + "provider": "anthropic", + "usage_metadata": { + "input_tokens": 9, + "output_tokens": 21, + "total_tokens": 30, + "input_cached_tokens": 4, + }, + } + ) + assert u.extras == {"cached": False, "step": 0, "log_id": "log_3"} + assert "usage_metadata" not in u.extras + + +@pytest.mark.parametrize( + "key", + [ + "input_cached_tokens", + "inputCachedTokens", + "cachedContentTokenCount", + "cache_read_input_tokens", + "input_cache_creation_tokens", + "inputCacheCreationTokens", + "cache_creation_input_tokens", + "reasoningTokens", + "reasoning_tokens", + "thoughtsTokenCount", + "input_tokens", + "output_tokens", + "total_tokens", + ], +) +def test_drift_every_mapped_spelling_stays_out_of_the_sweep(key: str) -> None: + """A key that IS consumed must not also show up as drift — that would read as an + unhandled counter in reconciliation and invite double-counting.""" + u = extract_cloudflare_log( + { + "tokens_in": 100, + "tokens_out": 10, + "provider": "anthropic", + "usage_metadata": {key: 7}, + } + ) + assert "usage_metadata" not in u.extras, f"{key} should be mapped, not swept" + + +def test_drift_no_captured_fixture_loses_a_counter() -> None: + """The sweep against real data, not constructed entries. + + Before the sweep existed this dropped `neurons` in 4 of the 14 captured entries and + `input_text_tokens` in 1 — measured, not hypothesised. Iterating the fixtures rather + than asserting a fixed key list is what makes this test survive a recapture: a new + counter Cloudflare starts sending is caught by the next `capture` run, with no test + edit and no re-audit of the module docstring. + """ + fixtures = sorted(FIX.glob("*.json")) + # Absent fixtures read as "not covered", never as a pass — same rule as the sweeps. + assert fixtures, "no captured Cloudflare fixtures found" + for path in fixtures: + entry = _load(path.name) + meta = entry.get("usage_metadata") + if not isinstance(meta, dict): + continue + u = extract_cloudflare_log(entry) + swept = u.extras.get("usage_metadata", {}) + for key in meta: + assert key in _MAPPED_USAGE_KEYS or key in swept, ( + f"{path.name}: {key!r} is neither mapped nor swept into extras" + ) diff --git a/tests/unit/test_pricing.py b/tests/unit/test_pricing.py index b8d68e3..7ae0663 100644 --- a/tests/unit/test_pricing.py +++ b/tests/unit/test_pricing.py @@ -5,6 +5,7 @@ import json import pathlib import re +import time import uuid from decimal import Decimal from typing import Any @@ -2038,3 +2039,113 @@ def test_workers_ai_compat_prefix_is_defined_exactly_once() -> None: f"expected one definition, found {definitions}" ) assert WORKERS_AI_COMPAT_PREFIX == "workers-ai/" + + +# ---------------------------------------------------------------------- +# prime() vs the TTL, and per-source backoff after a failed fetch. +# +# Both are about the same thing: pricing refresh runs on the queue's background +# thread, AHEAD of the drain, so wasted work there delays real billing events. +# ---------------------------------------------------------------------- +class _FailingFetcher(StubFetcher): + """StubFetcher whose Cloudflare fetch raises, like a bad CLOUDFLARE_API_TOKEN.""" + + def fetch_cloudflare_workers_ai(self) -> dict[str, ModelPrice]: + self.cloudflare_workers_ai_calls += 1 + raise RuntimeError("HTTP 401 Unauthorized") + + +def test_prime_respects_the_ttl_instead_of_refetching_every_tick() -> None: + """`prime()` used to set the stale flag unconditionally, so the TTL never applied + on this path at all — and `prime()` is reached from `_auto_prime_pricing_for` on a + matching `wrap()` and from `warm_pricing()`, both of which a server can run per + request. That re-downloaded the ~400-model OpenRouter catalogue on essentially + every flush tick, on the thread the queue drains events from.""" + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + for _ in range(4): + p.prime() + p.maybe_refresh() + assert fetcher.openrouter_calls == 1 + + +def test_prime_still_warms_a_table_that_has_aged_past_the_ttl() -> None: + """The TTL gate must not become a permanent lock-out: once a table is genuinely + stale, priming has to flag it again or prices freeze at the first fetch forever.""" + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=0) + for _ in range(3): + p.prime() + p.maybe_refresh() + assert fetcher.openrouter_calls == 3 + + +def test_a_failed_pricing_fetch_backs_off_instead_of_retrying_every_tick() -> None: + """Only the success path used to clear a source's stale flag, so a bad credential + re-attempted on every single tick — each attempt costing up to the 10s `_get_json` + timeout, and each one reported, ahead of the drain.""" + errors: list[str] = [] + fetcher = _FailingFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600, on_error=lambda e, w: errors.append(w)) + for _ in range(5): + p.prime(providers=["workers-ai"]) + p.maybe_refresh() + assert fetcher.cloudflare_workers_ai_calls == 1 + assert errors == ["pricing.fetch_cloudflare_workers_ai"] + assert p._in_backoff("cloudflare_workers_ai") + + +def test_a_failed_pricing_source_recovers_after_its_backoff_window() -> None: + """This is a backoff, not a permanent give-up. A rotated-back credential has to + start working again on its own — the same reasoning that makes a 401 transient in + the event queue.""" + fetcher = _FailingFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.prime(providers=["workers-ai"]) + p.maybe_refresh() + assert fetcher.cloudflare_workers_ai_calls == 1 + + # Expire the window rather than sleeping through it. + p._failure_backoff_until["cloudflare_workers_ai"] = time.time() - 0.001 + p.prime(providers=["workers-ai"]) + p.maybe_refresh() + assert fetcher.cloudflare_workers_ai_calls == 2 + + +def test_one_broken_pricing_source_does_not_delay_the_healthy_ones() -> None: + """Backoff is tracked per source precisely so a single bad credential cannot hold + up the other tables — the reason this is a dict and not one global timestamp.""" + fetcher = _FailingFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p.prime(providers=["workers-ai"]) + p.maybe_refresh() + assert fetcher.cloudflare_workers_ai_calls == 1 + assert fetcher.openrouter_calls == 1 + assert not p._in_backoff("openrouter") + + +def test_a_successful_fetch_clears_an_earlier_backoff() -> None: + """Otherwise the window keeps doubling across unrelated failures and a healthy + source inherits a minute-long delay from an outage that already ended.""" + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + p._note_failure("openrouter") + p._note_failure("openrouter") + assert p._in_backoff("openrouter") + p._note_success("openrouter") + assert not p._in_backoff("openrouter") + # And the next failure restarts at 1s rather than resuming at 4s. + p._note_failure("openrouter") + assert p._failure_backoff_seconds["openrouter"] == 1.0 + + +def test_pricing_backoff_growth_matches_the_queues_and_is_capped() -> None: + """Same 1→2→4→…→60s shape as the event queue's send backoff. The cap is what makes + a permanently-dead source settle into one attempt a minute instead of growing + without bound and effectively never retrying.""" + p = PricingProvider(fetcher=StubFetcher(), ttl_seconds=3600) + seq = [] + for _ in range(8): + p._note_failure("openrouter") + seq.append(p._failure_backoff_seconds["openrouter"]) + assert seq == [1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 60.0, 60.0] diff --git a/tests/unit/test_sdk.py b/tests/unit/test_sdk.py index cebd3d3..7811be7 100644 --- a/tests/unit/test_sdk.py +++ b/tests/unit/test_sdk.py @@ -128,6 +128,43 @@ def test_empty_or_absent_api_url_keeps_the_production_default(empty): sdk.shutdown(timeout=1.0) +def test_a_discarded_empty_api_url_is_reported_not_swallowed(): + """Falling back is right; falling back SILENTLY is the dangerous part. + + `LagoConfig`'s default is PRODUCTION, so the fallback above means + `api_url=os.environ.get("LAGO_API_URL", "")` with the var unset points a CI job or a + developer holding a real production key at production Lago, which accepts every + event — and ingested events cannot be un-ingested.""" + errors: list[tuple[Exception, str]] = [] + sdk = LagoSDK( + api_key="k", + api_url="", + config=LagoConfig(api_key="k", on_error=lambda e, w: errors.append((e, w))), + ) + try: + assert [w for _, w in errors] == ["config.api_url"] + # The message has to name where the events are actually going, or the report + # tells the reader nothing they can act on. + assert "api.getlago.com" in str(errors[0][0]) + finally: + sdk.shutdown(timeout=1.0) + + +def test_an_unpassed_api_url_is_not_reported(): + """`None` means "the caller never mentioned api_url", which is the overwhelmingly + common case and not a mistake. Reporting it would train customers to ignore + `on_error` — the one channel the empty-string case above depends on.""" + errors: list[tuple[Exception, str]] = [] + sdk = LagoSDK( + api_key="k", + config=LagoConfig(api_key="k", on_error=lambda e, w: errors.append((e, w))), + ) + try: + assert errors == [] + finally: + sdk.shutdown(timeout=1.0) + + def test_default_api_url_is_still_production_when_nothing_is_passed(): """Changing the parameter default to None must not change this.""" sdk = LagoSDK(api_key="k")