From 0017d5f1dbc316cb198a49d8cf46af5ad87bd35e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 May 2025 22:03:18 -0700 Subject: [PATCH] build(ui/): Allow empty values in daily agg table + reintroduce 'unassigned' teams in spend tracking --- .../migration.sql | 9 +++ .../litellm_proxy_extras/schema.prisma | 6 +- litellm/proxy/db/db_spend_update_writer.py | 4 +- .../common_daily_activity.py | 22 +++--- litellm/proxy/schema.prisma | 6 +- schema.prisma | 12 ++-- .../proxy/db/test_db_spend_update_writer.py | 70 +++++++++++++++++++ .../src/components/new_usage.tsx | 4 +- 8 files changed, 107 insertions(+), 26 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20250526154401_allow_null_entity_id/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250526154401_allow_null_entity_id/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250526154401_allow_null_entity_id/migration.sql new file mode 100644 index 0000000000..0746656a26 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250526154401_allow_null_entity_id/migration.sql @@ -0,0 +1,9 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DailyTagSpend" ALTER COLUMN "tag" DROP NOT NULL; + +-- AlterTable +ALTER TABLE "LiteLLM_DailyTeamSpend" ALTER COLUMN "team_id" DROP NOT NULL; + +-- AlterTable +ALTER TABLE "LiteLLM_DailyUserSpend" ALTER COLUMN "user_id" DROP NOT NULL; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 58064abd1d..71f10f9473 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -356,7 +356,7 @@ model LiteLLM_AuditLog { // Track daily user spend metrics per model and key model LiteLLM_DailyUserSpend { id String @id @default(uuid()) - user_id String + user_id String? date String api_key String model String @@ -383,7 +383,7 @@ model LiteLLM_DailyUserSpend { // Track daily team spend metrics per model and key model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) - team_id String + team_id String? date String api_key String model String @@ -410,7 +410,7 @@ model LiteLLM_DailyTeamSpend { // Track daily team spend metrics per model and key model LiteLLM_DailyTagSpend { id String @id @default(uuid()) - tag String + tag String? date String api_key String model String diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 57b5280b50..9a7b14272d 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -444,7 +444,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_team_spend_update_transactions, ) - + daily_tag_spend_update_transactions = ( await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() ) @@ -859,8 +859,6 @@ class DBSpendUpdateWriter: async with prisma_client.db.batch_() as batcher: for _, transaction in transactions_to_process.items(): entity_id = transaction.get(entity_id_field) - if not entity_id: - continue # Construct the where clause dynamically where_clause = { diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index d782adf85c..5efd31b232 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -86,17 +86,19 @@ def update_breakdown_metrics( # Update entity-specific metrics if entity_id_field is provided if entity_id_field: entity_value = getattr(record, entity_id_field, None) - if entity_value: - if entity_value not in breakdown.entities: - breakdown.entities[entity_value] = MetricWithMetadata( - metrics=SpendMetrics(), - metadata=entity_metadata_field.get(entity_value, {}) - if entity_metadata_field - else {}, - ) - breakdown.entities[entity_value].metrics = update_metrics( - breakdown.entities[entity_value].metrics, record + entity_value = ( + entity_value if entity_value else "Unassigned" + ) # allow for null entity_id_field + if entity_value not in breakdown.entities: + breakdown.entities[entity_value] = MetricWithMetadata( + metrics=SpendMetrics(), + metadata=entity_metadata_field.get(entity_value, {}) + if entity_metadata_field + else {}, ) + breakdown.entities[entity_value].metrics = update_metrics( + breakdown.entities[entity_value].metrics, record + ) return breakdown diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 58064abd1d..71f10f9473 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -356,7 +356,7 @@ model LiteLLM_AuditLog { // Track daily user spend metrics per model and key model LiteLLM_DailyUserSpend { id String @id @default(uuid()) - user_id String + user_id String? date String api_key String model String @@ -383,7 +383,7 @@ model LiteLLM_DailyUserSpend { // Track daily team spend metrics per model and key model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) - team_id String + team_id String? date String api_key String model String @@ -410,7 +410,7 @@ model LiteLLM_DailyTeamSpend { // Track daily team spend metrics per model and key model LiteLLM_DailyTagSpend { id String @id @default(uuid()) - tag String + tag String? date String api_key String model String diff --git a/schema.prisma b/schema.prisma index 9e0e763803..71f10f9473 100644 --- a/schema.prisma +++ b/schema.prisma @@ -204,7 +204,7 @@ model LiteLLM_VerificationToken { model_spend Json @default("{}") model_max_budget Json @default("{}") budget_id String? - organization_id String? // deprecated param, as user can be part of multiple organizations. Check LiteLLM_OrganizationMembership instead. + organization_id String? object_permission_id String? created_at DateTime? @default(now()) @map("created_at") created_by String? @@ -356,7 +356,7 @@ model LiteLLM_AuditLog { // Track daily user spend metrics per model and key model LiteLLM_DailyUserSpend { id String @id @default(uuid()) - user_id String + user_id String? date String api_key String model String @@ -383,7 +383,7 @@ model LiteLLM_DailyUserSpend { // Track daily team spend metrics per model and key model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) - team_id String + team_id String? date String api_key String model String @@ -410,7 +410,7 @@ model LiteLLM_DailyTeamSpend { // Track daily team spend metrics per model and key model LiteLLM_DailyTagSpend { id String @id @default(uuid()) - tag String + tag String? date String api_key String model String @@ -463,9 +463,9 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) } -model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use managed files +model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the id String @id @default(uuid()) - unified_object_id String @unique // The base64 encoded unified object ID + unified_object_id String @unique // The base64 encoded unified file ID model_object_id String @unique // the id returned by the backend API provider file_object Json // Stores the OpenAIFileObject file_purpose String // either 'batch' or 'fine-tune' 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 02b94c44ec..dfa075cbfc 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 @@ -65,3 +65,73 @@ async def test_daily_spend_tracking_with_disabled_spend_logs(): assert call_args["payload"]["spend"] == 0.1 assert call_args["payload"]["model"] == "gpt-4" assert call_args["payload"]["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_update_daily_spend_with_null_entity_id(): + """ + Test that table.upsert is called even when entity_id is null + + Ensures 'global view' has all daily spend transactions + """ + # Setup + mock_prisma_client = MagicMock() + mock_batcher = MagicMock() + mock_table = MagicMock() + mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher + mock_batcher.litellm_dailyuserspend = mock_table + + # Create a transaction with null entity_id + daily_spend_transactions = { + "test_key": { + "user_id": None, # null entity_id + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + } + + # Call the method + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=1, + prisma_client=mock_prisma_client, + proxy_logging_obj=MagicMock(), + daily_spend_transactions=daily_spend_transactions, + entity_type="user", + entity_id_field="user_id", + table_name="litellm_dailyuserspend", + unique_constraint_name="user_id_date_api_key_model_custom_llm_provider", + ) + + # Verify that table.upsert was called + mock_table.upsert.assert_called_once() + + # Verify the where clause contains null entity_id + call_args = mock_table.upsert.call_args[1] + where_clause = call_args["where"]["user_id_date_api_key_model_custom_llm_provider"] + assert where_clause["user_id"] is None + assert where_clause["date"] == "2024-01-01" + assert where_clause["api_key"] == "test-api-key" + assert where_clause["model"] == "gpt-4" + assert where_clause["custom_llm_provider"] == "openai" + + # Verify the create data contains null entity_id + create_data = call_args["data"]["create"] + assert create_data["user_id"] is None + assert create_data["date"] == "2024-01-01" + assert create_data["api_key"] == "test-api-key" + assert create_data["model"] == "gpt-4" + assert create_data["custom_llm_provider"] == "openai" + assert create_data["prompt_tokens"] == 10 + assert create_data["completion_tokens"] == 20 + assert create_data["spend"] == 0.1 + assert create_data["api_requests"] == 1 + assert create_data["successful_requests"] == 1 + assert create_data["failed_requests"] == 0 diff --git a/ui/litellm-dashboard/src/components/new_usage.tsx b/ui/litellm-dashboard/src/components/new_usage.tsx index b6378cd9cc..c5917f96b2 100644 --- a/ui/litellm-dashboard/src/components/new_usage.tsx +++ b/ui/litellm-dashboard/src/components/new_usage.tsx @@ -255,9 +255,11 @@ const NewUsagePage: React.FC = ({ return (
+ {all_admin_roles.includes(userRole || "") ? - This is the new usage dashboard.
You may see empty data, as these use new aggregate tables to allow UI to work at 1M+ spend logs. To access the old dashboard, go to Experimental {'>'} Old Usage. + Note: If you see key/model-level inconsistencies between Global View and Team Usage, it's because the Global View was missing spend when user_id = null, prior to v1.71.2. Learn more here. + : null} {all_admin_roles.includes(userRole || "") ? Global Usage : Your Usage}