diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 6c9289e3ff6..7ae1a1eff83 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1054,6 +1054,7 @@ async def _update_daily_spend( x[1].get("api_key") or "", x[1].get("model") or "", x[1].get("custom_llm_provider") or "", + x[1].get("mcp_namespaced_tool_name") or "", ), )[:BATCH_SIZE] ) 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 181d21b44f6..0ce9f625f2e 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 @@ -306,6 +306,20 @@ async def test_update_daily_spend_with_none_values_in_sorting_fields(): "successful_requests": 1, "failed_requests": 0, }, + "key6": { + "user_id": "user6", + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": None, # None mcp_namespaced_tool_name + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + }, } # Call the method - this should not raise TypeError @@ -320,8 +334,8 @@ async def test_update_daily_spend_with_none_values_in_sorting_fields(): unique_constraint_name="user_id_date_api_key_model_custom_llm_provider", ) - # Verify that table.upsert was called (should be called 5 times, once for each transaction) - assert mock_table.upsert.call_count == 5 + # Verify that table.upsert was called (should be called 6 times, once for each transaction) + assert mock_table.upsert.call_count == 6 # Tag Spend Tracking Tests