From 9c3fab24adb3a36fcfa02dc2dfa4f87d7893d0f4 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 12 Mar 2026 20:55:56 +0100 Subject: [PATCH] fix(proxy): restore per-entity breakdown in aggregated daily activity endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #21613 optimized the /user/daily/activity/aggregated endpoint by replacing find_many with a SQL GROUP BY query, but omitted entity_id from the SELECT/GROUP BY clauses and hardcoded entity_id_field=None in the call to _aggregate_spend_records. This caused breakdown.entities to always be empty in the response. Restore entity_id in the SQL query and forward entity_id_field and entity_metadata_field to the aggregation step. The GROUP BY performance benefit is preserved — the query still aggregates at the database level instead of fetching all individual rows into Python. --- .../common_daily_activity.py | 18 ++- .../test_common_daily_activity.py | 119 ++++++++++++++++++ 2 files changed, 126 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 011d2f7485..b05a183df0 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -475,10 +475,8 @@ def _build_aggregated_sql_query( ) -> Tuple[str, List[Any]]: """Build a parameterized SQL GROUP BY query for aggregated daily activity. - Groups by (date, api_key, model, model_group, custom_llm_provider, + Groups by (entity_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint) with SUMs on all metric columns. - The entity_id column is intentionally omitted from GROUP BY to collapse - rows across entities — this is where the biggest row reduction comes from. Returns: Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw(). @@ -539,6 +537,7 @@ def _build_aggregated_sql_query( sql_query = f""" SELECT + "{entity_id_field}", date, api_key, model, @@ -556,8 +555,8 @@ def _build_aggregated_sql_query( SUM(failed_requests)::bigint AS failed_requests FROM "{pg_table}" WHERE {where_clause} - GROUP BY date, api_key, model, model_group, custom_llm_provider, - mcp_namespaced_tool_name, endpoint + GROUP BY "{entity_id_field}", date, api_key, model, model_group, + custom_llm_provider, mcp_namespaced_tool_name, endpoint ORDER BY date DESC """ @@ -735,8 +734,7 @@ async def get_daily_activity_aggregated( """Aggregated variant that returns the full result set (no pagination). Uses SQL GROUP BY to aggregate rows in the database rather than fetching - all individual rows into Python. This collapses rows across entities - (users/teams/orgs), reducing ~150k rows to ~2-3k grouped rows. + all individual rows into Python, preserving per-entity granularity. Matches the response model of the paginated endpoint so the UI does not need to transform. """ @@ -773,13 +771,11 @@ async def get_daily_activity_aggregated( # Convert dicts to objects for compatibility with _aggregate_spend_records records = [SimpleNamespace(**row) for row in rows] - # entity_id_field=None skips entity breakdown (entity dimension was - # collapsed by the GROUP BY, so per-entity data is not available) aggregated = await _aggregate_spend_records( prisma_client=prisma_client, records=records, - entity_id_field=None, - entity_metadata_field=None, + entity_id_field=entity_id_field, + entity_metadata_field=entity_metadata_field, ) return SpendAnalyticsPaginatedResponse( diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 54fbac1264..c81792b0b2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -86,6 +86,7 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): # query_raw returns list of dicts (pre-aggregated by GROUP BY) mock_rows = [ { + "user_id": "user-1", "date": "2024-01-01", "endpoint": "/v1/chat/completions", "api_key": "key-1", @@ -103,6 +104,7 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): "failed_requests": 0, }, { + "user_id": "user-1", "date": "2024-01-01", "endpoint": "/v1/embeddings", "api_key": "key-2", @@ -452,6 +454,7 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): # query_raw returns list of dicts (pre-aggregated by GROUP BY) mock_rows = [ { + "user_id": "user-1", "date": "2024-01-01", "endpoint": "/v1/chat/completions", "api_key": "deleted-key-hash", @@ -507,3 +510,119 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): assert key_data.metadata.key_alias == "toto-test-2" assert key_data.metadata.team_id == "69cd4b77-b095-4489-8c46-4f2f31d840a2" assert key_data.metrics.spend == 10.0 + + +@pytest.mark.asyncio +async def test_get_daily_activity_aggregated_preserves_entity_breakdown(): + """Test that aggregated daily activity preserves per-entity breakdown. + + Regression test for PR #21613: the GROUP BY optimization omitted entity_id + from the query and hardcoded entity_id_field=None, causing breakdown.entities + to always be empty. + """ + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + + # query_raw returns rows WITH user_id (entity_id_field included in GROUP BY) + mock_rows = [ + { + "user_id": "user-alice", + "date": "2024-01-01", + "api_key": "key-1", + "model": "gpt-4", + "model_group": None, + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": None, + "endpoint": "/v1/chat/completions", + "spend": 20.0, + "prompt_tokens": 200, + "completion_tokens": 100, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "api_requests": 5, + "successful_requests": 5, + "failed_requests": 0, + }, + { + "user_id": "user-bob", + "date": "2024-01-01", + "api_key": "key-2", + "model": "gpt-4", + "model_group": None, + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": None, + "endpoint": "/v1/chat/completions", + "spend": 8.0, + "prompt_tokens": 80, + "completion_tokens": 40, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "api_requests": 2, + "successful_requests": 2, + "failed_requests": 0, + }, + { + "user_id": None, + "date": "2024-01-01", + "api_key": "key-3", + "model": "gpt-4", + "model_group": None, + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": None, + "endpoint": "/v1/chat/completions", + "spend": 2.0, + "prompt_tokens": 20, + "completion_tokens": 10, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + }, + ] + + mock_prisma.db.query_raw = AsyncMock(return_value=mock_rows) + mock_prisma.db.litellm_verificationtoken = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + result = await get_daily_activity_aggregated( + prisma_client=mock_prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + ) + + # Verify per-entity breakdown is populated (not empty) + daily_data = result.results[0] + assert len(daily_data.breakdown.entities) == 3, ( + "breakdown.entities should contain per-user entries" + ) + + # Verify individual entity metrics + assert "user-alice" in daily_data.breakdown.entities + assert daily_data.breakdown.entities["user-alice"].metrics.spend == 20.0 + assert daily_data.breakdown.entities["user-alice"].metrics.prompt_tokens == 200 + assert daily_data.breakdown.entities["user-alice"].metrics.api_requests == 5 + + assert "user-bob" in daily_data.breakdown.entities + assert daily_data.breakdown.entities["user-bob"].metrics.spend == 8.0 + assert daily_data.breakdown.entities["user-bob"].metrics.prompt_tokens == 80 + assert daily_data.breakdown.entities["user-bob"].metrics.api_requests == 2 + + # Verify NULL entity_id is mapped to "Unassigned" (line 294-296) + assert "Unassigned" in daily_data.breakdown.entities + assert daily_data.breakdown.entities["Unassigned"].metrics.spend == 2.0 + + # Verify per-entity API key breakdown + assert "key-1" in daily_data.breakdown.entities["user-alice"].api_key_breakdown + assert "key-2" in daily_data.breakdown.entities["user-bob"].api_key_breakdown + assert "key-3" in daily_data.breakdown.entities["Unassigned"].api_key_breakdown + + # Verify totals still correct + assert result.metadata.total_spend == 30.0 + assert result.metadata.total_api_requests == 8