tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-07 04:24:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(proxy): restore per-entity breakdown in aggregated daily activity endpoint

Browse Source 
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.
This commit is contained in:
michelligabriele
2026-03-12 20:55:56 +01:00
parent 7f0ec1e3d8
commit 9c3fab24ad
2 changed files with 126 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/management_endpoints/common_daily_activity.py
+7 -11
View File
@@ -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(
tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py
+119
View File
@@ -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