scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 359

refactor(proxy): fetch entity rollups with a companion query instead of extending the main one

Browse source 
The entity-as-extra-GROUPING-bit approach made the bitmask layout
mode-dependent: the same constant meant (date) for normal rows and
(date, entity) for entity rows, disambiguated by masking. Split it out:
the shared WHERE builder feeds both the untouched main query and a small
per-entity rollup query keyed by GROUPING(api_key), run concurrently, and
a fold writes breakdown.entities onto the built response.
This commit is contained in:
ryan-crabbe-berri 2026-08-11 13:16:19 -07:00
parent 67ac9edf67
commit fa677f395d
2 changed files with 262 additions and 126 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

303
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -121,7 +121,6 @@ class _AggregatedSpendData(TypedDict):
class _GroupingSetsRow(SimpleNamespace):
date: str
entity_id: str | None = None
api_key: str | None
model: str | None
model_group: str | None
@ -143,6 +142,11 @@ class _GroupingSetsRow(SimpleNamespace):
failed_requests: int | None
class _EntityRollupRow(_GroupingSetsRow):
entity_id: str | None
api_key_rolled: int
def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float:
"""Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled.
@ -556,35 +560,17 @@ def _build_where_conditions(
return where_conditions
def _build_aggregated_sql_query(
def _build_aggregated_where_clause(
*,
table_name: str,
entity_id_field: str,
entity_id: str | list[str] | None,
start_date: str,
end_date: str,
adjusted_start: str,
adjusted_end: str,
model: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
exclude_entity_ids: list[str] | None = None,
timezone_offset_minutes: int | None = None,
entity_breakdown_field: str | None = None,
exclude_entity_ids: list[str] | None, # mutable-ok: filter union shared with the paginated path
) -> tuple[str, list[str]]:
"""Build a parameterized SQL GROUP BY query for aggregated daily activity.
Groups by (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().
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes)
"""Build the WHERE clause and $N params shared by the aggregated queries."""
sql_conditions: Final[list[str]] = []
sql_params: Final[list[str]] = []
p = 1 # parameter index (1-based for PostgreSQL $N placeholders)
@ -637,7 +623,54 @@ def _build_aggregated_sql_query(
sql_params.append(api_key)
p += 1
where_clause: Final = " AND ".join(sql_conditions)
return " AND ".join(sql_conditions), sql_params
def _ptu_flat_cost_select(table_name: str) -> str:
"""Only LiteLLM_DailyTeamSpend carries ptu_flat_cost; other daily tables emit a
constant zero so the SpendMetrics.flat_cost response shape stays uniform."""
if table_name == "litellm_dailyteamspend":
return "SUM(ptu_flat_cost)::float AS ptu_flat_cost"
return "0::float AS ptu_flat_cost"
def _build_aggregated_sql_query(
*,
table_name: str,
entity_id_field: str,
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
start_date: str,
end_date: str,
model: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
timezone_offset_minutes: int | None = None,
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
"""Build a parameterized SQL GROUP BY query for aggregated daily activity.
Groups by (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().
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes)
where_clause, sql_params = _build_aggregated_where_clause(
entity_id_field=entity_id_field,
entity_id=entity_id,
adjusted_start=adjusted_start,
adjusted_end=adjusted_end,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
)
# Postgres computes every rollup level the response needs — per-date
# totals, per-(date, model), per-(date, model, api_key), per-provider,
@ -651,24 +684,8 @@ def _build_aggregated_sql_query(
# total_successful_requests metadata they feed) once the admin UI reads SGR
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
# api_requests rollups are still served from here.
#
# Only LiteLLM_DailyTeamSpend carries ptu_flat_cost; other daily tables emit a
# constant zero so the SpendMetrics.flat_cost response shape stays uniform.
ptu_flat_cost_select: Final = (
"SUM(ptu_flat_cost)::float AS ptu_flat_cost"
if table_name == "litellm_dailyteamspend"
else "0::float AS ptu_flat_cost"
)
entity_select: Final = f'"{entity_breakdown_field}" AS entity_id,' if entity_breakdown_field else ""
entity_grouping_arg: Final = f'"{entity_breakdown_field}", ' if entity_breakdown_field else ""
entity_grouping_sets: Final = (
f'(date, "{entity_breakdown_field}"), (date, "{entity_breakdown_field}", api_key),'
if entity_breakdown_field
else ""
)
sql_query: Final = f"""
SELECT
{entity_select}
date,
api_key,
model,
@ -676,11 +693,11 @@ def _build_aggregated_sql_query(
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
GROUPING({entity_grouping_arg}date, api_key, model, COALESCE(NULLIF(model_group, ''), model),
GROUPING(date, api_key, model, COALESCE(NULLIF(model_group, ''), model),
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
SUM(spend)::float AS spend,
{ptu_flat_cost_select},
{_ptu_flat_cost_select(table_name)},
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens,
@ -695,7 +712,6 @@ def _build_aggregated_sql_query(
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY GROUPING SETS (
{entity_grouping_sets}
(date),
(date, api_key),
(date, model),
@ -715,6 +731,70 @@ def _build_aggregated_sql_query(
return sql_query, sql_params
def _build_entity_rollup_sql_query(
*,
table_name: str,
entity_id_field: str,
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
start_date: str,
end_date: str,
model: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
timezone_offset_minutes: int | None = None,
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
"""Per-entity companion to _build_aggregated_sql_query.
Two rollup levels over the same WHERE clause — (date, entity) and
(date, entity, api_key) — told apart by GROUPING(api_key): 1 when the
api_key column is rolled up, 0 when it is part of the key.
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes)
where_clause, sql_params = _build_aggregated_where_clause(
entity_id_field=entity_id_field,
entity_id=entity_id,
adjusted_start=adjusted_start,
adjusted_end=adjusted_end,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
)
sql_query: Final = f"""
SELECT
"{entity_id_field}" AS entity_id,
date,
api_key,
GROUPING(api_key) AS api_key_rolled,
SUM(spend)::float AS spend,
{_ptu_flat_cost_select(table_name)},
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens,
SUM(cache_creation_input_tokens)::bigint AS cache_creation_input_tokens,
SUM(compression_saved_tokens)::bigint AS compression_saved_tokens,
SUM(compression_savings_spend)::float AS compression_savings_spend,
SUM(prompt_caching_savings_spend)::float AS prompt_caching_savings_spend,
SUM(autorouter_savings_spend)::float AS autorouter_savings_spend,
SUM(api_requests)::bigint AS api_requests,
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date, "{entity_id_field}"),
(date, "{entity_id_field}", api_key)
)
"""
return sql_query, sql_params
def _aggregate_spend_records_sync(
*,
records: Sequence[DailySpendRecord],
@ -815,10 +895,6 @@ _GROUP_DATE_MCP_API_KEY: Final = 29 # 0b0011101
_GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110
_GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110
# 0b10000000 — set when the entity column (leftmost GROUPING argument in
# entity-breakdown mode) is rolled up; clear only on the (date, entity[, api_key]) sets
_ENTITY_ROLLED_BIT: Final = 128
def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
"""Build a SpendMetrics directly from one already-aggregated rollup row.
@ -855,8 +931,6 @@ def _aggregate_grouping_sets_records_sync(
*,
records: Sequence[_GroupingSetsRow],
api_key_metadata: Mapping[str, _KeyMetadataDict],
entity_breakdown: bool = False,
entity_metadata_field: Mapping[str, dict[str, object]] | None = None, # mutable-ok: shared field shape
) -> _AggregatedSpendData:
"""Build the response from rollup rows produced by the GROUPING SETS query.
@ -896,41 +970,11 @@ def _aggregate_grouping_sets_records_sync(
metrics=metrics, metadata=_key_metadata(api_key_metadata, api_key)
)
def ensure_entity(
entities: dict[str, MetricWithMetadata], # mutable-ok: writes into breakdown.entities
entity_id: str,
) -> MetricWithMetadata:
existing: MetricWithMetadata | None = entities.get(entity_id)
if existing is not None:
return existing
created: Final = MetricWithMetadata(
metrics=SpendMetrics(),
metadata=(entity_metadata_field.get(entity_id, {}) if entity_metadata_field else {}),
)
entities[entity_id] = created
return created
def assign_entity_rollup(record: _GroupingSetsRow, metrics: SpendMetrics) -> None:
entity_id: Final = record.entity_id or "Unassigned"
entity_bucket: Final = ensure_entity(ensure_date(record.date)["breakdown"].entities, entity_id)
if record.group_level == _GROUP_DATE:
entity_bucket.metrics = metrics
return
if record.api_key and record.api_key != PTU_SENTINEL_API_KEY:
entity_bucket.api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
metrics=metrics, metadata=_key_metadata(api_key_metadata, record.api_key)
)
for record in records:
level = record.group_level & ~_ENTITY_ROLLED_BIT if entity_breakdown else record.group_level
level = record.group_level
metrics = _record_to_spend_metrics(record)
is_ptu_sentinel = record.api_key == PTU_SENTINEL_API_KEY
if entity_breakdown and not (record.group_level & _ENTITY_ROLLED_BIT):
if level in (_GROUP_DATE, _GROUP_DATE_API_KEY):
assign_entity_rollup(record, metrics)
continue
if level == _GROUP_GRAND_TOTAL:
total_metrics = metrics
continue
@ -1013,8 +1057,6 @@ async def _aggregate_grouping_sets_records(
*,
prisma_client: PrismaClient,
records: Sequence[_GroupingSetsRow],
entity_breakdown: bool = False,
entity_metadata_field: Mapping[str, dict[str, object]] | None = None, # mutable-ok: shared field shape
) -> _AggregatedSpendData:
"""Async wrapper: fetch api_key_metadata, then dispatch on a worker thread."""
api_keys: Final[set[str]] = {r.api_key for r in records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY}
@ -1027,8 +1069,6 @@ async def _aggregate_grouping_sets_records(
_aggregate_grouping_sets_records_sync,
records=records,
api_key_metadata=api_key_metadata,
entity_breakdown=entity_breakdown,
entity_metadata_field=entity_metadata_field,
)
@ -1156,6 +1196,42 @@ async def get_daily_activity(
)
def _fold_entity_rollups_sync(
*,
results: Sequence[DailySpendData],
entity_rows: Sequence[_EntityRollupRow],
api_key_metadata: Mapping[str, _KeyMetadataDict],
entity_metadata_field: Mapping[str, dict[str, object]] | None, # mutable-ok: shared field shape
) -> None:
"""Write breakdown.entities onto the already-built per-day results."""
by_date: Final = {day.date.strftime("%Y-%m-%d"): day for day in results} # mutable-ok: local fold index
for row in entity_rows:
day = by_date.get(row.date)
if day is None:
continue
entities = day.breakdown.entities
entity_id = row.entity_id or "Unassigned"
bucket = entities.get(entity_id)
if bucket is None:
bucket = MetricWithMetadata(
metrics=SpendMetrics(),
metadata=(
entity_metadata_field.get(entity_id, {}) if entity_metadata_field else {}
), # mutable-ok: pydantic metadata payload
)
entities[entity_id] = bucket
metrics = _record_to_spend_metrics(row)
if row.api_key_rolled:
bucket.metrics = metrics
elif row.api_key and row.api_key != PTU_SENTINEL_API_KEY:
bucket.api_key_breakdown[row.api_key] = KeyMetricWithMetadata(
metrics=metrics, metadata=_key_metadata(api_key_metadata, row.api_key)
)
async def get_daily_activity_aggregated(
prisma_client: PrismaClient | None,
table_name: str,
@ -1176,8 +1252,8 @@ async def get_daily_activity_aggregated(
all individual rows into Python. This collapses rows across entities
(users/teams/orgs), reducing ~150k rows to ~2-3k grouped rows.
include_entity_breakdown adds per-entity rollup levels so
`breakdown.entities` is populated, as entity-scoped views like Team Usage need.
include_entity_breakdown runs a small companion rollup query and folds
`breakdown.entities` onto the response, as entity-scoped views like Team Usage need.
Matches the response model of the paginated endpoint so the UI does not need to transform.
"""
@ -1204,25 +1280,62 @@ async def get_daily_activity_aggregated(
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
timezone_offset_minutes=timezone_offset_minutes,
entity_breakdown_field=entity_id_field if include_entity_breakdown else None,
)
# Execute GROUPING SETS query — returns one row per rollup level.
rows = await prisma_client.db.query_raw(sql_query, *sql_params)
if rows is None:
rows = []
entity_query: Final = (
_build_entity_rollup_sql_query(
table_name=table_name,
entity_id_field=entity_id_field,
entity_id=entity_id,
start_date=start_date,
end_date=end_date,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
timezone_offset_minutes=timezone_offset_minutes,
)
if include_entity_breakdown
else None
)
records: Final = [_GroupingSetsRow(**row) for row in rows]
# Execute the GROUPING SETS query (one row per rollup level), alongside
# the per-entity companion rollup when the caller wants entities.
raw_rows, raw_entity_rows = (
await asyncio.gather(
prisma_client.db.query_raw(sql_query, *sql_params),
prisma_client.db.query_raw(entity_query[0], *entity_query[1]),
)
if entity_query is not None
else (await prisma_client.db.query_raw(sql_query, *sql_params), None)
)
records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or [])]
# The grouping-sets dispatcher places each row directly in its bucket
# using the row's GROUPING() bitmask. No Python-side summing needed.
aggregated: Final = await _aggregate_grouping_sets_records(
prisma_client=prisma_client,
records=records,
entity_breakdown=include_entity_breakdown,
entity_metadata_field=entity_metadata_field,
)
if raw_entity_rows:
entity_records: Final = tuple(_EntityRollupRow(**row) for row in raw_entity_rows)
entity_api_keys: Final = frozenset(
r.api_key for r in entity_records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY
)
entity_key_metadata: Final = (
await get_api_key_metadata(prisma_client, set(entity_api_keys))
if entity_api_keys
else {} # mutable-ok: helper takes set; {} matches its return
)
await asyncio.to_thread(
_fold_entity_rollups_sync,
results=aggregated["results"],
entity_rows=entity_records,
api_key_metadata=entity_key_metadata,
entity_metadata_field=entity_metadata_field,
)
return SpendAnalyticsPaginatedResponse(
results=aggregated["results"],
metadata=DailySpendMetadata(

85
tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py
View file

@ -1721,11 +1721,16 @@ class TestFlagIsNotReadOnTheHotPath:
assert reads > 0
def test_build_aggregated_sql_query_entity_breakdown_and_api_key_list():
"""entity_breakdown_field adds the two entity rollup sets with the entity
column as the most-significant GROUPING bit; a list api_key filter becomes
a parameterized IN, and an empty list must match nothing (not everything)."""
sql, params = _build_aggregated_sql_query(
def test_entity_rollup_sql_query_and_api_key_list_filter():
"""The entity rollup companion query keeps its own two grouping sets keyed
by GROUPING(api_key), shares the WHERE builder (list api_key becomes a
parameterized IN, an empty list must match nothing), and the main
aggregated query stays entity-free."""
from litellm.proxy.management_endpoints.common_daily_activity import (
_build_entity_rollup_sql_query,
)
sql, params = _build_entity_rollup_sql_query(
table_name="litellm_dailyteamspend",
entity_id_field="team_id",
entity_id=None,
@ -1733,12 +1738,13 @@ def test_build_aggregated_sql_query_entity_breakdown_and_api_key_list():
end_date="2024-01-31",
model=None,
api_key=["key-1", "key-2"],
entity_breakdown_field="team_id",
)
assert '"team_id" AS entity_id' in sql
assert 'GROUPING("team_id", date' in sql
assert '(date, "team_id"), (date, "team_id", api_key),' in sql
assert "GROUPING(api_key) AS api_key_rolled" in sql
assert '(date, "team_id"),' in sql
assert '(date, "team_id", api_key)' in sql
assert "api_key IN ($3, $4)" in sql
assert "SUM(ptu_flat_cost)::float" in sql
assert params == ["2024-01-01", "2024-01-31", "key-1", "key-2"]
plain_sql, _ = _build_aggregated_sql_query(
@ -1768,14 +1774,13 @@ def test_build_aggregated_sql_query_entity_breakdown_and_api_key_list():
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_with_entity_breakdown():
"""include_entity_breakdown must populate breakdown.entities from the
(date, entity) / (date, entity, api_key) rollups while every pre-existing
rollup still dispatches correctly with the entity bit (128) set."""
"""include_entity_breakdown must run the companion entity rollup query and
fold breakdown.entities onto the response, without disturbing the main
query's rollup dispatch."""
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
base = {
"entity_id": None,
"model": None,
"model_group": None,
"custom_llm_provider": None,
@ -1794,24 +1799,39 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown():
"api_requests": 0,
"successful_requests": 0,
}
mock_rows = [
# () grand total: 127 + entity bit
{**base, "date": None, "group_level": 255, "spend": 18.0},
# (date): 63 + entity bit
{**base, "date": "2024-01-01", "group_level": 191, "spend": 18.0},
# (date, entity): entity bit clear
{**base, "date": "2024-01-01", "entity_id": "team-a", "group_level": 63, "spend": 12.0},
{**base, "date": "2024-01-01", "entity_id": "team-b", "group_level": 63, "spend": 6.0},
# (date, entity, api_key): entity bit clear
{**base, "date": "2024-01-01", "entity_id": "team-a", "api_key": "key-1", "group_level": 31, "spend": 12.0},
{**base, "date": "2024-01-01", "entity_id": "team-b", "api_key": "key-2", "group_level": 31, "spend": 6.0},
# (date, model): 47 + entity bit
{**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 175, "spend": 18.0},
# (date, api_key): 31 + entity bit
{**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 159, "spend": 12.0},
main_rows = [
{**base, "date": None, "group_level": 127, "spend": 18.0},
{**base, "date": "2024-01-01", "group_level": 63, "spend": 18.0},
{**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 47, "spend": 18.0},
{**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 31, "spend": 12.0},
]
entity_base = {
key: value
for key, value in base.items()
if key not in ("model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint")
}
entity_rows = [
{**entity_base, "date": "2024-01-01", "entity_id": "team-a", "api_key_rolled": 1, "spend": 12.0},
{**entity_base, "date": "2024-01-01", "entity_id": "team-b", "api_key_rolled": 1, "spend": 6.0},
{
**entity_base,
"date": "2024-01-01",
"entity_id": "team-a",
"api_key": "key-1",
"api_key_rolled": 0,
"spend": 12.0,
},
{
**entity_base,
"date": "2024-01-01",
"entity_id": "team-b",
"api_key": "key-2",
"api_key_rolled": 0,
"spend": 6.0,
},
]
mock_prisma.db.query_raw = AsyncMock(return_value=mock_rows)
mock_prisma.db.query_raw = AsyncMock(side_effect=[main_rows, entity_rows])
mock_prisma.db.litellm_verificationtoken = MagicMock()
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
@ -1828,9 +1848,12 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown():
include_entity_breakdown=True,
)
sql = mock_prisma.db.query_raw.call_args[0][0]
assert '"team_id" AS entity_id' in sql
assert '(date, "team_id"), (date, "team_id", api_key),' in sql
assert mock_prisma.db.query_raw.call_count == 2
main_sql = mock_prisma.db.query_raw.call_args_list[0][0][0]
entity_sql = mock_prisma.db.query_raw.call_args_list[1][0][0]
assert "entity_id" not in main_sql
assert '"team_id" AS entity_id' in entity_sql
assert '(date, "team_id"),' in entity_sql
assert result.metadata.total_spend == 18.0
assert len(result.results) == 1