From c3ddccd256983f22eb97412f509945010146b84b Mon Sep 17 00:00:00 2001 From: sathvikvittal Date: Thu, 6 Aug 2026 18:34:14 +0530 Subject: [PATCH 1/4] fix(proxy): include model group in daily usage aggregation --- .../migration.sql | 53 ++++++++++++++ litellm/proxy/db/daily_spend_bulk_upsert.py | 12 +++- litellm/proxy/db/db_spend_update_writer.py | 70 +++++++++++++++++-- litellm/proxy/schema.prisma | 12 ++-- schema.prisma | 12 ++-- .../proxy/db/test_daily_spend_bulk_upsert.py | 14 +++- .../proxy/db/test_db_spend_update_writer.py | 1 + 7 files changed, 151 insertions(+), 23 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260806000000_add_model_group_to_daily_usage_unique_keys/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260806000000_add_model_group_to_daily_usage_unique_keys/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260806000000_add_model_group_to_daily_usage_unique_keys/migration.sql new file mode 100644 index 00000000000..1166de79267 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260806000000_add_model_group_to_daily_usage_unique_keys/migration.sql @@ -0,0 +1,53 @@ +-- DropIndex +DROP INDEX IF EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key"; + +-- DropIndex +DROP INDEX IF EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key"; + +-- DropIndex +DROP INDEX IF EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key"; + +-- DropIndex +DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; + +-- DropIndex +DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key"; + +-- DropIndex +DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key"; + +-- Backfill +UPDATE "LiteLLM_DailyAgentSpend" SET "model_group" = '' WHERE "model_group" IS NULL; + +-- Backfill +UPDATE "LiteLLM_DailyEndUserSpend" SET "model_group" = '' WHERE "model_group" IS NULL; + +-- Backfill +UPDATE "LiteLLM_DailyOrganizationSpend" SET "model_group" = '' WHERE "model_group" IS NULL; + +-- Backfill +UPDATE "LiteLLM_DailyTagSpend" SET "model_group" = '' WHERE "model_group" IS NULL; + +-- Backfill +UPDATE "LiteLLM_DailyTeamSpend" SET "model_group" = '' WHERE "model_group" IS NULL; + +-- Backfill +UPDATE "LiteLLM_DailyUserSpend" SET "model_group" = '' WHERE "model_group" IS NULL; + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index 55d325177c6..6dbd19319e0 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -46,7 +46,15 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP # 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", + "model_group", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", +) _COUNTER_COLUMNS: Final = ( "prompt_tokens", @@ -129,7 +137,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 ()), @@ -141,7 +148,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 b0130db232a..956466b42eb 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1423,6 +1423,10 @@ class DBSpendUpdateWriter: e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj ) + @staticmethod + def _daily_transaction_key(*parts: object | None) -> str: + return json.dumps(tuple(str(part or "") for part in parts), separators=(",", ":")) + # fmt: off @overload @@ -1855,7 +1859,16 @@ class DBSpendUpdateWriter: 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_key: Final = self._daily_transaction_key( + payload["user"], + base_daily_transaction["date"], + payload["api_key"], + payload["model"], + base_daily_transaction.get("model_group"), + payload["custom_llm_provider"], + base_daily_transaction.get("mcp_namespaced_tool_name"), + endpoint_str, + ) daily_transaction: Final = DailyUserSpendTransaction(user_id=payload["user"], **base_daily_transaction) await self.daily_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) @@ -1878,7 +1891,16 @@ class DBSpendUpdateWriter: 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_key: Final = self._daily_transaction_key( + payload["team_id"], + base_daily_transaction["date"], + payload["api_key"], + payload["model"], + base_daily_transaction.get("model_group"), + payload["custom_llm_provider"], + base_daily_transaction.get("mcp_namespaced_tool_name"), + endpoint_str, + ) daily_transaction: Final = DailyTeamSpendTransaction(team_id=payload["team_id"], **base_daily_transaction) await self.daily_team_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) @@ -1911,7 +1933,16 @@ class DBSpendUpdateWriter: 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_key: Final = self._daily_transaction_key( + org_id, + base_daily_transaction["date"], + payload_with_org["api_key"], + payload_with_org["model"], + base_daily_transaction.get("model_group"), + payload_with_org["custom_llm_provider"], + base_daily_transaction.get("mcp_namespaced_tool_name"), + endpoint_str, + ) daily_transaction: Final = DailyOrganizationSpendTransaction(organization_id=org_id, **base_daily_transaction) await self.daily_org_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) @@ -1944,7 +1975,16 @@ class DBSpendUpdateWriter: 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_key: Final = self._daily_transaction_key( + end_user_id, + base_daily_transaction["date"], + payload_with_end_user_id["api_key"], + payload_with_end_user_id["model"], + base_daily_transaction.get("model_group"), + payload_with_end_user_id["custom_llm_provider"], + base_daily_transaction.get("mcp_namespaced_tool_name"), + endpoint_str, + ) daily_transaction: Final = DailyEndUserSpendTransaction(end_user_id=end_user_id, **base_daily_transaction) await self.daily_end_user_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) @@ -1971,7 +2011,16 @@ 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_key: Final = self._daily_transaction_key( + payload["agent_id"], + base_daily_transaction["date"], + payload_with_agent_id["api_key"], + payload_with_agent_id["model"], + base_daily_transaction.get("model_group"), + payload_with_agent_id["custom_llm_provider"], + base_daily_transaction.get("mcp_namespaced_tool_name"), + endpoint_str, + ) daily_transaction: Final = DailyAgentSpendTransaction(agent_id=payload["agent_id"], **base_daily_transaction) await self.daily_agent_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction}) @@ -2002,7 +2051,16 @@ class DBSpendUpdateWriter: raise ValueError(f"Invalid request_tags: {payload['request_tags']}") for tag in request_tags: 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_key = self._daily_transaction_key( + tag, + base_daily_transaction["date"], + payload["api_key"], + payload["model"], + base_daily_transaction.get("model_group"), + payload["custom_llm_provider"], + base_daily_transaction.get("mcp_namespaced_tool_name"), + endpoint_str, + ) daily_transaction = DailyTagSpendTransaction( tag=tag, **base_daily_transaction, request_id=payload["request_id"] ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 33fd9389b63..6c781b9b6a2 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -758,7 +758,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([user_id, date]) @@index([api_key]) @@ -793,7 +793,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([organization_id, date]) @@index([api_key]) @@ -827,7 +827,7 @@ model LiteLLM_DailyEndUserSpend { failed_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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([end_user_id, date]) @@index([api_key]) @@ -861,7 +861,7 @@ model LiteLLM_DailyAgentSpend { failed_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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([agent_id, date]) @@index([api_key]) @@ -897,7 +897,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([team_id, date]) @@index([api_key]) @@ -933,7 +933,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([tag, date]) @@index([api_key]) diff --git a/schema.prisma b/schema.prisma index 33fd9389b63..6c781b9b6a2 100644 --- a/schema.prisma +++ b/schema.prisma @@ -758,7 +758,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([user_id, date]) @@index([api_key]) @@ -793,7 +793,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([organization_id, date]) @@index([api_key]) @@ -827,7 +827,7 @@ model LiteLLM_DailyEndUserSpend { failed_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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([end_user_id, date]) @@index([api_key]) @@ -861,7 +861,7 @@ model LiteLLM_DailyAgentSpend { failed_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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([agent_id, date]) @@index([api_key]) @@ -897,7 +897,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([team_id, date]) @@index([api_key]) @@ -933,7 +933,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([tag, date]) @@index([api_key]) diff --git a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py index c2d0f64461a..05bebb61e86 100644 --- a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py @@ -17,7 +17,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", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") def tag_txn(**overrides): @@ -78,6 +78,15 @@ def test_distinct_keys_are_not_merged_and_are_ordered_deterministically(): assert merged == merge_by_conflict_key(TAG_TABLE, tuple(reversed(unordered))) +def test_distinct_model_groups_are_not_merged(): + merged = merge_by_conflict_key( + TAG_TABLE, + (tag_txn(model_group="public-a"), tag_txn(model_group="public-b")), + ) + + assert len(merged) == 2 + + def test_one_statement_carries_every_row_in_the_batch(): batch = merge_by_conflict_key(TAG_TABLE, tuple(tag_txn(tag=f"team-{i}") for i in range(100))) @@ -98,7 +107,8 @@ 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", "model_group", "custom_llm_provider", ' + '"mcp_namespaced_tool_name", "endpoint"' ) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 51810a28cdd..3970a4327df 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -308,6 +308,7 @@ async def test_update_daily_spend_with_null_entity_id(): assert _row_values(statement, "date") == ["2024-01-01"] assert _row_values(statement, "api_key") == ["test-api-key"] assert _row_values(statement, "model") == ["gpt-4"] + assert _row_values(statement, "model_group") == [""] assert _row_values(statement, "custom_llm_provider") == ["openai"] assert _row_values(statement, "mcp_namespaced_tool_name") == [""] assert _row_values(statement, "endpoint") == [""] From 282e882625b032143e4069c443a57f0dfb39589c Mon Sep 17 00:00:00 2001 From: sathvikvittal Date: Fri, 7 Aug 2026 12:53:56 +0530 Subject: [PATCH 2/4] fix(schema): include model group in daily spend unique constraints --- .../litellm_proxy_extras/schema.prisma | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 33fd9389b63..6c781b9b6a2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -758,7 +758,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([user_id, date]) @@index([api_key]) @@ -793,7 +793,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([organization_id, date]) @@index([api_key]) @@ -827,7 +827,7 @@ model LiteLLM_DailyEndUserSpend { failed_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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([end_user_id, date]) @@index([api_key]) @@ -861,7 +861,7 @@ model LiteLLM_DailyAgentSpend { failed_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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([agent_id, date]) @@index([api_key]) @@ -897,7 +897,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([team_id, date]) @@index([api_key]) @@ -933,7 +933,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, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) @@index([date]) @@index([tag, date]) @@index([api_key]) From 4f99f0a505c4690e8a040ec7f4f5b265913c6fb2 Mon Sep 17 00:00:00 2001 From: sathvikvittal Date: Fri, 7 Aug 2026 15:40:35 +0530 Subject: [PATCH 3/4] fix(proxy): disambiguate daily spend queue keys --- .../migration.sql | 18 ------- .../proxy/db/test_db_spend_update_writer.py | 53 +++++++++++++++++-- 2 files changed, 49 insertions(+), 22 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260806000000_add_model_group_to_daily_usage_unique_keys/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260806000000_add_model_group_to_daily_usage_unique_keys/migration.sql index 1166de79267..126852825c0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260806000000_add_model_group_to_daily_usage_unique_keys/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260806000000_add_model_group_to_daily_usage_unique_keys/migration.sql @@ -1,53 +1,35 @@ --- DropIndex DROP INDEX IF EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key"; --- DropIndex DROP INDEX IF EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key"; --- DropIndex DROP INDEX IF EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key"; --- DropIndex DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key"; --- DropIndex DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key"; --- DropIndex DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key"; --- Backfill UPDATE "LiteLLM_DailyAgentSpend" SET "model_group" = '' WHERE "model_group" IS NULL; --- Backfill UPDATE "LiteLLM_DailyEndUserSpend" SET "model_group" = '' WHERE "model_group" IS NULL; --- Backfill UPDATE "LiteLLM_DailyOrganizationSpend" SET "model_group" = '' WHERE "model_group" IS NULL; --- Backfill UPDATE "LiteLLM_DailyTagSpend" SET "model_group" = '' WHERE "model_group" IS NULL; --- Backfill UPDATE "LiteLLM_DailyTeamSpend" SET "model_group" = '' WHERE "model_group" IS NULL; --- Backfill UPDATE "LiteLLM_DailyUserSpend" SET "model_group" = '' WHERE "model_group" IS NULL; --- CreateIndex CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); --- CreateIndex CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); --- CreateIndex CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); --- CreateIndex CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); --- CreateIndex CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); --- CreateIndex CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 3970a4327df..03e28a879fc 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -336,6 +336,51 @@ def _daily_txn(user_id: str = "user1") -> dict: } +@pytest.mark.asyncio +async def test_daily_user_transaction_key_includes_model_group(): + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + writer.daily_spend_update_queue.add_update = AsyncMock() + base_payload = { + "request_id": "req-model-group-key", + "user": "test-user", + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "openai/gpt-4o-mini", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.2, + "metadata": '{"usage_object": {}}', + } + + await writer.add_spend_log_transaction_to_daily_user_transaction( + payload={**base_payload, "model_group": "a"}, + prisma_client=mock_prisma, + ) + await writer.add_spend_log_transaction_to_daily_user_transaction( + payload={**base_payload, "model_group": "b"}, + prisma_client=mock_prisma, + ) + + queued_keys = [ + next(iter(call_kwargs.kwargs["update"].keys())) + for call_kwargs in writer.daily_spend_update_queue.add_update.call_args_list + ] + + assert queued_keys == [ + '["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","",""]', + '["test-user","2024-01-01","test-key","openai/gpt-4o-mini","b","openai","",""]', + ] + + +def test_daily_transaction_key_is_not_delimiter_ambiguous(): + assert DBSpendUpdateWriter._daily_transaction_key( + "user", "model_a", "group" + ) != DBSpendUpdateWriter._daily_transaction_key("user", "model", "a_group") + + @pytest.mark.asyncio async def test_update_daily_spend_does_not_retry_post_send_ambiguous_errors(): # Regression for the double-apply hazard: a ReadTimeout means the batch was @@ -1036,7 +1081,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 key == f'["{org_id}","2024-01-01","test-key","gpt-4","gpt-4-group","openai","",""]' assert transaction["organization_id"] == org_id assert transaction["date"] == "2024-01-01" assert transaction["api_key"] == "test-key" @@ -1113,7 +1158,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 key == f'["{end_user_id}","2024-01-01","test-key","gpt-4","gpt-4-group","openai","",""]' assert transaction["end_user_id"] == end_user_id assert transaction["date"] == "2024-01-01" assert transaction["api_key"] == "test-key" @@ -1189,7 +1234,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 key == f'["{agent_id}","2024-01-01","test-key","gpt-4","gpt-4-group","openai","",""]' assert transaction["agent_id"] == agent_id assert transaction["date"] == "2024-01-01" assert transaction["api_key"] == "test-key" @@ -1310,7 +1355,7 @@ 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 key == '["test-user","2024-01-01","test-key","gpt-4","gpt-4-group","openai","","/chat/completions"]' # Verify endpoint is set in the transaction assert transaction["endpoint"] == "/chat/completions" From cc3ffbe89c707a768ccf49e7888e31b3446a7c67 Mon Sep 17 00:00:00 2001 From: sathvikvittal Date: Mon, 10 Aug 2026 11:49:14 +0530 Subject: [PATCH 4/4] fix(proxy): separate daily spend by MCP tool --- .../proxy/db/test_db_spend_update_writer.py | 51 ++++++++++++++++--- 1 file changed, 45 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 03e28a879fc..1bc08cc8118 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -12,7 +12,7 @@ sys.path.insert( from collections.abc import Callable from datetime import datetime, timezone -from unittest.mock import AsyncMock, MagicMock, call, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from redis.exceptions import DataError @@ -375,6 +375,47 @@ async def test_daily_user_transaction_key_includes_model_group(): ] +@pytest.mark.asyncio +async def test_daily_user_transaction_key_includes_mcp_tool_name(): + writer = DBSpendUpdateWriter() + writer.daily_spend_update_queue = MagicMock() + writer.daily_spend_update_queue.add_update = AsyncMock() + mock_prisma = MagicMock() + + base_payload = { + "request_id": "req-mcp-tool-key", + "user": "test-user", + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "openai/gpt-4o-mini", + "model_group": "a", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.2, + "metadata": '{"usage_object": {}}', + } + + await writer.add_spend_log_transaction_to_daily_user_transaction( + payload={**base_payload, "mcp_namespaced_tool_name": "server_a.tool"}, + prisma_client=mock_prisma, + ) + await writer.add_spend_log_transaction_to_daily_user_transaction( + payload={**base_payload, "mcp_namespaced_tool_name": "server_b.tool"}, + prisma_client=mock_prisma, + ) + + queued_keys = [ + next(iter(call_kwargs.kwargs["update"].keys())) + for call_kwargs in writer.daily_spend_update_queue.add_update.call_args_list + ] + + assert queued_keys == [ + '["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","server_a.tool",""]', + '["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","server_b.tool",""]', + ] + + def test_daily_transaction_key_is_not_delimiter_ambiguous(): assert DBSpendUpdateWriter._daily_transaction_key( "user", "model_a", "group" @@ -677,7 +718,7 @@ async def test_update_tag_db_with_valid_tags(): """ Test that _update_tag_db correctly processes valid tags and adds them to the spend update queue. """ - from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem + from litellm.proxy._types import Litellm_EntityType writer = DBSpendUpdateWriter() mock_prisma = MagicMock() @@ -1019,8 +1060,6 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i "metadata": '{"usage_object": {}}', } - # Mock the add_update method to capture what's being added - original_add_update = writer.daily_tag_spend_update_queue.add_update writer.daily_tag_spend_update_queue.add_update = AsyncMock() await writer.add_spend_log_transaction_to_daily_tag_transaction( @@ -1032,8 +1071,8 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i assert writer.daily_tag_spend_update_queue.add_update.call_count == 2 # Check that request_id is included in both transactions - for call in writer.daily_tag_spend_update_queue.add_update.call_args_list: - transaction_dict = call[1]["update"] + for update_call in writer.daily_tag_spend_update_queue.add_update.call_args_list: + transaction_dict = update_call[1]["update"] # Each transaction should have one key with the format tag_date_api_key_model_provider for key, transaction in transaction_dict.items(): assert (