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 83a0e6f9e05..c84b9f810fe 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 dc7383569dc..5d09280dc85 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5159,6 +5159,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 166a42adfc0..70106ab6ba7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -918,6 +918,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 @@ -3034,13 +3035,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 83a0e6f9e05..c84b9f810fe 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 83a0e6f9e05..c84b9f810fe 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"' )