scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-30 01:52:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

feat(usage): cursor pagination for the aggregated key breakdown

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-24 02:00:37 +00:00
parent 846a0c1133
commit 529c56adee
18 changed files with 777 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,11 +1,16 @@
import asyncio
import base64
import binascii
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from types import MappingProxyType, SimpleNamespace
from typing import TYPE_CHECKING, Final, Protocol
from fastapi import HTTPException, status
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
@ -178,6 +183,8 @@ class _GroupingSetsRow(_RollupMetricsRow):
endpoint: str | None
group_level: int
distinct_api_keys: int | None
key_spend: float | None
remaining_api_keys: int | None
class _EntityRollupRow(_RollupMetricsRow):
@ -815,6 +822,29 @@ def _key_free_source(pg_table: str, where_clause: str, marker_param: str | None)
) AS key_free_source"""
@dataclass(frozen=True, slots=True)
class KeyPageCursor:
spend: float
api_key: str
def encode(self) -> str:
return base64.urlsafe_b64encode(json.dumps({"spend": self.spend, "api_key": self.api_key}).encode()).decode()
_KEY_PAGE_CURSOR_ADAPTER: Final = TypeAdapter(KeyPageCursor)
def decode_key_page_cursor(raw: str) -> KeyPageCursor | None:
try:
decoded: Final = base64.b64decode(raw.encode(), altchars=b"-_", validate=True)
except (ValueError, binascii.Error):
return None
try:
return _KEY_PAGE_CURSOR_ADAPTER.validate_json(decoded)
except ValidationError:
return None
def _build_aggregated_sql_query(
*,
table_name: str,
@ -828,7 +858,8 @@ def _build_aggregated_sql_query(
timezone_offset_minutes: int | None = None,
include_current_utc_day: bool = False,
global_rollup_through: str | None = None,
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
cursor: KeyPageCursor | None = None,
) -> tuple[str, list[str | float]]: # mutable-ok: SQL text plus its ordered $N params
"""Build the GROUPING SETS query for aggregated daily activity.
Returns:
@ -853,6 +884,16 @@ def _build_aggregated_sql_query(
)
sentinel_param: Final = f"${len(where_params) + 1}"
marker_param: Final = None if global_rollup_through is None else f"${len(where_params) + 2}"
cursor_spend_param: Final = f"${len(where_params) + 2 + (1 if marker_param is not None else 0)}"
cursor_key_param: Final = f"${len(where_params) + 3 + (1 if marker_param is not None else 0)}"
cursor_clause: Final = (
""
if cursor is None
else (
f"WHERE key_spend < {cursor_spend_param}::float8 "
f"OR (key_spend = {cursor_spend_param}::float8 AND api_key > {cursor_key_param})"
)
)
metric_select: Final = _rollup_metric_select(table_name)
# TODO: drop the successful_requests/failed_requests aggregates (and the
@ -872,7 +913,8 @@ def _build_aggregated_sql_query(
| GROUPING(model, {_MODEL_GROUP_EXPR},
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
NULL::bigint AS distinct_api_keys,{metric_select}
NULL::bigint AS distinct_api_keys, NULL::float AS key_spend,
NULL::bigint AS remaining_api_keys,{metric_select}
FROM {_key_free_source(pg_table, where_clause, marker_param)}
GROUP BY GROUPING SETS (
(date),
@ -884,12 +926,18 @@ def _build_aggregated_sql_query(
()
))
UNION ALL
(WITH top_api_keys AS (
SELECT api_key, COUNT(*) OVER () AS distinct_api_keys
(WITH ranked_api_keys AS (
SELECT api_key, SUM(spend) AS key_spend, COUNT(*) OVER () AS distinct_api_keys
FROM "{pg_table}"
WHERE {where_clause} AND api_key <> {sentinel_param}
GROUP BY api_key
ORDER BY SUM(spend) DESC, api_key
),
top_api_keys AS (
SELECT api_key, key_spend, distinct_api_keys,
COUNT(*) OVER () AS remaining_api_keys
FROM ranked_api_keys
{cursor_clause}
ORDER BY key_spend DESC, api_key
LIMIT {USAGE_TOP_API_KEYS_LIMIT}
)
SELECT
@ -903,7 +951,9 @@ def _build_aggregated_sql_query(
GROUPING(date, api_key, model, {_MODEL_GROUP_EXPR},
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
MAX(top_api_keys.distinct_api_keys) AS distinct_api_keys,{metric_select}
MAX(top_api_keys.distinct_api_keys) AS distinct_api_keys,
MAX(top_api_keys.key_spend) AS key_spend,
MAX(top_api_keys.remaining_api_keys) AS remaining_api_keys,{metric_select}
FROM "{pg_table}" JOIN top_api_keys USING (api_key)
WHERE {where_clause}
GROUP BY GROUPING SETS (
@ -917,7 +967,8 @@ def _build_aggregated_sql_query(
"""
marker_params: Final = () if global_rollup_through is None else (global_rollup_through,)
return sql_query, [*where_params, PTU_SENTINEL_API_KEY, *marker_params]
cursor_params: Final = () if cursor is None else (cursor.spend, cursor.api_key)
return sql_query, [*where_params, PTU_SENTINEL_API_KEY, *marker_params, *cursor_params]
def _build_entity_rollup_sql_query(
@ -1431,6 +1482,7 @@ async def get_daily_activity_aggregated(
timezone_offset_minutes: int | None = None,
include_entity_breakdown: bool = False,
include_current_utc_day: bool = False,
cursor: KeyPageCursor | None = None,
) -> SpendAnalyticsPaginatedResponse:
"""Aggregated variant that returns the full result set (no pagination).
@ -1467,6 +1519,7 @@ async def get_daily_activity_aggregated(
sql_query, sql_params = _build_aggregated_sql_query(
**query_kwargs,
global_rollup_through=await global_rollup_reconciled_through(prisma_client, query_kwargs),
cursor=cursor,
)
entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None
@ -1477,6 +1530,20 @@ async def get_daily_activity_aggregated(
records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or ())]
total_api_keys: Final = next((r.distinct_api_keys for r in records if r.distinct_api_keys is not None), 0)
remaining_api_keys: Final = next(
(r.remaining_api_keys or 0 for r in records if getattr(r, "remaining_api_keys", None) is not None),
0,
)
page_last: Final = max(
(r for r in records if r.api_key is not None and getattr(r, "key_spend", None) is not None),
key=lambda r: (-(r.key_spend or 0.0), r.api_key or ""),
default=None,
)
next_cursor: Final = (
KeyPageCursor(page_last.key_spend or 0.0, page_last.api_key or "").encode()
if remaining_api_keys > USAGE_TOP_API_KEYS_LIMIT and page_last is not None
else None
)
# The grouping-sets dispatcher places each row directly in its bucket
# using the row's GROUPING() bitmask. No Python-side summing needed.
@ -1532,6 +1599,7 @@ async def get_daily_activity_aggregated(
has_more=False,
api_key_limit=USAGE_TOP_API_KEYS_LIMIT,
total_api_keys=total_api_keys,
next_cursor=next_cursor,
),
)

17
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -18,7 +18,7 @@ import traceback
from collections.abc import Awaitable, Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Any, Final, Literal, Protocol, cast, overload
from typing import Annotated, Any, Final, Literal, Protocol, cast, overload
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@ -54,6 +54,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
from litellm.proxy.management_endpoints.common_daily_activity import (
DailySpendRecord,
decode_key_page_cursor,
get_daily_activity,
get_daily_activity_aggregated,
)
@ -3055,6 +3056,12 @@ async def get_user_daily_activity_aggregated(
"terms) is included. Requires the timezone parameter. Historical ranges are "
"never extended.",
),
cursor: Annotated[
str | None,
fastapi.Query(
description="Opaque cursor from a previous response's metadata.next_cursor; loads the next page of api_keys"
),
] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> SpendAnalyticsPaginatedResponse:
"""
@ -3080,6 +3087,13 @@ async def get_user_daily_activity_aggregated(
detail={"error": "Please provide start_date and end_date"},
)
page_cursor: Final = None if cursor is None else decode_key_page_cursor(cursor)
if cursor is not None and page_cursor is None:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "Invalid cursor"},
)
try:
entity_id: Final = _resolve_user_daily_activity_entity_id(user_api_key_dict, user_id)
@ -3095,6 +3109,7 @@ async def get_user_daily_activity_aggregated(
api_key=api_key,
timezone_offset_minutes=timezone,
include_current_utc_day=include_current_utc_day,
cursor=page_cursor,
)
except HTTPException:

8
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -124,6 +124,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import (
resolve_model_budget,
)
from litellm.proxy.management_endpoints.common_daily_activity import (
decode_key_page_cursor,
get_daily_activity_aggregated,
)
from litellm.proxy.management_endpoints.common_utils import (
@ -6735,6 +6736,7 @@ async def get_team_daily_activity_aggregated(
api_key: str | None = None,
exclude_team_ids: str | None = None,
timezone: int | None = None,
cursor: str | None = None,
):
"""
Aggregated daily activity for teams without pagination, including per-team breakdown.
@ -6751,6 +6753,7 @@ async def get_team_daily_activity_aggregated(
api_key (Optional[str]): Filter by API key.
exclude_team_ids (Optional[str]): Comma-separated list of team IDs to exclude.
timezone (Optional[int]): Timezone offset in minutes from UTC, matching JavaScript's Date.getTimezoneOffset() convention.
cursor (Optional[str]): Opaque cursor from a previous response's metadata.next_cursor; loads the next page of api_keys.
Returns:
SpendAnalyticsPaginatedResponse: Response containing all daily activity data for the range.
"""
@ -6767,6 +6770,10 @@ async def get_team_daily_activity_aggregated(
if range_error is not None:
raise _daily_activity_error(status_code=400, message=range_error)
page_cursor: Final = None if cursor is None else decode_key_page_cursor(cursor)
if cursor is not None and page_cursor is None:
raise _daily_activity_error(status_code=400, message="Invalid cursor")
scope: Final = await _resolve_team_daily_activity_scope(
team_ids=team_ids,
exclude_team_ids=exclude_team_ids,
@ -6790,6 +6797,7 @@ async def get_team_daily_activity_aggregated(
exclude_entity_ids=scope.exclude_team_ids,
timezone_offset_minutes=timezone,
include_entity_breakdown=True,
cursor=page_cursor,
)

4
litellm/types/proxy/management_endpoints/common_daily_activity.py
View file

@ -111,6 +111,10 @@ class DailySpendMetadata(BaseModel):
description="Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key "
"lists are truncated to the highest-spend keys.",
)
next_cursor: str | None = Field(
default=None,
description="Opaque cursor for the next page of api_keys; None when every key is loaded",
)
class SpendAnalyticsPaginatedResponse(BaseModel):

57
tests/integration/spend/test_user_daily_activity_key_search.py
View file

@ -125,3 +125,60 @@ def test_key_search_by_exact_hash_only_returns_keys_the_caller_may_see(gateway:
)
assert as_admin.status_code == 200, as_admin.text
assert set(_key_breakdown(object_value(as_admin.json()))) == {stranger_digest}, as_admin.text
def test_aggregated_key_pages_reach_every_key_through_cursor(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
expensive: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
cheap: Final = scenario.model(input_cost_per_token=0.00001, output_cost_per_token=0.00002)
user: Final = scenario.user()
needle: Final = scenario.key(user_id=user, key_alias=f"needle-{uuid.uuid4().hex}", models=[cheap])
needle_digest: Final = sha256(needle.encode()).hexdigest()
hay: Final = tuple(scenario.key(user_id=user, models=[expensive]) for _ in range(USAGE_TOP_API_KEYS_LIMIT))
for key in (needle, *hay):
reply: Final = gateway.chat(
expensive if key in hay else cheap, key=key, text=f"key page {uuid.uuid4().hex}"
)
assert reply["usage"]["total_tokens"] == 40, reply
daily: Final = eventually(
lambda: read_rows('SELECT api_key FROM "LiteLLM_DailyUserSpend" WHERE user_id=%s', (user,)),
lambda values: len(values) == USAGE_TOP_API_KEYS_LIMIT + 1,
seconds=70,
)
assert len(daily) == USAGE_TOP_API_KEYS_LIMIT + 1, daily
first: Final = gateway.request(
"GET", "/user/daily/activity/aggregated", params={**_range_params(), "user_id": user}
)
assert first.status_code == 200, first.text
first_body: Final = object_value(first.json())
first_metadata: Final = object_value(first_body["metadata"])
first_keys: Final = _key_breakdown(first_body)
assert len(first_keys) == USAGE_TOP_API_KEYS_LIMIT, sorted(first_keys)
assert needle_digest not in first_keys, "Precondition: the cheap key must fall outside page 1"
next_cursor: Final = first_metadata["next_cursor"]
assert isinstance(next_cursor, str) and next_cursor, first.text
assert int(str(first_metadata["total_api_keys"])) == USAGE_TOP_API_KEYS_LIMIT + 1, first.text
second: Final = gateway.request(
"GET",
"/user/daily/activity/aggregated",
params={**_range_params(), "user_id": user, "cursor": next_cursor},
)
assert second.status_code == 200, second.text
second_body: Final = object_value(second.json())
second_metadata: Final = object_value(second_body["metadata"])
assert set(_key_breakdown(second_body)) == {needle_digest}, second.text
assert second_metadata["next_cursor"] is None, second.text
assert int(str(second_metadata["total_api_keys"])) == USAGE_TOP_API_KEYS_LIMIT + 1, second.text
assert float(str(second_metadata["total_spend"])) == pytest.approx(float(str(first_metadata["total_spend"]))), (
first.text,
second.text,
)
garbage: Final = gateway.request(
"GET",
"/user/daily/activity/aggregated",
params={**_range_params(), "user_id": user, "cursor": "not-a-cursor"},
)
assert garbage.status_code == 400, garbage.text

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

@ -17,11 +17,13 @@ from litellm.constants import (
USAGE_TOP_API_KEYS_LIMIT,
)
from litellm.proxy.management_endpoints.common_daily_activity import (
KeyPageCursor,
_adjust_dates_for_timezone,
_build_aggregated_sql_query,
_build_entity_rollup_sql_query,
_is_user_agent_tag,
_record_to_spend_metrics,
decode_key_page_cursor,
get_api_key_metadata,
get_daily_activity,
get_daily_activity_aggregated,
@ -2868,3 +2870,142 @@ def test_spend_logs_window_is_none_when_no_date_parses():
from litellm.proxy.management_endpoints.common_daily_activity import _spend_logs_window
assert _spend_logs_window({"garbage", ""}) is None
class TestKeyPageCursorCodec:
def test_round_trip(self):
cursor: Final = KeyPageCursor(spend=6.0, api_key="key-004")
decoded: Final = decode_key_page_cursor(cursor.encode())
assert decoded is not None
assert decoded.spend == 6.0
assert decoded.api_key == "key-004"
@pytest.mark.parametrize(
"raw",
[
"not-base64!!!",
"aGVsbG8=", # valid base64, not json
"eyJmb28iOiAxfQ==", # json object missing both fields
"eyJzcGVuZCI6ICIxLjUiLCAiYXBpX2tleSI6IDd9", # wrong types
],
)
def test_malformed_returns_none(self, raw):
assert decode_key_page_cursor(raw) is None
class TestBuildAggregatedSqlQueryCursor:
def test_no_cursor_keeps_sql_and_params_unchanged(self):
sql, params = _build_aggregated_sql_query(
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id="user-1",
start_date="2026-05-29",
end_date="2026-06-02",
model=None,
api_key=None,
timezone_offset_minutes=None,
)
assert "ranked_api_keys" in sql
assert "key_spend < $" not in sql
assert params == ["2026-05-29", "2026-06-02", "user-1", PTU_SENTINEL_API_KEY]
def test_cursor_appends_params_after_sentinel(self):
sql, params = _build_aggregated_sql_query(
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id="user-1",
start_date="2026-05-29",
end_date="2026-06-02",
model=None,
api_key=None,
timezone_offset_minutes=None,
cursor=KeyPageCursor(spend=6.0, api_key="key-004"),
)
assert params[-2:] == [6.0, "key-004"]
assert "key_spend < $5::float8" in sql
assert "key_spend = $5::float8 AND api_key > $6" in sql
def test_cursor_placeholders_follow_the_marker_param(self):
sql, params = _build_aggregated_sql_query(
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
start_date="2026-05-29",
end_date="2026-06-02",
model=None,
api_key=None,
timezone_offset_minutes=None,
global_rollup_through="2026-05-30",
cursor=KeyPageCursor(spend=6.0, api_key="key-004"),
)
assert params[-3:] == ["2026-05-30", 6.0, "key-004"]
assert "key_spend < $5::float8" in sql
assert "key_spend = $5::float8 AND api_key > $6" in sql
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_pages_api_keys_through_cursor(
_aggregated_postgresql: psycopg.Connection,
):
"""Same 105-key seed as the bounds test. Page 1 must hand back a cursor that
decodes to the cutoff key (the 6.0 tie at key-004), and page 2 must return
exactly the five keys below the cutoff with no further cursor."""
n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 5
key_rows: Final = [
(
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(n_keys)
]
_seed_daily_user_spend(_aggregated_postgresql, key_rows)
key_spend: Final = sum(6.0 if i == 4 else float(i + 1) for i in range(n_keys))
call_kwargs: Final = {
"prisma_client": None,
"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,
}
row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts)
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
first = await get_daily_activity_aggregated(**{**call_kwargs, "prisma_client": mock_prisma})
assert first.metadata.next_cursor is not None
cutoff: Final = decode_key_page_cursor(first.metadata.next_cursor)
assert cutoff is not None
assert (cutoff.spend, cutoff.api_key) == (6.0, "key-004")
first_day: Final = first.results[0]
assert len(first_day.breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT
row_counts.clear()
second = await get_daily_activity_aggregated(**{**call_kwargs, "prisma_client": mock_prisma}, cursor=cutoff)
second_day: Final = second.results[0]
assert set(second_day.breakdown.api_keys) == {"key-005", "key-003", "key-002", "key-001", "key-000"}
assert second.metadata.next_cursor is None
assert second.metadata.total_api_keys == n_keys
assert second.metadata.total_spend == pytest.approx(first.metadata.total_spend)
assert second.metadata.total_spend == pytest.approx(key_spend)
assert row_counts == [7 + 6 * 5]

75
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -2587,9 +2587,84 @@ async def test_get_user_daily_activity_aggregated_admin_global_view(monkeypatch,
api_key=None,
timezone_offset_minutes=480,
include_current_utc_day=include_current_utc_day,
cursor=None,
)
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_rejects_garbage_cursor(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
from fastapi import HTTPException
from litellm.proxy.management_endpoints.internal_user_endpoints import (
get_user_daily_activity_aggregated,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
mock_get_daily_agg = AsyncMock()
monkeypatch.setattr(
"litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated",
mock_get_daily_agg,
)
with pytest.raises(HTTPException) as exc_info:
await get_user_daily_activity_aggregated(
start_date="2025-01-01",
end_date="2025-01-31",
model=None,
api_key=None,
user_id=None,
timezone=None,
cursor="not-a-cursor",
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-user-001",
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
assert exc_info.value.status_code == 400
mock_get_daily_agg.assert_not_called()
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_decodes_cursor_to_key_page_cursor(monkeypatch):
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy.management_endpoints.common_daily_activity import KeyPageCursor
from litellm.proxy.management_endpoints.internal_user_endpoints import (
get_user_daily_activity_aggregated,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
mock_response = MagicMock()
mock_get_daily_agg = AsyncMock(return_value=mock_response)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated",
mock_get_daily_agg,
)
encoded: Final = KeyPageCursor(spend=6.0, api_key="key-004").encode()
result = await get_user_daily_activity_aggregated(
start_date="2025-01-01",
end_date="2025-01-31",
model=None,
api_key=None,
user_id=None,
timezone=None,
cursor=encoded,
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-user-001",
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
assert result is mock_response
forwarded: Final = mock_get_daily_agg.call_args.kwargs["cursor"]
assert isinstance(forwarded, KeyPageCursor)
assert (forwarded.spend, forwarded.api_key) == (6.0, "key-004")
@pytest.mark.asyncio
async def test_get_user_daily_activity_aggregated_non_admin_cannot_view_other_users(
monkeypatch,

59
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -14569,6 +14569,65 @@ async def test_get_team_daily_activity_aggregated_rejects_bad_ranges(
mock_aggregated.assert_not_called()
@pytest.mark.asyncio
async def test_get_team_daily_activity_aggregated_rejects_garbage_cursor(mock_db_client):
from litellm.proxy.management_endpoints.team_endpoints import (
get_team_daily_activity_aggregated,
)
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated",
new_callable=AsyncMock,
) as mock_aggregated:
with pytest.raises(HTTPException) as exc_info:
await get_team_daily_activity_aggregated(
team_ids=None,
start_date="2024-01-01",
end_date="2024-01-31",
model=None,
api_key=None,
exclude_team_ids=None,
timezone=None,
cursor="not-a-cursor",
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert exc_info.value.status_code == 400
mock_aggregated.assert_not_called()
@pytest.mark.asyncio
async def test_get_team_daily_activity_aggregated_decodes_cursor_to_key_page_cursor(mock_db_client):
from litellm.proxy.management_endpoints.common_daily_activity import KeyPageCursor
from litellm.proxy.management_endpoints.team_endpoints import (
get_team_daily_activity_aggregated,
)
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated",
new_callable=AsyncMock,
) as mock_aggregated:
mock_aggregated.return_value = MagicMock()
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
encoded = KeyPageCursor(spend=6.0, api_key="key-004").encode()
await get_team_daily_activity_aggregated(
team_ids=None,
start_date="2024-01-01",
end_date="2024-01-31",
model=None,
api_key=None,
exclude_team_ids=None,
timezone=None,
cursor=encoded,
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
forwarded = mock_aggregated.call_args.kwargs["cursor"]
assert isinstance(forwarded, KeyPageCursor)
assert (forwarded.spend, forwarded.api_key) == (6.0, "key-004")
def _wire_new_team_prisma(mock_db_client):
mock_db_client.jsonify_team_object = lambda db_data: db_data
mock_db_client.get_data = AsyncMock(return_value=None)

4
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -1786,7 +1786,7 @@
"count": 1
},
"max-params": {
"count": 21
"count": 22
},
"no-nested-ternary": {
"count": 5
@ -2432,4 +2432,4 @@
"count": 1
}
}
}
}

54
ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
View file

@ -664,7 +664,59 @@ describe("EntityUsage", () => {
fireEvent.click(screen.getByText("Key Activity"));
});
expect(await screen.findByRole("note")).toHaveTextContent("Only the 100 highest-spend keys of 3,000 are loaded");
// processActivityData is mocked to a single-entry record here, so the loaded count reads 1
expect(await screen.findByRole("note")).toHaveTextContent("Only the 1 highest-spend keys of 3,000 are loaded");
});
it("loads the next key page through the metadata cursor and clears the button when the cursor runs out", async () => {
mockTeamDailyActivityAggregatedCall.mockResolvedValue({
...mockSpendData,
results: [
{
...mockSpendData.results[0],
breakdown: {
...mockSpendData.results[0].breakdown,
api_keys: { "hash-a": createKeyMetrics(0.5, { key_alias: "loaded-a", team_id: null }) },
},
},
],
metadata: { ...mockSpendData.metadata, api_key_limit: 1, total_api_keys: 2, next_cursor: "cursor-1" },
});
render(<EntityUsage {...defaultProps} entityType="team" />);
await waitFor(() => {
expect(mockTeamDailyActivityAggregatedCall).toHaveBeenCalled();
});
act(() => {
fireEvent.click(screen.getByText("Key Activity"));
});
mockTeamDailyActivityAggregatedCall.mockResolvedValue({
...mockSpendData,
results: [
{
...mockSpendData.results[0],
breakdown: {
...mockSpendData.results[0].breakdown,
api_keys: { "hash-b": createKeyMetrics(0.01, { key_alias: "needle-b", team_id: null }) },
},
},
],
metadata: { ...mockSpendData.metadata, next_cursor: null },
});
fireEvent.click(await screen.findByRole("button", { name: "Load more keys" }));
expect(mockTeamDailyActivityAggregatedCall).toHaveBeenLastCalledWith(
"test-token",
expect.anything(),
expect.anything(),
null,
"cursor-1",
);
await waitFor(() => {
expect(screen.queryByRole("button", { name: "Load more keys" })).not.toBeInTheDocument();
});
});
// An inactive tab panel is marked aria-selected="false" by one tab library and hidden by the

24
ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
View file

@ -44,6 +44,7 @@ import EndpointUsage from "../EndpointUsage/EndpointUsage";
import ModelViewToggle, { ModelViewType } from "../ModelViewToggle";
import TopKeyView from "@/components/UsagePage/components/EntityUsage/TopKeyView";
import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel";
import { mergeKeyPage } from "@/components/UsagePage/keyPageMerge";
import TopModelView from "./TopModelView";
import TeamUserSpendCard from "./TeamUserSpendCard";
@ -73,6 +74,7 @@ interface EntitySpendData {
total_tokens: number;
api_key_limit?: number | null;
total_api_keys?: number | null;
next_cursor?: string | null;
};
}
@ -161,9 +163,28 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
aggregatedFetchFn,
});
const spendData = spendDataRaw as unknown as EntitySpendData;
const spendDataBase = spendDataRaw as unknown as EntitySpendData;
const [loadedPages, setLoadedPages] = useState<{ base: EntitySpendData; merged: EntitySpendData } | null>(null);
const spendData = loadedPages?.base === spendDataBase ? loadedPages.merged : spendDataBase;
const apiKeyTruncation = getApiKeyTruncation(spendData.metadata?.api_key_limit, spendData.metadata?.total_api_keys);
const loadMoreKeys =
aggregatedFetchFn && accessToken && startTime && endTime && spendData.metadata?.next_cursor
? () =>
teamDailyActivityAggregatedCall(
accessToken,
startTime,
endTime,
Array.isArray(entityFilterArg) ? entityFilterArg : null,
spendData.metadata?.next_cursor ?? null,
).then((page: EntitySpendData) => {
setLoadedPages((prev) => {
const base = prev?.base === spendDataBase ? prev.merged : spendDataBase;
return { base: spendDataBase, merged: mergeKeyPage(base, page) };
});
})
: undefined;
const {
data: agentSpendDataRaw,
isFetchingMore: agentIsFetchingMore,
@ -667,6 +688,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
keyMetrics={keyMetrics}
hidePromptCachingMetrics={entityType === "agent"}
apiKeyTruncation={apiKeyTruncation}
loadMoreKeys={loadMoreKeys}
/>
),
},

20
ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
View file

@ -32,6 +32,7 @@ import UserDropdown from "@/components/common_components/UserDropdown";
import EntityUsageExportModal from "@/components/EntityUsageExport";
import { getApiKeyTruncation, getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason";
import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel";
import { mergeKeyPage } from "@/components/UsagePage/keyPageMerge";
import { Team } from "@/components/key_team_helpers/key_list";
import {
gatewayDailyActivityCall,
@ -438,6 +439,24 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
[userSpendData, modelViewType, teams],
);
const keyMetrics = useMemo(() => processActivityData(userSpendData, "api_keys", teams), [userSpendData, teams]);
const aggregatedCursor = activeAggregated?.metadata?.next_cursor;
const loadMoreKeys = useCallback(() => {
if (!accessToken || !startTime || !endTime || !aggregatedCursor) return Promise.resolve();
const rangeKey = currentAggregatedRangeKey;
return userDailyActivityAggregatedCall(
accessToken,
startTime,
endTime,
effectiveUserId,
false,
null,
aggregatedCursor,
).then((page) => {
setAggregatedData((prev) =>
prev && prev.rangeKey === rangeKey ? { rangeKey, value: mergeKeyPage(prev.value, page) } : prev,
);
});
}, [accessToken, startTime, endTime, effectiveUserId, aggregatedCursor, currentAggregatedRangeKey]);
const searchKeys = useCallback(
(q: string) => {
if (!accessToken || !startTime || !endTime) return Promise.resolve({});
@ -879,6 +898,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
keyMetrics={keyMetrics}
apiKeyTruncation={spendFetchState.apiKeyTruncation}
searchKeys={searchKeys}
loadMoreKeys={aggregatedCursor ? loadMoreKeys : undefined}
/>
</TabsContent>
<TabsContent value="mcp" keepMounted>

40
ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx
View file

@ -142,4 +142,44 @@ describe("KeyActivityPanel", () => {
expect(await screen.findByRole("alert")).toHaveTextContent("Key search failed");
expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alice");
});
it("shows no Load more keys button without the prop", () => {
render(<KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={{ limit: 2, total: 3 }} />);
expect(screen.queryByRole("button", { name: "Load more keys" })).not.toBeInTheDocument();
});
it("calls loadMoreKeys when the button is clicked", async () => {
const loadMoreKeys = vi.fn<() => Promise<void>>().mockResolvedValue(undefined);
render(
<KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={{ limit: 2, total: 3 }} loadMoreKeys={loadMoreKeys} />,
);
fireEvent.click(screen.getByRole("button", { name: "Load more keys" }));
expect(loadMoreKeys).toHaveBeenCalledTimes(1);
expect(await screen.findByRole("button", { name: "Load more keys" })).toBeEnabled();
});
it("shows a pending status while loadMoreKeys is in flight", async () => {
const loadMoreKeys = vi.fn<() => Promise<void>>().mockReturnValue(new Promise(() => {}));
render(
<KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={{ limit: 2, total: 3 }} loadMoreKeys={loadMoreKeys} />,
);
fireEvent.click(screen.getByRole("button", { name: "Load more keys" }));
expect(await screen.findByRole("status")).toHaveTextContent("Loading more keys...");
expect(screen.getByRole("button", { name: "Load more keys" })).toBeDisabled();
});
it("reports a failed key page load", async () => {
const loadMoreKeys = vi.fn<() => Promise<void>>().mockRejectedValue(new Error("boom"));
render(
<KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={{ limit: 2, total: 3 }} loadMoreKeys={loadMoreKeys} />,
);
fireEvent.click(screen.getByRole("button", { name: "Load more keys" }));
expect(await screen.findByRole("alert")).toHaveTextContent("Loading more keys failed");
});
});

32
ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx
View file

@ -13,6 +13,7 @@ interface KeyActivityPanelProps {
hidePromptCachingMetrics?: boolean;
apiKeyTruncation?: ApiKeyTruncation;
searchKeys?: SearchKeys;
loadMoreKeys?: () => Promise<void>;
}
type SearchKeys = (query: string) => Promise<Record<string, ModelActivityData>>;
@ -23,6 +24,8 @@ type RemoteSearch =
| { status: "done"; query: string; searchKeys: SearchKeys; keys: Record<string, ModelActivityData> }
| { status: "error"; query: string; searchKeys: SearchKeys };
type LoadMoreState = { status: "idle" } | { status: "loading" } | { status: "error" };
const REMOTE_SEARCH_DEBOUNCE_MS = 300;
const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({
@ -30,9 +33,11 @@ const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({
hidePromptCachingMetrics = false,
apiKeyTruncation,
searchKeys,
loadMoreKeys,
}) => {
const [query, setQuery] = useState("");
const [remote, setRemote] = useState<RemoteSearch>({ status: "idle" });
const [loadMore, setLoadMore] = useState<LoadMoreState>({ status: "idle" });
const filtered = useMemo(() => filterKeyActivity(keyMetrics, query), [keyMetrics, query]);
const trimmedQuery = query.trim();
const remoteEnabled = searchKeys !== undefined && apiKeyTruncation !== undefined && trimmedQuery !== "";
@ -110,10 +115,35 @@ const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({
)}
{apiKeyTruncation !== undefined && (
<span className="text-sm text-muted-foreground" role="note">
Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "}
Only the {totalKeys.toLocaleString()} highest-spend keys of{" "}
{apiKeyTruncation.total.toLocaleString()} are loaded
</span>
)}
{loadMoreKeys !== undefined && (
<button
type="button"
className="text-sm text-muted-foreground underline disabled:no-underline disabled:opacity-50"
disabled={loadMore.status === "loading"}
onClick={() => {
setLoadMore({ status: "loading" });
loadMoreKeys()
.then(() => setLoadMore({ status: "idle" }))
.catch(() => setLoadMore({ status: "error" }));
}}
>
Load more keys
</button>
)}
{loadMore.status === "loading" && (
<span role="status" className="text-sm text-muted-foreground">
Loading more keys...
</span>
)}
{loadMore.status === "error" && (
<span role="alert" className="text-sm text-muted-foreground">
Loading more keys failed
</span>
)}
</div>
{noMatches ? (
<p className="rounded-lg border p-6 text-center text-sm text-muted-foreground">

101
ui/litellm-dashboard/src/components/UsagePage/keyPageMerge.test.ts Normal file
View file

@ -0,0 +1,101 @@
import { describe, expect, it } from "vitest";
import { mergeKeyPage } from "./keyPageMerge";
import type { DailyData, KeyMetricWithMetadata, MetricWithMetadata, SpendMetrics } from "./types";
const metrics = (spend: number): SpendMetrics => ({
spend,
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
api_requests: 0,
successful_requests: 0,
failed_requests: 0,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
});
const keyEntry = (spend: number, alias: string): KeyMetricWithMetadata => ({
metrics: metrics(spend),
metadata: { key_alias: alias, team_id: null, user_id: null, user_email: null },
});
const modelEntry = (apiKeys: Record<string, KeyMetricWithMetadata>): MetricWithMetadata => ({
metrics: metrics(0),
metadata: {},
api_key_breakdown: apiKeys,
});
const day = (date: string, apiKeys: Record<string, KeyMetricWithMetadata>, model = "gpt-4"): DailyData => ({
date,
metrics: metrics(1),
breakdown: {
models: { [model]: modelEntry(apiKeys) },
model_groups: { "group-a": modelEntry(apiKeys) },
providers: { openai: modelEntry(apiKeys) },
mcp_servers: {},
api_keys: apiKeys,
entities: {},
endpoints: { "/v1/chat": modelEntry(apiKeys) },
},
});
const page = (results: DailyData[], next_cursor: string | null) => ({
results,
metadata: { total_spend: 42, total_api_keys: 3, next_cursor },
});
describe("mergeKeyPage", () => {
it("adds keys from the next page to the same day's breakdown", () => {
const prev = page([day("2026-01-01", { "hash-a": keyEntry(1, "a") })], "c1");
const next = page([day("2026-01-01", { "hash-b": keyEntry(0.5, "b") })], null);
const merged = mergeKeyPage(prev, next);
expect(Object.keys(merged.results[0].breakdown.api_keys)).toEqual(["hash-a", "hash-b"]);
expect(merged.results[0].breakdown.models["gpt-4"].api_key_breakdown).toHaveProperty("hash-b");
expect(merged.results[0].breakdown.providers.openai.api_key_breakdown).toHaveProperty("hash-b");
expect(merged.results[0].breakdown.endpoints?.["/v1/chat"].api_key_breakdown).toHaveProperty("hash-b");
});
it("lets the next page win when a key appears on both pages", () => {
const prev = page([day("2026-01-01", { "hash-a": keyEntry(1, "a") })], "c1");
const next = page([day("2026-01-01", { "hash-a": keyEntry(9, "a-new") })], null);
const merged = mergeKeyPage(prev, next);
expect(merged.results[0].breakdown.api_keys["hash-a"].metadata.key_alias).toBe("a-new");
expect(merged.results[0].breakdown.models["gpt-4"].api_key_breakdown["hash-a"].metrics.spend).toBe(9);
});
it("keeps prev metadata totals and takes next_cursor from next", () => {
const prev = page([day("2026-01-01", { "hash-a": keyEntry(1, "a") })], "c1");
const next = page([day("2026-01-01", { "hash-b": keyEntry(0.5, "b") })], null);
const merged = mergeKeyPage(prev, next);
expect(merged.metadata.total_spend).toBe(42);
expect(merged.metadata.total_api_keys).toBe(3);
expect(merged.metadata.next_cursor).toBeNull();
});
it("appends a day that only exists in the next page", () => {
const prev = page([day("2026-01-01", { "hash-a": keyEntry(1, "a") })], "c1");
const next = page([day("2026-01-02", { "hash-b": keyEntry(0.5, "b") })], "c2");
const merged = mergeKeyPage(prev, next);
expect(merged.results.map((d) => d.date)).toEqual(["2026-01-01", "2026-01-02"]);
expect(merged.metadata.next_cursor).toBe("c2");
});
it("does not mutate prev", () => {
const prev = page([day("2026-01-01", { "hash-a": keyEntry(1, "a") })], "c1");
const next = page([day("2026-01-01", { "hash-b": keyEntry(0.5, "b") })], null);
mergeKeyPage(prev, next);
expect(Object.keys(prev.results[0].breakdown.api_keys)).toEqual(["hash-a"]);
expect(prev.results[0].breakdown.models["gpt-4"].api_key_breakdown).not.toHaveProperty("hash-b");
});
});

53
ui/litellm-dashboard/src/components/UsagePage/keyPageMerge.ts Normal file
View file

@ -0,0 +1,53 @@
import type { DailyData, MetricWithMetadata } from "./types";
type AggregatedResponse<M> = {
results: DailyData[];
metadata: M;
};
const BREAKDOWN_LISTS = ["models", "model_groups", "providers", "mcp_servers"] as const;
function mergeEntryMap(
prev: { [key: string]: MetricWithMetadata } | undefined,
next: { [key: string]: MetricWithMetadata } | undefined,
): { [key: string]: MetricWithMetadata } {
const merged: { [key: string]: MetricWithMetadata } = { ...prev };
for (const [name, nextEntry] of Object.entries(next ?? {})) {
const prevEntry = merged[name];
merged[name] = prevEntry
? { ...prevEntry, api_key_breakdown: { ...prevEntry.api_key_breakdown, ...nextEntry.api_key_breakdown } }
: nextEntry;
}
return merged;
}
function mergeDay(prev: DailyData, next: DailyData): DailyData {
const breakdown = { ...prev.breakdown };
breakdown.api_keys = { ...prev.breakdown.api_keys, ...next.breakdown.api_keys };
for (const list of BREAKDOWN_LISTS) {
breakdown[list] = mergeEntryMap(prev.breakdown[list], next.breakdown[list]);
}
if (prev.breakdown.endpoints !== undefined || next.breakdown.endpoints !== undefined) {
breakdown.endpoints = mergeEntryMap(prev.breakdown.endpoints, next.breakdown.endpoints);
}
return { ...prev, breakdown };
}
export function mergeKeyPage<M extends { next_cursor?: string | null }>(
prev: AggregatedResponse<M>,
next: AggregatedResponse<M>,
): AggregatedResponse<M> {
const mergedResults = [...prev.results];
for (const nextDay of next.results) {
const index = mergedResults.findIndex((day) => day.date === nextDay.date);
if (index === -1) {
mergedResults.push(nextDay);
} else {
mergedResults[index] = mergeDay(mergedResults[index], nextDay);
}
}
return {
results: mergedResults,
metadata: { ...prev.metadata, next_cursor: next.metadata?.next_cursor ?? null },
};
}

12
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -1446,6 +1446,7 @@ export const teamDailyActivityAggregatedCall = async (
startTime: Date,
endTime: Date,
teamIds: string[] | null = null,
cursor: string | null = null,
) => {
/**
* Get aggregated daily team activity with per-team breakdown (no pagination)
@ -1459,6 +1460,7 @@ export const teamDailyActivityAggregatedCall = async (
timezone: new Date().getTimezoneOffset().toString(),
team_ids: teamIds && teamIds.length > 0 ? teamIds.join(",") : undefined,
exclude_team_ids: "litellm-dashboard",
cursor: cursor ?? undefined,
},
});
} catch (error) {
@ -2523,12 +2525,17 @@ export const userDailyActivityAggregatedCall = async (
accessToken: string,
startTime: Date,
endTime: Date,
...options: [userId?: string | null, includeCurrentUtcDay?: boolean, apiKey?: string | null]
...options: [
userId?: string | null,
includeCurrentUtcDay?: boolean,
apiKey?: string | null,
cursor?: string | null,
]
) => {
/**
* Get aggregated daily user activity (no pagination)
*/
const [userId = null, includeCurrentUtcDay = false, apiKey = null] = options;
const [userId = null, includeCurrentUtcDay = false, apiKey = null, cursor = null] = options;
try {
const formatDate = (date: Date) => {
const year = date.getFullYear();
@ -2548,6 +2555,7 @@ export const userDailyActivityAggregatedCall = async (
user_id: userId,
include_current_utc_day: includeCurrentUtcDay ? "true" : undefined,
api_key: apiKey,
cursor: cursor ?? undefined,
},
});
} catch (error) {

9
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -15633,6 +15633,7 @@ export interface paths {
* api_key (Optional[str]): Filter by API key.
* exclude_team_ids (Optional[str]): Comma-separated list of team IDs to exclude.
* timezone (Optional[int]): Timezone offset in minutes from UTC, matching JavaScript's Date.getTimezoneOffset() convention.
* cursor (Optional[str]): Opaque cursor from a previous response's metadata.next_cursor; loads the next page of api_keys.
* Returns:
* SpendAnalyticsPaginatedResponse: Response containing all daily activity data for the range.
*/
@ -28383,6 +28384,11 @@ export interface components {
* @default false
*/
has_more: boolean;
/**
* Next Cursor
* @description Opaque cursor for the next page of api_keys; None when every key is loaded
*/
next_cursor?: string | null;
/**
* Page
* @default 1
@ -62309,6 +62315,7 @@ export interface operations {
api_key?: string | null;
exclude_team_ids?: string | null;
timezone?: number | null;
cursor?: string | null;
};
header?: never;
path?: never;
@ -64439,6 +64446,8 @@ export interface operations {
timezone?: number | null;
/** @description When the range ends on the caller's current local day, extend it to today's UTC bucket so spend written after the caller's local midnight (in UTC terms) is included. Requires the timezone parameter. Historical ranges are never extended. */
include_current_utc_day?: boolean;
/** @description Opaque cursor from a previous response's metadata.next_cursor; loads the next page of api_keys */
cursor?: string | null;
};
header?: never;
path?: never;