From c1ad5fd4145593bef1078d5bf216c276a6f4f3af Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Mon, 19 Oct 2026 12:55:30 -0700 Subject: [PATCH] feat(proxy): record the sending team on daily tag spend rows LiteLLM_DailyTagSpend gains team_id TEXT NOT NULL DEFAULT '' ("" = the request had no team) and the unique key becomes (tag, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint), so one tag used by several teams is kept as one row per team. A (team_id, date, tag) index serves team-scoped tag reads. Migrations add the column, build the new unique index and the team index CONCURRENTLY, then drop the old unique index CONCURRENTLY. The spend writer stamps the request's team_id on each daily tag transaction and keys the in-memory/Redis merge on it; the bulk upsert adds team_id to the tag table's insert columns and ON CONFLICT target. Autorouter baseline adjustments carry team_id too. Entries queued by an older pod, and baseline records stored before this change, carry no team_id and land under "". Upgrade: an old-version writer targets the dropped unique index and fails every flush, so drain and stop old instances before running these migrations. LIT-8516 --- .../migration.sql | 8 ++++++++ .../migration.sql | 14 ++++++++++++++ .../migration.sql | 12 ++++++++++++ .../migration.sql | 15 +++++++++++++++ .../litellm_proxy_extras/schema.prisma | 4 +++- litellm/proxy/_types.py | 4 ++++ litellm/proxy/db/baseline_accounting.py | 4 ++++ litellm/proxy/db/daily_spend_bulk_upsert.py | 19 ++++++++++++++----- litellm/proxy/db/db_spend_update_writer.py | 8 ++++++-- litellm/proxy/schema.prisma | 4 +++- schema.prisma | 4 +++- .../proxy/db/test_daily_spend_bulk_upsert.py | 10 +++++----- 12 files changed, 91 insertions(+), 15 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql new file mode 100644 index 00000000000..1a8322fa254 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql @@ -0,0 +1,8 @@ +-- Records the team that sent each request on its daily tag spend rows, so one tag used by +-- several teams lands in one row per team. "" means the request had no team, and every row +-- written before this migration keeps "" (attribution is not retroactive). +-- +-- NOT NULL DEFAULT '' rather than nullable: the daily spend upsert normalizes every conflict +-- column to '' because a NULL never matches itself in a unique index. On PostgreSQL 11+ adding +-- a column with a constant default is a metadata-only change, so this does not rewrite the table. +ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "team_id" TEXT NOT NULL DEFAULT ''; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql new file mode 100644 index 00000000000..1816ff63ce6 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000001_daily_tag_spend_team_id_unique_index/migration.sql @@ -0,0 +1,14 @@ +-- CreateIndex (CONCURRENTLY) +-- +-- The new conflict target for the daily tag spend upsert: the old unique key plus team_id. +-- Built before the old unique index is dropped (next migrations), so the table is never +-- without a unique index the writer can arbitrate on. +-- +-- Disclaimer: +-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a +-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. +-- - If the build is interrupted, Postgres may leave an INVALID index that must be dropped and +-- recreated. +-- - Do not edit this file after it has been applied to any database: Prisma checksums +-- migrations; add a new migration instead. +CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key" ON "LiteLLM_DailyTagSpend"("tag", "team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql new file mode 100644 index 00000000000..474d5311305 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000002_daily_tag_spend_team_id_date_tag_index/migration.sql @@ -0,0 +1,12 @@ +-- CreateIndex (CONCURRENTLY) +-- +-- Serves team-scoped tag reads: WHERE team_id = ANY(...) AND date BETWEEN ... grouped by tag. +-- +-- Disclaimer: +-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a +-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. +-- - If the build is interrupted, Postgres may leave an INVALID index that must be dropped and +-- recreated. +-- - Do not edit this file after it has been applied to any database: Prisma checksums +-- migrations; add a new migration instead. +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_DailyTagSpend_team_id_date_tag_idx" ON "LiteLLM_DailyTagSpend"("team_id", "date", "tag"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql new file mode 100644 index 00000000000..7db51879c01 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql @@ -0,0 +1,15 @@ +-- DropIndex (CONCURRENTLY) +-- +-- Drops the unique key without team_id. With it in place, two teams sending the same tag on +-- the same day for the same key and model would collide on it, so it has to go once the +-- team-aware unique index exists. +-- +-- Upgrade note: an old-version writer still targets this index in its ON CONFLICT clause and +-- fails every flush once it is gone. Drain and stop old instances before running migrations. +-- +-- Disclaimer: +-- - DROP INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a +-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. +-- - Do not edit this file after it has been applied to any database: Prisma checksums +-- migrations; add a new migration instead. +DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 179b741fe7c..b018cf4788c 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? + team_id String @default("") // "" = the request had no team date String api_key String model String? @@ -1111,9 +1112,10 @@ 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, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") @@index([date]) @@index([tag, date]) + @@index([team_id, date, tag]) @@index([api_key]) @@index([model]) @@index([mcp_namespaced_tool_name]) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7b7bfc55ade..1d8d80b2e6f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5270,6 +5270,10 @@ class DailyEndUserSpendTransaction(BaseDailySpendTransaction): class DailyTagSpendTransaction(BaseDailySpendTransaction): request_id: str | None tag: str + # "" when the request had no team. Not required: entries queued by a pod running the + # previous release, or replayed from the Redis buffer across an upgrade, carry no such key, + # and the upsert reads a missing value as "". + team_id: NotRequired[str] class DailyAgentSpendTransaction(BaseDailySpendTransaction): diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index f7a2be654c8..5aca1ebaa50 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -58,6 +58,9 @@ class DailyBaselineAttribution(LiteLLMBaseModel): model_group: str | None = None endpoint: str | None = None mcp_namespaced_tool_name: str | None = None + # The request's team, for the team-scoped tag rows. Defaulted so records stored before + # the column existed still validate; they land under "" like any pre-upgrade tag spend. + team_id: str = "" targets: tuple[DailyBaselineTarget, ...] = () def adjustment(self, target: DailyBaselineTarget, savings_delta: float, request_id: str) -> SpendRow: @@ -71,6 +74,7 @@ class DailyBaselineAttribution(LiteLLMBaseModel): "model_group": self.model_group, "endpoint": self.endpoint, "mcp_namespaced_tool_name": self.mcp_namespaced_tool_name, + **({"team_id": self.team_id} if table.team_scoped else {}), table.entity_id_column: target.entity_id, "request_id": request_id, "autorouter_savings_spend": savings_delta, diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index eb130a5196f..b66d4cbc014 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -32,6 +32,9 @@ class DailySpendTable: name: str entity_id_column: str carries_request_id: bool = False + # Rows also keyed by the team that sent the request, so one entity id shared by several + # teams is kept as one row per team. + team_scoped: bool = False DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( @@ -41,7 +44,9 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), - "tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True), + "tag": DailySpendTable( + name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True, team_scoped=True + ), } ) @@ -124,9 +129,14 @@ def _as_float(value: object) -> float: return float(value) if isinstance(value, (int, float)) else 0.0 +def _conflict_columns(table: DailySpendTable) -> tuple[str, ...]: + """The table's unique constraint, in constraint order.""" + return (table.entity_id_column, *(("team_id",) if table.team_scoped else ()), *_KEY_COLUMNS) + + def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" - return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) + return tuple(_as_text(transaction.get(column)) for column in _conflict_columns(table)) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -174,8 +184,7 @@ def _row_params( def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", - table.entity_id_column, - *_KEY_COLUMNS, + *_conflict_columns(table), "model_group", *_COUNTER_COLUMNS, *_SPEND_COLUMNS, @@ -213,7 +222,7 @@ def build_bulk_upsert( sql: Final = ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted(_conflict_columns(table))}) DO UPDATE SET\n" f" {increments}{request_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d9fb236632b..4b46d2eeaad 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1040,6 +1040,7 @@ class DBSpendUpdateWriter: model_group=base.get("model_group"), endpoint=base.get("endpoint"), mcp_namespaced_tool_name=base.get("mcp_namespaced_tool_name"), + team_id=payload.get("team_id") or "", targets=tuple( DailyBaselineTarget(entity=entity, entity_id=identity) for entity, values in bases @@ -3094,13 +3095,16 @@ class DBSpendUpdateWriter: return request_tags: Final = daily_spend_entity_ids(payload, "tag") + team_id: Final = payload.get("team_id") or "" 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}" + # team_id is part of the key so the queue never folds two teams' spend on one tag + # into a single entry before the upsert can keep them apart. + daily_transaction_key = f"{tag}_{team_id}_{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"] + tag=tag, team_id=team_id, **base_daily_transaction, request_id=payload["request_id"] ) await self.daily_tag_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 179b741fe7c..b018cf4788c 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? + team_id String @default("") // "" = the request had no team date String api_key String model String? @@ -1111,9 +1112,10 @@ 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, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") @@index([date]) @@index([tag, date]) + @@index([team_id, date, tag]) @@index([api_key]) @@index([model]) @@index([mcp_namespaced_tool_name]) diff --git a/schema.prisma b/schema.prisma index 179b741fe7c..b018cf4788c 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1086,6 +1086,7 @@ model LiteLLM_DailyTagSpend { id String @id @default(uuid()) request_id String? tag String? + team_id String @default("") // "" = the request had no team date String api_key String model String? @@ -1111,9 +1112,10 @@ 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, team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyTagSpend_tag_team_id_dimensions_key") @@index([date]) @@index([tag, date]) + @@index([team_id, date, tag]) @@index([api_key]) @@index([model]) @@index([mcp_namespaced_tool_name]) 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 5bce9137666..0912c0c406d 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -88,10 +88,10 @@ def test_one_statement_carries_every_row_in_the_batch(): assert sql.count("INSERT INTO") == 1 assert len(re.findall(r"ON CONFLICT", sql)) == 1 - # 25 bound columns per row plus the inlined updated_at, so the row count is what - # separates one multi-row statement from a hundred single-row ones. - assert len(params) == 100 * 25 - assert "$2500::text" in sql + # 26 bound columns per row (team_id included) plus the inlined updated_at, so the row + # count is what separates one multi-row statement from a hundred single-row ones. + assert len(params) == 100 * 26 + assert "$2600::text" in sql assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1 @@ -101,7 +101,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", "team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' )