scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge d43af0ece0 into 1c61c2606e

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-09-12 16:50:34 +00:00 • committed by GitHub
commit cec422ae67
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 154 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -742,15 +742,23 @@ def _build_aggregated_sql_query(
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
# api_requests rollups are still served from here.
sql_query: Final = f"""
WITH top_api_keys AS (
SELECT api_key AS top_api_key
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY api_key
ORDER BY SUM(spend) DESC, api_key
LIMIT {_MAX_API_KEYS_IN_BREAKDOWN}
)
SELECT
date,
api_key,
tk.top_api_key AS api_key,
model,
COALESCE(NULLIF(model_group, ''), model) AS model_group,
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
GROUPING(date, api_key, model, COALESCE(NULLIF(model_group, ''), model),
GROUPING(date, tk.top_api_key, model, COALESCE(NULLIF(model_group, ''), model),
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
SUM(spend)::float AS spend,
@ -767,21 +775,22 @@ def _build_aggregated_sql_query(
SUM(api_requests)::bigint AS api_requests,
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests
FROM "{pg_table}"
FROM "{pg_table}" t
LEFT JOIN top_api_keys tk ON tk.top_api_key = t.api_key
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date),
(date, api_key),
(date, tk.top_api_key),
(date, model),
(date, model, api_key),
(date, model, tk.top_api_key),
(date, COALESCE(NULLIF(model_group, ''), model)),
(date, COALESCE(NULLIF(model_group, ''), model), api_key),
(date, COALESCE(NULLIF(model_group, ''), model), tk.top_api_key),
(date, custom_llm_provider),
(date, custom_llm_provider, api_key),
(date, custom_llm_provider, tk.top_api_key),
(date, mcp_namespaced_tool_name),
(date, mcp_namespaced_tool_name, api_key),
(date, mcp_namespaced_tool_name, tk.top_api_key),
(date, endpoint),
(date, endpoint, api_key),
(date, endpoint, tk.top_api_key),
()
)
"""
@ -959,6 +968,11 @@ _GROUP_DATE_MCP_API_KEY: Final = 29 # 0b0011101
_GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110
_GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110
# Six of the thirteen grouping sets in _build_aggregated_sql_query are keyed on api_key, so
# without a cap the result grows with every distinct key and the prisma query engine OOMs
# buffering it. Keys outside the top N fall into a NULL api_key bucket the dispatcher skips.
_MAX_API_KEYS_IN_BREAKDOWN: Final = 100
def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
"""Build a SpendMetrics directly from one already-aggregated rollup row.

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

@ -1,14 +1,16 @@
import re
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import psycopg
import pytest
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
from psycopg.rows import dict_row
from pytest_postgresql import factories
from litellm.proxy.management_endpoints.common_daily_activity import (
_MAX_API_KEYS_IN_BREAKDOWN,
_adjust_dates_for_timezone,
_build_aggregated_sql_query,
_build_entity_rollup_sql_query,
@ -19,6 +21,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
get_daily_activity_aggregated,
update_metrics,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
from litellm.types.proxy.management_endpoints.common_daily_activity import (
DailySpendMetadata,
SpendMetrics,
@ -472,9 +475,7 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash(
return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")]
)
mock_prisma.db.query_raw = AsyncMock(
return_value=[
{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}
]
return_value=[{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}]
)
result = await get_api_key_metadata(
@ -1236,10 +1237,10 @@ class TestBuildAggregatedSqlQuery:
fallback = "COALESCE(NULLIF(model_group, ''), model)"
assert f"{fallback} AS model_group" in normalized
assert (
f"GROUPING(date, api_key, model, {fallback}, "
f"GROUPING(date, tk.top_api_key, model, {fallback}, "
"custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level" in normalized
)
assert f"(date, {fallback}), (date, {fallback}, api_key)," in normalized
assert f"(date, {fallback}), (date, {fallback}, tk.top_api_key)," in normalized
assert "(date, model_group)" not in normalized
assert "COALESCE(model_group, model)" not in normalized
@ -1360,6 +1361,128 @@ async def test_get_daily_activity_aggregated_empty_result_set():
assert result.metadata.total_compression_saved_tokens == 0
_aggregated_postgresql_proc: Final = factories.postgresql_proc()
_aggregated_postgresql: Final = factories.postgresql("_aggregated_postgresql_proc")
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_bounds_api_key_rollups(
_aggregated_postgresql: psycopg.Connection,
):
"""Exercise the GROUPING SETS query against real Postgres with more keys than the cap.
key-004 and key-005 tie on spend exactly at the _MAX_API_KEYS_IN_BREAKDOWN
cutoff; the deterministic api_key tiebreaker must keep key-004 and drop
key-005. Excluded keys still count toward the totals via the grand-total
and date-level rollup rows.
"""
conn: Final = _aggregated_postgresql
with conn.cursor() as cur:
cur.execute(
"""
CREATE TABLE "LiteLLM_DailyUserSpend" (
id TEXT PRIMARY KEY,
user_id TEXT,
date TEXT NOT NULL,
api_key TEXT NOT NULL,
model TEXT,
model_group TEXT,
custom_llm_provider TEXT,
mcp_namespaced_tool_name TEXT,
endpoint TEXT,
prompt_tokens BIGINT DEFAULT 0,
completion_tokens BIGINT DEFAULT 0,
cache_read_input_tokens BIGINT DEFAULT 0,
cache_creation_input_tokens BIGINT DEFAULT 0,
compression_saved_tokens BIGINT DEFAULT 0,
compression_savings_spend DOUBLE PRECISION DEFAULT 0,
prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0,
gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0,
autorouter_savings_spend DOUBLE PRECISION DEFAULT 0,
spend DOUBLE PRECISION DEFAULT 0,
api_requests BIGINT DEFAULT 0,
successful_requests BIGINT DEFAULT 0,
failed_requests BIGINT DEFAULT 0
)
"""
)
cur.executemany(
"""
INSERT INTO "LiteLLM_DailyUserSpend"
(id, user_id, date, api_key, model, model_group, custom_llm_provider,
endpoint, prompt_tokens, spend, api_requests, successful_requests)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""",
[
(
f"row-{i:03d}",
f"user-{i:03d}",
"2026-06-01",
f"key-{i:03d}",
"gpt-5",
"",
"openai",
"/v1/chat/completions",
10,
6.0 if i == 4 else float(i + 1),
1,
1,
)
for i in range(_MAX_API_KEYS_IN_BREAKDOWN + 5)
],
)
conn.commit()
row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim
async def query_raw(sql: str, *params: str) -> list[dict[str, object]]:
converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql)
with conn.cursor(row_factory=dict_row) as cur:
cur.execute(
converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query
{f"p{i}": v for i, v in enumerate(params, start=1)},
)
rows: Final = cur.fetchall()
row_counts.append(len(rows))
return rows
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = query_raw
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.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="2026-06-01",
end_date="2026-06-01",
model=None,
api_key=None,
)
# 13 grouping sets: 7 key-less sets emit one row each; the 6 api_key sets
# emit the top N keys plus one NULL bucket for the excluded keys.
assert row_counts == [7 + 6 * (_MAX_API_KEYS_IN_BREAKDOWN + 1)]
# Totals still cover every key, capped or not.
assert result.metadata.total_spend == pytest.approx(5566.0)
assert result.metadata.total_api_requests == _MAX_API_KEYS_IN_BREAKDOWN + 5
expected_top: Final = {f"key-{i:03d}" for i in range(6, 105)} | {"key-004"}
day: Final = result.results[0]
assert len(day.breakdown.api_keys) == _MAX_API_KEYS_IN_BREAKDOWN
assert set(day.breakdown.api_keys) == expected_top
assert day.breakdown.api_keys["key-004"].metrics.spend == 6.0
assert "key-005" not in day.breakdown.api_keys
assert day.breakdown.models["gpt-5"].metrics.spend == pytest.approx(5566.0)
assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == expected_top
def _no_spend_record():
"""A rollup row for a key with no spend, where SUM() returns NULL (None)."""
return SimpleNamespace(