From 09105585e09079214b343b6342abe4b4a9db711c Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 13 Aug 2026 17:18:07 +0000 Subject: [PATCH] Show estimated per-model cost in `ucode usage` Co-authored-by: Isaac --- src/ucode/databricks.py | 41 ++++ src/ucode/ui.py | 11 + src/ucode/usage.py | 514 ++++++++++++++++++++++++++++------------ tests/test_usage.py | 354 +++++++++++++++++++-------- 4 files changed, 663 insertions(+), 257 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 705b5519..c26c169e 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1651,6 +1651,47 @@ def fetch_model_recommendation(workspace: str, token: str) -> tuple[dict, str | return payload, None +# The gateway's per-model price catalog (USD per million tokens), sourced from the same Zippy data +# the server uses to bill external-model spend. We read it to estimate per-model cost from token +# counts, since no API returns per-model dollars directly. +_EXTERNAL_PROVIDER_MODELS_API_PATH = "/api/ai-gateway/v2/external-provider-models" +_EXTERNAL_PROVIDER_MODELS_PAGE_SIZE = 1000 +_EXTERNAL_PROVIDER_MODELS_MAX_PAGES = 20 + + +def fetch_external_model_prices(workspace: str, token: str) -> tuple[list[dict], str | None]: + """List external-provider models and their `base_pricing` (USD per million tokens) via the gateway. + + Returns ``(models, reason)`` with each model the raw API entry; ``reason`` is non-None on failure + (callers omit cost rather than fail). + """ + hostname = workspace_hostname(workspace) + base_url = f"https://{hostname}{_EXTERNAL_PROVIDER_MODELS_API_PATH}" + models: list[dict] = [] + page_token: str | None = None + seen_tokens: set[str] = set() + for _ in range(_EXTERNAL_PROVIDER_MODELS_MAX_PAGES): + params: dict[str, str] = {"page_size": str(_EXTERNAL_PROVIDER_MODELS_PAGE_SIZE)} + if page_token: + params["page_token"] = page_token + payload, reason = _http_get_json(f"{base_url}?{urlencode(params)}", token, timeout=30) + if payload is None: + # Return what we have if a later page blips; only the first-page failure is fatal. + return (models, None) if models else ([], reason or "unknown error") + if not isinstance(payload, dict): + return [], "external-provider-models returned an unexpected response shape" + for entry in payload.get("models") or []: + if isinstance(entry, dict) and entry.get("model_name"): + models.append(entry) + page_token = payload.get("next_page_token") or None + if not page_token or page_token in seen_tokens: + break + seen_tokens.add(page_token) + if not models: + return [], "external-provider-models listing returned no models" + return models, None + + # Every field ucode's manifest can set, as `update_mask` paths for a PATCH. The server rejects a # missing or empty mask, and rejects paths outside its own mutable set — this is that set minus the # fields ucode doesn't author: `budget_id` (deprecated in favour of `budget_policy.budget_id`, and diff --git a/src/ucode/ui.py b/src/ucode/ui.py index 3747d6bc..d660b7b7 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -254,6 +254,17 @@ def format_usd(amount: Decimal) -> str: return f"${amount.quantize(Decimal('0.01'), rounding=ROUND_HALF_UP):,}" +def format_cost_usd(amount: Decimal) -> str: + """Like `format_usd`, but keeps more precision for sub-cent per-model costs. + + A model that cost a fraction of a cent would round to ``$0.00`` at two + decimals and read as free, so amounts under a cent show four decimals. + """ + if Decimal(0) < amount < Decimal("0.01"): + return f"${amount.quantize(Decimal('0.0001'), rounding=ROUND_HALF_UP)}" + return format_usd(amount) + + def format_meter(fraction: float, width: int = 30) -> str: """Text meter for `fraction` of a whole, clamped to [0, 1].""" clamped = min(max(fraction, 0.0), 1.0) diff --git a/src/ucode/usage.py b/src/ucode/usage.py index 8d690a06..014b7991 100644 --- a/src/ucode/usage.py +++ b/src/ucode/usage.py @@ -6,16 +6,18 @@ from __future__ import annotations import json +import re from collections.abc import Mapping from datetime import date, datetime, timedelta -from decimal import Decimal -from typing import cast +from decimal import Decimal, InvalidOperation +from typing import NamedTuple, cast from ucode.databricks import ( SqlWarehouse, apply_pat_environment, discover_sql_warehouses, ensure_databricks_auth, + fetch_external_model_prices, get_databricks_token, resolve_current_budget_spend, run_usage_query, @@ -23,7 +25,7 @@ from ucode.state import load_state from ucode.ui import ( console, - format_duration, + format_cost_usd, format_meter, format_token_count, format_usd, @@ -43,9 +45,139 @@ QUERY_MESSAGE = "Querying system.ai_gateway.usage..." STARTUP_MESSAGE = "Starting up warehouse..." +PRICES_MESSAGE = "Fetching model prices..." # `REQUESTED` is an explicit --warehouse-id, whose state we never looked up. WARM_WAREHOUSE_STATES = ("RUNNING", "REQUESTED") +MILLION = Decimal(1_000_000) + +# Region/provider prefixes the price catalog puts on some model ids (e.g. `eu/gpt-5.6-sol`, +# `us.anthropic.claude-opus-4-8`). They're stripped before matching so a regional id and its base id +# collide onto the same price key. +_PRICE_STRIP_PREFIXES = ( + "eu/", + "us/", + "global/", + "apac.", + "au.", + "us.", + "eu.", + "anthropic.", + "databricks-", + "ap-northeast-1/", + "ap-northeast-2/", + "ap-southeast-1/", + "ap-southeast-2/", + "ap-south-1/", + "ca-central-1/", + "sa-east-1/", +) + + +class ModelPrice(NamedTuple): + """USD rates per million tokens. Any field may be None when the catalog omits it.""" + + input: Decimal | None + output: Decimal | None + cache_read: Decimal | None + + +def _price_decimal(raw: object) -> Decimal | None: + if not isinstance(raw, (int, float, str)) or not str(raw).strip(): + return None + try: + return Decimal(str(raw).strip()) + except InvalidOperation: + return None + + +def normalize_price_key(model_name: str) -> str: + """Separator-insensitive key that collapses a model's spellings for price matching. + + Usage `gpt-5-6-sol` / catalog `gpt-5.6-sol` / `us.anthropic.claude-opus-4-8` all reduce to one + key by stripping region/provider prefixes and version/date suffixes, then all non-alphanumerics. + """ + name = (model_name or "").strip().lower() + # Loop so stacked prefixes collapse too, e.g. `us.` then `anthropic.`. + changed = True + while changed: + changed = False + for prefix in _PRICE_STRIP_PREFIXES: + if name.startswith(prefix): + name = name[len(prefix) :] + changed = True + break + name = name.split("@", 1)[0] + name = re.sub(r"-v\d+(:\d+)?$", "", name) # Bedrock version suffix, e.g. `-v1:0` + name = re.sub(r"-20\d{6}$", "", name) # date stamp, e.g. `-20251001` + return re.sub(r"[^a-z0-9]+", "", name) + + +def _is_bare_model_name(model_name: str) -> bool: + """True when a catalog id carries no region/provider prefix (its base price).""" + lowered = model_name.strip().lower() + return not any(lowered.startswith(prefix) for prefix in _PRICE_STRIP_PREFIXES) + + +def build_price_lookup(raw_models: list[dict]) -> dict[str, ModelPrice]: + """Map normalized price keys to `ModelPrice` from the raw catalog listing. + + When several catalog ids collapse to one key (a regional id and its base), the priced, un-prefixed + entry wins so we bill at the base rate rather than a regional markup. + """ + lookup: dict[str, ModelPrice] = {} + scores: dict[str, tuple[bool, bool]] = {} + for entry in raw_models: + if not isinstance(entry, dict): + continue + name = entry.get("model_name") + if not isinstance(name, str) or not name.strip(): + continue + key = normalize_price_key(name) + if not key: + continue + pricing = entry.get("base_pricing") + pricing = pricing if isinstance(pricing, Mapping) else {} + price = ModelPrice( + input=_price_decimal(pricing.get("input_per_million_tokens")), + output=_price_decimal(pricing.get("output_per_million_tokens")), + cache_read=_price_decimal(pricing.get("cache_read_per_million_tokens")), + ) + score = (price.input is not None, _is_bare_model_name(name)) + if key not in scores or score > scores[key]: + scores[key] = score + lookup[key] = price + return lookup + + +def estimate_model_cost( + price: ModelPrice | None, + input_tokens: int, + cached_tokens: int, + output_tokens: int, +) -> Decimal | None: + """Estimate USD cost for one model's token usage, or None when it can't be priced. + + Cached tokens are a subset of input tokens billed at the (cheaper) cache-read rate, so the + uncached remainder is priced at the input rate. Returns None when the catalog has no input rate + or when no token breakdown is available (input/output/cached all zero), so the caller shows tokens + without a dollar figure rather than a misleading $0. + """ + if price is None or price.input is None: + return None + if input_tokens <= 0 and output_tokens <= 0 and cached_tokens <= 0: + return None + cached = min(max(cached_tokens, 0), max(input_tokens, 0)) + uncached = max(input_tokens, 0) - cached + cache_rate = price.cache_read if price.cache_read is not None else price.input + output_rate = price.output if price.output is not None else Decimal(0) + total = ( + Decimal(uncached) * price.input + + Decimal(cached) * cache_rate + + Decimal(max(output_tokens, 0)) * output_rate + ) + return total / MILLION + def build_usage_report_query() -> str: return f""" @@ -63,7 +195,10 @@ def build_usage_report_query() -> str: request_id, event_time, destination_model, - COALESCE(total_tokens, 0) AS total_tokens_used + COALESCE(total_tokens, 0) AS total_tokens_used, + COALESCE(input_tokens, 0) AS input_tokens_used, + COALESCE(token_details.cache_read_input_tokens, 0) AS cached_tokens_used, + COALESCE(output_tokens, 0) AS output_tokens_used FROM system.ai_gateway.usage WHERE event_time >= current_timestamp() - interval {USAGE_SUMMARY_DAYS} days AND requester = current_user() @@ -80,9 +215,7 @@ def build_usage_report_query() -> str: tool, usage_day, SUM(total_tokens_used) AS total_tokens_used, - COUNT(DISTINCT request_id) AS sessions, - MIN(event_time) AS first_event_time, - MAX(event_time) AS last_event_time + COUNT(DISTINCT request_id) AS sessions FROM usage_events GROUP BY 1, 2, 3 ), @@ -92,9 +225,12 @@ def build_usage_report_query() -> str: tool, usage_day, destination_model, - SUM(total_tokens_used) AS model_tokens_used + COUNT(DISTINCT request_id) AS model_requests, + SUM(total_tokens_used) AS model_tokens_used, + SUM(input_tokens_used) AS model_input_used, + SUM(cached_tokens_used) AS model_cached_used, + SUM(output_tokens_used) AS model_output_used FROM usage_events - WHERE destination_model IS NOT NULL AND destination_model != '' GROUP BY 1, 2, 3, 4 ), model_rollup AS ( @@ -102,11 +238,17 @@ def build_usage_report_query() -> str: requester_name, tool, usage_day, - CONCAT_WS(', ', SORT_ARRAY(COLLECT_SET(destination_model))) AS models, TO_JSON( SORT_ARRAY( COLLECT_LIST( - NAMED_STRUCT('model', destination_model, 'tokens', model_tokens_used) + NAMED_STRUCT( + 'model', destination_model, + 'requests', model_requests, + 'tokens', model_tokens_used, + 'input', model_input_used, + 'cached', model_cached_used, + 'output', model_output_used + ) ) ) ) AS model_tokens @@ -119,9 +261,6 @@ def build_usage_report_query() -> str: daily_usage.usage_day, daily_usage.total_tokens_used, daily_usage.sessions, - daily_usage.first_event_time, - daily_usage.last_event_time, - COALESCE(model_rollup.models, '') AS models, COALESCE(model_rollup.model_tokens, '[]') AS model_tokens FROM daily_usage LEFT JOIN model_rollup @@ -168,18 +307,6 @@ def coerce_date(value_obj: object) -> date | None: return None -def coerce_datetime(value_obj: object) -> datetime | None: - if isinstance(value_obj, datetime): - return value_obj - if isinstance(value_obj, str): - candidate = value_obj.replace("Z", "+00:00") - try: - return datetime.fromisoformat(candidate) - except ValueError: - return None - return None - - def simplify_model_name(tool: str, model_name: str) -> str: normalized = (model_name or "").strip() if not normalized: @@ -219,91 +346,136 @@ def summarize_models(tool: str, raw_models: object) -> str: return ", ".join(parts) if parts else "-" -def _coerce_model_token_item(tool: str, item: object) -> tuple[str, int] | None: - if not isinstance(item, Mapping): - return None - item_mapping = cast(Mapping[str, object], item) +class ModelUsage(NamedTuple): + """One model's usage for a day/window. - raw_model = item_mapping.get("model") - if not isinstance(raw_model, str) or not raw_model.strip(): - return None + `key` is a separator-insensitive identity used to merge spellings of the same model (id form vs. + display form vs. regional variant); `name` is the display name derived from `raw_names` (the raw + `destination_model` ids). `input` includes cached tokens (`cached` is its cache-read subset). + """ - raw_tokens = item_mapping.get("tokens") - try: - token_total = int(cast(int | float | str, raw_tokens or 0)) - except (TypeError, ValueError): - token_total = 0 + name: str + key: str + requests: int + total: int + input: int + cached: int + output: int + raw_names: tuple[str, ...] - model_name = simplify_model_name(tool, raw_model) - if model_name == "-": - return None - return model_name, token_total +UNKNOWN_MODEL = "" +_MODEL_DATABRICKS_PREFIX = "databricks-" -def extract_model_token_breakdown( - tool: str, - raw_model_tokens: object, - raw_models: object = None, - total_tokens: int = 0, -) -> list[tuple[str, int]]: - items: object - if isinstance(raw_model_tokens, str) and raw_model_tokens.strip(): - try: - items = json.loads(raw_model_tokens) - except json.JSONDecodeError: - items = [] - else: - items = raw_model_tokens - model_tokens: dict[str, int] = {} - if isinstance(items, list): - for item in items: - coerced = _coerce_model_token_item(tool, item) - if not coerced: - continue - model_name, token_total = coerced - model_tokens[model_name] = model_tokens.get(model_name, 0) + token_total +def model_identity_key(name: str) -> str: + """Identity used to merge spellings of the same model, falling back to the lowercased name.""" + return normalize_price_key(name) or name.strip().lower() or UNKNOWN_MODEL - if model_tokens: - return sorted(model_tokens.items(), key=lambda item: (-item[1], item[0].lower())) - models = extract_model_names(tool, raw_models) - if len(models) == 1 and total_tokens: - return [(models[0], total_tokens)] - return [(model_name, 0) for model_name in models] +def canonical_model_name(raw_names: tuple[str, ...]) -> str: + """Display name for a merged model: prefer the lowercase, space-free id spelling over a display + label (e.g. `claude-opus-4-8` over `Claude Opus 4.8`), keeping the family prefix and dropping + only `databricks-`.""" + cleaned: list[str] = [] + for raw in raw_names: + name = raw.strip() + if name.lower().startswith(_MODEL_DATABRICKS_PREFIX): + name = name[len(_MODEL_DATABRICKS_PREFIX) :] + if name: + cleaned.append(name) + if not cleaned: + return UNKNOWN_MODEL + id_forms = [name for name in cleaned if " " not in name and name == name.lower()] + return sorted(id_forms or cleaned, key=lambda name: (len(name), name))[0] -def summarize_model_tokens( - tool: str, - raw_model_tokens: object, - raw_models: object, - total_tokens: int, -) -> str: - model_tokens = extract_model_token_breakdown( - tool, - raw_model_tokens, - raw_models, - total_tokens, +def _coerce_int(raw: object) -> int: + try: + return int(cast(int | float | str, raw or 0)) + except (TypeError, ValueError): + return 0 + + +def _coerce_model_usage_item(item: object) -> ModelUsage | None: + if not isinstance(item, Mapping): + return None + item_mapping = cast(Mapping[str, object], item) + + raw_model = item_mapping.get("model") + if isinstance(raw_model, str) and raw_model.strip(): + raw_names: tuple[str, ...] = (raw_model.strip(),) + key = model_identity_key(raw_model) + name = canonical_model_name(raw_names) + else: + # Requests the gateway logged without a destination model still count toward the total. + raw_names = () + key = UNKNOWN_MODEL + name = UNKNOWN_MODEL + return ModelUsage( + name=name, + key=key, + requests=_coerce_int(item_mapping.get("requests")), + total=_coerce_int(item_mapping.get("tokens")), + input=_coerce_int(item_mapping.get("input")), + cached=_coerce_int(item_mapping.get("cached")), + output=_coerce_int(item_mapping.get("output")), + raw_names=raw_names, ) - if not model_tokens: - return "-" - return ", ".join( - f"{model_name} ({format_token_count(token_total)})" if token_total else model_name - for model_name, token_total in model_tokens + + +def extract_model_usage(raw_model_tokens: object) -> list[ModelUsage]: + """Per-model usage from a row's `model_tokens` JSON, merged by identity, highest tokens first.""" + try: + items = ( + json.loads(raw_model_tokens) if isinstance(raw_model_tokens, str) else raw_model_tokens + ) + except json.JSONDecodeError: + items = [] + + merged: dict[str, ModelUsage] = {} + if isinstance(items, list): + for item in items: + coerced = _coerce_model_usage_item(item) + if coerced: + _merge_model_usage(merged, coerced) + return sorted(merged.values(), key=lambda u: (-u.total, u.name.lower())) + + +def _merge_model_usage(merged: dict[str, ModelUsage], usage: ModelUsage) -> None: + """Accumulate `usage` into `merged` by identity key; the display name is recomputed from the + unioned raw names so a model's id and display spellings fold into one row across days.""" + existing = merged.get(usage.key) + if existing is None: + merged[usage.key] = usage + return + raw_names = tuple(dict.fromkeys(existing.raw_names + usage.raw_names)) + merged[usage.key] = existing._replace( + name=canonical_model_name(raw_names) if raw_names else existing.name, + requests=existing.requests + usage.requests, + total=existing.total + usage.total, + input=existing.input + usage.input, + cached=existing.cached + usage.cached, + output=existing.output + usage.output, + raw_names=raw_names, ) -def empty_tool_day(tool: str, usage_day: date) -> dict[str, object]: - return { - "tool": tool, - "usage_day": usage_day, - "total_tokens_used": 0, - "sessions": 0, - "first_event_time": None, - "last_event_time": None, - "models": "-", - "model_tokens": "[]", - } +def model_usage_cost( + usage: ModelUsage, + price_lookup: dict[str, ModelPrice] | None, +) -> Decimal | None: + """Estimated USD cost for one model's usage, or None when it can't be priced.""" + if not price_lookup: + return None + for raw_name in usage.raw_names: + price = price_lookup.get(normalize_price_key(raw_name)) + if price is None: + continue + cost = estimate_model_cost(price, usage.input, usage.cached, usage.output) + if cost is not None: + return cost + return None def has_tool_usage_last_week(records: list[dict[str, object]], tool: str) -> bool: @@ -322,44 +494,71 @@ def has_tool_usage_last_week(records: list[dict[str, object]], tool: str) -> boo return False -def build_tool_breakdown_rows(records: list[dict[str, object]], tool: str) -> list[list[str]]: +class ToolUsageTotals(NamedTuple): + """Tool-level rollup for the week: total requests, tokens, and cost (None when unpriceable).""" + + requests: int + tokens: int + cost: Decimal | None + + +TOOL_MODEL_TABLE_HEADERS = ["Model", "Requests", "Input (incl. cache)", "Output", "Cost (USD)"] + + +def aggregate_tool_model_usage(records: list[dict[str, object]], tool: str) -> list[ModelUsage]: + """Per-model usage for `tool` over the last 7 days, merged across days, highest tokens first.""" today = date.today() - rows_by_day: dict[date, dict[str, object]] = {} + week_start = today - timedelta(days=USAGE_BREAKDOWN_DAYS - 1) + merged: dict[str, ModelUsage] = {} for record in records: if record.get("tool") != tool: continue usage_day = coerce_date(record.get("usage_day")) - if usage_day: - rows_by_day[usage_day] = record - - rendered_rows: list[list[str]] = [] - for day_offset in range(USAGE_BREAKDOWN_DAYS): - usage_day = today - timedelta(days=day_offset) - record = rows_by_day.get(usage_day) or empty_tool_day(tool, usage_day) - first_event_time = coerce_datetime(record.get("first_event_time")) - last_event_time = coerce_datetime(record.get("last_event_time")) - duration = None - if first_event_time and last_event_time: - duration = last_event_time - first_event_time - token_total = int(cast(int, record.get("total_tokens_used") or 0)) - session_total = int(cast(int, record.get("sessions") or 0)) - rendered_rows.append( + if not usage_day or usage_day < week_start: + continue + for model_usage in extract_model_usage(record.get("model_tokens")): + _merge_model_usage(merged, model_usage) + return sorted(merged.values(), key=lambda u: (-u.total, u.name.lower())) + + +def build_tool_model_rows( + records: list[dict[str, object]], + tool: str, + price_lookup: dict[str, ModelPrice] | None = None, +) -> tuple[list[list[str]], ToolUsageTotals]: + """Per-model table rows + tool totals for the week. + + Columns match `TOOL_MODEL_TABLE_HEADERS`: model, requests, input (incl. cache), output, cost. + Cost is a dash for models the price catalog doesn't cover; the totals' cost is None when nothing + in the tool could be priced. + """ + rows: list[list[str]] = [] + total_requests = 0 + total_tokens = 0 + total_cost = Decimal(0) + any_priced = False + for usage in aggregate_tool_model_usage(records, tool): + cost = model_usage_cost(usage, price_lookup) + rows.append( [ - usage_day.strftime("%m-%d"), - usage_day.strftime("%a"), - format_token_count(token_total) if token_total else "-", - str(session_total) if session_total else "-", - format_duration(duration), - summarize_model_tokens( - tool, - record.get("model_tokens"), - record.get("models"), - token_total, - ), + usage.name, + f"{usage.requests:,}", + format_token_count(usage.input), + format_token_count(usage.output), + format_cost_usd(cost) if cost is not None else "-", ] ) - - return rendered_rows + total_requests += usage.requests + total_tokens += usage.input + usage.output + if cost is not None: + total_cost += cost + any_priced = True + totals = ToolUsageTotals( + requests=total_requests, + tokens=total_tokens, + cost=total_cost if any_priced else None, + ) + return rows, totals def find_requester_name( @@ -403,6 +602,7 @@ def render_usage_summary( requester_name: str, tool_displays: dict[str, str], budget_spend: tuple[Decimal, Decimal] | None = None, + price_lookup: dict[str, ModelPrice] | None = None, ) -> str: today = date.today() week_start = today - timedelta(days=USAGE_BREAKDOWN_DAYS - 1) @@ -412,7 +612,7 @@ def render_usage_summary( weekly_total = 0 monthly_total = 0 active_tools_last_week: list[str] = [] - weekly_model_tokens: dict[str, int] = {} + weekly_model_usage: dict[str, ModelUsage] = {} for record in records: usage_day = coerce_date(record.get("usage_day")) if not usage_day: @@ -430,15 +630,8 @@ def render_usage_summary( ): active_tools_last_week.append(tool) if isinstance(tool, str): - for model_name, model_token_total in extract_model_token_breakdown( - tool, - record.get("model_tokens"), - record.get("models"), - token_total, - ): - weekly_model_tokens[model_name] = ( - weekly_model_tokens.get(model_name, 0) + model_token_total - ) + for model_usage in extract_model_usage(record.get("model_tokens")): + _merge_model_usage(weekly_model_usage, model_usage) if usage_day == today: daily_total += token_total @@ -453,16 +646,22 @@ def render_usage_summary( if active_tools_last_week: tool_text = ", ".join(tool_displays[tool] for tool in active_tools_last_week) lines.append(f"{label('Active tools:')} {value(tool_text)}") - if weekly_model_tokens: + if weekly_model_usage: top_models = sorted( - weekly_model_tokens.items(), - key=lambda item: (-item[1], item[0].lower()), + weekly_model_usage.values(), + key=lambda u: (-u.total, u.name.lower()), )[:3] - models_text = ", ".join( - f"{model_name} ({format_token_count(token_total)})" - for model_name, token_total in top_models - ) + models_text = ", ".join(usage.name for usage in top_models) lines.append(f"{label('Top models this week:')} {value(models_text)}") + weekly_cost = sum( + ( + model_usage_cost(usage, price_lookup) or Decimal(0) + for usage in weekly_model_usage.values() + ), + Decimal(0), + ) + if weekly_cost > 0: + lines.append(f"{label('Est. cost (7 days):')} {value(format_cost_usd(weekly_cost))}") lines.extend(render_budget_lines(budget_spend)) return "\n".join(lines) @@ -541,6 +740,12 @@ def usage(warehouse_id: str | None = None) -> int: with spinner("Checking budget spend..."): budget_spend, _ = resolve_current_budget_spend(workspace, token) + # Per-model dollar cost is estimated from tokens × catalog prices; omit cost rather than fail + # when the price catalog is unreachable. + with spinner(PRICES_MESSAGE): + raw_prices, _ = fetch_external_model_prices(workspace, token) + price_lookup = build_price_lookup(raw_prices) + tool_displays = {tool: spec["display"] for tool, spec in TOOL_SPECS.items()} configured_tools = configured_usage_tools(state, tool_displays) configured_tool_displays = {tool: tool_displays[tool] for tool in configured_tools} @@ -552,11 +757,14 @@ def usage(warehouse_id: str | None = None) -> int: requester_name, configured_tool_displays, budget_spend=budget_spend, + price_lookup=price_lookup, ) ) - table_headers = ["Date", "Day", "Tokens", "Sessions", "Duration", "Models"] - table_widths = [8, 5, 10, 8, 8, 24] + table_widths = [24, 10, 20, 10, 12] + today = date.today() + week_start = today - timedelta(days=USAGE_BREAKDOWN_DAYS - 1) + date_range = f"{week_start:%b %d}–{today:%b %d, %Y}" if not configured_tools: print_note("No coding agents configured. Run `ucode configure` to set up agents.") @@ -568,11 +776,11 @@ def usage(warehouse_id: str | None = None) -> int: if not has_tool_usage_last_week(records, tool): print_note(f"No usage for {display} in the last {USAGE_BREAKDOWN_DAYS} days.") continue - console.print( - render_box_table( - table_headers, - build_tool_breakdown_rows(records, tool), - max_widths=table_widths, - ) - ) + rows, totals = build_tool_model_rows(records, tool, price_lookup) + console.print(muted(date_range)) + console.print(f"{label('Requests:')} {value(f'{totals.requests:,}')}") + console.print(f"{label('Total tokens:')} {value(f'{totals.tokens:,}')}") + if totals.cost is not None: + console.print(f"{label('Cost (USD):')} {value(format_cost_usd(totals.cost))}") + console.print(render_box_table(TOOL_MODEL_TABLE_HEADERS, rows, max_widths=table_widths)) return 0 diff --git a/tests/test_usage.py b/tests/test_usage.py index d6ff2a32..f0e97199 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -14,23 +14,26 @@ from ucode.usage import ( USAGE_BREAKDOWN_DAYS, USAGE_SUMMARY_DAYS, + ModelPrice, + aggregate_tool_model_usage, build_current_user_query, - build_tool_breakdown_rows, + build_price_lookup, + build_tool_model_rows, build_usage_report_query, coerce_date, - coerce_datetime, configured_usage_tools, - empty_tool_day, + estimate_model_cost, extract_model_names, - extract_model_token_breakdown, + extract_model_usage, filter_records_for_tools, has_tool_usage_last_week, + model_usage_cost, + normalize_price_key, parse_usage_rows, render_budget_lines, render_usage_summary, run_query_on_first_working_warehouse, simplify_model_name, - summarize_model_tokens, summarize_models, usage, ) @@ -54,7 +57,23 @@ def test_includes_per_model_token_rollup(self): q = build_usage_report_query() assert "model_tokens" in q assert "SUM(total_tokens_used) AS model_tokens_used" in q - assert "NAMED_STRUCT('model', destination_model, 'tokens', model_tokens_used)" in q + assert "'model', destination_model" in q + assert "'tokens', model_tokens_used" in q + + def test_includes_per_model_cost_token_breakdown(self): + q = build_usage_report_query() + # input/cached/output tokens are needed to price a model's usage at distinct rates. + assert "COALESCE(input_tokens, 0) AS input_tokens_used" in q + assert "COALESCE(token_details.cache_read_input_tokens, 0) AS cached_tokens_used" in q + assert "COALESCE(output_tokens, 0) AS output_tokens_used" in q + assert "'input', model_input_used" in q + assert "'cached', model_cached_used" in q + assert "'output', model_output_used" in q + + def test_includes_per_model_request_count(self): + q = build_usage_report_query() + assert "COUNT(DISTINCT request_id) AS model_requests" in q + assert "'requests', model_requests" in q class TestBuildCurrentUserQuery: @@ -169,27 +188,6 @@ def test_none_returns_none(self): assert coerce_date(None) is None -class TestCoerceDatetime: - def test_datetime_passthrough(self): - dt = datetime(2024, 6, 1, 0, 0, 0) - assert coerce_datetime(dt) == dt - - def test_iso_string(self): - result = coerce_datetime("2024-06-01T12:00:00") - assert isinstance(result, datetime) - assert result.date() == date(2024, 6, 1) - - def test_z_suffix(self): - result = coerce_datetime("2024-06-01T12:00:00Z") - assert isinstance(result, datetime) - - def test_invalid_string_returns_none(self): - assert coerce_datetime("bad") is None - - def test_none_returns_none(self): - assert coerce_datetime(None) is None - - class TestSimplifyModelName: def test_strips_databricks_and_tool_prefix(self): # databricks- stripped first, then claude- stripped → "sonnet-4" @@ -257,56 +255,214 @@ def test_none_returns_dash(self): assert summarize_models("claude", None) == "-" -class TestModelTokenBreakdown: - def test_extracts_json_model_tokens(self): +class TestExtractModelUsage: + def test_extracts_json_model_tokens_full_names_ordered(self): raw = ( '[{"model":"databricks-claude-opus-4", "tokens":236000}, ' '{"model":"databricks-claude-haiku-4.5", "tokens":920}]' ) - result = extract_model_token_breakdown("claude", raw) - assert result == [("opus-4", 236000), ("haiku-4.5", 920)] + # Full model names (family prefix kept), highest tokens first. + usages = extract_model_usage(raw) + assert [(u.name, u.total) for u in usages] == [ + ("claude-opus-4", 236000), + ("claude-haiku-4.5", 920), + ] - def test_merges_simplified_duplicate_model_names(self): + def test_merges_duplicate_model_spellings(self): raw = [ {"model": "databricks-claude-opus-4", "tokens": 100}, {"model": "claude-opus-4", "tokens": 50}, ] - result = extract_model_token_breakdown("claude", raw) - assert result == [("opus-4", 150)] - - def test_single_model_legacy_fallback_uses_total_tokens(self): - result = extract_model_token_breakdown( - "codex", - None, - "databricks-gpt-5", - 13300, + usages = extract_model_usage(raw) + assert [(u.name, u.total) for u in usages] == [("claude-opus-4", 150)] + + def test_empty_or_unparseable_yields_nothing(self): + assert extract_model_usage("[]") == [] + assert extract_model_usage("not json") == [] + assert extract_model_usage(None) == [] + + +class TestNormalizePriceKey: + def test_dash_and_dot_versions_collapse(self): + # Usage table uses dashes (`gpt-5-6-sol`); the catalog uses dots (`gpt-5.6-sol`). + assert normalize_price_key("gpt-5-6-sol") == normalize_price_key("gpt-5.6-sol") + + def test_display_name_matches_id(self): + assert normalize_price_key("Claude Opus 4.8") == normalize_price_key( + "anthropic.claude-opus-4-8" ) - assert result == [("5", 13300)] - - def test_multi_model_legacy_fallback_does_not_assign_total_to_each_model(self): - result = extract_model_token_breakdown( - "claude", - None, - "databricks-claude-haiku-4.5, databricks-claude-opus-4", - 237000, + + def test_region_prefix_stripped(self): + assert normalize_price_key("eu/gpt-5.6-sol") == normalize_price_key("gpt-5.6-sol") + + def test_bedrock_version_and_date_suffix_stripped(self): + assert normalize_price_key( + "us.anthropic.claude-haiku-4-5-20251001-v1:0" + ) == normalize_price_key("claude-haiku-4-5") + + def test_distinct_models_stay_distinct(self): + assert normalize_price_key("gpt-5-nano") != normalize_price_key("gpt-5") + + +class TestBuildPriceLookup: + def _catalog(self): + return [ + { + "model_name": "gpt-5.6-sol", + "base_pricing": { + "input_per_million_tokens": 5.0, + "output_per_million_tokens": 30.0, + "cache_read_per_million_tokens": 0.5, + }, + }, + { + "model_name": "eu/gpt-5.6-sol", + "base_pricing": { + "input_per_million_tokens": 6.0, + "output_per_million_tokens": 36.0, + }, + }, + ] + + def test_prefers_bare_name_over_region_prefixed(self): + lookup = build_price_lookup(self._catalog()) + price = lookup[normalize_price_key("gpt-5.6-sol")] + assert price.input == Decimal("5.0") + assert price.cache_read == Decimal("0.5") + + def test_ignores_entries_without_model_name(self): + assert build_price_lookup([{"base_pricing": {"input_per_million_tokens": 1.0}}]) == {} + + +class TestEstimateModelCost: + def test_prices_uncached_cached_and_output_separately(self): + price = ModelPrice(input=Decimal("5"), output=Decimal("30"), cache_read=Decimal("0.5")) + # 1M input of which 200k cached, 100k output. + cost = estimate_model_cost(price, 1_000_000, 200_000, 100_000) + # 800k*5 + 200k*0.5 + 100k*30, all /1e6 = 4.0 + 0.1 + 3.0 + assert cost == Decimal("7.1") + + def test_cached_falls_back_to_input_rate_when_no_cache_price(self): + price = ModelPrice(input=Decimal("5"), output=Decimal("30"), cache_read=None) + cost = estimate_model_cost(price, 1_000_000, 200_000, 0) + assert cost == Decimal("5") + + def test_none_when_no_input_rate(self): + price = ModelPrice(input=None, output=Decimal("30"), cache_read=None) + assert estimate_model_cost(price, 1_000_000, 0, 0) is None + + def test_none_when_no_token_breakdown(self): + price = ModelPrice(input=Decimal("5"), output=Decimal("30"), cache_read=None) + assert estimate_model_cost(price, 0, 0, 0) is None + + +class TestModelUsageCostAndRendering: + def _lookup(self): + return build_price_lookup( + [ + { + "model_name": "anthropic.claude-opus-4-8", + "base_pricing": { + "input_per_million_tokens": 5.0, + "output_per_million_tokens": 25.0, + "cache_read_per_million_tokens": 0.5, + }, + } + ] ) - assert result == [("haiku-4.5", 0), ("opus-4", 0)] - def test_summarizes_tokens_next_to_each_model(self): - raw = '[{"model":"databricks-claude-opus-4", "tokens":236000}]' - result = summarize_model_tokens("claude", raw, "", 0) - assert result == "opus-4 (236.0K)" + def test_extract_model_usage_carries_token_breakdown(self): + raw = ( + '[{"model":"claude-opus-4-8","requests":3,' + '"tokens":1000,"input":800,"cached":200,"output":100}]' + ) + usages = extract_model_usage(raw) + assert len(usages) == 1 + u = usages[0] + assert (u.name, u.requests, u.total, u.input, u.cached, u.output) == ( + "claude-opus-4-8", + 3, + 1000, + 800, + 200, + 100, + ) + assert u.raw_names == ("claude-opus-4-8",) + def test_null_model_bucketed_as_unknown(self): + raw = '[{"model":null,"requests":1,"tokens":500,"input":400,"cached":0,"output":100}]' + (u,) = extract_model_usage(raw) + assert u.name == "" + assert u.raw_names == () -class TestEmptyToolDay: - def test_structure(self): - d = date(2024, 6, 1) - row = empty_tool_day("claude", d) - assert row["tool"] == "claude" - assert row["usage_day"] == d - assert row["total_tokens_used"] == 0 - assert row["sessions"] == 0 - assert row["models"] == "-" + def test_cost_matched_via_raw_name(self): + raw = ( + '[{"model":"claude-opus-4-8","requests":1,"tokens":1100000,' + '"input":1000000,"cached":200000,"output":100000}]' + ) + (u,) = extract_model_usage(raw) + cost = model_usage_cost(u, self._lookup()) + # 800k*5 + 200k*0.5 + 100k*25 = 4 + 0.1 + 2.5 + assert cost == Decimal("6.6") + + +class TestBuildToolModelRows: + def _lookup(self): + return build_price_lookup( + [ + { + "model_name": "anthropic.claude-opus-4-8", + "base_pricing": { + "input_per_million_tokens": 5.0, + "output_per_million_tokens": 25.0, + "cache_read_per_million_tokens": 0.5, + }, + } + ] + ) + + def _records(self): + # Two days of the same tool; opus-4-8 appears on both and must aggregate. + return [ + { + "tool": "claude", + "usage_day": date.today(), + "total_tokens_used": 1_100_000, + "model_tokens": ( + '[{"model":"claude-opus-4-8","requests":2,"tokens":1100000,' + '"input":1000000,"cached":200000,"output":100000}]' + ), + }, + { + "tool": "claude", + "usage_day": date.today() - timedelta(days=1), + "total_tokens_used": 500, + "model_tokens": ( + '[{"model":"mystery-model","requests":1,"tokens":500,' + '"input":400,"cached":0,"output":100}]' + ), + }, + ] + + def test_aggregates_per_model_over_week(self): + usages = aggregate_tool_model_usage(self._records(), "claude") + # opus-4-8 (1.1M) sorts ahead of mystery-model (500). + assert [u.name for u in usages] == ["claude-opus-4-8", "mystery-model"] + assert usages[0].requests == 2 + + def test_rows_and_totals(self): + rows, totals = build_tool_model_rows(self._records(), "claude", self._lookup()) + # columns: model, requests, input (incl cache), output, cost + assert rows[0] == ["claude-opus-4-8", "2", "1.0M", "100.0K", "$6.60"] + assert rows[1] == ["mystery-model", "1", "400", "100", "-"] + assert totals.requests == 3 + assert totals.tokens == 1_100_000 + 500 + assert totals.cost == Decimal("6.6") + + def test_totals_cost_none_when_nothing_priced(self): + records = [self._records()[1]] # only the unpriced mystery-model + _, totals = build_tool_model_rows(records, "claude", self._lookup()) + assert totals.cost is None class TestRenderBudgetLines: @@ -341,13 +497,11 @@ def test_thousands_separator(self): class TestRenderUsageSummary: - def _make_record(self, days_ago: int, tool: str, tokens: int, model: str = "") -> dict: - d = date.today() - timedelta(days=days_ago) + def _make_record(self, days_ago: int, tool: str, tokens: int) -> dict: return { "tool": tool, - "usage_day": d, + "usage_day": date.today() - timedelta(days=days_ago), "total_tokens_used": tokens, - "models": model, } def test_contains_requester_name(self): @@ -376,9 +530,16 @@ def test_active_tools_listed(self): assert "Claude Code" in result def test_top_models_listed(self): - records = [self._make_record(0, "claude", 5000, "databricks-claude-sonnet-4")] + records = [ + { + "tool": "claude", + "usage_day": date.today(), + "total_tokens_used": 5000, + "model_tokens": '[{"model":"databricks-claude-sonnet-4","tokens":5000}]', + } + ] result = render_usage_summary(records, "user", {"claude": "Claude Code"}) - assert "sonnet-4" in result + assert "claude-sonnet-4" in result def test_includes_budget_spend_when_available(self): records = [self._make_record(0, "claude", 1000)] @@ -420,27 +581,12 @@ def test_top_models_uses_per_model_token_totals(self): "user", {"claude": "Claude Code", "codex": "Codex"}, ) - assert "opus-4 (236.1K)" in result - assert "5 (13.3K)" in result - assert "haiku-4.5 (920)" in result - assert "haiku-4.5 (237.0K)" not in result - - def test_daily_table_shows_per_model_token_totals(self): - records = [ - { - "tool": "claude", - "usage_day": date.today(), - "total_tokens_used": 237000, - "sessions": 2, - "models": "databricks-claude-haiku-4.5, databricks-claude-opus-4", - "model_tokens": ( - '[{"model":"databricks-claude-haiku-4.5", "tokens":920}, ' - '{"model":"databricks-claude-opus-4", "tokens":236080}]' - ), - } - ] - rows = build_tool_breakdown_rows(records, "claude") - assert rows[0][5] == "opus-4 (236.1K), haiku-4.5 (920)" + # Top-models line is full names only, ranked by per-model token totals + # (claude-opus-4 236.1K > gpt-5 13.3K > claude-haiku-4.5 920). + assert "Top models this week:" in result + assert "claude-opus-4, gpt-5, claude-haiku-4.5" in result + # No token counts in this line — those live in the per-model table. + assert "236.1K" not in result def test_empty_records(self): result = render_usage_summary([], "user", {"claude": "Claude Code"}) @@ -457,9 +603,6 @@ def test_filters_to_configured_agents_and_skips_inactive_tables(self, monkeypatc "usage_day", "total_tokens_used", "sessions", - "first_event_time", - "last_event_time", - "models", "model_tokens", ] rows = [ @@ -469,10 +612,7 @@ def test_filters_to_configured_agents_and_skips_inactive_tables(self, monkeypatc today, 100, 1, - None, - None, - "databricks-gpt-5", - '[{"model":"databricks-gpt-5", "tokens":100}]', + '[{"model":"databricks-gpt-5","requests":1,"tokens":100,"input":80,"output":20}]', ), ( "user@example.com", @@ -480,9 +620,6 @@ def test_filters_to_configured_agents_and_skips_inactive_tables(self, monkeypatc old_day, 200, 1, - None, - None, - "databricks-claude-opus-4", '[{"model":"databricks-claude-opus-4", "tokens":200}]', ), ( @@ -491,9 +628,6 @@ def test_filters_to_configured_agents_and_skips_inactive_tables(self, monkeypatc today, 900, 1, - None, - None, - "databricks-gemini-2.0-flash", '[{"model":"databricks-gemini-2.0-flash", "tokens":900}]', ), ] @@ -527,6 +661,9 @@ def fake_render_box_table(headers, table_rows, max_widths=None): monkeypatch.setattr( usage_mod, "resolve_current_budget_spend", lambda *args, **kwargs: (None, "disabled") ) + monkeypatch.setattr( + usage_mod, "fetch_external_model_prices", lambda *args, **kwargs: ([], "disabled") + ) monkeypatch.setattr(usage_mod, "console", DummyConsole()) monkeypatch.setattr(usage_mod, "print_heading", headings.append) monkeypatch.setattr(usage_mod, "print_note", notes.append) @@ -542,7 +679,10 @@ def fake_render_box_table(headers, table_rows, max_widths=None): f"No usage for Claude Code in the last {USAGE_BREAKDOWN_DAYS} days.", ] assert len(rendered_tables) == 1 - assert rendered_tables[0][0][2] == "100" + # One per-model row for codex: full model name "gpt-5", 1 request, 80 input / 20 output. + assert rendered_tables[0][0][0] == "gpt-5" + assert rendered_tables[0][0][1] == "1" + assert rendered_tables[0][0][2] == "80" assert "gemini" not in "\n".join(printed).lower() assert "900" not in "\n".join(printed) @@ -620,6 +760,12 @@ def fake_discover(workspace, token, *, warehouse_id=None): monkeypatch.setattr(usage_mod, "get_databricks_token", lambda *a, **k: "token") monkeypatch.setattr(usage_mod, "discover_sql_warehouses", fake_discover) monkeypatch.setattr(usage_mod, "run_usage_query", lambda *a, **k: (["c"], [])) + monkeypatch.setattr( + usage_mod, "resolve_current_budget_spend", lambda *a, **k: (None, "disabled") + ) + monkeypatch.setattr( + usage_mod, "fetch_external_model_prices", lambda *a, **k: ([], "disabled") + ) monkeypatch.setattr(usage_mod, "print_note", lambda *a: None) monkeypatch.setattr(usage_mod, "console", type("C", (), {"print": lambda *a: None})())