fix(proxy): include model group in daily usage aggregation

2026-09-11 22:51:28 +00:00 · 2026-08-06 18:34:14 +05:30 · 2026-08-06 18:34:14 +05:30 · c3ddccd256
commit c3ddccd256
parent b0fac57fe4
7 changed files with 151 additions and 23 deletions
--- 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
@ -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");
--- 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 ()),
--- 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"]
            )
--- 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])
--- 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])
--- 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"'
    )


--- 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") == [""]