fix(proxy): disambiguate daily spend queue keys

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");

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"