From 8934ca807f10db9feb4772a4789c011d8d569acb Mon Sep 17 00:00:00 2001 From: Anass Date: Wed, 10 Jun 2026 18:08:49 +0200 Subject: [PATCH] Add price mode: emit computed dollar cost instead of token counts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New optional pricing_mode ("tokens" default | "price"). In price mode the SDK emits one llm_cost event per call whose value is Σ(unit_price × tokens) × markup, with a full per-field breakdown (tokens, unit_price, cost, base_cost, markup, price_source) in properties. Pricing sources (public, no API keys): - OpenRouter /api/v1/models for native anthropic/openai/mistral/gemini (USD/token) - AWS Bedrock Price List Bulk API (per-region offer files) for Bedrock Design: - pricing.py: PricingProvider with a TTL cache + injectable fetcher; lookup() is pure in-memory and never blocks the call; maybe_refresh() does the HTTP on the queue's background thread. Fork-safe via a PID self-heal (no register_at_fork — that tripped macOS's objc fork-safety abort). - Conservative, vendor-gated model matching; Bedrock parser keys on inferenceType and rejects priority/flex/batch tiers, scales per-1K units. - Money in Decimal floored to 12 dp; identical output to the JS BigInt impl (locked by a cross-repo golden fixture). - Fallback: unavailable price -> emit token events + on_error (never under-bill). - mode + markup overridable per-call via extra_lago; global via LagoConfig. Default mode is "tokens" -> zero behavior change. New config: pricing_mode, markup, cost_metric_code, pricing_ttl_seconds, bedrock_default_region, pricing_provider. 28 new pricing unit tests + env-gated live test. Gate: ruff + format + mypy clean; 346 unit tests; coverage 88.27%. --- CHANGELOG.md | 3 + README.md | 51 ++ src/lago_agent_sdk/__init__.py | 10 +- src/lago_agent_sdk/config.py | 30 +- src/lago_agent_sdk/exceptions.py | 11 + src/lago_agent_sdk/pricing.py | 568 ++++++++++++++++++ src/lago_agent_sdk/queue.py | 15 + src/lago_agent_sdk/sdk.py | 148 ++++- src/lago_agent_sdk/wrappers/anthropic.py | 43 +- src/lago_agent_sdk/wrappers/boto3_bedrock.py | 8 + src/lago_agent_sdk/wrappers/gemini.py | 33 +- src/lago_agent_sdk/wrappers/mistral.py | 27 +- src/lago_agent_sdk/wrappers/openai.py | 29 +- tests/integration/test_live_pricing.py | 58 ++ tests/unit/fixtures/pricing/money_golden.json | 59 ++ tests/unit/test_pricing.py | 525 ++++++++++++++++ 16 files changed, 1526 insertions(+), 92 deletions(-) create mode 100644 src/lago_agent_sdk/pricing.py create mode 100644 tests/integration/test_live_pricing.py create mode 100644 tests/unit/fixtures/pricing/money_golden.json create mode 100644 tests/unit/test_pricing.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a03d61d..679920e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ All notable changes to this project will be documented here. Format follows [Kee ## [Unreleased] +### Added +- **Price mode — emit computed dollar cost instead of token counts.** New `pricing_mode` config (`"tokens"` default | `"price"`), plus `markup`, `cost_metric_code` (default `llm_cost`), `pricing_ttl_seconds`, and `bedrock_default_region`. In price mode the SDK emits one `llm_cost` event per call carrying a top-level `precise_total_amount_cents` (cost in cents, after markup) for Lago's **dynamic charge model**, with a full per-field breakdown in `properties` (value in USD, base, markup, source, per-field tokens/unit_price/cost). Live unit prices come from public, no-auth sources: OpenRouter (`/api/v1/models`) for native anthropic/openai/mistral/gemini, and the AWS Bedrock Price List **Bulk** API for Bedrock. Prices are fetched + cached on the background queue thread (never blocking the customer's call); a missing price falls back to token events and calls `on_error` (never silently under-bills). Mode and markup are overridable per-call via `extra_lago={"mode": "price", "markup": 1.5}`. Money is computed with `Decimal` floored to 12 dp, identical to the JS implementation (cross-repo golden fixture). New `pricing.py` module + `PricingProvider`; default `pricing_mode="tokens"` keeps existing behavior unchanged. + ### Fixed - **Anthropic `messages.create(stream=True)` under-billed input tokens.** The stream wrapper read only top-level `usage`, which on a basic stream appears only on `message_delta` as `{output_tokens: N}` — the authoritative `input_tokens` / `cache_*` counts arrive nested under `message.usage` on the `message_start` event and were ignored, so input billed 0. The wrapper now merges usage from `message_start` (input/cache) and `message_delta` (cumulative output). Sync + async paths; regression tests use the realistic wire shape (delta carries no input echo). - **Legacy `google-generativeai` SDK silently emitted no events.** The detector matched both the new `google-genai` and the deprecated `google-generativeai` SDKs, but the wrapper only instruments the unified `Client.models` / `.aio` surface — a legacy `GenerativeModel` routed through and wrapped nothing. `wrap()` now rejects legacy clients with a clear pointer to migrate to `google-genai`. diff --git a/README.md b/README.md index ee0d425..70562a2 100644 --- a/README.md +++ b/README.md @@ -189,6 +189,57 @@ For both OpenAI and Gemini, `cache_read`, `audio_input`, and `image_input` are * OpenAI's Predicted Outputs tokens (`accepted_prediction_tokens`, `rejected_prediction_tokens`) are not surfaced — see the OpenAI adapter docstring for details on this intentional gap. +## Pricing mode — send dollar cost instead of tokens + +By default the SDK emits **token counts** (`pricing_mode="tokens"`). You can instead have it +compute and emit the **dollar cost** of each call: `Σ(unit_price_per_token × tokens) × markup`. + +```python +from lago_agent_sdk import LagoSDK, LagoConfig + +sdk = LagoSDK(api_key="...", config=LagoConfig( + api_key="...", + default_subscription_id="sub_123", + pricing_mode="price", # "tokens" (default) | "price" + markup=1.2, # optional cost multiplier (1.2 = +20%) +)) +client = sdk.wrap(anthropic_client) +# ... use the client normally ... +``` + +In **price mode** the SDK emits **one event per call** with code `llm_cost`. The event carries a +top-level `precise_total_amount_cents` (the total cost in cents, after markup) for Lago's +**dynamic charge model**, plus a breakdown in `properties`: `unit` (total tokens), `value` (USD +total), `base_cost` (pre-markup), `markup`, `price_source`, and per-field `*_tokens` / +`*_unit_price` / `*_cost`. Set up in Lago a `sum`-aggregation billable metric `llm_cost` on +`field_name: "unit"` and a **dynamic** charge on it — Lago sums each event's +`precise_total_amount_cents` into a single fee (`unit` is the displayed usage quantity). See +`testing/lago_setup_pricing_plan.py` for a script that creates this. + +Per-call override via `extra_lago` (mode and markup, in addition to subscription/dimensions): + +```python +client.messages.create(model="claude-...", messages=[...], + extra_lago={"mode": "price", "markup": 1.5}) +``` + +**Live, public pricing sources (no API keys):** +- **OpenRouter** (`/api/v1/models`) for native `anthropic` / `openai` / `mistral` / `gemini` + clients — USD per token. +- **AWS Bedrock Price List Bulk API** (public) for Bedrock — parsed per region. + +Prices are fetched and cached in the background (TTL `pricing_ttl_seconds`, default 1h); the +refresh runs on the SDK's background thread, so **your LLM call is never blocked on pricing**. + +**Fallback (never under-bill):** if a price is unavailable (table not warm on the first call, +or the model isn't found in the source), the SDK **falls back to emitting token-count events** +and calls `on_error` so it's visible — it never silently drops the usage. + +**Bedrock note:** AWS's public bulk data lists many models (Titan, Llama, Mistral, Cohere, and +older Claude) but, at time of writing, **not the current Claude 3.5/3.7/4 models**. Bedrock +calls for models absent from AWS's data fall back to token events. Native Anthropic clients are +priced via OpenRouter and unaffected. + ## Error policy The SDK never breaks your LLM call. If anything in instrumentation fails (adapter bug, Lago down, network error), the SDK swallows it, logs a warning, and your call returns normally. diff --git a/src/lago_agent_sdk/__init__.py b/src/lago_agent_sdk/__init__.py index 6e3347b..42837be 100644 --- a/src/lago_agent_sdk/__init__.py +++ b/src/lago_agent_sdk/__init__.py @@ -1,13 +1,15 @@ """Lago Agent SDK — Python.""" from .canonical import CanonicalUsage -from .config import DEFAULT_METRIC_CODES, LagoConfig +from .config import DEFAULT_COST_METRIC_CODE, DEFAULT_METRIC_CODES, LagoConfig from .exceptions import ( LagoApiError, LagoConfigError, LagoSDKError, + PricingUnavailableError, UnknownClientError, ) +from .pricing import HttpPricingFetcher, ModelPrice, PricingProvider, compute_cost from .sdk import LagoSDK __all__ = [ @@ -17,7 +19,13 @@ "LagoApiError", "LagoConfigError", "LagoSDKError", + "PricingUnavailableError", "UnknownClientError", "DEFAULT_METRIC_CODES", + "DEFAULT_COST_METRIC_CODE", + "PricingProvider", + "HttpPricingFetcher", + "ModelPrice", + "compute_cost", ] __version__ = "0.1.0" diff --git a/src/lago_agent_sdk/config.py b/src/lago_agent_sdk/config.py index 28bb1c5..ea81f64 100644 --- a/src/lago_agent_sdk/config.py +++ b/src/lago_agent_sdk/config.py @@ -4,6 +4,7 @@ from collections.abc import Callable from dataclasses import dataclass, field +from typing import Any, Literal DEFAULT_METRIC_CODES: dict[str, str] = { "input": "llm_input_tokens", @@ -19,6 +20,13 @@ "audio_output": "llm_audio_output_tokens", } +# Metric code for the single per-call dollar-cost event emitted in price mode. +DEFAULT_COST_METRIC_CODE = "llm_cost" + +# Pricing mode: emit raw token counts (default, backward-compatible) or a single +# computed dollar-cost event per call. +PricingMode = Literal["tokens", "price"] + def _mask_api_key(api_key: str) -> str: """Render an api key safe for logs/repr: keeps a 4-char tail for debuggability.""" @@ -42,6 +50,21 @@ class LagoConfig: max_retry_seconds: float = 60.0 on_error: Callable[[Exception, str], None] | None = None + # --- pricing (price mode) --- + # Global default mode. "tokens" preserves the existing behavior exactly. + pricing_mode: PricingMode = "tokens" + # Multiplier applied to the computed cost (1.0 = no markup, 1.2 = +20%). + markup: float = 1.0 + # Metric code for the single dollar-cost event emitted in price mode. + cost_metric_code: str = DEFAULT_COST_METRIC_CODE + # How long a fetched pricing table stays fresh before a background refresh. + pricing_ttl_seconds: float = 3600.0 + # Region used for Bedrock pricing when the model id carries no region prefix. + bedrock_default_region: str = "us-east-1" + # Optional injected PricingProvider (or a stub) — primarily for tests/overrides. + # Typed Any to avoid a config→pricing import cycle. + pricing_provider: Any | None = field(default=None, repr=False) + def __repr__(self) -> str: return ( f"LagoConfig(api_key={_mask_api_key(self.api_key)!r}, " @@ -51,5 +74,10 @@ def __repr__(self) -> str: f"max_batch_size={self.max_batch_size}, " f"max_buffer_size={self.max_buffer_size}, " f"request_timeout_seconds={self.request_timeout_seconds}, " - f"max_retry_seconds={self.max_retry_seconds})" + f"max_retry_seconds={self.max_retry_seconds}, " + f"pricing_mode={self.pricing_mode!r}, " + f"markup={self.markup}, " + f"cost_metric_code={self.cost_metric_code!r}, " + f"pricing_ttl_seconds={self.pricing_ttl_seconds}, " + f"bedrock_default_region={self.bedrock_default_region!r})" ) diff --git a/src/lago_agent_sdk/exceptions.py b/src/lago_agent_sdk/exceptions.py index 4e60af9..1a2004d 100644 --- a/src/lago_agent_sdk/exceptions.py +++ b/src/lago_agent_sdk/exceptions.py @@ -22,3 +22,14 @@ def __init__(self, status: int, body: str) -> None: class UnknownClientError(LagoConfigError): """`wrap()` received a client kind the SDK does not recognize.""" + + +class PricingUnavailableError(LagoSDKError): + """Price mode could not resolve a price (table not warm yet, or model not + matched). Surfaced via on_error; the SDK falls back to emitting token events.""" + + def __init__(self, provider: str, model: str, api: str) -> None: + super().__init__(f"no price for provider={provider!r} model={model!r} api={api!r}") + self.provider = provider + self.model = model + self.api = api diff --git a/src/lago_agent_sdk/pricing.py b/src/lago_agent_sdk/pricing.py new file mode 100644 index 0000000..3b3dfe7 --- /dev/null +++ b/src/lago_agent_sdk/pricing.py @@ -0,0 +1,568 @@ +"""Pricing — optional dollar-cost computation for price mode. + +Fetches live, public, no-auth per-token unit prices and computes the cost of a +call as ``Σ(unit_price × token_count) × markup``. + +Sources: + - OpenRouter (``https://openrouter.ai/api/v1/models``) for native providers + (anthropic / openai / mistral / gemini). Prices are USD per token. + - AWS Bedrock Price List **Bulk** API (public, no credentials) for Bedrock. + +Design constraints (mirror the queue's non-blocking guarantee): + - ``lookup()`` is pure in-memory and O(1); it NEVER does network I/O, so the + customer's LLM call is never blocked on pricing. + - All HTTP happens in ``maybe_refresh()``, which the EventQueue's background + worker calls on its flush tick. Tables are swapped atomically under a lock. + - A cold/missing table returns ``None`` from ``lookup`` → the caller falls back + to emitting token events (see sdk.emit), so we never silently under-bill. + +Money is computed with ``decimal.Decimal`` and floored to 12 decimal places +(ROUND_DOWN) so results are deterministic and match the JS implementation +byte-for-byte. +""" + +from __future__ import annotations + +import logging +import os +import re +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass +from decimal import ROUND_DOWN, Decimal, InvalidOperation +from typing import Any, Protocol + +from .canonical import CanonicalUsage + +logger = logging.getLogger("lago_agent_sdk.pricing") + +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" + +# Canonical usage fields we know how to price. +PRICED_FIELDS = ("input", "output", "cache_read", "cache_write", "reasoning") + +# Providers whose reported `input` token count ALREADY includes the cached +# (`cache_read`) tokens — i.e. cache_read is a subset of input, not additive. +# For these, the cached portion must be billed at the cache-read rate, not the +# full prompt rate, so compute_cost moves it out of `input`. Anthropic reports +# input EXCLUSIVE of cache (cache_read/cache_write are additive), so it's absent. +_INPUT_INCLUDES_CACHE_READ = frozenset({"openai", "gemini"}) + +# Providers whose reported `output` token count ALREADY includes the reasoning +# tokens (reasoning is a subset of output). For these, reasoning is billed as +# part of output and must NOT be billed again separately. (Gemini's `thoughts` +# are additive to output, so it's absent here.) +_OUTPUT_INCLUDES_REASONING = frozenset({"openai"}) + +# Canonical field -> OpenRouter pricing key. +_OPENROUTER_FIELD_MAP = { + "input": "prompt", + "output": "completion", + "cache_read": "input_cache_read", + "cache_write": "input_cache_write", + "reasoning": "internal_reasoning", +} + +# Our provider name -> OpenRouter vendor prefix. +_VENDOR_MAP = { + "anthropic": "anthropic", + "openai": "openai", + "mistral": "mistralai", + "gemini": "google", + "google": "google", +} + +# Bedrock cross-region inference prefix -> a representative AWS region. +_BEDROCK_REGION_PREFIX = { + "us": "us-east-1", + "eu": "eu-west-1", + "apac": "ap-southeast-1", +} + +# Vendor words that may lead an AWS Bedrock product's model name. +_BEDROCK_VENDOR_WORDS = { + "anthropic", + "mistral", + "mistralai", + "ai21", + "cohere", + "meta", + "amazon", + "stability", + "stabilityai", + "google", +} + +_SCALE = 12 +_Q = Decimal(1).scaleb(-_SCALE) # Decimal("1E-12") +_VERSION_DATE_SUFFIX = re.compile(r"-(?:\d{8}|v\d+)$") + + +# ---------------------------------------------------------------------- +# Money helpers (kept in lock-step with the JS implementation) +# ---------------------------------------------------------------------- +def _parse_price(value: Any) -> Decimal | None: + """Parse a price into a Decimal floored to 12 dp. None on invalid/negative.""" + try: + d = Decimal(str(value)) + except (InvalidOperation, ValueError, TypeError): + return None + if d.is_nan() or d.is_infinite() or d < 0: + return None + return d.quantize(_Q, rounding=ROUND_DOWN) + + +def _fmt_money(d: Decimal) -> str: + """Floor to 12 dp, render as a plain decimal string, trim trailing zeros.""" + q = d.quantize(_Q, rounding=ROUND_DOWN) + s = format(q, "f") + if "." in s: + s = s.rstrip("0").rstrip(".") + return s or "0" + + +def _norm(s: str) -> str: + """Lowercase + unify '.'/'-' so 'claude-opus-4.8' == 'claude-opus-4-8'.""" + return s.lower().replace(".", "-") + + +def _alnum(s: str) -> str: + """Lowercase, keep only [a-z0-9] — for cross-format (AWS) matching.""" + return re.sub(r"[^a-z0-9]", "", s.lower()) + + +def _strip_version(model: str) -> str: + """Drop a trailing -YYYYMMDD date or -vN version tag.""" + return _VERSION_DATE_SUFFIX.sub("", model) + + +# ---------------------------------------------------------------------- +# Price tables +# ---------------------------------------------------------------------- +@dataclass(frozen=True) +class ModelPrice: + """Per-token USD prices for one model. None = no price for that field.""" + + source: str + input: Decimal | None = None + output: Decimal | None = None + cache_read: Decimal | None = None + cache_write: Decimal | None = None + reasoning: Decimal | None = None + + def get(self, field_name: str) -> Decimal | None: + return getattr(self, field_name, None) + + +@dataclass +class CostBreakdown: + """Result of compute_cost — all amounts are money strings ready for an event.""" + + total: str # after-markup total in USD (billable value) + total_cents: str # same total in CENTS — Lago dynamic charge `precise_total_amount_cents` + base: str # pre-markup + markup: str + source: str + fields: dict[str, dict[str, str]] # field -> {tokens, unit_price, cost} + + +def compute_cost(usage: CanonicalUsage, price: ModelPrice, markup: Decimal) -> CostBreakdown: + """Compute ``Σ(unit_price × count) × markup`` for the priced fields present. + + Fields without a unit price are excluded from the sum (recorded nowhere); a + 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): + # • 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: + counts["reasoning"] = 0 + if provider in _INPUT_INCLUDES_CACHE_READ and price.get("cache_read") is not None: + counts["input"] = max(0, counts["input"] - counts["cache_read"]) + + base = Decimal(0) + fields: dict[str, dict[str, str]] = {} + for f in PRICED_FIELDS: + count = counts[f] + if not count: + continue + unit = price.get(f) + if unit is None: + continue + cost = unit * count + base += cost + fields[f] = { + "tokens": str(count), + "unit_price": _fmt_money(unit), + "cost": _fmt_money(cost), + } + # Floor the USD total to 12 dp FIRST, then derive cents from it, so cents == + # billed-USD × 100 exactly (matches the JS integer-division implementation). + total = (base * markup).quantize(_Q, rounding=ROUND_DOWN) + return CostBreakdown( + total=_fmt_money(total), + total_cents=_fmt_money(total * 100), + base=_fmt_money(base), + markup=_fmt_money(markup), + source=price.source, + fields=fields, + ) + + +def coerce_markup(markup: Any) -> tuple[Decimal, bool]: + """Return (markup_decimal, ok). Falls back to 1.0 when invalid/non-positive.""" + d = _parse_price(markup) + if d is None or d <= 0: + return Decimal(1), False + return d, True + + +# ---------------------------------------------------------------------- +# OpenRouter parsing + matching +# ---------------------------------------------------------------------- +def parse_openrouter(data: Any) -> dict[str, Any]: + """Parse the /models response into {'exact': {...}, 'norm': {...}} tables.""" + exact: dict[str, ModelPrice] = {} + norm: dict[tuple[str, str], ModelPrice] = {} + models = data.get("data") if isinstance(data, dict) else None + if not isinstance(models, list): + return {"exact": exact, "norm": norm} + for m in models: + if not isinstance(m, dict): + continue + mid = m.get("id") + pricing = m.get("pricing") + if not isinstance(mid, str) or not isinstance(pricing, dict): + continue + mp = ModelPrice( + source="openrouter", + input=_parse_price(pricing.get(_OPENROUTER_FIELD_MAP["input"])), + output=_parse_price(pricing.get(_OPENROUTER_FIELD_MAP["output"])), + cache_read=_parse_price(pricing.get(_OPENROUTER_FIELD_MAP["cache_read"])), + cache_write=_parse_price(pricing.get(_OPENROUTER_FIELD_MAP["cache_write"])), + reasoning=_parse_price(pricing.get(_OPENROUTER_FIELD_MAP["reasoning"])), + ) + exact[mid] = mp + if "/" in mid: + vendor, _, suffix = mid.partition("/") + norm[(vendor.lower(), _norm(suffix))] = mp + return {"exact": exact, "norm": norm} + + +def lookup_openrouter(table: dict[str, Any], provider: str, model: str) -> ModelPrice | None: + """Match (provider, model) to an OpenRouter price. Conservative: vendor-gated.""" + vendor = _VENDOR_MAP.get((provider or "").lower(), (provider or "").lower()) + exact: dict[str, ModelPrice] = table.get("exact", {}) + norm: dict[tuple[str, str], ModelPrice] = table.get("norm", {}) + # 1. exact id + hit = exact.get(f"{vendor}/{model}") + if hit is not None: + return hit + # 2. normalized suffix (. <-> -) + hit = norm.get((vendor, _norm(model))) + if hit is not None: + return hit + # 3. date/version-stripped, normalized + hit = norm.get((vendor, _norm(_strip_version(model)))) + if hit is not None: + return hit + return None + + +# ---------------------------------------------------------------------- +# Bedrock parsing + matching +# +# The AWS Price List offer schema is large and its attribute keys vary by +# product; this parser is deliberately defensive and is validated end-to-end by +# the env-gated live test. A miss returns None → safe token fallback. +# ---------------------------------------------------------------------- +def parse_bedrock_region(model: str, default_region: str) -> str: + head = model.split(".", 1)[0].lower() if "." in model else "" + return _BEDROCK_REGION_PREFIX.get(head, default_region) + + +def bedrock_model_key(model: str) -> str: + """Reduce a Bedrock model id to the alnum key used to index AWS prices. + + e.g. 'eu.anthropic.claude-sonnet-4-6' -> 'claudesonnet46'; + 'anthropic.claude-haiku-4-5-20251001-v1:0' -> 'claudehaiku45'; + 'mistral.mixtral-8x7b-instruct-v0:1' -> 'mixtral8x7binstruct'. + """ + parts = model.split(".") + if parts and parts[0].lower() in _BEDROCK_REGION_PREFIX: + parts = parts[1:] + if len(parts) > 1: + model_part = ".".join(parts[1:]) # drop vendor + else: + model_part = parts[0] if parts else "" + model_part = re.sub(r":\d+$", "", model_part) # ':0' + model_part = re.sub(r"-v\d+$", "", model_part) # '-v1' + model_part = _strip_version(model_part) + return _alnum(model_part) + + +def _aws_model_keys(name: str) -> list[str]: + """Candidate alnum keys for an AWS model name (with/without vendor prefix).""" + base = _strip_version(_norm(name)) + keys = {_alnum(base)} + words = name.split() + if words and words[0].lower() in _BEDROCK_VENDOR_WORDS: + keys.add(_alnum(_strip_version(_norm(" ".join(words[1:]))))) + return [k for k in keys if k] + + +def _usd_per_token(term: Any) -> Decimal | None: + """Extract a USD-per-token price from a terms.OnDemand[sku] entry.""" + if not isinstance(term, dict): + return None + for offer in term.values(): + dims = offer.get("priceDimensions") if isinstance(offer, dict) else None + if not isinstance(dims, dict): + continue + for dim in dims.values(): + if not isinstance(dim, dict): + continue + ppu = dim.get("pricePerUnit") + usd = ppu.get("USD") if isinstance(ppu, dict) else None + price = _parse_price(usd) + if price is None: + continue + unit = str(dim.get("unit", "")).lower() + # AWS sometimes prices per 1K tokens. + if "1k" in unit or "1000" in unit or "thousand" in unit: + price = (price / Decimal(1000)).quantize(_Q, rounding=ROUND_DOWN) + return price + return None + + +def parse_bedrock_offer(offer: Any, region: str) -> dict[str, ModelPrice]: + """Build {alnum_model_key: ModelPrice(input/output)} from an AWS offer file.""" + if not isinstance(offer, dict): + return {} + products = offer.get("products") + terms = offer.get("terms") + on_demand = terms.get("OnDemand") if isinstance(terms, dict) else None + if not isinstance(products, dict) or not isinstance(on_demand, dict): + return {} + + table: dict[str, dict[str, Decimal]] = {} + for sku, product in products.items(): + if not isinstance(product, dict): + continue + attrs = product.get("attributes") + if not isinstance(attrs, dict): + continue + name = attrs.get("model") or attrs.get("titleModelId") or attrs.get("modelName") + if not isinstance(name, str) or not name: + continue + direction = _bedrock_direction(attrs) + if direction is None: + continue + price = _usd_per_token(on_demand.get(sku)) + if price is None: + continue + for key in _aws_model_keys(name): + table.setdefault(key, {})[direction] = price + + return { + key: ModelPrice(source="aws_bedrock", input=v.get("input"), output=v.get("output")) + for key, v in table.items() + } + + +def _bedrock_direction(attrs: dict[str, Any]) -> str | None: + """Classify a Bedrock product as standard on-demand 'input'/'output' tokens. + + Prefers the explicit ``inferenceType`` ("Input tokens" / "Output tokens"). + Rejects tiered variants ("... priority/flex/batch") so we capture the + standard on-demand price, not a discounted/surge tier. Falls back to a + usagetype scan only when inferenceType is absent. + """ + it = str(attrs.get("inferenceType", "")).strip().lower() + if it == "input tokens": + return "input" + if it == "output tokens": + return "output" + if it: + # Present but a tier variant (priority/flex/batch) or non-token → skip. + return None + # inferenceType absent: fall back to usagetype, excluding batch/non-token. + blob = " ".join(str(attrs.get(k, "")) for k in ("usagetype", "operation", "feature")).lower() + if "batch" in blob or "token" not in blob: + return None + if "input" in blob: + return "input" + if "output" in blob: + return "output" + return None + + +def lookup_bedrock(region_table: dict[str, ModelPrice], model: str) -> ModelPrice | None: + return region_table.get(bedrock_model_key(model)) + + +# ---------------------------------------------------------------------- +# Fetcher (real HTTP; injectable for tests) +# ---------------------------------------------------------------------- +class PricingFetcher(Protocol): + def fetch_openrouter(self) -> dict[str, Any]: ... + def fetch_bedrock(self, region: str) -> dict[str, ModelPrice]: ... + + +class HttpPricingFetcher: + """Default fetcher using ``requests`` (already a core dependency).""" + + def __init__(self, timeout: float = 10.0) -> None: + self._timeout = timeout + + def fetch_openrouter(self) -> dict[str, Any]: + import requests + + resp = requests.get(OPENROUTER_URL, timeout=self._timeout) + resp.raise_for_status() + return parse_openrouter(resp.json()) + + def fetch_bedrock(self, region: str) -> dict[str, ModelPrice]: + import requests + + idx = requests.get(AWS_BEDROCK_REGION_INDEX, timeout=self._timeout) + idx.raise_for_status() + regions = idx.json().get("regions", {}) + entry = regions.get(region) + if not isinstance(entry, dict) or not entry.get("currentVersionUrl"): + return {} + offer = requests.get(AWS_PRICING_HOST + entry["currentVersionUrl"], timeout=self._timeout) + offer.raise_for_status() + return parse_bedrock_offer(offer.json(), region) + + +# ---------------------------------------------------------------------- +# PricingProvider — cache + background refresh + non-blocking lookup +# ---------------------------------------------------------------------- +class PricingProvider: + def __init__( + self, + fetcher: PricingFetcher | None = None, + ttl_seconds: float = 3600.0, + default_region: str = "us-east-1", + on_error: Callable[[Exception, str], None] | None = None, + ) -> None: + self._fetcher: PricingFetcher = fetcher or HttpPricingFetcher() + self._ttl = ttl_seconds + self._default_region = default_region + self._on_error = on_error + self._lock = threading.Lock() + self._pid = os.getpid() + self._openrouter: dict[str, Any] | None = None + self._openrouter_fetched = 0.0 + # Not stale by default: token-mode SDKs never trigger a pricing fetch. + # A price-mode lookup flags the relevant source stale on first use. + self._openrouter_stale = False + self._bedrock: dict[str, dict[str, ModelPrice]] = {} + self._bedrock_fetched: dict[str, float] = {} + self._bedrock_stale: set[str] = set() + self._refreshing: set[str] = set() + + def _heal_fork(self) -> None: + """Self-heal after a fork: a lock copied from the parent may be held by a + thread that doesn't exist in the child. Detect a PID change and replace + the lock + mark tables stale so the child's queue thread refetches. Cheap + PID read on the hot path; avoids os.register_at_fork (whose extra + fork-time work trips macOS's objc fork-safety abort).""" + if os.getpid() != self._pid: + self._lock = threading.Lock() + self._pid = os.getpid() + self._openrouter_stale = self._openrouter is not None or self._openrouter_stale + self._bedrock_stale = set(self._bedrock.keys()) + self._refreshing = set() + + def prime(self) -> None: + """Flag the OpenRouter table for an eager background warm (used when + price mode is the global default) to shrink the cold-start window.""" + with self._lock: + self._openrouter_stale = True + + # ---- non-blocking lookup (customer thread) ---- + def lookup(self, provider: str, model: str, api: str) -> ModelPrice | None: + try: + self._heal_fork() + if (api or "").startswith("bedrock"): + region = parse_bedrock_region(model, self._default_region) + with self._lock: + table = self._bedrock.get(region) + fresh = ( + table is not None + and (time.time() - self._bedrock_fetched.get(region, 0.0)) < self._ttl + ) + if not fresh: + self._bedrock_stale.add(region) + return lookup_bedrock(table, model) if table is not None else None + with self._lock: + table_or = self._openrouter + fresh = table_or is not None and (time.time() - self._openrouter_fetched) < self._ttl + if not fresh: + self._openrouter_stale = True + return lookup_openrouter(table_or, provider, model) if table_or is not None else None + except Exception: # noqa: BLE001 — lookup must never raise + return None + + # ---- background refresh (queue worker thread) ---- + def maybe_refresh(self) -> None: + self._heal_fork() + # Lock-free fast path: when nothing is stale (the common case, and always + # in token mode), do no work at all — not even acquire the lock. This + # keeps the queue's background tick essentially free and avoids extra + # cross-thread lock churn. The reads are racy but harmless: a missed flag + # just defers a refresh by one tick. + if not self._openrouter_stale and not self._bedrock_stale: + return + with self._lock: + do_openrouter = self._openrouter_stale and "openrouter" not in self._refreshing + if do_openrouter: + self._refreshing.add("openrouter") + regions = [r for r in self._bedrock_stale if f"bedrock:{r}" not in self._refreshing] + for r in regions: + self._refreshing.add(f"bedrock:{r}") + + if do_openrouter: + try: + table = self._fetcher.fetch_openrouter() + with self._lock: + self._openrouter = table + self._openrouter_fetched = time.time() + self._openrouter_stale = False + except Exception as exc: # noqa: BLE001 + self._report(exc, "pricing.fetch_openrouter") + finally: + with self._lock: + self._refreshing.discard("openrouter") + + for r in regions: + try: + table = self._fetcher.fetch_bedrock(r) + with self._lock: + self._bedrock[r] = table + self._bedrock_fetched[r] = time.time() + self._bedrock_stale.discard(r) + except Exception as exc: # noqa: BLE001 + self._report(exc, "pricing.fetch_bedrock") + finally: + with self._lock: + self._refreshing.discard(f"bedrock:{r}") + + def _report(self, exc: Exception, where: str) -> None: + if self._on_error: + try: + self._on_error(exc, where) + except Exception: # noqa: BLE001 + pass + logger.warning("lago %s failed: %s", where, exc) diff --git a/src/lago_agent_sdk/queue.py b/src/lago_agent_sdk/queue.py index 8c7f401..ac62274 100644 --- a/src/lago_agent_sdk/queue.py +++ b/src/lago_agent_sdk/queue.py @@ -29,6 +29,7 @@ def __init__( max_buffer_size: int = 10_000, max_retry_seconds: float = 60.0, on_error: Callable[[Exception, str], None] | None = None, + pricing: Any | None = None, ) -> None: self._sender = sender self._flush_interval = flush_interval @@ -36,6 +37,9 @@ def __init__( self._max_buffer_size = max_buffer_size self._max_retry_seconds = max_retry_seconds self._on_error = on_error + # Optional PricingProvider — its (blocking) HTTP refresh runs on this + # background thread so the customer's call is never blocked on pricing. + self._pricing = pricing self._buffer: deque[dict[str, Any]] = deque() self._lock = threading.Lock() @@ -62,6 +66,10 @@ def _after_in_child(self) -> None: self._buffer = deque() # don't replay parent's events from the child self._backoff_seconds = 0.0 self._http_calls = 0 + # Note: the PricingProvider self-heals on fork via a PID check inside + # lookup()/maybe_refresh(); we deliberately do NOT call into it from this + # fork handler (touching it here changes thread timing enough to trip + # macOS's objc fork-safety abort). self._thread = threading.Thread(target=self._run, name="lago-queue", daemon=True) self._thread.start() @@ -115,6 +123,13 @@ def _run(self) -> None: self._wake.wait(timeout=self._flush_interval) self._wake.clear() + # Refresh pricing tables on this background thread (off the hot path). + if self._pricing is not None: + 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: diff --git a/src/lago_agent_sdk/sdk.py b/src/lago_agent_sdk/sdk.py index d9dd8fa..dfe8d06 100644 --- a/src/lago_agent_sdk/sdk.py +++ b/src/lago_agent_sdk/sdk.py @@ -11,8 +11,9 @@ from .canonical import CanonicalUsage from .config import LagoConfig from .detector import detect_client_kind -from .exceptions import UnknownClientError +from .exceptions import PricingUnavailableError, UnknownClientError from .lago_client import LagoClient +from .pricing import PricingProvider, coerce_markup, compute_cost from .queue import EventQueue logger = logging.getLogger("lago_agent_sdk") @@ -47,6 +48,16 @@ def __init__( api_url=self.config.api_url, timeout=self.config.request_timeout_seconds, ) + # Pricing provider (price mode). Default does no network until a + # price-mode lookup flags a source stale; refreshes run on the queue + # thread, never on the customer's call. + self._pricing: PricingProvider = self.config.pricing_provider or PricingProvider( + ttl_seconds=self.config.pricing_ttl_seconds, + default_region=self.config.bedrock_default_region, + on_error=self.config.on_error, + ) + if self.config.pricing_mode == "price": + self._pricing.prime() # eager warm when price mode is the global default self._queue = EventQueue( sender=self._lago_client.send_batch, flush_interval=self.config.flush_interval_seconds, @@ -54,6 +65,7 @@ def __init__( max_buffer_size=self.config.max_buffer_size, max_retry_seconds=self.config.max_retry_seconds, on_error=self.config.on_error, + pricing=self._pricing, ) # ------------------------------------------------------------------ @@ -124,7 +136,16 @@ def emit( usage: CanonicalUsage, subscription: str | None = None, dimensions: dict[str, Any] | None = None, + mode: str | None = None, + markup: float | None = None, ) -> None: + """Emit usage to Lago. + + In ``tokens`` mode (default), pushes one event per nonzero token field. + In ``price`` mode, pushes a single dollar-cost event; if no price is + available it falls back to token events and reports via on_error. + Precedence for mode/markup: per-call arg > config default. + """ try: sub = self._resolve_subscription(subscription) if not sub: @@ -134,37 +155,106 @@ def emit( ) return - nonzero = usage.nonzero_numeric() - if not nonzero: - # Mistral legacy / empty — nothing to bill + effective_mode = mode or self.config.pricing_mode + if effective_mode != "price": + self._emit_token_events(usage, sub, dimensions) return - now = int(time.time()) - for field_name, value in nonzero.items(): - code = self.config.metric_codes.get(field_name) - if not code: - continue - event = { - "transaction_id": str(uuid.uuid4()), - "external_subscription_id": sub, - "code": code, - "timestamp": now, - "properties": { - "value": str(value), - "model": usage.model, - "provider": usage.provider, - "api": usage.api, - **(dimensions or {}), - }, - } - self._queue.push(event) + price = self._pricing.lookup(usage.provider, usage.model, usage.api) + if price is None: + # Don't silently under-bill: fall back to token events + report. + self._report_error(PricingUnavailableError(usage.provider, usage.model, usage.api), "pricing") + self._emit_token_events(usage, sub, dimensions) + return + + markup_value, ok = coerce_markup(markup if markup is not None else self.config.markup) + if not ok: + self._report_error( + ValueError( + f"invalid markup {markup if markup is not None else self.config.markup!r}; using 1.0" + ), + "pricing", + ) + self._emit_cost_event(usage, price, markup_value, sub, dimensions) except Exception as exc: # noqa: BLE001 — never raise from emit - if self.config.on_error: - try: - self.config.on_error(exc, "emit") - except Exception: # noqa: BLE001 - pass - logger.warning("lago emit failed: %s", exc) + self._report_error(exc, "emit") + + def _emit_token_events(self, usage: CanonicalUsage, sub: str, dimensions: dict[str, Any] | None) -> None: + nonzero = usage.nonzero_numeric() + if not nonzero: + # Mistral legacy / empty — nothing to bill + return + now = int(time.time()) + for field_name, value in nonzero.items(): + code = self.config.metric_codes.get(field_name) + if not code: + continue + event = { + "transaction_id": str(uuid.uuid4()), + "external_subscription_id": sub, + "code": code, + "timestamp": now, + "properties": { + "value": str(value), + "model": usage.model, + "provider": usage.provider, + "api": usage.api, + **(dimensions or {}), + }, + } + self._queue.push(event) + + def _emit_cost_event( + self, + usage: CanonicalUsage, + price: Any, + markup: Any, + sub: str, + dimensions: dict[str, Any] | None, + ) -> None: + breakdown = compute_cost(usage, price, markup) + # `unit` = total tokens for the call — the quantity the sum-aggregation + # billable metric sums (the dynamic charge's fee comes from + # precise_total_amount_cents; unit is the displayed usage quantity). + # Sum the *billed* per-field counts from the breakdown, which compute_cost + # has already de-overlapped (e.g. cache_read carved out of input), so + # subset fields aren't double-counted in the displayed total. + unit = sum(int(parts["tokens"]) for parts in breakdown.fields.values()) + properties: dict[str, Any] = { + "unit": str(unit), + "value": breakdown.total, + "base_cost": breakdown.base, + "markup": breakdown.markup, + "model": usage.model, + "provider": usage.provider, + "api": usage.api, + "price_source": breakdown.source, + } + for field_name, parts in breakdown.fields.items(): + properties[f"{field_name}_tokens"] = parts["tokens"] + properties[f"{field_name}_unit_price"] = parts["unit_price"] + properties[f"{field_name}_cost"] = parts["cost"] + properties.update(dimensions or {}) + self._queue.push( + { + "transaction_id": str(uuid.uuid4()), + "external_subscription_id": sub, + "code": self.config.cost_metric_code, + "timestamp": int(time.time()), + # Top-level amount (in cents) for Lago's dynamic charge model — + # the charge sums these into a single fee. + "precise_total_amount_cents": breakdown.total_cents, + "properties": properties, + } + ) + + def _report_error(self, exc: Exception, where: str) -> None: + if self.config.on_error: + try: + self.config.on_error(exc, where) + except Exception: # noqa: BLE001 + pass + logger.warning("lago %s failed: %s", where, exc) def flush(self, timeout: float = 5.0) -> bool: return self._queue.flush(timeout=timeout) diff --git a/src/lago_agent_sdk/wrappers/anthropic.py b/src/lago_agent_sdk/wrappers/anthropic.py index 333e191..ded2252 100644 --- a/src/lago_agent_sdk/wrappers/anthropic.py +++ b/src/lago_agent_sdk/wrappers/anthropic.py @@ -98,15 +98,18 @@ def wrap_anthropic_client( original_stream = getattr(messages, "stream", None) is_async = type(client).__name__.startswith("Async") - def _resolve_opts(lago_opts: dict[str, Any]) -> tuple[str | None, dict[str, Any]]: - sub = lago_opts.get("subscription") or base_sub - dims = {**base_dims, **(lago_opts.get("dimensions") or {})} - return sub, dims - - def _emit_from(payload: Any, model_id: str, sub: str | None, dims: dict[str, Any]) -> None: + def _resolve_opts(lago_opts: dict[str, Any]) -> dict[str, Any]: + return { + "subscription": lago_opts.get("subscription") or base_sub, + "dimensions": {**base_dims, **(lago_opts.get("dimensions") or {})}, + "mode": lago_opts.get("mode"), + "markup": lago_opts.get("markup"), + } + + def _emit_from(payload: Any, model_id: str, opts: dict[str, Any]) -> None: try: usage = extract_anthropic_native(payload, model_id=model_id) - sdk.emit(usage, subscription=sub, dimensions=dims) + sdk.emit(usage, **opts) except Exception as exc: # noqa: BLE001 logger.warning("lago: anthropic emit failed: %s", exc) @@ -117,11 +120,11 @@ def _create(*args: Any, **kwargs: Any) -> Any: assert original_create is not None lago_opts = _pop_lago_kwarg(kwargs) model_id = kwargs.get("model", "") - sub, dims = _resolve_opts(lago_opts) + opts = _resolve_opts(lago_opts) response = original_create(*args, **kwargs) if _is_message_like(response): - _emit_from(response, model_id, sub, dims) + _emit_from(response, model_id, opts) return response # Streaming — wrap the iterator, merging usage across message_start @@ -135,7 +138,7 @@ def _wrap_stream(src: Iterator[Any]) -> Iterator[Any]: yield event finally: if accumulated: - _emit_from({"usage": accumulated}, model_id, sub, dims) + _emit_from({"usage": accumulated}, model_id, opts) return _wrap_stream(response) @@ -146,11 +149,11 @@ async def _create_async(*args: Any, **kwargs: Any) -> Any: assert original_create is not None lago_opts = _pop_lago_kwarg(kwargs) model_id = kwargs.get("model", "") - sub, dims = _resolve_opts(lago_opts) + opts = _resolve_opts(lago_opts) response = await original_create(*args, **kwargs) if _is_message_like(response): - _emit_from(response, model_id, sub, dims) + _emit_from(response, model_id, opts) return response async def _wrap_async_stream(src: AsyncIterator[Any]) -> AsyncIterator[Any]: @@ -162,7 +165,7 @@ async def _wrap_async_stream(src: AsyncIterator[Any]) -> AsyncIterator[Any]: yield event finally: if accumulated: - _emit_from({"usage": accumulated}, model_id, sub, dims) + _emit_from({"usage": accumulated}, model_id, opts) return _wrap_async_stream(response) @@ -177,9 +180,9 @@ def _wrap_stream_manager(*args: Any, **kwargs: Any) -> Any: assert original_stream is not None lago_opts = _pop_lago_kwarg(kwargs) model_id = kwargs.get("model", "") - sub, dims = _resolve_opts(lago_opts) + opts = _resolve_opts(lago_opts) inner = original_stream(*args, **kwargs) - return _LagoStreamManager(inner, sdk, model_id, sub, dims, is_async=is_async) + return _LagoStreamManager(inner, sdk, model_id, opts, is_async=is_async) if original_create is not None: messages.create = _create_async if is_async else _create @@ -202,16 +205,14 @@ def __init__( inner: Any, sdk: Any, model_id: str, - sub: str | None, - dims: dict[str, Any], + opts: dict[str, Any], *, is_async: bool, ) -> None: self._inner = inner self._sdk = sdk self._model_id = model_id - self._sub = sub - self._dims = dims + self._opts = opts self._stream: Any = None self._is_async = is_async @@ -249,7 +250,7 @@ def _emit_final(self) -> None: from ..adapters import extract_anthropic_native usage = extract_anthropic_native(final, model_id=self._model_id) - self._sdk.emit(usage, subscription=self._sub, dimensions=self._dims) + self._sdk.emit(usage, **self._opts) except Exception as exc: # noqa: BLE001 logger.warning("lago: anthropic stream-manager emit failed: %s", exc) @@ -268,6 +269,6 @@ async def _emit_final_async(self) -> None: from ..adapters import extract_anthropic_native usage = extract_anthropic_native(final, model_id=self._model_id) - self._sdk.emit(usage, subscription=self._sub, dimensions=self._dims) + self._sdk.emit(usage, **self._opts) except Exception as exc: # noqa: BLE001 logger.warning("lago: anthropic async stream-manager emit failed: %s", exc) diff --git a/src/lago_agent_sdk/wrappers/boto3_bedrock.py b/src/lago_agent_sdk/wrappers/boto3_bedrock.py index d465942..2e07f10 100644 --- a/src/lago_agent_sdk/wrappers/boto3_bedrock.py +++ b/src/lago_agent_sdk/wrappers/boto3_bedrock.py @@ -73,6 +73,8 @@ def _converse(*args: Any, **kwargs: Any) -> Any: usage, subscription=lago_opts.get("subscription") or base_sub, dimensions={**base_dims, **(lago_opts.get("dimensions") or {})}, + mode=lago_opts.get("mode"), + markup=lago_opts.get("markup"), ) except Exception as exc: # noqa: BLE001 — never break the call logger.warning("lago: converse instrumentation failed: %s", exc) @@ -108,6 +110,8 @@ def _wrap_stream() -> Iterator[Any]: usage, subscription=lago_opts.get("subscription") or base_sub, dimensions={**base_dims, **(lago_opts.get("dimensions") or {})}, + mode=lago_opts.get("mode"), + markup=lago_opts.get("markup"), ) except Exception as exc: # noqa: BLE001 logger.warning("lago: converse_stream instrumentation failed: %s", exc) @@ -135,6 +139,8 @@ def _invoke_model(*args: Any, **kwargs: Any) -> Any: usage, subscription=lago_opts.get("subscription") or base_sub, dimensions={**base_dims, **(lago_opts.get("dimensions") or {})}, + mode=lago_opts.get("mode"), + markup=lago_opts.get("markup"), ) except Exception as exc: # noqa: BLE001 logger.warning("lago: invoke_model parse/emit failed: %s", exc) @@ -203,6 +209,8 @@ def _wrap_invoke_stream() -> Iterator[Any]: usage, subscription=lago_opts.get("subscription") or base_sub, dimensions={**base_dims, **(lago_opts.get("dimensions") or {})}, + mode=lago_opts.get("mode"), + markup=lago_opts.get("markup"), ) except Exception as exc: # noqa: BLE001 logger.warning("lago: invoke_model_with_response_stream instrumentation failed: %s", exc) diff --git a/src/lago_agent_sdk/wrappers/gemini.py b/src/lago_agent_sdk/wrappers/gemini.py index f53ec51..ccf4144 100644 --- a/src/lago_agent_sdk/wrappers/gemini.py +++ b/src/lago_agent_sdk/wrappers/gemini.py @@ -45,15 +45,18 @@ def wrap_gemini_client( base_dims = dict(dimensions or {}) base_sub = subscription - def _resolve_opts(lago_opts: dict[str, Any]) -> tuple[str | None, dict[str, Any]]: - sub = lago_opts.get("subscription") or base_sub - dims = {**base_dims, **(lago_opts.get("dimensions") or {})} - return sub, dims - - def _emit_from(payload: Any, model_id: str, sub: str | None, dims: dict[str, Any]) -> None: + def _resolve_opts(lago_opts: dict[str, Any]) -> dict[str, Any]: + return { + "subscription": lago_opts.get("subscription") or base_sub, + "dimensions": {**base_dims, **(lago_opts.get("dimensions") or {})}, + "mode": lago_opts.get("mode"), + "markup": lago_opts.get("markup"), + } + + def _emit_from(payload: Any, model_id: str, opts: dict[str, Any]) -> None: try: usage = extract_gemini_native(payload, model_id=model_id) - sdk.emit(usage, subscription=sub, dimensions=dims) + sdk.emit(usage, **opts) except Exception as exc: # noqa: BLE001 logger.warning("lago: gemini emit failed: %s", exc) @@ -61,9 +64,9 @@ def _make_sync_generate(original: Any) -> Any: def _generate(*args: Any, **kwargs: Any) -> Any: lago_opts = _pop_lago_kwarg(kwargs) model_id = kwargs.get("model") or (args[0] if args else "") - sub, dims = _resolve_opts(lago_opts) + opts = _resolve_opts(lago_opts) response = original(*args, **kwargs) - _emit_from(response, str(model_id), sub, dims) + _emit_from(response, str(model_id), opts) return response return _generate @@ -72,9 +75,9 @@ def _make_async_generate(original: Any) -> Any: async def _generate_async(*args: Any, **kwargs: Any) -> Any: lago_opts = _pop_lago_kwarg(kwargs) model_id = kwargs.get("model") or (args[0] if args else "") - sub, dims = _resolve_opts(lago_opts) + opts = _resolve_opts(lago_opts) response = await original(*args, **kwargs) - _emit_from(response, str(model_id), sub, dims) + _emit_from(response, str(model_id), opts) return response return _generate_async @@ -83,7 +86,7 @@ def _make_sync_stream(original: Any) -> Any: def _stream(*args: Any, **kwargs: Any) -> Iterator[Any]: lago_opts = _pop_lago_kwarg(kwargs) model_id = kwargs.get("model") or (args[0] if args else "") - sub, dims = _resolve_opts(lago_opts) + opts = _resolve_opts(lago_opts) src = original(*args, **kwargs) def _iter() -> Iterator[Any]: @@ -96,7 +99,7 @@ def _iter() -> Iterator[Any]: yield chunk finally: if last_with_usage is not None: - _emit_from(last_with_usage, str(model_id), sub, dims) + _emit_from(last_with_usage, str(model_id), opts) return _iter() @@ -106,7 +109,7 @@ def _make_async_stream(original: Any) -> Any: async def _stream_async(*args: Any, **kwargs: Any) -> AsyncIterator[Any]: lago_opts = _pop_lago_kwarg(kwargs) model_id = kwargs.get("model") or (args[0] if args else "") - sub, dims = _resolve_opts(lago_opts) + opts = _resolve_opts(lago_opts) src = await original(*args, **kwargs) async def _aiter() -> AsyncIterator[Any]: @@ -119,7 +122,7 @@ async def _aiter() -> AsyncIterator[Any]: yield chunk finally: if last_with_usage is not None: - _emit_from(last_with_usage, str(model_id), sub, dims) + _emit_from(last_with_usage, str(model_id), opts) return _aiter() diff --git a/src/lago_agent_sdk/wrappers/mistral.py b/src/lago_agent_sdk/wrappers/mistral.py index 8ce5d2a..729d6b1 100644 --- a/src/lago_agent_sdk/wrappers/mistral.py +++ b/src/lago_agent_sdk/wrappers/mistral.py @@ -58,10 +58,13 @@ def wrap_mistral_client( original_complete_async = getattr(chat, "complete_async", None) original_stream_async = getattr(chat, "stream_async", None) - def _resolve_opts(lago_opts: dict[str, Any]) -> tuple[str | None, dict[str, Any]]: - sub = lago_opts.get("subscription") or base_sub - dims = {**base_dims, **(lago_opts.get("dimensions") or {})} - return sub, dims + def _resolve_opts(lago_opts: dict[str, Any]) -> dict[str, Any]: + return { + "subscription": lago_opts.get("subscription") or base_sub, + "dimensions": {**base_dims, **(lago_opts.get("dimensions") or {})}, + "mode": lago_opts.get("mode"), + "markup": lago_opts.get("markup"), + } # ------------------------------------------------------------------ # chat.complete — non-streaming @@ -73,8 +76,8 @@ def _complete(*args: Any, **kwargs: Any) -> Any: response = original_complete(*args, **kwargs) try: usage = extract_mistral_native(_to_dict(response), model_id=model_id) - sub, dims = _resolve_opts(lago_opts) - sdk.emit(usage, subscription=sub, dimensions=dims) + opts = _resolve_opts(lago_opts) + sdk.emit(usage, **opts) except Exception as exc: # noqa: BLE001 — never break the call logger.warning("lago: mistral.chat.complete instrumentation failed: %s", exc) return response @@ -101,8 +104,8 @@ def _wrap_iter() -> Iterator[Any]: if last_usage is not None: try: usage = extract_mistral_native(last_usage, model_id=model_id) - sub, dims = _resolve_opts(lago_opts) - sdk.emit(usage, subscription=sub, dimensions=dims) + opts = _resolve_opts(lago_opts) + sdk.emit(usage, **opts) except Exception as exc: # noqa: BLE001 logger.warning("lago: mistral.chat.stream instrumentation failed: %s", exc) @@ -118,8 +121,8 @@ async def _complete_async(*args: Any, **kwargs: Any) -> Any: response = await original_complete_async(*args, **kwargs) try: usage = extract_mistral_native(_to_dict(response), model_id=model_id) - sub, dims = _resolve_opts(lago_opts) - sdk.emit(usage, subscription=sub, dimensions=dims) + opts = _resolve_opts(lago_opts) + sdk.emit(usage, **opts) except Exception as exc: # noqa: BLE001 logger.warning("lago: mistral.chat.complete_async instrumentation failed: %s", exc) return response @@ -147,8 +150,8 @@ async def _agen() -> AsyncIterator[Any]: if last_usage is not None: try: usage = extract_mistral_native(last_usage, model_id=model_id) - sub, dims = _resolve_opts(lago_opts) - sdk.emit(usage, subscription=sub, dimensions=dims) + opts = _resolve_opts(lago_opts) + sdk.emit(usage, **opts) except Exception as exc: # noqa: BLE001 logger.warning("lago: mistral.chat.stream_async instrumentation failed: %s", exc) diff --git a/src/lago_agent_sdk/wrappers/openai.py b/src/lago_agent_sdk/wrappers/openai.py index 4ae271d..90ccb2f 100644 --- a/src/lago_agent_sdk/wrappers/openai.py +++ b/src/lago_agent_sdk/wrappers/openai.py @@ -83,15 +83,18 @@ def wrap_openai_client( base_sub = subscription is_async = type(client).__name__.startswith("Async") - def _resolve_opts(lago_opts: dict[str, Any]) -> tuple[str | None, dict[str, Any]]: - sub = lago_opts.get("subscription") or base_sub - dims = {**base_dims, **(lago_opts.get("dimensions") or {})} - return sub, dims - - def _emit_from(payload: Any, model_id: str, sub: str | None, dims: dict[str, Any]) -> None: + def _resolve_opts(lago_opts: dict[str, Any]) -> dict[str, Any]: + return { + "subscription": lago_opts.get("subscription") or base_sub, + "dimensions": {**base_dims, **(lago_opts.get("dimensions") or {})}, + "mode": lago_opts.get("mode"), + "markup": lago_opts.get("markup"), + } + + def _emit_from(payload: Any, model_id: str, opts: dict[str, Any]) -> None: try: usage = extract_openai_native(payload, model_id=model_id) - sdk.emit(usage, subscription=sub, dimensions=dims) + sdk.emit(usage, **opts) except Exception as exc: # noqa: BLE001 logger.warning("lago: openai emit failed: %s", exc) @@ -125,11 +128,11 @@ def _create(*args: Any, **kwargs: Any) -> Any: if not is_responses_api: _ensure_stream_options_include_usage(kwargs) model_id = kwargs.get("model", "") - sub, dims = _resolve_opts(lago_opts) + opts = _resolve_opts(lago_opts) response = original(*args, **kwargs) if _is_response_like(response): - _emit_from(response, model_id, sub, dims) + _emit_from(response, model_id, opts) return response # Streaming — wrap the iterator to capture the final usage on close. @@ -144,7 +147,7 @@ def _wrap_stream(src: Iterator[Any]) -> Iterator[Any]: yield event finally: if last_usage is not None: - _emit_from(last_usage, model_id, sub, dims) + _emit_from(last_usage, model_id, opts) return _wrap_stream(response) @@ -156,11 +159,11 @@ async def _create_async(*args: Any, **kwargs: Any) -> Any: if not is_responses_api: _ensure_stream_options_include_usage(kwargs) model_id = kwargs.get("model", "") - sub, dims = _resolve_opts(lago_opts) + opts = _resolve_opts(lago_opts) response = await original(*args, **kwargs) if _is_response_like(response): - _emit_from(response, model_id, sub, dims) + _emit_from(response, model_id, opts) return response async def _wrap_async_stream(src: AsyncIterator[Any]) -> AsyncIterator[Any]: @@ -174,7 +177,7 @@ async def _wrap_async_stream(src: AsyncIterator[Any]) -> AsyncIterator[Any]: yield event finally: if last_usage is not None: - _emit_from(last_usage, model_id, sub, dims) + _emit_from(last_usage, model_id, opts) return _wrap_async_stream(response) diff --git a/tests/integration/test_live_pricing.py b/tests/integration/test_live_pricing.py new file mode 100644 index 0000000..c03b104 --- /dev/null +++ b/tests/integration/test_live_pricing.py @@ -0,0 +1,58 @@ +"""Live pricing test — hits the real OpenRouter + AWS Bedrock bulk APIs. + +Skipped unless LAGO_LIVE_PRICING=1 (it makes real network calls, no keys needed +since both sources are public). Validates that the real fetchers build tables +and that known models resolve to sane USD-per-token prices — in particular it +exercises the AWS Bedrock offer-file parser against the live schema. +""" + +from __future__ import annotations + +import os +from decimal import Decimal + +import pytest + +from lago_agent_sdk.pricing import HttpPricingFetcher, lookup_bedrock, lookup_openrouter + +pytestmark = pytest.mark.skipif( + os.environ.get("LAGO_LIVE_PRICING") != "1", + reason="LAGO_LIVE_PRICING != 1 (live network test)", +) + + +def test_openrouter_live_table_and_known_models() -> None: + table = HttpPricingFetcher(timeout=30).fetch_openrouter() + exact = table["exact"] + assert len(exact) > 50, "expected a substantial OpenRouter model list" + + # A few well-known models should resolve with a positive input price. + resolved = 0 + for provider, model in [ + ("openai", "gpt-4o"), + ("anthropic", "claude-3.5-sonnet"), + ("google", "gemini-2.5-flash"), + ]: + mp = lookup_openrouter(table, provider, model) + if mp is not None and mp.input is not None and mp.input >= Decimal(0): + resolved += 1 + assert resolved >= 1, "expected at least one well-known OpenRouter model to resolve" + + +def test_bedrock_live_table_builds_and_resolves() -> None: + region = "us-east-1" + table = HttpPricingFetcher(timeout=30).fetch_bedrock(region) + # The parser should extract at least some priced models from the live offer. + assert table, "AWS Bedrock offer parsed to an empty table — schema may have changed" + priced = [mp for mp in table.values() if mp.input is not None or mp.output is not None] + assert priced, "no Bedrock models had input/output token prices" + + # A common Bedrock model should resolve (best-effort; logs the key on miss). + for model in [ + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "anthropic.claude-3-haiku-20240307-v1:0", + ]: + mp = lookup_bedrock(table, model) + if mp is not None and (mp.input or mp.output): + return + pytest.skip(f"no probed Bedrock model matched; {len(table)} keys built — refine matcher if needed") diff --git a/tests/unit/fixtures/pricing/money_golden.json b/tests/unit/fixtures/pricing/money_golden.json new file mode 100644 index 0000000..5ac4612 --- /dev/null +++ b/tests/unit/fixtures/pricing/money_golden.json @@ -0,0 +1,59 @@ +{ + "_comment": "Cross-repo golden money cases. Python (Decimal) and JS (scaled BigInt) must produce identical base/total/total_cents strings. Money is floored to 12 decimal places (ROUND_DOWN). Prices are USD per token (strings); markup is a string multiplier. total_cents = total USD x 100 (Lago dynamic charge precise_total_amount_cents).", + "cases": [ + { + "name": "input+output, no markup", + "prices": { "input": "0.000003", "output": "0.000015" }, + "counts": { "input": 1000, "output": 500 }, + "markup": "1", + "base": "0.0105", + "total": "0.0105", + "total_cents": "1.05" + }, + { + "name": "input+output, 1.2x markup", + "prices": { "input": "0.000003", "output": "0.000015" }, + "counts": { "input": 1000, "output": 500 }, + "markup": "1.2", + "base": "0.0105", + "total": "0.0126", + "total_cents": "1.26" + }, + { + "name": "cache_read + reasoning, 1.5x markup", + "prices": { "cache_read": "0.0000005", "reasoning": "0.00001" }, + "counts": { "cache_read": 2000, "reasoning": 100 }, + "markup": "1.5", + "base": "0.002", + "total": "0.003", + "total_cents": "0.3" + }, + { + "name": "free model -> zero", + "prices": { "input": "0", "output": "0" }, + "counts": { "input": 1000, "output": 200 }, + "markup": "1", + "base": "0", + "total": "0", + "total_cents": "0" + }, + { + "name": "12dp exact precision", + "prices": { "input": "0.000000333333" }, + "counts": { "input": 3 }, + "markup": "1", + "base": "0.000000999999", + "total": "0.000000999999", + "total_cents": "0.0000999999" + }, + { + "name": "floor-12dp truncation on markup", + "prices": { "input": "0.000001" }, + "counts": { "input": 1000 }, + "markup": "1.333333333333", + "base": "0.001", + "total": "0.001333333333", + "total_cents": "0.1333333333" + } + ] +} diff --git a/tests/unit/test_pricing.py b/tests/unit/test_pricing.py new file mode 100644 index 0000000..8b83926 --- /dev/null +++ b/tests/unit/test_pricing.py @@ -0,0 +1,525 @@ +"""Pricing tests — matching, money math, provider cache, and SDK price mode.""" + +from __future__ import annotations + +import json +import pathlib +from decimal import Decimal +from typing import Any + +import pytest + +from lago_agent_sdk import CanonicalUsage, LagoConfig, LagoSDK, ModelPrice +from lago_agent_sdk.pricing import ( + PricingProvider, + bedrock_model_key, + coerce_markup, + compute_cost, + lookup_bedrock, + lookup_openrouter, + parse_bedrock_offer, + parse_bedrock_region, + parse_openrouter, +) + +FIXTURES = pathlib.Path(__file__).parent / "fixtures" / "pricing" + + +# ---------------------------------------------------------------------- +# Stub fetcher (no network) — mirrors the queue's injectable sender pattern +# ---------------------------------------------------------------------- +class StubFetcher: + def __init__(self, openrouter: dict | None = None, bedrock: dict | None = None) -> None: + self._openrouter = openrouter or {"exact": {}, "norm": {}} + self._bedrock = bedrock or {} + self.openrouter_calls = 0 + self.bedrock_calls: list[str] = [] + + def fetch_openrouter(self) -> dict[str, Any]: + self.openrouter_calls += 1 + return self._openrouter + + def fetch_bedrock(self, region: str) -> dict[str, ModelPrice]: + self.bedrock_calls.append(region) + return self._bedrock.get(region, {}) + + +_OPENROUTER_RAW = { + "data": [ + { + "id": "anthropic/claude-opus-4.8", + "pricing": { + "prompt": "0.000005", + "completion": "0.000025", + "input_cache_read": "0.0000005", + "input_cache_write": "0.00000625", + "internal_reasoning": "0.000025", + }, + }, + { + "id": "openai/gpt-4o", + "pricing": { + "prompt": "0.0000025", + "completion": "0.00001", + "input_cache_read": "0.00000125", + "internal_reasoning": "0.00001", + }, + }, + {"id": "mistralai/mistral-large", "pricing": {"prompt": "0.000002", "completion": "0.000006"}}, + { + "id": "google/gemini-2.5-flash", + "pricing": { + "prompt": "0.0000003", + "completion": "0.0000025", + "input_cache_read": "0.000000075", + "internal_reasoning": "0.0000025", + }, + }, + ] +} + + +# ---------------------------------------------------------------------- +# OpenRouter parsing + matching +# ---------------------------------------------------------------------- +def test_openrouter_exact_and_normalized_match() -> None: + table = parse_openrouter(_OPENROUTER_RAW) + # normalized: our "claude-opus-4-8" matches OpenRouter "claude-opus-4.8" + mp = lookup_openrouter(table, "anthropic", "claude-opus-4-8") + assert mp is not None + assert mp.input == Decimal("0.000005") + assert mp.output == Decimal("0.000025") + assert mp.cache_read == Decimal("0.0000005") + assert mp.reasoning == Decimal("0.000025") + assert mp.source == "openrouter" + + +def test_openrouter_vendor_map_mistral_and_gemini() -> None: + table = parse_openrouter(_OPENROUTER_RAW) + # provider "mistral" -> vendor "mistralai" + assert lookup_openrouter(table, "mistral", "mistral-large") is not None + # provider "gemini" -> vendor "google" + assert lookup_openrouter(table, "gemini", "gemini-2.5-flash") is not None + + +def test_openrouter_date_version_stripped_match() -> None: + table = parse_openrouter( + {"data": [{"id": "anthropic/claude-haiku-4.5", "pricing": {"prompt": "0.000001"}}]} + ) + # our id carries a date suffix; matcher strips it + mp = lookup_openrouter(table, "anthropic", "claude-haiku-4-5-20251001") + assert mp is not None + assert mp.input == Decimal("0.000001") + + +def test_openrouter_miss_returns_none() -> None: + table = parse_openrouter(_OPENROUTER_RAW) + assert lookup_openrouter(table, "anthropic", "totally-made-up-model") is None + # vendor-gated: right model name, wrong vendor -> miss + assert lookup_openrouter(table, "openai", "claude-opus-4-8") is None + + +# ---------------------------------------------------------------------- +# Bedrock region + key + offer parsing +# ---------------------------------------------------------------------- +@pytest.mark.parametrize( + "model,expected", + [ + ("eu.anthropic.claude-sonnet-4-6", "eu-west-1"), + ("us.anthropic.claude-sonnet-4-6", "us-east-1"), + ("apac.anthropic.claude-sonnet-4-6", "ap-southeast-1"), + ("anthropic.claude-haiku-4-5-20251001-v1:0", "us-east-1"), # no prefix -> default + ], +) +def test_bedrock_region_detection(model: str, expected: str) -> None: + assert parse_bedrock_region(model, "us-east-1") == expected + + +@pytest.mark.parametrize( + "model,expected_key", + [ + ("eu.anthropic.claude-sonnet-4-6", "claudesonnet46"), + ("anthropic.claude-haiku-4-5-20251001-v1:0", "claudehaiku45"), + ("mistral.mixtral-8x7b-instruct-v0:1", "mixtral8x7binstruct"), + ], +) +def test_bedrock_model_key(model: str, expected_key: str) -> None: + assert bedrock_model_key(model) == expected_key + + +def _aws_product(model: str, inference_type: str, usd: str, unit: str = "1K tokens") -> tuple[dict, dict]: + """Build one (product, term) pair matching the real AWS Bedrock offer schema.""" + sku = f"{model}:{inference_type}".replace(" ", "") + product = { + sku: { + "productFamily": "...", + "attributes": { + "model": model, + "usagetype": f"USE1-{model.replace(' ', '')}-{inference_type.replace(' ', '-')}", + "inferenceType": inference_type, + "feature": "On-demand Inference", + "provider": "Anthropic", + }, + } + } + term = {sku: {"off": {"priceDimensions": {"d": {"pricePerUnit": {"USD": usd}, "unit": unit}}}}} + return product, term + + +def test_bedrock_offer_parse_and_lookup() -> None: + # Real AWS schema: inferenceType distinguishes direction; unit is "1K tokens". + p_in, t_in = _aws_product("Claude Sonnet 4.6", "Input tokens", "0.003") # $3/M + p_out, t_out = _aws_product("Claude Sonnet 4.6", "Output tokens", "0.015") # $15/M + offer = {"products": {**p_in, **p_out}, "terms": {"OnDemand": {**t_in, **t_out}}} + table = parse_bedrock_offer(offer, "us-east-1") + mp = lookup_bedrock(table, "us.anthropic.claude-sonnet-4-6") + assert mp is not None + assert mp.input == Decimal("0.000003") # 0.003 per 1K -> 3e-6 per token + assert mp.output == Decimal("0.000015") + assert mp.source == "aws_bedrock" + + +def test_bedrock_offer_rejects_tier_variants() -> None: + # Standard on-demand tier must win over priority/flex/batch variants. + p_std, t_std = _aws_product("Claude Sonnet 4.6", "Input tokens", "0.003") + p_pri, t_pri = _aws_product("Claude Sonnet 4.6", "Input tokens priority", "0.006") + p_flex, t_flex = _aws_product("Claude Sonnet 4.6", "Input tokens flex", "0.0015") + offer = { + "products": {**p_std, **p_pri, **p_flex}, + "terms": {"OnDemand": {**t_std, **t_pri, **t_flex}}, + } + table = parse_bedrock_offer(offer, "us-east-1") + mp = lookup_bedrock(table, "anthropic.claude-sonnet-4-6") + assert mp is not None + assert mp.input == Decimal("0.000003") # the standard tier, not priority/flex + + +def test_bedrock_usagetype_fallback_when_no_inference_type() -> None: + # Resilience: if AWS ever drops inferenceType, fall back to usagetype scan. + offer = { + "products": { + "S": {"attributes": {"model": "Mixtral 8x7B Instruct", "usagetype": "USE1-Input-Tokens"}} + }, + "terms": { + "OnDemand": { + "S": { + "o": {"priceDimensions": {"d": {"pricePerUnit": {"USD": "0.0005"}, "unit": "1K tokens"}}} + } + } + }, + } + table = parse_bedrock_offer(offer, "us-east-1") + mp = lookup_bedrock(table, "mistral.mixtral-8x7b-instruct-v0:1") + assert mp is not None + assert mp.input == Decimal("0.0000005") # 0.0005 per 1K -> 5e-7 per token + + +# ---------------------------------------------------------------------- +# compute_cost + golden money parity +# ---------------------------------------------------------------------- +def test_compute_cost_excludes_unpriced_fields() -> None: + price = ModelPrice(source="openrouter", input=Decimal("0.000003"), output=Decimal("0.000015")) + # tool_calls is not a priced field; image_input has no unit price + usage = CanonicalUsage( + input=1000, output=500, tool_calls=3, image_input=50, model="m", provider="p", api="native" + ) + b = compute_cost(usage, price, Decimal("1")) + assert set(b.fields) == {"input", "output"} + assert b.base == "0.0105" + assert b.total == "0.0105" + + +def test_compute_cost_only_unpriced_fields_yields_zero() -> None: + # model priced but the call's only count is an unpriced field + price = ModelPrice(source="openrouter", input=Decimal("0.000003")) + usage = CanonicalUsage(tool_calls=5, model="m", provider="p", api="native") + b = compute_cost(usage, price, Decimal("1")) + assert b.total == "0" + assert b.fields == {} + + +def test_money_golden_cases() -> None: + cases = json.loads((FIXTURES / "money_golden.json").read_text())["cases"] + for c in cases: + prices = {k: Decimal(v) for k, v in c["prices"].items()} + price = ModelPrice(source="openrouter", **prices) + usage = CanonicalUsage(model="m", provider="p", api="native", **c["counts"]) + b = compute_cost(usage, price, Decimal(c["markup"])) + assert b.base == c["base"], f"{c['name']}: base {b.base} != {c['base']}" + assert b.total == c["total"], f"{c['name']}: total {b.total} != {c['total']}" + assert b.total_cents == c["total_cents"], f"{c['name']}: cents {b.total_cents} != {c['total_cents']}" + + +def test_coerce_markup() -> None: + assert coerce_markup(1.2) == (Decimal("1.2"), True) + assert coerce_markup("2") == (Decimal("2"), True) + assert coerce_markup(0) == (Decimal("1"), False) + assert coerce_markup(-1) == (Decimal("1"), False) + assert coerce_markup("nonsense") == (Decimal("1"), False) + + +# ---------------------------------------------------------------------- +# PricingProvider — cache + refresh + non-blocking lookup +# ---------------------------------------------------------------------- +def test_provider_cold_lookup_flags_stale_then_refresh_warms() -> None: + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + # cold: no table yet -> None, and source flagged for refresh + assert p.lookup("anthropic", "claude-opus-4-8", "native") is None + assert fetcher.openrouter_calls == 0 + # background worker would call this; we call it directly + p.maybe_refresh() + assert fetcher.openrouter_calls == 1 + # now warm + mp = p.lookup("anthropic", "claude-opus-4-8", "native") + assert mp is not None and mp.input == Decimal("0.000005") + + +def test_provider_token_mode_does_no_fetch() -> None: + fetcher = StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600) + # No lookups performed -> nothing flagged stale -> refresh is a no-op. + p.maybe_refresh() + assert fetcher.openrouter_calls == 0 + + +def test_provider_bedrock_region_routing() -> None: + bedrock_table = parse_bedrock_offer( + { + "products": { + "S": { + "attributes": { + "model": "Claude Sonnet 4.6", + "usagetype": "Input-Tokens", + "unit": "tokens", + } + } + }, + "terms": { + "OnDemand": { + "S": { + "o": { + "priceDimensions": {"d": {"pricePerUnit": {"USD": "0.000003"}, "unit": "tokens"}} + } + } + } + }, + }, + "eu-west-1", + ) + fetcher = StubFetcher(bedrock={"eu-west-1": bedrock_table}) + p = PricingProvider(fetcher=fetcher, ttl_seconds=3600, default_region="us-east-1") + assert p.lookup("anthropic", "eu.anthropic.claude-sonnet-4-6", "bedrock_converse") is None + p.maybe_refresh() + assert fetcher.bedrock_calls == ["eu-west-1"] + mp = p.lookup("anthropic", "eu.anthropic.claude-sonnet-4-6", "bedrock_converse") + assert mp is not None and mp.input == Decimal("0.000003") + + +# ---------------------------------------------------------------------- +# SDK price mode integration +# ---------------------------------------------------------------------- +def _warm_provider() -> PricingProvider: + p = PricingProvider(fetcher=StubFetcher(openrouter=parse_openrouter(_OPENROUTER_RAW)), ttl_seconds=3600) + p.prime() + p.maybe_refresh() # warm synchronously for deterministic tests + return p + + +def _price_sdk( + provider: PricingProvider, default_sub: str = "sub_default", on_error=None, markup: float = 1.0 +): + received: list = [] + cfg = LagoConfig( + api_key="dummy", + default_subscription_id=default_sub, + pricing_mode="price", + markup=markup, + pricing_provider=provider, + on_error=on_error, + ) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: received.append(list(b)) # type: ignore[attr-defined] + return sdk, received + + +def test_price_mode_emits_single_cost_event() -> None: + sdk, received = _price_sdk(_warm_provider()) + u = CanonicalUsage(input=1000, output=500, model="claude-opus-4-8", provider="anthropic", api="native") + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert len(flat) == 1 + ev = flat[0] + assert ev["code"] == "llm_cost" + # Lago dynamic charge: top-level cents amount = 0.0175 USD * 100 = 1.75 + assert ev["precise_total_amount_cents"] == "1.75" + props = ev["properties"] + # `unit` = total tokens (1000 + 500) — the sum-aggregation quantity + assert props["unit"] == "1500" + # 1000*0.000005 + 500*0.000025 = 0.005 + 0.0125 = 0.0175 + assert props["value"] == "0.0175" + assert props["base_cost"] == "0.0175" + assert props["price_source"] == "openrouter" + assert props["input_tokens"] == "1000" + assert props["input_unit_price"] == "0.000005" + assert props["output_cost"] == "0.0125" + + +def test_price_mode_markup_scales_value() -> None: + sdk, received = _price_sdk(_warm_provider(), markup=2.0) + u = CanonicalUsage(input=1000, output=500, model="claude-opus-4-8", provider="anthropic", api="native") + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + ev = [e for batch in received for e in batch][0] + assert ev["properties"]["base_cost"] == "0.0175" + assert ev["properties"]["value"] == "0.035" # 0.0175 * 2 + assert ev["properties"]["markup"] == "2" + + +def test_per_call_markup_overrides_global() -> None: + sdk, received = _price_sdk(_warm_provider(), markup=1.0) + u = CanonicalUsage(input=1000, output=500, model="claude-opus-4-8", provider="anthropic", api="native") + sdk.emit(u, markup=3.0) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + ev = [e for batch in received for e in batch][0] + assert ev["properties"]["value"] == "0.0525" # 0.0175 * 3 + + +# ---------------------------------------------------------------------- +# Subset semantics: some providers report `input` INCLUSIVE of cache_read and +# `output` INCLUSIVE of reasoning. Pricing the parent at full count AND the +# subset separately would double-bill — these tests lock the de-overlap. +# ---------------------------------------------------------------------- +def test_price_mode_openai_cache_read_subset_not_double_billed() -> None: + sdk, received = _price_sdk(_warm_provider()) + # OpenAI: input (prompt_tokens)=1000 ALREADY includes cache_read=800. + u = CanonicalUsage( + input=1000, output=500, cache_read=800, model="gpt-4o", provider="openai", api="native" + ) + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + props = [e for batch in received for e in batch][0]["properties"] + # input billed for only the non-cached portion (1000 - 800); cache billed at cache rate + assert props["input_tokens"] == "200" + assert props["cache_read_tokens"] == "800" + # 200*0.0000025 + 800*0.00000125 + 500*0.00001 = 0.0005 + 0.001 + 0.005 = 0.0065 + # (the bug would bill input at full 1000 -> 0.0085) + assert props["value"] == "0.0065" + # unit = billed tokens 200 + 800 + 500 = 1500 = prompt(1000) + completion(500) + assert props["unit"] == "1500" + + +def test_price_mode_gemini_cache_subset_and_reasoning_additive() -> None: + sdk, received = _price_sdk(_warm_provider()) + # Gemini: input=1000 INCLUDES cache_read=300; reasoning(thoughts)=100 is ADDITIVE. + u = CanonicalUsage( + input=1000, + output=400, + cache_read=300, + reasoning=100, + model="gemini-2.5-flash", + provider="gemini", + api="native", + ) + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + props = [e for batch in received for e in batch][0]["properties"] + assert props["input_tokens"] == "700" # 1000 - 300 cached + assert props["cache_read_tokens"] == "300" + assert props["output_tokens"] == "400" + assert props["reasoning_tokens"] == "100" # billed separately (additive for Gemini) + # 700*3e-7 + 300*7.5e-8 + 400*2.5e-6 + 100*2.5e-6 = 0.00021+0.0000225+0.001+0.00025 = 0.0014825 + assert props["value"] == "0.0014825" + # unit = 700+300+400+100 = 1500 = prompt(1000)+candidates(400)+thoughts(100) + assert props["unit"] == "1500" + + +def test_price_mode_openai_reasoning_in_output_not_double_billed() -> None: + sdk, received = _price_sdk(_warm_provider()) + # OpenAI o-series: output (completion_tokens)=500 ALREADY includes reasoning=200. + u = CanonicalUsage(input=100, output=500, reasoning=200, model="gpt-4o", provider="openai", api="native") + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + props = [e for batch in received for e in batch][0]["properties"] + # reasoning folded into output — no separate reasoning line, output billed in full + assert "reasoning_tokens" not in props + assert props["output_tokens"] == "500" + # 100*0.0000025 + 500*0.00001 = 0.00025 + 0.005 = 0.00525 (bug would add 200*1e-5=0.002) + assert props["value"] == "0.00525" + assert props["unit"] == "600" # 100 + 500; reasoning not double-counted + + +def test_price_mode_anthropic_cache_is_additive() -> None: + sdk, received = _price_sdk(_warm_provider()) + # Anthropic: input EXCLUDES cache; cache_read/cache_write are additive (no subtraction). + u = CanonicalUsage( + input=1000, + output=500, + cache_read=400, + cache_write=200, + model="claude-opus-4-8", + provider="anthropic", + api="native", + ) + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + props = [e for batch in received for e in batch][0]["properties"] + assert props["input_tokens"] == "1000" # unchanged — additive provider + assert props["cache_read_tokens"] == "400" + assert props["cache_write_tokens"] == "200" + # 1000*5e-6 + 500*25e-6 + 400*5e-7 + 200*6.25e-6 = 0.005+0.0125+0.0002+0.00125 = 0.01895 + assert props["value"] == "0.01895" + assert props["unit"] == "2100" # 1000+500+400+200, all additive + + +def test_price_unavailable_falls_back_to_token_events_and_reports() -> None: + errors: list = [] + # warm provider but ask for an unknown model -> price None -> fallback + sdk, received = _price_sdk( + _warm_provider(), on_error=lambda exc, where: errors.append((type(exc).__name__, where)) + ) + u = CanonicalUsage(input=10, output=20, model="unknown-model-xyz", provider="anthropic", api="native") + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + codes = {e["code"] for e in flat} + assert codes == {"llm_input_tokens", "llm_output_tokens"} # token fallback + assert any(name == "PricingUnavailableError" and where == "pricing" for name, where in errors) + + +def test_per_call_price_mode_overrides_global_tokens() -> None: + # global mode is tokens (default); per-call asks for price + provider = _warm_provider() + received: list = [] + cfg = LagoConfig(api_key="dummy", default_subscription_id="sub_default", pricing_provider=provider) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: received.append(list(b)) # type: ignore[attr-defined] + u = CanonicalUsage(input=1000, output=500, model="claude-opus-4-8", provider="anthropic", api="native") + sdk.emit(u, mode="price") + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert len(flat) == 1 and flat[0]["code"] == "llm_cost" + + +def test_default_mode_is_tokens_unchanged() -> None: + provider = _warm_provider() + received: list = [] + cfg = LagoConfig(api_key="dummy", default_subscription_id="sub_default", pricing_provider=provider) + sdk = LagoSDK(api_key="dummy", config=cfg) + sdk._queue._sender = lambda b: received.append(list(b)) # type: ignore[attr-defined] + u = CanonicalUsage(input=1000, output=500, model="claude-opus-4-8", provider="anthropic", api="native") + sdk.emit(u) + assert sdk.flush(timeout=2.0) + sdk.shutdown(timeout=1.0) + flat = [e for batch in received for e in batch] + assert {e["code"] for e in flat} == {"llm_input_tokens", "llm_output_tokens"}