From c6669babf482e3aa18f43ab7ff74826d0a4ec6d9 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 3 Oct 2026 00:32:55 -0700 Subject: [PATCH] fix(usage): preserve daily spend by public model --- .../migration.sql | 27 ++++++++++ .../litellm_proxy_extras/schema.prisma | 12 ++--- litellm-proxy-extras/migration_runbook.md | 8 +++ litellm/proxy/db/daily_spend_bulk_upsert.py | 12 +++-- litellm/proxy/db/db_spend_update_writer.py | 25 ++++----- litellm/proxy/schema.prisma | 12 ++--- .../spend_tracking/ptu_flat_cost_rollup.py | 45 ++++++++-------- schema.prisma | 12 ++--- .../proxy/db/test_daily_spend_bulk_upsert.py | 21 +++++--- .../proxy/db/test_db_spend_update_writer.py | 43 ++++++++++++--- .../test_ptu_flat_cost_rollup.py | 54 +++++++++++++------ .../src/components/activity_metrics.test.tsx | 10 ++-- .../src/components/activity_metrics.tsx | 15 ++++-- 13 files changed, 199 insertions(+), 97 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260922000000_daily_spend_model_group_identity/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260922000000_daily_spend_model_group_identity/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260922000000_daily_spend_model_group_identity/migration.sql new file mode 100644 index 00000000000..bc803c0a096 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260922000000_daily_spend_model_group_identity/migration.sql @@ -0,0 +1,27 @@ +BEGIN; + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_group_key" ON "LiteLLM_DailyUserSpend" ("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_group_key" ON "LiteLLM_DailyTeamSpend" ("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_group_key" ON "LiteLLM_DailyOrganizationSpend" ("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_group_key" ON "LiteLLM_DailyEndUserSpend" ("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_group_key" ON "LiteLLM_DailyAgentSpend" ("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_group_key" ON "LiteLLM_DailyTagSpend" ("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"); + +DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key"; + +DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key"; + +DROP INDEX IF EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key"; + +DROP INDEX IF EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key"; + +DROP INDEX IF EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key"; + +DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; + +COMMIT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index cf76b764350..cbd31d98814 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -873,7 +873,7 @@ model LiteLLM_DailyUserSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyUserSpend_group_key") @@index([date]) @@index([user_id, date]) @@index([api_key]) @@ -942,7 +942,7 @@ model LiteLLM_DailyOrganizationSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyOrganizationSpend_group_key") @@index([date]) @@index([organization_id, date]) @@index([api_key]) @@ -979,7 +979,7 @@ model LiteLLM_DailyEndUserSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyEndUserSpend_group_key") @@index([date]) @@index([end_user_id, date]) @@index([api_key]) @@ -1016,7 +1016,7 @@ model LiteLLM_DailyAgentSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyAgentSpend_group_key") @@index([date]) @@index([agent_id, date]) @@index([api_key]) @@ -1055,7 +1055,7 @@ model LiteLLM_DailyTeamSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTeamSpend_group_key") @@index([date]) @@index([team_id, date]) @@index([api_key]) @@ -1094,7 +1094,7 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTagSpend_group_key") @@index([date]) @@index([tag, date]) @@index([api_key]) diff --git a/litellm-proxy-extras/migration_runbook.md b/litellm-proxy-extras/migration_runbook.md index b1e9236e520..164952d3edf 100644 --- a/litellm-proxy-extras/migration_runbook.md +++ b/litellm-proxy-extras/migration_runbook.md @@ -94,6 +94,14 @@ When the guard fires: > **AI AGENTS / ASSISTANTS:** Do **not** automatically re-run the command with `--allow-destructive`. If the guard fires while you are driving the runbook for a human, stop, show them the error, and wait for their explicit confirmation before passing the flag. Auto-passing `--allow-destructive` is the exact failure mode this guard exists to prevent. +## Daily spend model-group identity cutover + +`20260922000000_daily_spend_model_group_identity` separates daily usage by the requested public model group. Schedule a maintenance window: the six replacement unique indexes scan existing daily spend tables and block writes while building + +Drain incoming requests and all in-memory and shared Redis spend buffers, then stop every old proxy and spend collector before applying this migration. Start only upgraded writers after it completes. Old writers require the removed conflict target and cannot run alongside the expanded identity. The migration changes indexes only; previously merged attribution and legacy NULL groups remain as recorded + +PTU flat charges retain their deployment/day identity across public-model renames. New inference rows include the requested group in their identity. Rolling back requires stopping writers and reconciling group-separated rows before restoring the old unique indexes; simply starting an older image is unsupported + ## Common Fixes **Missing testing module:** diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index eb130a5196f..c99410e3dd7 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -78,7 +78,15 @@ def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEnti # The unique constraint's columns after the entity id, in constraint order. A NULL can # never match itself in a unique index, so every one of these is normalized to '': the # conflict target has to be NULL-free or the row is re-inserted on every single flush. -_KEY_COLUMNS: Final = ("date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") +_KEY_COLUMNS: Final = ( + "date", + "api_key", + "model", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + "model_group", +) _COUNTER_COLUMNS: Final = ( "prompt_tokens", @@ -164,7 +172,6 @@ def _row_params( return ( str(uuid.uuid4()), *key, - None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")), *(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS), *(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS), *((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()), @@ -176,7 +183,6 @@ def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: "id", table.entity_id_column, *_KEY_COLUMNS, - "model_group", *_COUNTER_COLUMNS, *_SPEND_COLUMNS, *(("request_id",) if table.carries_request_id else ()), diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 26a21069c83..1d8dfc97938 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -54,6 +54,7 @@ from litellm.proxy.common_utils.user_api_key_cache import project_cache_key from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, + conflict_key, daily_spend_entity_ids, merge_by_conflict_key, ) @@ -2811,7 +2812,7 @@ class DBSpendUpdateWriter: """ Add a spend log transaction to the `daily_spend_update_queue` - Key = @@unique([user_id, date, api_key, model, custom_llm_provider]) ) + Key matches the daily user spend database identity, including the requested model group. If key exists, update the transaction with the new spend and usage """ @@ -2825,9 +2826,8 @@ class DBSpendUpdateWriter: if base_daily_transaction is None: return - endpoint_str: Final = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{payload['user']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" daily_transaction: Final = DailyUserSpendTransaction(user_id=payload["user"], **base_daily_transaction) + daily_transaction_key: Final = json.dumps(conflict_key(DAILY_SPEND_TABLES["user"], daily_transaction)) await self.daily_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_team_transaction( @@ -2848,9 +2848,8 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("team_id is None for request. Skipping incrementing team spend.") return - endpoint_str: Final = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{payload['team_id']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" daily_transaction: Final = DailyTeamSpendTransaction(team_id=payload["team_id"], **base_daily_transaction) + daily_transaction_key: Final = json.dumps(conflict_key(DAILY_SPEND_TABLES["team"], daily_transaction)) await self.daily_team_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_org_transaction( @@ -2881,9 +2880,8 @@ class DBSpendUpdateWriter: if base_daily_transaction is None: return - endpoint_str: Final = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{org_id}_{base_daily_transaction['date']}_{payload_with_org['api_key']}_{payload_with_org['model']}_{payload_with_org['custom_llm_provider']}_{endpoint_str}" daily_transaction: Final = DailyOrganizationSpendTransaction(organization_id=org_id, **base_daily_transaction) + daily_transaction_key: Final = json.dumps(conflict_key(DAILY_SPEND_TABLES["org"], daily_transaction)) await self.daily_org_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_end_user_transaction( @@ -2914,9 +2912,8 @@ class DBSpendUpdateWriter: if base_daily_transaction is None: return - endpoint_str: Final = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{end_user_id}_{base_daily_transaction['date']}_{payload_with_end_user_id['api_key']}_{payload_with_end_user_id['model']}_{payload_with_end_user_id['custom_llm_provider']}_{endpoint_str}" daily_transaction: Final = DailyEndUserSpendTransaction(end_user_id=end_user_id, **base_daily_transaction) + daily_transaction_key: Final = json.dumps(conflict_key(DAILY_SPEND_TABLES["end_user"], daily_transaction)) await self.daily_end_user_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_agent_transaction( @@ -2941,9 +2938,8 @@ class DBSpendUpdateWriter: ) if base_daily_transaction is None: return - endpoint_str: Final = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}_{endpoint_str}" daily_transaction: Final = DailyAgentSpendTransaction(agent_id=payload["agent_id"], **base_daily_transaction) + daily_transaction_key: Final = json.dumps(conflict_key(DAILY_SPEND_TABLES["agent"], daily_transaction)) await self.daily_agent_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) async def add_spend_log_transaction_to_daily_tag_transaction( @@ -2968,10 +2964,9 @@ class DBSpendUpdateWriter: for tag in request_tags: if tag is None: continue - endpoint_str = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" daily_transaction = DailyTagSpendTransaction( tag=tag, **base_daily_transaction, request_id=payload["request_id"] ) - - await self.daily_tag_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) + await self.daily_tag_spend_update_queue.add_update( + update={json.dumps(conflict_key(DAILY_SPEND_TABLES["tag"], daily_transaction)): daily_transaction} + ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index cf76b764350..cbd31d98814 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -873,7 +873,7 @@ model LiteLLM_DailyUserSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyUserSpend_group_key") @@index([date]) @@index([user_id, date]) @@index([api_key]) @@ -942,7 +942,7 @@ model LiteLLM_DailyOrganizationSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyOrganizationSpend_group_key") @@index([date]) @@index([organization_id, date]) @@index([api_key]) @@ -979,7 +979,7 @@ model LiteLLM_DailyEndUserSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyEndUserSpend_group_key") @@index([date]) @@index([end_user_id, date]) @@index([api_key]) @@ -1016,7 +1016,7 @@ model LiteLLM_DailyAgentSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyAgentSpend_group_key") @@index([date]) @@index([agent_id, date]) @@index([api_key]) @@ -1055,7 +1055,7 @@ model LiteLLM_DailyTeamSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTeamSpend_group_key") @@index([date]) @@index([team_id, date]) @@index([api_key]) @@ -1094,7 +1094,7 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTagSpend_group_key") @@index([date]) @@index([tag, date]) @@index([api_key]) diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 99e5c35ae14..2e0f91d9d8d 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -9,7 +9,7 @@ and ``cost_per_ptu_per_hour`` accrues flat cost of ``[ptu_effective_from, ptu_effective_to)`` window (a window opening at 23:00 charges one hour that day). The amount is written to ``LiteLLM_DailyTeamSpend`` under a sentinel api_key so the rows are distinguishable from per-request rows -and share the existing unique constraint. +and keep their identity across display-name changes. """ import asyncio @@ -19,6 +19,7 @@ from dataclasses import dataclass from datetime import date, datetime, time, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final +from uuid import NAMESPACE_URL, uuid5 from litellm._logging import verbose_proxy_logger from litellm.constants import ( @@ -30,6 +31,7 @@ from litellm.constants import ( PTU_SENTINEL_API_KEY, ) from litellm.litellm_core_utils.ptu_pricing import ptu_terms +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient, writer_wrapper from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.repositories.model_repository import ModelRepository from litellm.repositories.prisma_protocols import TableActions @@ -242,32 +244,27 @@ async def _upsert_ptu_daily_row( date_str: str, flat_cost: float, ) -> None: - """Idempotent upsert of a sentinel-api_key row on LiteLLM_DailyTeamSpend. - - ``model`` holds the deployment id because it is part of the table's unique key and a - rename must not move the row. ``model_group`` carries the operator-facing name, which - is outside the key and is what the usage views display. - """ - where: Final = { - "team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": { - "team_id": team_id, - "date": date_str, - "api_key": PTU_SENTINEL_API_KEY, - "model": model_id, - "custom_llm_provider": "", - "mcp_namespaced_tool_name": "", - "endpoint": "", - } + """Upsert one deployment/day charge, retaining legacy IDs and a mutable display name.""" + natural_key: Final = { + "team_id": team_id, + "date": date_str, + "api_key": PTU_SENTINEL_API_KEY, + "model": model_id, } + table: Final = _DailyTeamSpendRepository(WriterPinnedClient(writer_wrapper(prisma_client.db))).table + existing: Final = await table.find_first(where=natural_key) + row_id: Final = ( + existing.id + if existing is not None + else str(uuid5(NAMESPACE_URL, "litellm:ptu:" + json.dumps(tuple(natural_key.values()), separators=(",", ":")))) + ) now: Final = datetime.now(timezone.utc) - await _daily_team_spend_table(prisma_client).upsert( - where=where, + await table.upsert( + where={"id": row_id}, data={ "create": { - "team_id": team_id, - "date": date_str, - "api_key": PTU_SENTINEL_API_KEY, - "model": model_id, + **natural_key, + "id": row_id, "model_group": model_name, "custom_llm_provider": "", "mcp_namespaced_tool_name": "", @@ -291,7 +288,7 @@ async def _upsert_charge_with_retry( ) -> bool: """Write one charge, retrying transient failures. Returns False once attempts are spent. - The upsert is idempotent on the sentinel unique key, so a retry can only rewrite the + The upsert is idempotent on the stable row ID, so a retry can only rewrite the same amount for the same day. Retrying in-run matters because the scheduled job moves on to the next date: a write lost here is a day of PTU cost that no later run replays. """ diff --git a/schema.prisma b/schema.prisma index cf76b764350..cbd31d98814 100644 --- a/schema.prisma +++ b/schema.prisma @@ -873,7 +873,7 @@ model LiteLLM_DailyUserSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyUserSpend_group_key") @@index([date]) @@index([user_id, date]) @@index([api_key]) @@ -942,7 +942,7 @@ model LiteLLM_DailyOrganizationSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyOrganizationSpend_group_key") @@index([date]) @@index([organization_id, date]) @@index([api_key]) @@ -979,7 +979,7 @@ model LiteLLM_DailyEndUserSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyEndUserSpend_group_key") @@index([date]) @@index([end_user_id, date]) @@index([api_key]) @@ -1016,7 +1016,7 @@ model LiteLLM_DailyAgentSpend { timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyAgentSpend_group_key") @@index([date]) @@index([agent_id, date]) @@index([api_key]) @@ -1055,7 +1055,7 @@ model LiteLLM_DailyTeamSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTeamSpend_group_key") @@index([date]) @@index([team_id, date]) @@index([api_key]) @@ -1094,7 +1094,7 @@ model LiteLLM_DailyTagSpend { created_at DateTime @default(now()) updated_at DateTime @updatedAt - @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTagSpend_group_key") @@index([date]) @@index([tag, date]) @@index([api_key]) diff --git a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py index 7893fb82281..e470f93b16e 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -3,11 +3,13 @@ import re from collections.abc import AsyncIterator from contextlib import AbstractAsyncContextManager, asynccontextmanager +from typing import Final import pytest from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, + DailySpendTable, build_bulk_upsert, conflict_key, merge_by_conflict_key, @@ -19,7 +21,7 @@ USER_TABLE = DAILY_SPEND_TABLES["user"] # Every nullable member of the unique constraint, so a test that only varied the provider # cannot pass while a sibling column still leaks a NULL into the conflict target. -NULLABLE_KEY_COLUMNS = ("model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") +NULLABLE_KEY_COLUMNS = ("model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group") def tag_txn(**overrides): @@ -71,13 +73,16 @@ def test_null_and_empty_provider_merge_into_one_row(order): assert folded["api_requests"] == 4 -def test_distinct_keys_are_not_merged_and_are_ordered_deterministically(): - unordered = (tag_txn(tag="z-team"), tag_txn(tag="a-team"), tag_txn(tag="m-team")) +@pytest.mark.parametrize("table", tuple(DAILY_SPEND_TABLES.values()), ids=tuple(DAILY_SPEND_TABLES)) +@pytest.mark.parametrize("group_only", (False, True)) +def test_distinct_keys_are_not_merged_and_are_ordered_deterministically(table: DailySpendTable, group_only: bool) -> None: + column: Final = "model_group" if group_only else table.entity_id_column + unordered: Final = tuple(tag_txn(**{column: name}) for name in ("z-team", "a-team", "m-team")) - merged = merge_by_conflict_key(TAG_TABLE, unordered) + merged: Final = merge_by_conflict_key(table, unordered) - assert [txn["tag"] for _, txn in merged] == ["a-team", "m-team", "z-team"] - assert merged == merge_by_conflict_key(TAG_TABLE, tuple(reversed(unordered))) + assert [txn[column] for _, txn in merged] == ["a-team", "m-team", "z-team"] + assert merged == merge_by_conflict_key(table, tuple(reversed(unordered))) def test_one_statement_carries_every_row_in_the_batch(): @@ -100,7 +105,7 @@ def test_conflict_target_is_the_full_unique_constraint(): conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql) assert conflict_target is not None assert conflict_target.group(1) == ( - '"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' + '"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"' ) @@ -188,7 +193,7 @@ async def test_writer_survives_a_transaction_whose_key_columns_are_null(): """A NULL key column used to raise out of prisma and drop the whole batch's spend.""" prisma_client = _RecordingPrismaClient() transactions = { - "mcp": tag_txn(model=None, custom_llm_provider=None, mcp_namespaced_tool_name="server/tool"), + "mcp": tag_txn(model=None, model_group=None, custom_llm_provider=None, mcp_namespaced_tool_name="server/tool"), "chat": tag_txn(), } diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 4de90d5d7f7..c4809c6a631 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -8,6 +8,7 @@ import re from collections.abc import AsyncIterator, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager from datetime import datetime, timedelta, timezone +from functools import partial from types import SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, call, patch @@ -1418,7 +1419,7 @@ async def test_add_spend_log_transaction_to_daily_org_transaction_injects_org_id update_dict = call_args["update"] assert len(update_dict) == 1 for key, transaction in update_dict.items(): - assert key == f"{org_id}_2024-01-01_test-key_gpt-4_openai_" + assert json.loads(key) == [org_id, "2024-01-01", "test-key", "gpt-4", "openai", "", "", "gpt-4-group"] assert transaction["organization_id"] == org_id assert transaction["date"] == "2024-01-01" assert transaction["api_key"] == "test-key" @@ -1495,7 +1496,7 @@ async def test_add_spend_log_transaction_to_daily_end_user_transaction_injects_e update_dict = call_args["update"] assert len(update_dict) == 1 for key, transaction in update_dict.items(): - assert key == f"{end_user_id}_2024-01-01_test-key_gpt-4_openai_" + assert json.loads(key) == [end_user_id, "2024-01-01", "test-key", "gpt-4", "openai", "", "", "gpt-4-group"] assert transaction["end_user_id"] == end_user_id assert transaction["date"] == "2024-01-01" assert transaction["api_key"] == "test-key" @@ -1571,7 +1572,7 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_injects_agen update_dict = call_args["update"] assert len(update_dict) == 1 for key, transaction in update_dict.items(): - assert key == f"{agent_id}_2024-01-01_test-key_gpt-4_openai_" + assert json.loads(key) == [agent_id, "2024-01-01", "test-key", "gpt-4", "openai", "", "", "gpt-4-group"] assert transaction["agent_id"] == agent_id assert transaction["date"] == "2024-01-01" assert transaction["api_key"] == "test-key" @@ -1686,7 +1687,9 @@ async def test_endpoint_field_is_correctly_mapped_from_call_type(): for key, transaction in update_dict.items(): # Verify endpoint is included in the key - assert key == f"test-user_2024-01-01_test-key_gpt-4_openai_/chat/completions" + assert json.loads(key) == [ + "test-user", "2024-01-01", "test-key", "gpt-4", "openai", "", "/chat/completions", "gpt-4-group" + ] # Verify endpoint is set in the transaction assert transaction["endpoint"] == "/chat/completions" @@ -1732,7 +1735,9 @@ async def test_endpoint_field_maps_retrieve_batch_spend_row_to_batches_endpoint( assert len(update_dict) == 1 for key, transaction in update_dict.items(): - assert key == "test-user_2024-01-01_test-key_gpt-4_openai_/batches" + assert json.loads(key) == [ + "test-user", "2024-01-01", "test-key", "gpt-4", "openai", "", "/batches", "gpt-4-group" + ] assert transaction["endpoint"] == "/batches" @@ -3632,7 +3637,8 @@ async def test_commit_spend_updates_to_db_reports_table_committed_before_cache_i @pytest.mark.asyncio -async def test_daily_transaction_internal_call_keeps_spend_but_not_request_counts(): +@pytest.mark.parametrize("internal_first", (True, False)) +async def test_daily_transaction_internal_call_keeps_spend_but_not_request_counts(internal_first: bool) -> None: """Internal sub-calls (auto-router classifier, shadow eval's shadow and judge) bill spend and tokens to the key but are not requests the caller made: api_requests, successful_requests, and autorouter_savings_spend must all stay zero for them.""" @@ -3644,6 +3650,10 @@ async def test_daily_transaction_internal_call_keeps_spend_but_not_request_count return { "request_id": "req-internal-1", "user": "test-user", + "team_id": "test-team", + "end_user": "test-end-user", + "agent_id": "test-agent", + "request_tags": '["test-tag"]', "startTime": "2026-08-11T00:00:00", "api_key": "test-key", "model": "claude-sonnet-5", @@ -3677,6 +3687,27 @@ async def test_daily_transaction_internal_call_keeps_spend_but_not_request_count assert user_sent["api_requests"] == 1 assert user_sent["successful_requests"] == 1 + groups: Final = (("router/shadow", {"internal_call_origin": "shadow_eval_judge"}), ("production", {})) + for producer, queue in ( + (writer.add_spend_log_transaction_to_daily_user_transaction, writer.daily_spend_update_queue), + (writer.add_spend_log_transaction_to_daily_team_transaction, writer.daily_team_spend_update_queue), + ( + partial(writer.add_spend_log_transaction_to_daily_org_transaction, org_id="test-org"), + writer.daily_org_spend_update_queue, + ), + (writer.add_spend_log_transaction_to_daily_end_user_transaction, writer.daily_end_user_spend_update_queue), + (writer.add_spend_log_transaction_to_daily_agent_transaction, writer.daily_agent_spend_update_queue), + (writer.add_spend_log_transaction_to_daily_tag_transaction, writer.daily_tag_spend_update_queue), + ): + for group, metadata in groups if internal_first else reversed(groups): + payload: Final = {**_payload(metadata), "model_group": group, "spend": 0.05 if metadata else 0.15} + await producer(payload=payload, prisma_client=mock_prisma) + rows: Final = await queue.flush_and_get_aggregated_daily_spend_update_transactions() + assert {row["model_group"]: (row["spend"], row["api_requests"]) for row in rows.values()} == { + "router/shadow": (0.05, 0), + "production": (0.15, 1), + } + def _response_time_payload(request_duration_ms: object, metadata: dict | None = None) -> dict: return { diff --git a/tests/unit/proxy/spend_tracking/test_ptu_flat_cost_rollup.py b/tests/unit/proxy/spend_tracking/test_ptu_flat_cost_rollup.py index 8f25cffecf5..c466ddacb07 100644 --- a/tests/unit/proxy/spend_tracking/test_ptu_flat_cost_rollup.py +++ b/tests/unit/proxy/spend_tracking/test_ptu_flat_cost_rollup.py @@ -1,16 +1,18 @@ """Tests for the per-model PTU flat-cost daily rollup.""" +import asyncio import json import types from datetime import date, datetime, timedelta, timezone -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import ANY, AsyncMock, MagicMock +from uuid import uuid4 import pytest import litellm.proxy.spend_tracking.ptu_flat_cost_rollup as ptu_rollup from litellm.constants import PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR -from litellm.types.router import ModelInfo from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( PTUModel, _active_hours_on_day, @@ -20,6 +22,7 @@ from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( run_ptu_flat_cost_rollup, run_scheduled_ptu_rollup, ) +from litellm.types.router import ModelInfo DAY = date(2026, 7, 30) TODAY = date(2026, 7, 31) @@ -166,6 +169,7 @@ def _prisma_with_models(rows, existing_sentinel_rows=()): model_table = MagicMock() model_table.find_many = AsyncMock(return_value=rows) daily = MagicMock() + daily.find_first = AsyncMock(return_value=None) daily.find_many = AsyncMock(return_value=list(existing_sentinel_rows)) daily.upsert = AsyncMock() daily.delete_many = AsyncMock() @@ -174,25 +178,28 @@ def _prisma_with_models(rows, existing_sentinel_rows=()): @pytest.mark.asyncio -async def test_rollup_writes_sentinel_row_with_hourly_cost(): +@pytest.mark.parametrize("legacy_value", [None, ""]) +@pytest.mark.parametrize("writer_unavailable", [False, True]) +async def test_rollup_writes_hourly_cost_to_the_legacy_primary_row(legacy_value, writer_unavailable): rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "team_x"})] - prisma, table = _prisma_with_models(rows) + table = _FakeSentinelTable() + table.seed("team_x", DAY, "m1", 1.0) + legacy = table.rows[("team_x", DAY.isoformat(), PTU_SENTINEL_API_KEY, "m1")] + legacy.update( + {field: legacy_value for field in ("model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint")} + ) + prisma = _prisma_for(rows, table) + replica = _prisma_for(rows, _FakeSentinelTable()) + prisma.db = RoutingPrismaWrapper(writer=prisma.db, reader=replica.db) + prisma.db._writer_unavailable = writer_unavailable result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) assert result.models_processed == 1 assert result.rows_written == 1 - created = table.upsert.await_args.kwargs["data"]["create"] - assert created["api_key"] == PTU_SENTINEL_API_KEY - assert created["ptu_flat_cost"] == pytest.approx(240.0) - assert created["team_id"] == "team_x" - # identity in the key, display beside it, so a rename cannot move the row - assert created["model"] == "m1" - assert created["model_group"] == "gpt-4o-mini-ptu" - keyed = table.upsert.await_args.kwargs["where"][ - "team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" + assert list(table.rows.values()) == [ + {**legacy, "model_group": "gpt-4o-mini-ptu", "ptu_flat_cost": pytest.approx(240.0), "updated_at": ANY} ] - assert keyed["model"] == "m1" @pytest.mark.asyncio @@ -689,21 +696,32 @@ class _FakeSentinelTable: self.rows = {} self._upsert_gate = upsert_gate self.upsert_keys = [] + self.upsert_ids = [] self.delete_many_calls = [] self.find_many_calls = [] + self.create = AsyncMock() async def upsert(self, where, data): + await asyncio.sleep(0) if self._upsert_gate is not None: await self._upsert_gate.wait() - key = where["team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"] + key = data["create"] row_key = (key["team_id"], key["date"], key["api_key"], key["model"]) self.upsert_keys.append(row_key) + self.upsert_ids.append(where["id"]) + existing = self.rows.get(row_key) self.rows[row_key] = { - "ptu_flat_cost": data["create"]["ptu_flat_cost"], - "model_group": data["create"]["model_group"], "updated_at": datetime.now(timezone.utc), + **({**existing, **data["update"]} if existing and existing["id"] == where["id"] else key), } + async def find_first(self, where): + for key, value in self.rows.items(): + row = {**dict(zip(("team_id", "date", "api_key", "model"), key)), **value} + if all(row.get(field) == expected for field, expected in where.items()): + return types.SimpleNamespace(**row) + return None + async def delete_many(self, where): self.delete_many_calls.append(where) cutoff = where["updated_at"]["lt"] @@ -737,6 +755,7 @@ class _FakeSentinelTable: def seed(self, team_id, day, model_id, flat_cost, updated_at=None, model_group=None): """Seed a row the way the rollup writes one: keyed on the deployment id.""" self.rows[(team_id, day.isoformat(), PTU_SENTINEL_API_KEY, model_id)] = { + "id": str(uuid4()), "ptu_flat_cost": flat_cost, "model_group": model_group or model_id, "updated_at": updated_at or datetime.now(timezone.utc), @@ -1683,6 +1702,7 @@ async def test_concurrent_runs_straddling_a_rename_write_one_row(): per_day[day] = per_day.get(day, 0) + 1 assert set(per_day.values()) == {1}, f"a day carries more than one charge: {per_day}" assert {key[3] for key in table.rows} == {"dep-1"} + assert len(set(table.upsert_ids)) == len(table.rows) == len(set(zip(table.upsert_keys, table.upsert_ids))) @pytest.mark.asyncio diff --git a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx index 8ce18884718..159a5ddbbf8 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx @@ -455,19 +455,23 @@ describe("ActivityMetrics", () => { expect(avgSpendElements.some((el) => el.textContent?.includes("1.058"))).toBe(true); }); - it("should handle zero successful requests without division error", () => { + it.each([0, 10])("should handle zero successful requests with $%s spend", (spend) => { const modelWithZeroRequests: Record = { "gpt-4": { ...mockModelMetrics["gpt-4"], + total_requests: 0, total_successful_requests: 0, - total_tokens: 0, - total_spend: 0, + total_tokens: spend * 1000, + total_spend: spend, }, }; render(); const zeroElements = screen.getAllByText("0"); expect(zeroElements.length).toBeGreaterThan(0); + expect(screen.getByText("N/A avg per successful request")).toBeInTheDocument(); + expect(screen.getByText("N/A per successful request")).toBeInTheDocument(); + expect(screen.getByText(/Request counts exclude internal calls/)).toBeInTheDocument(); }); it("should display prompt caching token counts when visible", () => { diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index e7712a99c71..d5efe86d00c 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -201,7 +201,10 @@ export const ModelSection = ({

Total Tokens

{metrics.total_tokens.toLocaleString()}

- {Math.round(metrics.total_tokens / metrics.total_successful_requests)} avg per successful request + {metrics.total_successful_requests > 0 + ? Math.round(metrics.total_tokens / metrics.total_successful_requests) + : "N/A"}{" "} + avg per successful request

@@ -210,8 +213,10 @@ export const ModelSection = ({

Total Spend

${formatNumberWithCommas(metrics.total_spend, 2)}

- ${formatNumberWithCommas(metrics.total_spend / metrics.total_successful_requests, 3)} per successful - request + {metrics.total_successful_requests > 0 + ? `$${formatNumberWithCommas(metrics.total_spend / metrics.total_successful_requests, 3)}` + : "N/A"}{" "} + per successful request

@@ -227,6 +232,10 @@ export const ModelSection = ({ +

+ Request counts exclude internal calls, such as shadow evaluations. Spend and tokens include them. Daily usage is + grouped by UTC day. +

{fetchTopApiKeys && }