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"