scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-02 02:11:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

refactor(repositories): daily activity repository with centralized bounded usage queries (#43398)

Browse source 
Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-01 16:02:01 -07:00 • committed by GitHub
parent 4b1d9bf148
commit aa601ce4e8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 3736 additions and 1173 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/constants.py
View file

@ -2162,6 +2162,17 @@ MCP_SPEND_LOG_MODEL_PREFIX: Final[str] = "MCP: "
PTU_SENTINEL_API_KEY: Final[str] = "__ptu_flat_cost__"
PTU_ROLLUP_JOB_ID: Final[str] = "ptu_flat_cost_rollup_job"
PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900
USAGE_TOP_API_KEYS_DEFAULT: Final[int] = 100
USAGE_TOP_API_KEYS_MAX: Final[int] = 1000
USAGE_KEY_PAGE_DEFAULT: Final[int] = 50
USAGE_KEY_PAGE_MAX: Final[int] = 100
USAGE_KEY_SEARCH_DEFAULT: Final[int] = 100
USAGE_KEY_SEARCH_MAX: Final[int] = 100
USAGE_MODEL_TOP_KEYS_DEFAULT: Final[int] = 5
USAGE_MODEL_TOP_KEYS_MAX: Final[int] = 100
USAGE_CACHE_LEAKAGE_KEYS_DEFAULT: Final[int] = 20
USAGE_CACHE_LEAKAGE_KEYS_MAX: Final[int] = 100
USAGE_EXPORT_BATCH_SIZE: Final[int] = 1000
# Furthest back the catch-up pass looks for unpriced PTU days when a deployment
# declares no ptu_effective_from, bounding the scan for an open-ended window.
PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90

39
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -3455,6 +3455,33 @@
},
"DailySpendMetadata": {
"properties": {
"api_key_limit": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"description": "When set, api_keys and every api_key_breakdown list at most this many keys, ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.",
"title": "Api Key Limit"
},
"entity_total_api_keys": {
"anyOf": [
{
"additionalProperties": {
"type": "integer"
},
"type": "object"
},
{
"type": "null"
}
],
"description": "Distinct API keys per entity over the requested range, set when the entity breakdown is included. When an entity's count exceeds api_key_limit, its api_key_breakdown lists only its keys among the top api_key_limit keys overall.",
"title": "Entity Total Api Keys"
},
"has_more": {
"default": false,
"title": "Has More",
@ -3465,6 +3492,18 @@
"title": "Page",
"type": "integer"
},
"total_api_keys": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"description": "Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key lists are truncated to the highest-spend keys.",
"title": "Total Api Keys"
},
"total_api_requests": {
"default": 0,
"title": "Total Api Requests",

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

File diff suppressed because it is too large Load diff

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

@ -54,6 +54,8 @@ from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventH
from litellm.proxy.management.teams.access import is_team_admin
from litellm.proxy.management_endpoints.common_daily_activity import (
DailySpendRecord,
daily_activity_repository,
daily_activity_scope,
get_daily_activity,
get_daily_activity_aggregated,
)
@ -3079,18 +3081,22 @@ async def get_user_daily_activity_aggregated(
)
entity_id = user_id
repository: Final = daily_activity_repository(prisma_client)
scope: Final = daily_activity_scope(
"litellm_dailyuserspend",
"user_id",
entity_id,
None,
api_key,
start_date,
end_date,
model,
timezone,
include_current_utc_day,
)
return await get_daily_activity_aggregated(
prisma_client=prisma_client,
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=entity_id,
entity_metadata_field=None,
start_date=start_date,
end_date=end_date,
model=model,
api_key=api_key,
timezone_offset_minutes=timezone,
include_current_utc_day=include_current_utc_day,
repository,
scope,
)
except HTTPException:

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

@ -125,6 +125,8 @@ from litellm.proxy.hooks.model_max_budget_limiter import (
from litellm.proxy.management.teams.access import TEAM_OR_ORG_ADMIN, TeamRole, is_team_admin, team_access_denied
from litellm.proxy.management.teams.dependencies import get_team_access
from litellm.proxy.management_endpoints.common_daily_activity import (
daily_activity_repository,
daily_activity_scope,
get_daily_activity_aggregated,
)
from litellm.proxy.management_endpoints.common_utils import (
@ -6761,18 +6763,22 @@ async def get_team_daily_activity_aggregated(
proxy_logging_obj=proxy_logging_obj,
)
repository: Final = daily_activity_repository(prisma_client)
activity_scope: Final = daily_activity_scope(
"litellm_dailyteamspend",
"team_id",
scope.team_ids,
scope.exclude_team_ids,
scope.api_key_filter,
start_date,
end_date,
model,
timezone,
)
return await get_daily_activity_aggregated(
prisma_client=prisma_client,
table_name="litellm_dailyteamspend",
entity_id_field="team_id",
entity_id=scope.team_ids,
repository,
activity_scope,
entity_metadata_field=scope.team_alias_metadata,
start_date=start_date,
end_date=end_date,
model=model,
api_key=scope.api_key_filter,
exclude_entity_ids=scope.exclude_team_ids,
timezone_offset_minutes=timezone,
include_entity_breakdown=True,
)

32
litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py
View file

@ -8,11 +8,13 @@ from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable,
from datetime import date
from typing import Final, Literal, NamedTuple, Protocol, cast, overload
from fastapi import HTTPException
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL
from litellm.proxy._types import CommonProxyErrors
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
@ -259,22 +261,34 @@ async def _query_activity(
) -> SpendAnalyticsPaginatedResponse:
"""Shared helper that calls the daily activity query layer."""
from litellm.proxy.management_endpoints.common_daily_activity import (
daily_activity_repository,
daily_activity_scope,
get_daily_activity,
get_daily_activity_aggregated,
)
from litellm.proxy.proxy_server import prisma_client
if use_aggregated:
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
repository: Final = daily_activity_repository(prisma_client)
scope: Final = daily_activity_scope(
table_name,
entity_id_field,
entity_id,
None,
None,
start_date,
end_date,
None,
None,
)
return await get_daily_activity_aggregated(
prisma_client=prisma_client,
table_name=table_name,
entity_id_field=entity_id_field,
entity_id=entity_id,
entity_metadata_field=None,
start_date=start_date,
end_date=end_date,
model=None,
api_key=None,
repository,
scope,
)
return await get_daily_activity(
prisma_client=prisma_client,

304
litellm/repositories/daily_activity_repository.py Normal file
View file

@ -0,0 +1,304 @@
import asyncio
from collections.abc import AsyncIterator, Mapping, Sequence
from datetime import datetime
from itertools import groupby
from types import MappingProxyType
from typing import Final, Protocol
from pydantic import StrictStr, TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm import constants
from litellm._logging import verbose_proxy_logger
from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.daily_activity_sql import (
ExportCursor,
SqlQuery,
adjust_dates_for_timezone,
build_aggregated_sql,
build_cache_leakage_keys_sql,
build_entity_rollup_sql,
build_export_sql,
build_key_page_sql,
build_key_search_sql,
build_model_top_keys_sql,
)
from litellm.repositories.prisma_protocols import TableActions
from litellm.types.repositories.daily_activity import (
AggregatedRows,
DailyActivityProxyReads,
DailyActivityRow,
DailyActivityScope,
DailyActivityTable,
DailyRowsPage,
EntityRollupRow,
ExportRow,
ExportType,
GroupingSetsRow,
KeyMetadataRow,
KeyPage,
KeySpendRow,
SpendLogsWindow,
)
class _VerificationTokenRow(Protocol):
token: str
key_alias: str | None
team_id: str | None
user_id: str | None
metadata: Mapping[str, object] | None
class _DeletedVerificationTokenRow(_VerificationTokenRow, Protocol):
deleted_at: datetime
class _QueryRaw(Protocol):
async def __call__(self, query: str, *values: object) -> Sequence[Mapping[str, object]] | None: ...
class _DailyActivityDatabase(Protocol):
query_raw: _QueryRaw
litellm_verificationtoken: TableActions[_VerificationTokenRow]
litellm_deletedverificationtoken: TableActions[_DeletedVerificationTokenRow]
@property
def litellm_dailyuserspend(self) -> TableActions[DailyActivityRow]: ...
@property
def litellm_dailyteamspend(self) -> TableActions[DailyActivityRow]: ...
@property
def litellm_dailytagspend(self) -> TableActions[DailyActivityRow]: ...
@property
def litellm_dailyorganizationspend(self) -> TableActions[DailyActivityRow]: ...
@property
def litellm_dailyenduserspend(self) -> TableActions[DailyActivityRow]: ...
@property
def litellm_dailyagentspend(self) -> TableActions[DailyActivityRow]: ...
class DailyActivityDatabase(Protocol):
@property
def db(self) -> _DailyActivityDatabase: ...
_GROUPING_ADAPTER: Final = TypeAdapter(tuple[GroupingSetsRow, ...])
_ENTITY_ADAPTER: Final = TypeAdapter(tuple[EntityRollupRow, ...])
_KEY_SPEND_ADAPTER: Final = TypeAdapter(tuple[KeySpendRow, ...])
_KEY_PAGE_TOTAL_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(int)
_EXPORT_ADAPTER: Final = TypeAdapter(tuple[ExportRow, ...])
_METADATA_TAGS_ADAPTER: Final = TypeAdapter(list[StrictStr])
def _metadata_tags(value: object) -> tuple[str, ...]:
stable_value: Final = value
if not isinstance(value, list):
return ()
try:
return tuple(_METADATA_TAGS_ADAPTER.validate_python(stable_value))
except ValidationError:
return ()
def _daily_rows_table(
prisma_client: DailyActivityDatabase, table: DailyActivityTable
) -> TableActions[DailyActivityRow]:
if table is DailyActivityTable.USER:
return prisma_client.db.litellm_dailyuserspend
if table is DailyActivityTable.TEAM:
return prisma_client.db.litellm_dailyteamspend
if table is DailyActivityTable.TAG:
return prisma_client.db.litellm_dailytagspend
if table is DailyActivityTable.ORGANIZATION:
return prisma_client.db.litellm_dailyorganizationspend
if table is DailyActivityTable.CUSTOMER:
return prisma_client.db.litellm_dailyenduserspend
if table is DailyActivityTable.AGENT:
return prisma_client.db.litellm_dailyagentspend
assert_never(table)
raise AssertionError("unreachable")
def _next_export_cursor(batch: tuple[ExportRow, ...], export_type: ExportType) -> ExportCursor:
last: Final = batch[-1]
cursor_key: Final = (
last.api_key
if export_type is ExportType.DAILY_WITH_KEYS
else last.model
if export_type is ExportType.DAILY_WITH_MODELS
else last.user_id
if export_type is ExportType.DAILY_WITH_USERS
else ""
)
return ExportCursor(date=last.date, entity_id=last.entity_id, group_key=cursor_key or "")
class DailyActivityRepository:
def __init__(self, prisma_client: DailyActivityDatabase, *, proxy_reads: DailyActivityProxyReads) -> None:
self._prisma_client = prisma_client
self._proxy_reads = proxy_reads
async def _query(self, query: SqlQuery) -> tuple[Mapping[str, object], ...]:
first_line: Final = query.sql.lstrip().splitlines()[0].lstrip("(").strip()
verbose_proxy_logger.debug("DailyActivityRepository query: %s", first_line)
result: Sequence[Mapping[str, object]] | None = await self._prisma_client.db.query_raw(query.sql, *query.params)
if result is None:
return ()
return tuple(result)
async def aggregated(
self, scope: DailyActivityScope, *, include_entity_breakdown: bool, api_key_limit: int
) -> AggregatedRows:
grouping_query: Final = build_aggregated_sql(scope, api_key_limit=api_key_limit)
entity_query: Final = (
build_entity_rollup_sql(scope, api_key_limit=api_key_limit) if include_entity_breakdown else None
)
grouping_result, entity_result = await asyncio.gather(
self._query(grouping_query),
self._query(entity_query) if entity_query is not None else asyncio.sleep(0, result=None),
)
grouping_rows: Final = _GROUPING_ADAPTER.validate_python(grouping_result)
entity_rows: Final = None if entity_result is None else _ENTITY_ADAPTER.validate_python(entity_result)
distinct_api_keys: Final = next(
(row.distinct_api_keys for row in grouping_rows if row.distinct_api_keys is not None), 0
)
return AggregatedRows(
grouping_rows=grouping_rows,
entity_rows=entity_rows,
distinct_api_keys=distinct_api_keys,
)
async def search_keys(self, scope: DailyActivityScope, *, search: str, limit: int) -> tuple[str, ...]:
if not 1 <= limit <= constants.USAGE_KEY_SEARCH_MAX:
raise ValueError(f"limit must be between 1 and {constants.USAGE_KEY_SEARCH_MAX}")
query: Final = build_key_search_sql(scope, search=search, limit=limit)
rows: Final = _KEY_SPEND_ADAPTER.validate_python(await self._query(query))
return tuple(row.api_key for row in rows)
async def key_page(self, scope: DailyActivityScope, *, offset: int, limit: int) -> KeyPage:
query: Final = build_key_page_sql(scope, offset=offset, limit=limit)
result: Final = await self._query(query)
total_api_keys_value: Final = result[0].get("total_api_keys") if result else 0
total_api_keys: Final = (
_KEY_PAGE_TOTAL_ADAPTER.validate_python(total_api_keys_value) if total_api_keys_value is not None else 0
)
rows: Final = _KEY_SPEND_ADAPTER.validate_python(tuple(row for row in result if row.get("api_key") is not None))
return KeyPage(rows=rows, total_api_keys=total_api_keys)
async def model_top_keys(
self, scope: DailyActivityScope, *, model_group: str, by_model_group: bool, limit: int
) -> tuple[KeySpendRow, ...]:
if not 1 <= limit <= constants.USAGE_MODEL_TOP_KEYS_MAX:
raise ValueError(f"limit must be between 1 and {constants.USAGE_MODEL_TOP_KEYS_MAX}")
query: Final = build_model_top_keys_sql(
scope,
model_group=model_group,
by_model_group=by_model_group,
limit=limit,
)
return _KEY_SPEND_ADAPTER.validate_python(await self._query(query))
async def cache_leakage_keys(self, scope: DailyActivityScope, *, limit: int) -> tuple[KeySpendRow, ...]:
if not 1 <= limit <= constants.USAGE_CACHE_LEAKAGE_KEYS_MAX:
raise ValueError(f"limit must be between 1 and {constants.USAGE_CACHE_LEAKAGE_KEYS_MAX}")
query: Final = build_cache_leakage_keys_sql(scope, limit=limit)
return _KEY_SPEND_ADAPTER.validate_python(await self._query(query))
async def export_rows(self, scope: DailyActivityScope, *, export_type: ExportType) -> AsyncIterator[ExportRow]:
batch_size: Final = constants.USAGE_EXPORT_BATCH_SIZE
cursor: ExportCursor | None = None # rebind-ok: each page advances the export keyset cursor
while True:
batch: tuple[ExportRow, ...] = _EXPORT_ADAPTER.validate_python(
await self._query(build_export_sql(scope, export_type=export_type, after=cursor, batch_size=batch_size))
)
for row in batch:
yield row
if len(batch) < batch_size:
return
cursor = _next_export_cursor(batch, export_type)
async def _active_token_rows(self, values: tuple[str, ...]) -> tuple[_VerificationTokenRow, ...]:
return await find_many_in(self._prisma_client.db.litellm_verificationtoken, "token", values)
async def _deleted_token_rows(self, values: tuple[str, ...]) -> tuple[_DeletedVerificationTokenRow, ...]:
try:
return await find_many_in(self._prisma_client.db.litellm_deletedverificationtoken, "token", values)
except Exception as exc:
verbose_proxy_logger.warning("Could not read deleted verification token metadata: %s", exc)
return ()
async def key_metadata(
self, api_keys: frozenset[str], window: SpendLogsWindow | None
) -> Mapping[str, KeyMetadataRow]:
if not api_keys:
return {}
values: Final = tuple(api_keys)
active_rows: Final = await self._active_token_rows(values)
active: Final = MappingProxyType({row.token: self._metadata_row(row, key_exists=True) for row in active_rows})
missing: Final = tuple(key for key in values if key not in active)
deleted_rows: Final = await self._deleted_token_rows(missing)
deleted_by_token: Final = MappingProxyType(
{
token: max(rows, key=lambda row: row.deleted_at)
for token, rows in groupby(
sorted(deleted_rows, key=lambda row: row.token),
key=lambda row: row.token,
)
}
)
deleted: Final = MappingProxyType(
{
key: self._metadata_row(deleted_by_token[key], key_exists=False)
for key in missing
if key in deleted_by_token
}
)
resolved: Final = MappingProxyType({**deleted, **active})
return await self._proxy_reads.recover_key_metadata(resolved, api_keys, window)
@staticmethod
def _metadata_row(row: _VerificationTokenRow, *, key_exists: bool) -> KeyMetadataRow:
tags: Final = _metadata_tags(row.metadata.get("tags") if row.metadata is not None else None)
return KeyMetadataRow(
api_key=row.token,
key_alias=row.key_alias,
team_id=row.team_id,
user_id=row.user_id,
user_email=None,
key_exists=key_exists,
tags=tags,
)
async def daily_rows(self, scope: DailyActivityScope, *, page: int, page_size: int) -> DailyRowsPage:
table: Final = _daily_rows_table(self._prisma_client, scope.table)
adjusted_start, adjusted_end = adjust_dates_for_timezone(
scope.start_date,
scope.end_date,
scope.timezone_offset_minutes,
include_current_utc_day=scope.include_current_utc_day,
)
entity_filter: Final = {
**({"in": list(scope.entity_ids)} if scope.entity_ids is not None else {}),
**({"not": {"in": list(scope.exclude_entity_ids)}} if scope.exclude_entity_ids else {}),
}
conditions: Final = {
"date": {"gte": adjusted_start, "lte": adjusted_end},
**({scope.entity_id_field: entity_filter} if entity_filter else {}),
**({"model": scope.model} if scope.model else {}),
**({"api_key": {"in": list(scope.api_keys)}} if scope.api_keys is not None else {}),
}
count, rows = await asyncio.gather(
table.count(where=conditions),
table.find_many(
where=conditions,
skip=(page - 1) * page_size,
take=page_size,
order=({"date": "desc"}, {"id": "asc"}),
),
)
return DailyRowsPage(total_count=count, rows=tuple(rows))

519
litellm/repositories/daily_activity_sql.py Normal file
View file

@ -0,0 +1,519 @@
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from itertools import count, islice
from types import MappingProxyType
from typing import Final
from typing_extensions import assert_never
from litellm import constants
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.types.repositories.daily_activity import DailyActivityScope, DailyActivityTable, ExportType
_API_KEY_ROLLED_UP_BIT: Final = 32
_MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)"
@dataclass(frozen=True, slots=True)
class SqlQuery:
sql: str
params: tuple[object, ...]
@dataclass(frozen=True, slots=True)
class ExportCursor:
date: str
entity_id: str
group_key: str
PRISMA_TO_PG_TABLE: Final[Mapping[DailyActivityTable, str]] = MappingProxyType(
{
DailyActivityTable.USER: "LiteLLM_DailyUserSpend",
DailyActivityTable.TEAM: "LiteLLM_DailyTeamSpend",
DailyActivityTable.TAG: "LiteLLM_DailyTagSpend",
DailyActivityTable.ORGANIZATION: "LiteLLM_DailyOrganizationSpend",
DailyActivityTable.CUSTOMER: "LiteLLM_DailyEndUserSpend",
DailyActivityTable.AGENT: "LiteLLM_DailyAgentSpend",
}
)
def adjust_dates_for_timezone(
start_date: str,
end_date: str,
timezone_offset_minutes: int | None,
include_current_utc_day: bool = False,
utc_now: datetime | None = None,
) -> tuple[str, str]:
if not include_current_utc_day or timezone_offset_minutes is None:
return start_date, end_date
now: Final = utc_now if utc_now is not None else datetime.now(timezone.utc)
caller_local_today: Final = (now - timedelta(minutes=timezone_offset_minutes)).date().isoformat()
if end_date < caller_local_today:
return start_date, end_date
return start_date, max(end_date, now.date().isoformat())
def build_where_clause(scope: DailyActivityScope, *, start_index: int = 1) -> tuple[str, tuple[object, ...]]:
adjusted_start, adjusted_end = adjust_dates_for_timezone(
scope.start_date,
scope.end_date,
scope.timezone_offset_minutes,
scope.include_current_utc_day,
)
entity_index: Final = start_index + 2
has_entity_array: Final = scope.entity_ids is not None and bool(scope.entity_ids)
exclusion_index: Final = entity_index + int(has_entity_array)
model_index: Final = exclusion_index + int(bool(scope.exclude_entity_ids))
api_keys_index: Final = model_index + int(bool(scope.model))
conditions: Final = (
f"date >= ${start_index}",
f"date <= ${start_index + 1}",
*(
("FALSE",)
if scope.entity_ids == ()
else (f'"{scope.entity_id_field}" = ANY(${entity_index}::text[])',)
if has_entity_array
else ()
),
*((f'NOT ("{scope.entity_id_field}" = ANY(${exclusion_index}::text[]))',) if scope.exclude_entity_ids else ()),
*((f"model = ${model_index}",) if scope.model else ()),
*(
("FALSE",)
if scope.api_keys == ()
else (f"api_key = ANY(${api_keys_index}::text[])",)
if scope.api_keys
else ()
),
)
params: Final = (
adjusted_start,
adjusted_end,
*((list(scope.entity_ids or ()),) if has_entity_array else ()),
*((list(scope.exclude_entity_ids),) if scope.exclude_entity_ids else ()),
*((scope.model,) if scope.model else ()),
*((list(scope.api_keys),) if scope.api_keys else ()),
)
return " AND ".join(conditions), params
def _ptu_flat_cost_select(table: DailyActivityTable, *, aggregate: bool = True) -> str:
if table is DailyActivityTable.TEAM:
return "SUM(ptu_flat_cost)::float AS ptu_flat_cost" if aggregate else "SUM(scoped.ptu_flat_cost)::float"
return "0::float AS ptu_flat_cost" if aggregate else "0::float"
def _rollup_metric_select(table: DailyActivityTable) -> str:
return f"""
SUM(spend)::float AS spend,
{_ptu_flat_cost_select(table)},
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(gateway_injected_caching_savings_spend)::float AS gateway_injected_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,
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
SUM(timed_requests)::bigint AS timed_requests"""
def _validate_api_key_limit(api_key_limit: int) -> None:
if not 1 <= api_key_limit <= constants.USAGE_TOP_API_KEYS_MAX:
raise ValueError(f"api_key_limit must be between 1 and {constants.USAGE_TOP_API_KEYS_MAX}")
def _top_api_keys_sql(pg_table: str, where_clause: str, *, sentinel_param: int, limit_param: int) -> str:
return f"""
SELECT api_key, 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::numeric) DESC, api_key
LIMIT ${limit_param}
"""
def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> SqlQuery:
pg_table: Final = PRISMA_TO_PG_TABLE[scope.table]
where_clause, where_params = build_where_clause(scope)
_validate_api_key_limit(api_key_limit)
sentinel_param: Final = len(where_params) + 1
top_keys_limit_param: Final = len(where_params) + 2
top_api_keys: Final = _top_api_keys_sql(
pg_table, where_clause, sentinel_param=sentinel_param, limit_param=top_keys_limit_param
)
metric_select: Final = _rollup_metric_select(scope.table)
sql: Final = f"""
(SELECT
date,
NULL::text AS api_key,
model,
{_MODEL_GROUP_EXPR} AS model_group,
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
(GROUPING(date) << 6) | {_API_KEY_ROLLED_UP_BIT}
| GROUPING(model, {_MODEL_GROUP_EXPR},
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
NULL::bigint AS distinct_api_keys,{metric_select}
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date),
(date, model),
(date, {_MODEL_GROUP_EXPR}),
(date, custom_llm_provider),
(date, mcp_namespaced_tool_name),
(date, endpoint),
()
))
UNION ALL
(WITH top_api_keys AS (
{top_api_keys}
)
SELECT
date,
api_key,
model,
{_MODEL_GROUP_EXPR} AS model_group,
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
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}
FROM "{pg_table}" JOIN top_api_keys USING (api_key)
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date, api_key),
(date, model, api_key),
(date, {_MODEL_GROUP_EXPR}, api_key),
(date, custom_llm_provider, api_key),
(date, mcp_namespaced_tool_name, api_key),
(date, endpoint, api_key)
))
"""
return SqlQuery(
sql=sql,
params=(*where_params, PTU_SENTINEL_API_KEY, api_key_limit),
)
def build_entity_rollup_sql(scope: DailyActivityScope, *, api_key_limit: int) -> SqlQuery:
pg_table: Final = PRISMA_TO_PG_TABLE[scope.table]
where_clause, where_params = build_where_clause(scope)
_validate_api_key_limit(api_key_limit)
sentinel_param: Final = len(where_params) + 1
top_keys_limit_param: Final = len(where_params) + 2
top_api_keys: Final = _top_api_keys_sql(
pg_table, where_clause, sentinel_param=sentinel_param, limit_param=top_keys_limit_param
)
metric_select: Final = _rollup_metric_select(scope.table)
sql: Final = f"""
WITH top_api_keys AS (
{top_api_keys}
),
entity_api_keys AS (
SELECT COALESCE("{scope.entity_id_field}", '') AS entity_id,
COUNT(DISTINCT api_key)::bigint AS distinct_api_keys
FROM "{pg_table}"
WHERE {where_clause} AND api_key <> ${sentinel_param}
GROUP BY COALESCE("{scope.entity_id_field}", '')
)
(SELECT e.*, COALESCE(k.distinct_api_keys, 0)::bigint AS distinct_api_keys
FROM (
SELECT COALESCE("{scope.entity_id_field}", '') AS entity_id,
date,
NULL::text AS api_key,
1 AS api_key_rolled,{metric_select}
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY date, COALESCE("{scope.entity_id_field}", '')
) e
LEFT JOIN entity_api_keys k ON k.entity_id = e.entity_id)
UNION ALL
(SELECT COALESCE("{scope.entity_id_field}", '') AS entity_id,
date,
api_key,
0 AS api_key_rolled,{metric_select},
NULL::bigint AS distinct_api_keys
FROM "{pg_table}" JOIN top_api_keys USING (api_key)
WHERE {where_clause}
GROUP BY date, COALESCE("{scope.entity_id_field}", ''), api_key)
"""
return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, api_key_limit))
def _key_spend_select() -> str:
return """
COALESCE(SUM(spend), 0)::float AS spend,
COALESCE(SUM(prompt_tokens), 0)::bigint AS prompt_tokens,
COALESCE(SUM(completion_tokens), 0)::bigint AS completion_tokens,
(COALESCE(SUM(prompt_tokens), 0) + COALESCE(SUM(completion_tokens), 0))::bigint AS total_tokens,
COALESCE(SUM(api_requests), 0)::bigint AS api_requests,
COALESCE(SUM(successful_requests), 0)::bigint AS successful_requests,
COALESCE(SUM(failed_requests), 0)::bigint AS failed_requests,
COALESCE(SUM(cache_read_input_tokens), 0)::bigint AS cache_read_input_tokens,
COALESCE(SUM(cache_creation_input_tokens), 0)::bigint AS cache_creation_input_tokens"""
def build_key_page_sql(scope: DailyActivityScope, *, offset: int, limit: int) -> SqlQuery:
if not 1 <= limit <= constants.USAGE_KEY_PAGE_MAX:
raise ValueError(f"limit must be between 1 and {constants.USAGE_KEY_PAGE_MAX}")
if offset < 0:
raise ValueError("offset must be non-negative")
where_clause, where_params = build_where_clause(scope)
sentinel_param: Final = len(where_params) + 1
limit_param: Final = sentinel_param + 1
offset_param: Final = limit_param + 1
sql: Final = f"""
WITH ranked AS (
SELECT api_key,{_key_spend_select()}, SUM(spend::numeric) AS rank_spend
FROM "{PRISMA_TO_PG_TABLE[scope.table]}"
WHERE {where_clause} AND api_key <> ${sentinel_param}
GROUP BY api_key
)
SELECT (SELECT COUNT(*) FROM ranked)::bigint AS total_api_keys, page.*
FROM (SELECT 1) AS one
LEFT JOIN LATERAL (
SELECT * FROM ranked
ORDER BY rank_spend DESC, api_key
LIMIT ${limit_param} OFFSET ${offset_param}
) AS page ON TRUE
"""
return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, limit, offset))
def _bounded_limit(limit: int, *, minimum: int = 1) -> None:
if limit < minimum:
raise ValueError(f"limit must be at least {minimum}")
def build_key_search_sql(scope: DailyActivityScope, *, search: str, limit: int) -> SqlQuery:
_bounded_limit(limit)
where_clause, where_params = build_where_clause(scope)
search_param: Final = len(where_params) + 1
sentinel_param: Final = search_param + 1
limit_param: Final = sentinel_param + 1
escaped: Final = search.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
sql: Final = f"""
SELECT api_key,{_key_spend_select()}
FROM "{PRISMA_TO_PG_TABLE[scope.table]}"
WHERE {where_clause}
AND api_key <> ${sentinel_param}
AND (
api_key ILIKE ${search_param} ESCAPE '\\'
OR api_key IN (
SELECT v.token FROM "LiteLLM_VerificationToken" v
LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = v.user_id
WHERE v.key_alias ILIKE ${search_param} ESCAPE '\\'
OR v.user_id ILIKE ${search_param} ESCAPE '\\'
OR u.user_email ILIKE ${search_param} ESCAPE '\\'
UNION
SELECT d.token FROM "LiteLLM_DeletedVerificationToken" d
LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = d.user_id
WHERE d.key_alias ILIKE ${search_param} ESCAPE '\\'
OR d.user_id ILIKE ${search_param} ESCAPE '\\'
OR u.user_email ILIKE ${search_param} ESCAPE '\\'
)
)
GROUP BY api_key
ORDER BY SUM(spend::numeric) DESC, api_key
LIMIT ${limit_param}
"""
return SqlQuery(sql=sql, params=(*where_params, f"%{escaped}%", PTU_SENTINEL_API_KEY, limit))
def build_model_top_keys_sql(
scope: DailyActivityScope, *, model_group: str, by_model_group: bool, limit: int
) -> SqlQuery:
_bounded_limit(limit)
where_clause, where_params = build_where_clause(scope)
model_param: Final = len(where_params) + 1
sentinel_param: Final = model_param + 1
limit_param: Final = sentinel_param + 1
model_clause: Final = (
f"COALESCE(NULLIF(model_group, ''), model) = ${model_param}" if by_model_group else f"model = ${model_param}"
)
sql: Final = f"""
SELECT api_key,{_key_spend_select()}
FROM "{PRISMA_TO_PG_TABLE[scope.table]}"
WHERE {where_clause} AND {model_clause} AND api_key <> ${sentinel_param}
GROUP BY api_key
ORDER BY SUM(spend::numeric) DESC, api_key
LIMIT ${limit_param}
"""
return SqlQuery(sql=sql, params=(*where_params, model_group, PTU_SENTINEL_API_KEY, limit))
def build_cache_leakage_keys_sql(scope: DailyActivityScope, *, limit: int) -> SqlQuery:
_bounded_limit(limit)
where_clause, where_params = build_where_clause(scope)
sentinel_param: Final = len(where_params) + 1
limit_param: Final = sentinel_param + 1
sql: Final = f"""
SELECT api_key,{_key_spend_select()}
FROM "{PRISMA_TO_PG_TABLE[scope.table]}"
WHERE {where_clause} AND api_key <> ${sentinel_param}
GROUP BY api_key
HAVING SUM(prompt_tokens) - SUM(cache_read_input_tokens) > 0
ORDER BY SUM(prompt_tokens) - SUM(cache_read_input_tokens) DESC, api_key
LIMIT ${limit_param}
"""
return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, limit))
def build_export_sql(
scope: DailyActivityScope, *, export_type: ExportType, after: ExportCursor | None, batch_size: int
) -> SqlQuery:
_bounded_limit(batch_size)
where_clause, where_params = build_where_clause(scope)
group_key, output_key, user_fields, type_joins = _export_grouping(export_type)
grouping_keys: Final = (
f"scoped.date, COALESCE(scoped.\"{scope.entity_id_field}\", '')",
*((group_key,) if export_type is not ExportType.DAILY else ()),
)
entity_joins: Final = (
('LEFT JOIN "LiteLLM_TeamTable" tt ON tt.team_id = scoped.team_id',)
if scope.table is DailyActivityTable.TEAM
else ('LEFT JOIN "LiteLLM_OrganizationTable" ot ON ot.organization_id = scoped.organization_id',)
if scope.table is DailyActivityTable.ORGANIZATION
else ()
)
joins: Final = (*type_joins, *entity_joins)
alias_expression: Final = (
"MAX(tt.team_alias)"
if scope.table is DailyActivityTable.TEAM
else "MAX(ot.organization_alias)"
if scope.table is DailyActivityTable.ORGANIZATION
else "NULL::text"
)
parameter_indexes: Final = count(len(where_params) + 1)
sentinel_param: Final = next(parameter_indexes) if export_type is not ExportType.DAILY else None
cursor_indexes: Final = tuple(islice(parameter_indexes, 3)) if after is not None else ()
limit_param: Final = next(parameter_indexes)
cursor_clause, cursor_params = _export_cursor_clause(
scope, after=after, cursor_indexes=cursor_indexes, group_key=group_key
)
sentinel_clause: Final = f" AND api_key <> ${sentinel_param}" if sentinel_param is not None else ""
table: Final = PRISMA_TO_PG_TABLE[scope.table]
flat_cost: Final = _ptu_flat_cost_select(scope.table, aggregate=False)
sql: Final = f"""
WITH scoped AS (
SELECT * FROM "{table}"
WHERE {where_clause}{sentinel_clause}
)
SELECT
scoped.date,
COALESCE(scoped."{scope.entity_id_field}", '') AS entity_id,
{alias_expression} AS entity_alias,
{output_key} AS api_key,
{user_fields},
{"NULLIF(COALESCE(scoped.model, ''), '')" if export_type is ExportType.DAILY_WITH_MODELS else "NULL::text"} AS model,
COALESCE(SUM(scoped.spend), 0)::float AS spend,
{flat_cost} AS flat_cost,
COALESCE(SUM(scoped.prompt_tokens), 0)::bigint AS prompt_tokens,
COALESCE(SUM(scoped.completion_tokens), 0)::bigint AS completion_tokens,
COALESCE(SUM(scoped.api_requests), 0)::bigint AS api_requests,
COALESCE(SUM(scoped.successful_requests), 0)::bigint AS successful_requests,
COALESCE(SUM(scoped.failed_requests), 0)::bigint AS failed_requests,
COALESCE(SUM(scoped.cache_read_input_tokens), 0)::bigint AS cache_read_input_tokens,
COALESCE(SUM(scoped.cache_creation_input_tokens), 0)::bigint AS cache_creation_input_tokens
FROM scoped
{" ".join(joins)}
WHERE TRUE{cursor_clause}
GROUP BY {", ".join(grouping_keys)}
ORDER BY {", ".join(grouping_keys)}
LIMIT ${limit_param}
"""
return SqlQuery(
sql=sql,
params=(
*where_params,
*((PTU_SENTINEL_API_KEY,) if export_type is not ExportType.DAILY else ()),
*cursor_params,
batch_size,
),
)
def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str, ...]]:
if export_type is ExportType.DAILY:
return (
"''",
"NULL::text",
"NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email",
(),
)
if export_type is ExportType.DAILY_WITH_KEYS:
return (
"scoped.api_key",
"NULLIF(scoped.api_key, '')",
"MAX(COALESCE(vt.key_alias, dvt.key_alias)) AS key_alias, "
"MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email",
(
'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key',
"""LEFT JOIN LATERAL (
SELECT key_alias, user_id
FROM "LiteLLM_DeletedVerificationToken"
WHERE token = scoped.api_key
ORDER BY deleted_at DESC
LIMIT 1
) dvt ON vt.token IS NULL""",
'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = COALESCE(vt.user_id, dvt.user_id)',
),
)
if export_type is ExportType.DAILY_WITH_MODELS:
return (
"COALESCE(scoped.model, '')",
"NULL::text",
"NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email",
(),
)
if export_type is ExportType.DAILY_WITH_USERS:
return (
"COALESCE(vt.user_id, dvt.user_id, '')",
"NULL::text",
"NULL::text AS key_alias, MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, "
"MAX(u.user_email) AS user_email",
(
'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key',
"""LEFT JOIN LATERAL (
SELECT key_alias, user_id
FROM "LiteLLM_DeletedVerificationToken"
WHERE token = scoped.api_key
ORDER BY deleted_at DESC
LIMIT 1
) dvt ON vt.token IS NULL""",
'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = COALESCE(vt.user_id, dvt.user_id)',
),
)
assert_never(export_type)
raise AssertionError("unreachable")
def _export_cursor_clause(
scope: DailyActivityScope,
*,
after: ExportCursor | None,
cursor_indexes: tuple[int, ...],
group_key: str,
) -> tuple[str, tuple[object, ...]]:
if after is None:
return "", ()
first_cursor_index: Final = cursor_indexes[0]
clause: Final = (
f""" AND (scoped.date, COALESCE(scoped."{scope.entity_id_field}", ''), {group_key}) """
f"> (${first_cursor_index}, ${cursor_indexes[1]}, ${cursor_indexes[2]})"
)
return clause, (after.date, after.entity_id, after.group_key)

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

@ -101,6 +101,22 @@ class DailySpendMetadata(BaseModel):
page: int = Field(default=1)
total_pages: int = Field(default=1)
has_more: bool = Field(default=False)
api_key_limit: int | None = Field(
default=None,
description="When set, api_keys and every api_key_breakdown list at most this many keys, "
"ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.",
)
total_api_keys: int | None = Field(
default=None,
description="Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key "
"lists are truncated to the highest-spend keys.",
)
entity_total_api_keys: dict[str, int] | None = Field(
default=None,
description="Distinct API keys per entity over the requested range, set when the entity breakdown is "
"included. When an entity's count exceeds api_key_limit, its api_key_breakdown lists only its keys "
"among the top api_key_limit keys overall.",
)
class SpendAnalyticsPaginatedResponse(BaseModel):

0
litellm/types/repositories/__init__.py Normal file
View file

192
litellm/types/repositories/daily_activity.py Normal file
View file

@ -0,0 +1,192 @@
from collections.abc import Mapping
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from types import MappingProxyType
from typing import Protocol, TypeAlias
class DailyActivityTable(str, Enum):
USER = "litellm_dailyuserspend"
TEAM = "litellm_dailyteamspend"
TAG = "litellm_dailytagspend"
ORGANIZATION = "litellm_dailyorganizationspend"
CUSTOMER = "litellm_dailyenduserspend"
AGENT = "litellm_dailyagentspend"
_ENTITY_FIELDS: Mapping[DailyActivityTable, frozenset[str]] = MappingProxyType(
{
DailyActivityTable.USER: frozenset(("user_id",)),
DailyActivityTable.TEAM: frozenset(("team_id",)),
DailyActivityTable.TAG: frozenset(("tag",)),
DailyActivityTable.ORGANIZATION: frozenset(("organization_id",)),
DailyActivityTable.CUSTOMER: frozenset(("end_user_id",)),
DailyActivityTable.AGENT: frozenset(("agent_id",)),
}
)
@dataclass(frozen=True, slots=True)
class DailyActivityScope:
table: DailyActivityTable
entity_id_field: str
entity_ids: tuple[str, ...] | None
exclude_entity_ids: tuple[str, ...]
api_keys: tuple[str, ...] | None
start_date: str
end_date: str
model: str | None
timezone_offset_minutes: int | None
include_current_utc_day: bool = False
def __post_init__(self) -> None:
if self.entity_id_field not in _ENTITY_FIELDS[self.table]:
raise ValueError(f"Invalid entity_id_field {self.entity_id_field!r} for {self.table.value}")
@dataclass(frozen=True, slots=True)
class KeySpendRow:
api_key: str
spend: float
prompt_tokens: int
completion_tokens: int
total_tokens: int
api_requests: int
successful_requests: int
failed_requests: int
cache_read_input_tokens: int
cache_creation_input_tokens: int
@dataclass(frozen=True, slots=True)
class KeyPage:
rows: tuple[KeySpendRow, ...]
total_api_keys: int
@dataclass(frozen=True, slots=True)
class KeyMetadataRow:
api_key: str
key_alias: str | None
team_id: str | None
user_id: str | None
user_email: str | None
key_exists: bool
tags: tuple[str, ...]
class ExportType(str, Enum):
DAILY = "daily"
DAILY_WITH_KEYS = "daily_with_keys"
DAILY_WITH_MODELS = "daily_with_models"
DAILY_WITH_USERS = "daily_with_users"
@dataclass(frozen=True, slots=True)
class ExportRow:
date: str
entity_id: str
entity_alias: str | None
api_key: str | None
key_alias: str | None
user_id: str | None
user_email: str | None
model: str | None
spend: float
flat_cost: float
prompt_tokens: int
completion_tokens: int
api_requests: int
successful_requests: int
failed_requests: int
cache_read_input_tokens: int
cache_creation_input_tokens: int
@dataclass(frozen=True, slots=True)
class RollupMetricsRow:
date: str | None
api_key: str | None
spend: float | None
ptu_flat_cost: float | None = field(default=None, kw_only=True)
prompt_tokens: int | None
completion_tokens: int | None
cache_read_input_tokens: int | None
cache_creation_input_tokens: int | None
compression_saved_tokens: int | None
compression_savings_spend: float | None
prompt_caching_savings_spend: float | None
gateway_injected_caching_savings_spend: float | None
autorouter_savings_spend: float | None
api_requests: int | None
successful_requests: int | None
failed_requests: int | None
total_response_time_ms: int | None
timed_requests: int | None
@dataclass(frozen=True, slots=True)
class GroupingSetsRow(RollupMetricsRow):
model: str | None
model_group: str | None
custom_llm_provider: str | None
mcp_namespaced_tool_name: str | None
endpoint: str | None
group_level: int
distinct_api_keys: int | None
@dataclass(frozen=True, slots=True)
class EntityRollupRow(RollupMetricsRow):
entity_id: str | None
api_key_rolled: int
distinct_api_keys: int | None
@dataclass(frozen=True, slots=True)
class AggregatedRows:
grouping_rows: tuple[GroupingSetsRow, ...]
entity_rows: tuple[EntityRollupRow, ...] | None
distinct_api_keys: int
SpendLogsWindow: TypeAlias = tuple[datetime, datetime]
class DailyActivityProxyReads(Protocol):
async def recover_key_metadata(
self, resolved: Mapping[str, KeyMetadataRow], api_keys: frozenset[str], window: SpendLogsWindow | None
) -> Mapping[str, KeyMetadataRow]: ...
class DailyActivityRow(Protocol):
id: str
date: str
api_key: str
model: str | None
model_group: str | None
custom_llm_provider: str | None
mcp_namespaced_tool_name: str | None
endpoint: str | None
prompt_tokens: int
completion_tokens: int
cache_read_input_tokens: int
cache_creation_input_tokens: int
compression_saved_tokens: int
compression_savings_spend: float
prompt_caching_savings_spend: float
gateway_injected_caching_savings_spend: float
autorouter_savings_spend: float
spend: float
api_requests: int
successful_requests: int
failed_requests: int
total_response_time_ms: int
timed_requests: int
@dataclass(frozen=True, slots=True)
class DailyRowsPage:
total_count: int
rows: tuple[DailyActivityRow, ...]

11
tests/code_coverage_tests/unbounded_in_baseline.txt
View file

@ -39,14 +39,6 @@ litellm/proxy/management_endpoints/auto_router_endpoints.py start_shadow_eval pr
litellm/proxy/management_endpoints/auto_router_endpoints.py start_shadow_eval prisma token.in `list(data.api_key_ids)` 0
litellm/proxy/management_endpoints/auto_router_endpoints.py start_shadow_eval prisma user_id.in `list(data.user_ids)` 0
litellm/proxy/management_endpoints/budget_management_endpoints.py info_budget prisma budget_id.in `data.budgets` 0
litellm/proxy/management_endpoints/common_daily_activity.py _build_aggregated_where_clause raw-sql api_key.IN `IN ({placeholders})` 0
litellm/proxy/management_endpoints/common_daily_activity.py _build_aggregated_where_clause raw-sql {entity_id_field}.IN `IN ({placeholders})` 0
litellm/proxy/management_endpoints/common_daily_activity.py _build_aggregated_where_clause raw-sql {entity_id_field}.IN `IN ({placeholders})` 1
litellm/proxy/management_endpoints/common_daily_activity.py _build_where_conditions prisma [entity_id_field].in `entity_id` 0
litellm/proxy/management_endpoints/common_daily_activity.py _build_where_conditions prisma api_key.in `api_key` 0
litellm/proxy/management_endpoints/common_daily_activity.py _build_where_conditions prisma not.in `exclude_entity_ids` 0
litellm/proxy/management_endpoints/common_daily_activity.py get_api_key_metadata prisma token.in `list(api_keys)` 0
litellm/proxy/management_endpoints/common_daily_activity.py get_api_key_metadata prisma token.in `list(missing_keys)` 0
litellm/proxy/management_endpoints/common_utils.py _team_admin_can_invite_user prisma team_id.in `admin_user_obj.teams` 0
litellm/proxy/management_endpoints/common_utils.py _user_has_admin_privileges prisma team_id.in `user_obj.teams` 0
litellm/proxy/management_endpoints/customer_endpoints.py delete_end_user prisma user_id.in `data.user_ids` 0
@ -149,4 +141,7 @@ litellm/proxy/utils.py PrismaClient.get_data prisma budget_id.in `budget_id_list
litellm/proxy/utils.py PrismaClient.get_data prisma team_id.in `team_id_list` 0
litellm/proxy/utils.py PrismaClient.get_data prisma user_id.in `user_id_list` 0
litellm/proxy/utils.py prefetch_config_params prisma param_name.in `param_names` 0
litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma ?.in `list(scope.entity_ids)` 0
litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma api_key.in `list(scope.api_keys)` 0
litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma not.in `list(scope.exclude_entity_ids)` 0
litellm/router_utils/auto_router_model_naming.py <module> raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0

319
tests/integration/spend/_daily_activity_fixtures.py Normal file
View file

@ -0,0 +1,319 @@
from itertools import product
from typing import Final
import psycopg
from psycopg import sql
_TABLE_NAMES: Final = (
"LiteLLM_DailyUserSpend",
"LiteLLM_DailyTeamSpend",
"LiteLLM_VerificationToken",
"LiteLLM_DeletedVerificationToken",
"LiteLLM_UserTable",
"LiteLLM_TeamTable",
)
_TAG_KEY_MEMBERSHIPS: Final = (
("tag-a", "entity-key-0"),
("tag-a", "entity-key-1"),
("tag-a", "entity-key-2"),
("tag-a", "entity-key-3"),
("tag-a", "entity-key-4"),
("tag-b", "entity-key-0"),
("tag-b", "entity-key-1"),
("tag-b", "entity-key-5"),
("tag-c", "entity-key-2"),
("tag-c", "entity-key-3"),
("tag-c", "entity-key-6"),
("tag-c", "entity-key-7"),
("tag-d", "entity-key-4"),
("tag-d", "entity-key-5"),
("tag-d", "entity-key-6"),
("tag-d", "entity-key-7"),
)
_TAG_ACTIVITY_DATES: Final = ("2026-06-01", "2026-06-02")
def seed_daily_activity_fixture(connection: psycopg.Connection, *, schema: str, ptu_sentinel_api_key: str) -> None:
daily_user_table: Final = sql.Identifier(schema, "LiteLLM_DailyUserSpend")
daily_team_table: Final = sql.Identifier(schema, "LiteLLM_DailyTeamSpend")
verification_token_table: Final = sql.Identifier(schema, "LiteLLM_VerificationToken")
deleted_token_table: Final = sql.Identifier(schema, "LiteLLM_DeletedVerificationToken")
user_table: Final = sql.Identifier(schema, "LiteLLM_UserTable")
team_table: Final = sql.Identifier(schema, "LiteLLM_TeamTable")
keys: Final = (
("key-a", "model-popular", 100.0, 2, 1),
("key-b", "model-popular", 90.0, 3, 1),
("key-c", "model-popular", 80.0, 4, 1),
("key-target", "model-target", 1.0, 5, 2),
("key-cache", "model-cache", 2.0, 1000, 1),
)
user_rows: Final = tuple(
(
f"user-row-{index}",
"user-1",
"2026-06-01",
api_key,
model,
"",
"provider-a",
None,
"/v1/chat/completions",
prompt_tokens,
2,
cache_read_tokens,
0,
spend,
1,
1,
0,
"2026-06-01 12:00:00",
)
for index, (api_key, model, spend, prompt_tokens, cache_read_tokens) in enumerate(keys)
)
team_rows: Final = tuple(
(
f"team-row-{index}",
"team-1",
"2026-06-01",
api_key,
model,
"",
"provider-a",
None,
"/v1/chat/completions",
prompt_tokens,
2,
cache_read_tokens,
0,
spend,
1,
1,
0,
0.0,
"2026-06-01 12:00:00",
)
for index, (api_key, model, spend, prompt_tokens, cache_read_tokens) in enumerate(keys)
)
sentinel_user_row: Final = (
"user-row-ptu",
"user-1",
"2026-06-01",
ptu_sentinel_api_key,
"model-ptu",
"",
"provider-a",
None,
"/v1/chat/completions",
0,
0,
0,
0,
1000.0,
0,
0,
0,
"2026-06-01 12:00:00",
)
sentinel_team_row: Final = (
"team-row-ptu",
"team-1",
"2026-06-01",
ptu_sentinel_api_key,
"model-ptu",
"",
"provider-a",
None,
"/v1/chat/completions",
0,
0,
0,
0,
1000.0,
0,
0,
0,
42.0,
"2026-06-01 12:00:00",
)
with connection.cursor() as cursor:
for table_name in _TABLE_NAMES:
cursor.execute(
sql.SQL("CREATE TABLE {} (LIKE {} INCLUDING DEFAULTS INCLUDING CONSTRAINTS)").format(
sql.Identifier(schema, table_name),
sql.Identifier(table_name),
)
)
cursor.executemany(
sql.SQL("""
INSERT INTO {}
(id, user_id, date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, prompt_tokens, completion_tokens,
cache_read_input_tokens, cache_creation_input_tokens, spend, api_requests,
successful_requests, failed_requests, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""").format(daily_user_table),
(*user_rows, sentinel_user_row),
)
cursor.executemany(
sql.SQL("""
INSERT INTO {}
(id, team_id, date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, prompt_tokens, completion_tokens,
cache_read_input_tokens, cache_creation_input_tokens, spend, api_requests,
successful_requests, failed_requests, ptu_flat_cost, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""").format(daily_team_table),
(*team_rows, sentinel_team_row),
)
cursor.executemany(
sql.SQL(
"INSERT INTO {} (token, key_alias, team_id, user_id, metadata, models) VALUES (%s, %s, %s, %s, %s, %s)"
).format(verification_token_table),
(
("key-a", "alias-a", "team-1", "user-1", '{"tags": ["blue", "gold"]}', []),
("key-b", "alias-b", "team-1", "user-1", '{"tags": []}', []),
("key-c", "alias-c", "team-1", "user-1", '{"tags": []}', []),
("key-cache", "alias-cache", "team-1", "user-1", '{"tags": []}', []),
),
)
cursor.executemany(
sql.SQL("""
INSERT INTO {}
(id, token, key_alias, team_id, user_id, metadata, models, deleted_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
""").format(deleted_token_table),
(
("deleted-old", "key-target", "older-target", "team-1", "user-1", '{"tags": []}', [], "2026-06-01"),
(
"deleted-new",
"key-target",
"deleted-target",
"team-1",
"user-1",
'{"tags": ["archived"]}',
[],
"2026-06-02",
),
),
)
cursor.execute(
sql.SQL("INSERT INTO {} (user_id, user_email, models) VALUES (%s, %s, %s)").format(user_table),
("user-1", "user@example.com", []),
)
cursor.execute(
sql.SQL("INSERT INTO {} (team_id, team_alias, admins, members, models) VALUES (%s, %s, %s, %s, %s)").format(
team_table
),
("team-1", "Usage Team", [], [], []),
)
connection.commit()
def seed_daily_tag_activity_fixture(connection: psycopg.Connection, *, schema: str) -> None:
tag_table: Final = sql.Identifier(schema, "LiteLLM_DailyTagSpend")
rows: Final = tuple(
(
f"tag-rollup-{row_index}",
tag,
date,
api_key,
"entity-rollup-model",
"",
"provider-a",
None,
"/v1/chat/completions",
row_index + 1,
float(row_index + 1),
row_index % 5 + 1,
f"{date} 12:00:00",
)
for row_index, (date, (tag, api_key)) in enumerate(product(_TAG_ACTIVITY_DATES, _TAG_KEY_MEMBERSHIPS))
)
with connection.cursor() as cursor:
cursor.execute(
sql.SQL("CREATE TABLE {} (LIKE {} INCLUDING DEFAULTS INCLUDING CONSTRAINTS)").format(
tag_table,
sql.Identifier("LiteLLM_DailyTagSpend"),
)
)
cursor.executemany(
sql.SQL("""
INSERT INTO {}
(id, tag, date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, prompt_tokens, spend, api_requests, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""").format(tag_table),
rows,
)
connection.commit()
def seed_daily_tag_float_tie_fixture(connection: psycopg.Connection, *, schema: str) -> None:
tag_table: Final = sql.Identifier(schema, "LiteLLM_DailyTagSpend")
key_spends: Final = (
("key-z", 0.1),
("key-z", 0.2),
("key-z", 0.3),
("key-a", 0.3),
("key-a", 0.2),
("key-a", 0.1),
)
rows: Final = (
(
f"float-tie-{row_index}",
"tag-float-tie",
"2026-06-01",
api_key,
"float-tie-model",
"",
"provider-a",
None,
"/v1/chat/completions",
1,
spend,
1,
"2026-06-01 12:00:00",
)
for row_index, (api_key, spend) in enumerate(key_spends, start=1)
)
with connection.cursor() as cursor:
cursor.execute(
sql.SQL("CREATE TABLE {} (LIKE {} INCLUDING DEFAULTS INCLUDING CONSTRAINTS)").format(
tag_table,
sql.Identifier("LiteLLM_DailyTagSpend"),
)
)
cursor.executemany(
sql.SQL("""
INSERT INTO {}
(id, tag, date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, prompt_tokens, spend, api_requests, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""").format(tag_table),
rows,
)
connection.commit()
def seed_daily_team_unassigned_fixture(
connection: psycopg.Connection, *, schema: str, ptu_sentinel_api_key: str
) -> None:
team_table: Final = sql.Identifier(schema, "LiteLLM_DailyTeamSpend")
rows: Final = (
("unassigned-null", None, "key-unassigned-null", 3.0, 0.0),
("unassigned-empty", "", "key-unassigned-empty", 7.0, 0.0),
("unassigned-ptu", None, ptu_sentinel_api_key, 13.0, 13.0),
)
with connection.cursor() as cursor:
cursor.executemany(
sql.SQL("""
INSERT INTO {}
(id, team_id, date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, prompt_tokens, spend, api_requests, ptu_flat_cost, updated_at)
VALUES (%s, %s, '2026-06-03', %s, 'model-a', '', 'provider-a', NULL, '/v1/chat/completions',
1, %s, 1, %s, '2026-06-03 12:00:00')
""").format(team_table),
rows,
)
connection.commit()

1
tests/integration/spend/fixtures/daily_activity_team.json Normal file
View file

File diff suppressed because one or more lines are too long

1
tests/integration/spend/fixtures/daily_activity_user.json Normal file
View file

File diff suppressed because one or more lines are too long

93
tests/integration/spend/test_daily_activity_aggregated_breakdowns.py
View file

@ -6,7 +6,7 @@ from typing import Final
import pytest
from pydantic import JsonValue, TypeAdapter
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_DEFAULT
from tests.integration._support.client import Gateway, object_value
from tests.integration._support.database import write_rows
@ -43,57 +43,86 @@ def _row_id() -> str:
return f"agg-{uuid.uuid4().hex}"
def _ranked_key_rows(day: str, count: int) -> list[tuple[object, ...]]:
return [
(
_row_id(),
f"user-{i:03d}",
day,
f"key-{i:03d}",
"gpt-5",
"",
"openai",
None,
"/v1/chat/completions",
10,
6.0 if i == 4 else float(i + 1),
1,
1,
)
for i in range(count)
]
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_returns_every_api_key(gateway: Gateway) -> None:
async def test_get_daily_activity_aggregated_bounds_api_key_rollups(gateway: Gateway) -> None:
"""key-004 and key-005 tie on spend exactly at the default api_key_limit cutoff; the api_key
tiebreaker keeps key-004 and drops key-005. The PTU sentinel outspends every key but takes no
slot. Dropped keys and the sentinel still count toward the totals and the model rollup."""
key_count: Final = USAGE_TOP_API_KEYS_DEFAULT + 5
day: Final = _unique_day()
_seed(
day,
[
*[
(
_row_id(),
f"user-{i:03d}",
day,
f"key-{i:03d}",
"gpt-5",
"",
"openai",
None,
"/v1/chat/completions",
10,
6.0 if i == 4 else float(i + 1),
1,
1,
)
for i in range(105)
],
*_ranked_key_rows(day, key_count),
(_row_id(), None, day, PTU_SENTINEL_API_KEY, "gpt-5", "", "azure", None, None, 0, 1000.0, 0, 0),
],
)
key_spend: Final = sum(6.0 if i == 4 else float(i + 1) for i in range(key_count))
try:
body: Final = _activity(gateway, day)
metadata: Final = object_value(body["metadata"])
assert metadata["total_spend"] == pytest.approx(6566.0)
assert metadata["total_api_requests"] == 105
assert metadata["total_spend"] == pytest.approx(key_spend + 1000.0)
assert metadata["total_api_requests"] == key_count
assert metadata["total_api_keys"] == key_count
assert metadata["api_key_limit"] == USAGE_TOP_API_KEYS_DEFAULT
results: Final = _RESULTS.validate_python(body["results"])
assert len(results) == 1
result_day: Final = object_value(results[0])
assert object_value(result_day["metrics"])["spend"] == pytest.approx(6566.0)
assert object_value(result_day["metrics"])["spend"] == pytest.approx(key_spend + 1000.0)
breakdown: Final = object_value(result_day["breakdown"])
expected_api_keys: Final = {f"key-{i:03d}" for i in range(105)}
expected_top: Final = {f"key-{i:03d}" for i in range(6, key_count)} | {"key-004"}
api_keys: Final = object_value(breakdown["api_keys"])
assert set(api_keys) == expected_api_keys
assert set(api_keys) == expected_top
assert object_value(object_value(api_keys["key-004"])["metrics"])["spend"] == 6.0
assert PTU_SENTINEL_API_KEY not in api_keys
models: Final = object_value(breakdown["models"])
gpt5: Final = object_value(models["gpt-5"])
assert object_value(gpt5["metrics"])["spend"] == pytest.approx(6566.0)
assert set(object_value(gpt5["api_key_breakdown"])) == expected_api_keys
assert object_value(gpt5["metrics"])["spend"] == pytest.approx(key_spend + 1000.0)
assert set(object_value(gpt5["api_key_breakdown"])) == expected_top
providers: Final = object_value(breakdown["providers"])
openai: Final = object_value(providers["openai"])
assert object_value(openai["metrics"])["spend"] == pytest.approx(5566.0)
assert set(object_value(openai["api_key_breakdown"])) == expected_api_keys
assert object_value(openai["metrics"])["spend"] == pytest.approx(key_spend)
assert set(object_value(openai["api_key_breakdown"])) == expected_top
endpoints: Final = object_value(breakdown["endpoints"])
assert object_value(object_value(endpoints["/v1/chat/completions"])["metrics"])["api_requests"] == 105
assert object_value(object_value(endpoints["/v1/chat/completions"])["metrics"])["api_requests"] == key_count
finally:
_clean(day)
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_reports_exact_limit_key_count_as_complete(gateway: Gateway) -> None:
"""With exactly USAGE_TOP_API_KEYS_DEFAULT keys nothing is dropped and total_api_keys equals the limit."""
day: Final = _unique_day()
_seed(day, _ranked_key_rows(day, USAGE_TOP_API_KEYS_DEFAULT))
try:
body: Final = _activity(gateway, day)
metadata: Final = object_value(body["metadata"])
assert metadata["total_api_keys"] == USAGE_TOP_API_KEYS_DEFAULT
assert metadata["api_key_limit"] == USAGE_TOP_API_KEYS_DEFAULT
results: Final = _RESULTS.validate_python(body["results"])
api_keys: Final = object_value(object_value(object_value(results[0])["breakdown"])["api_keys"])
assert set(api_keys) == {f"key-{i:03d}" for i in range(USAGE_TOP_API_KEYS_DEFAULT)}
finally:
_clean(day)
@ -126,7 +155,9 @@ async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_resu
)
try:
body: Final = _activity(gateway, day, api_key="key-1")
assert object_value(body["metadata"])["total_spend"] == 2.0
metadata: Final = object_value(body["metadata"])
assert metadata["total_spend"] == 2.0
assert metadata["total_api_keys"] == 1
results: Final = _RESULTS.validate_python(body["results"])
assert len(results) == 1
breakdown: Final = object_value(object_value(results[0])["breakdown"])

624
tests/integration/spend/test_daily_activity_repository.py Normal file
View file

@ -0,0 +1,624 @@
import os
import uuid
from collections.abc import AsyncIterator, Mapping
from contextlib import asynccontextmanager
from dataclasses import dataclass
from math import isclose
from pathlib import Path
from types import MappingProxyType
from typing import Final, cast
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import psycopg
import pytest
from integration.spend._daily_activity_fixtures import (
seed_daily_activity_fixture,
seed_daily_tag_activity_fixture,
seed_daily_tag_float_tie_fixture,
seed_daily_team_unassigned_fixture,
)
from prisma import Prisma
from psycopg import sql
from pydantic import TypeAdapter
from litellm import constants
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity_aggregated
from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.daily_activity_repository import DailyActivityDatabase, DailyActivityRepository
from litellm.types.repositories.daily_activity import (
DailyActivityProxyReads,
DailyActivityScope,
DailyActivityTable,
ExportType,
KeyMetadataRow,
SpendLogsWindow,
)
@dataclass(frozen=True, slots=True)
class _TagRollupMetrics:
tag: str | None
date: str
spend: float
api_requests: int
prompt_tokens: int
@dataclass(frozen=True, slots=True)
class _TagApiKeyCount:
tag: str | None
distinct_api_keys: int
@dataclass(frozen=True, slots=True)
class _TagKeyMembershipCount:
api_key: str
tag_count: int
@dataclass(frozen=True, slots=True)
class _TagFloatSpend:
api_key: str
spend: float
@dataclass(frozen=True, slots=True)
class _TagRankedKey:
api_key: str
@dataclass(frozen=True, slots=True)
class _TagDistinctKeyCount:
total_api_keys: int
_TAG_ROLLUP_METRICS_ADAPTER: Final = TypeAdapter(tuple[_TagRollupMetrics, ...])
_TAG_API_KEY_COUNT_ADAPTER: Final = TypeAdapter(tuple[_TagApiKeyCount, ...])
_TAG_KEY_MEMBERSHIP_COUNT_ADAPTER: Final = TypeAdapter(tuple[_TagKeyMembershipCount, ...])
_TAG_FLOAT_SPEND_ADAPTER: Final = TypeAdapter(tuple[_TagFloatSpend, ...])
_TAG_RANKED_KEY_ADAPTER: Final = TypeAdapter(tuple[_TagRankedKey, ...])
_TAG_DISTINCT_KEY_COUNT_ADAPTER: Final = TypeAdapter(tuple[_TagDistinctKeyCount, ...])
def _scoped_url(url: str, schema: str) -> str:
parsed: Final = urlsplit(url)
return urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})))
@asynccontextmanager
async def _daily_activity_database(
*,
include_tag_activity: bool = False,
include_tag_float_tie_activity: bool = False,
include_team_unassigned_activity: bool = False,
) -> AsyncIterator[Prisma]:
schema: Final = f"integration_{uuid.uuid4().hex}"
url: Final = os.environ["DATABASE_URL"]
with psycopg.connect(url, autocommit=True) as setup:
setup.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
try:
with psycopg.connect(url) as connection:
seed_daily_activity_fixture(
connection,
schema=schema,
ptu_sentinel_api_key=constants.PTU_SENTINEL_API_KEY,
)
if include_tag_activity:
seed_daily_tag_activity_fixture(connection, schema=schema)
if include_tag_float_tie_activity:
seed_daily_tag_float_tie_fixture(connection, schema=schema)
if include_team_unassigned_activity:
seed_daily_team_unassigned_fixture(
connection, schema=schema, ptu_sentinel_api_key=constants.PTU_SENTINEL_API_KEY
)
database: Final = Prisma(datasource={"url": _scoped_url(url, schema)})
await database.connect()
try:
yield database
finally:
await database.disconnect()
finally:
setup.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))
@dataclass(frozen=True, slots=True)
class _PrismaDatabase:
db: Prisma
@dataclass(frozen=True, slots=True)
class _ProxyReads(DailyActivityProxyReads):
database: Prisma
async def recover_key_metadata(
self, resolved: Mapping[str, KeyMetadataRow], api_keys: frozenset[str], window: SpendLogsWindow | None
) -> Mapping[str, KeyMetadataRow]:
user_ids: Final = frozenset(row.user_id for row in resolved.values() if row.user_id)
user_rows: Final = await find_many_in(self.database.litellm_usertable, "user_id", user_ids) if user_ids else ()
user_emails: Final = MappingProxyType({row.user_id: row.user_email for row in user_rows if row.user_email})
return MappingProxyType(
{
key: KeyMetadataRow(
api_key=row.api_key,
key_alias=row.key_alias,
team_id=row.team_id,
user_id=row.user_id,
user_email=row.user_email or user_emails.get(row.user_id),
key_exists=row.key_exists,
tags=row.tags,
)
for key, row in resolved.items()
}
)
def _repository(database: Prisma) -> DailyActivityRepository:
client: Final = cast(DailyActivityDatabase, _PrismaDatabase(database))
return DailyActivityRepository(client, proxy_reads=_ProxyReads(database))
def _scope(
table: DailyActivityTable,
entity_id_field: str,
entity_id: str,
api_keys: tuple[str, ...] | None = None,
) -> DailyActivityScope:
return DailyActivityScope(
table=table,
entity_id_field=entity_id_field,
entity_ids=(entity_id,),
exclude_entity_ids=(),
api_keys=api_keys,
start_date="2026-06-01",
end_date="2026-06-01",
model=None,
timezone_offset_minutes=None,
)
@pytest.mark.asyncio
async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch: pytest.MonkeyPatch) -> None:
async with _daily_activity_database() as database:
repository: Final = _repository(database)
team_scope: Final = _scope(DailyActivityTable.TEAM, "team_id", "team-1")
monkeypatch.setattr(constants, "USAGE_EXPORT_BATCH_SIZE", 2)
aggregate: Final = await repository.aggregated(team_scope, include_entity_breakdown=True, api_key_limit=3)
totals: Final = tuple(row for row in aggregate.grouping_rows if row.group_level == 127)
assert len(totals) == 1
assert totals[0].spend == 1273.0
assert totals[0].ptu_flat_cost == 42.0
assert aggregate.distinct_api_keys == 5
grouped_keys: Final = frozenset(
row.api_key for row in aggregate.grouping_rows if row.group_level == 31 and row.api_key
)
assert grouped_keys == frozenset(("key-a", "key-b", "key-c"))
entity_totals: Final = tuple(row for row in aggregate.entity_rows or () if row.api_key_rolled)
assert len(entity_totals) == 1
assert entity_totals[0].spend == 1273.0
assert entity_totals[0].ptu_flat_cost == 42.0
targeted_model_keys: Final = await repository.model_top_keys(
team_scope, model_group="model-target", by_model_group=False, limit=3
)
assert tuple(row.api_key for row in targeted_model_keys) == ("key-target",)
popular_model_keys: Final = await repository.model_top_keys(
team_scope, model_group="model-popular", by_model_group=False, limit=3
)
assert tuple(row.api_key for row in popular_model_keys) == ("key-a", "key-b", "key-c")
assert await repository.search_keys(team_scope, search="target", limit=10) == ("key-target",)
assert await repository.search_keys(team_scope, search="deleted-target", limit=20) == ("key-target",)
leakage_keys: Final = await repository.cache_leakage_keys(team_scope, limit=2)
assert tuple(row.api_key for row in leakage_keys) == ("key-cache", "key-c")
exports: Final = tuple(
[row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_KEYS)]
)
assert tuple(row.api_key for row in exports) == (
"key-a",
"key-b",
"key-c",
"key-cache",
"key-target",
)
assert sum(row.spend for row in exports) == 273.0
assert sum(row.flat_cost for row in exports) == 0.0
deleted_key_export: Final = next(row for row in exports if row.api_key == "key-target")
assert (deleted_key_export.key_alias, deleted_key_export.user_id, deleted_key_export.user_email) == (
"deleted-target",
"user-1",
"user@example.com",
)
user_exports: Final = tuple(
[row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_USERS)]
)
assert len(user_exports) == 1
assert (user_exports[0].user_id, user_exports[0].user_email, user_exports[0].spend) == (
"user-1",
"user@example.com",
273.0,
)
daily_export: Final = tuple(
[row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY)]
)
assert len(daily_export) == 1
assert daily_export[0].spend == 1273.0
assert daily_export[0].flat_cost == 42.0
metadata: Final = await repository.key_metadata(frozenset(("key-a", "key-target")), None)
assert metadata["key-a"].key_exists is True
assert metadata["key-a"].key_alias == "alias-a"
assert metadata["key-a"].tags == ("blue", "gold")
assert metadata["key-a"].user_email == "user@example.com"
assert metadata["key-target"].key_exists is False
assert metadata["key-target"].key_alias == "deleted-target"
assert metadata["key-target"].tags == ("archived",)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("search", "api_keys", "expected_keys"),
(
("needle-alias", None, ("needle-key",)),
("needle-user", None, ("needle-key",)),
("needle@example.com", None, ("needle-key",)),
("needle-alias", ("key-a",), ()),
),
)
async def test_search_keys_matches_token_metadata_outside_top_n_and_respects_scope(
search: str,
api_keys: tuple[str, ...] | None,
expected_keys: tuple[str, ...],
) -> None:
async with _daily_activity_database() as database:
await database.execute_raw(
"""
INSERT INTO "LiteLLM_DailyTeamSpend" (
id, team_id, date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, prompt_tokens, completion_tokens,
cache_read_input_tokens, cache_creation_input_tokens, spend, api_requests,
successful_requests, failed_requests, ptu_flat_cost, updated_at
) VALUES (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19::timestamp
)
""",
"needle-row",
"team-1",
"2026-06-01",
"needle-key",
"model-needle",
"",
"provider-a",
None,
"/v1/chat/completions",
1,
1,
0,
0,
0.5,
1,
1,
0,
0.0,
"2026-06-01 12:00:00",
)
await database.execute_raw(
"""
INSERT INTO "LiteLLM_VerificationToken" (token, key_alias, team_id, user_id, metadata, models)
VALUES ($1, $2, $3, $4, $5::jsonb, $6::text[])
""",
"needle-key",
"needle-alias",
"team-1",
"needle-user",
'{"tags": []}',
[],
)
await database.execute_raw(
"""
INSERT INTO "LiteLLM_UserTable" (user_id, user_email, models)
VALUES ($1, $2, $3::text[])
""",
"needle-user",
"needle@example.com",
[],
)
repository: Final = _repository(database)
team_scope: Final = _scope(DailyActivityTable.TEAM, "team_id", "team-1", api_keys=api_keys)
aggregate: Final = await repository.aggregated(team_scope, include_entity_breakdown=False, api_key_limit=1)
top_keys: Final = frozenset(
row.api_key for row in aggregate.grouping_rows if row.group_level == 31 and row.api_key
)
assert "needle-key" not in top_keys
assert await repository.search_keys(team_scope, search=search, limit=10) == expected_keys
@pytest.mark.asyncio
async def test_aggregated_returns_totals_with_a_one_key_limit() -> None:
async with _daily_activity_database() as database:
aggregate: Final = await _repository(database).aggregated(
_scope(DailyActivityTable.TEAM, "team_id", "team-1"),
include_entity_breakdown=False,
api_key_limit=1,
)
totals: Final = tuple(row for row in aggregate.grouping_rows if row.group_level == 127)
per_key_rows: Final = tuple(row for row in aggregate.grouping_rows if row.api_key is not None)
per_key_names: Final = frozenset(row.api_key for row in per_key_rows)
assert len(totals) == 1
assert totals[0].spend == 1273.0
assert len(per_key_rows) == 6
assert len(per_key_names) == 1
@pytest.mark.asyncio
async def test_tag_entity_rollups_bound_keys_and_preserve_full_scope_totals() -> None:
async with _daily_activity_database(include_tag_activity=True) as database:
scope: Final = DailyActivityScope(
table=DailyActivityTable.TAG,
entity_id_field="tag",
entity_ids=None,
exclude_entity_ids=(),
api_keys=None,
start_date="2026-06-01",
end_date="2026-06-02",
model=None,
timezone_offset_minutes=None,
)
repository: Final = _repository(database)
aggregate: Final = await repository.aggregated(scope, include_entity_breakdown=True, api_key_limit=3)
independent_metrics: Final = _TAG_ROLLUP_METRICS_ADAPTER.validate_python(
await database.query_raw(
"""
SELECT tag, date, SUM(spend)::float AS spend,
SUM(api_requests)::bigint AS api_requests,
SUM(prompt_tokens)::bigint AS prompt_tokens
FROM "LiteLLM_DailyTagSpend"
WHERE date >= $1 AND date <= $2
GROUP BY tag, date
""",
"2026-06-01",
"2026-06-02",
)
)
independent_key_counts: Final = _TAG_API_KEY_COUNT_ADAPTER.validate_python(
await database.query_raw(
"""
SELECT tag, COUNT(DISTINCT api_key)::bigint AS distinct_api_keys
FROM "LiteLLM_DailyTagSpend"
WHERE date >= $1 AND date <= $2 AND api_key <> $3
GROUP BY tag
""",
"2026-06-01",
"2026-06-02",
constants.PTU_SENTINEL_API_KEY,
)
)
independent_key_memberships: Final = _TAG_KEY_MEMBERSHIP_COUNT_ADAPTER.validate_python(
await database.query_raw(
"""
SELECT api_key, COUNT(DISTINCT tag)::bigint AS tag_count
FROM "LiteLLM_DailyTagSpend"
WHERE date >= $1 AND date <= $2
GROUP BY api_key
""",
"2026-06-01",
"2026-06-02",
)
)
expected_metrics: Final = MappingProxyType({(row.date, row.tag): row for row in independent_metrics})
expected_key_counts: Final = MappingProxyType(
{row.tag: row.distinct_api_keys for row in independent_key_counts}
)
assert {row.date for row in independent_metrics} == {"2026-06-01", "2026-06-02"}
assert len(expected_key_counts) == 4
assert len(independent_key_memberships) == 8
assert all(row.tag_count == 2 for row in independent_key_memberships)
assert max(expected_key_counts.values()) > 3
entity_rows: Final = aggregate.entity_rows or ()
rolled_rows: Final = tuple(row for row in entity_rows if row.api_key_rolled)
keyed_rows: Final = tuple(row for row in entity_rows if not row.api_key_rolled and row.api_key)
top_level_keys: Final = frozenset(
row.api_key for row in aggregate.grouping_rows if row.group_level == 31 and row.api_key
)
entity_day_keys: Final = MappingProxyType(
{
key: frozenset(row.api_key for row in keyed_rows if (row.date, row.entity_id) == key and row.api_key)
for key in frozenset((row.date, row.entity_id) for row in keyed_rows)
}
)
assert entity_day_keys
assert max(len(keys) for keys in entity_day_keys.values()) <= 3
assert frozenset(row.api_key for row in keyed_rows) <= top_level_keys
rolled_by_entity_day: Final = MappingProxyType(
{(row.date, row.entity_id): row for row in rolled_rows if row.date is not None}
)
assert set(rolled_by_entity_day) == set(expected_metrics)
for key, row in rolled_by_entity_day.items():
assert row.spend is not None
assert isclose(row.spend, expected_metrics[key].spend, rel_tol=1e-9, abs_tol=1e-9)
assert row.api_requests == expected_metrics[key].api_requests
assert row.prompt_tokens == expected_metrics[key].prompt_tokens
assert row.distinct_api_keys == expected_key_counts[row.entity_id]
response: Final = await get_daily_activity_aggregated(
repository,
scope,
include_entity_breakdown=True,
api_key_limit=3,
)
assert response.metadata.entity_total_api_keys == {
tag: count for tag, count in expected_key_counts.items() if tag is not None
}
assert all(
all(len(entity.api_key_breakdown) <= 3 for entity in day.breakdown.entities.values())
for day in response.results
)
@pytest.mark.asyncio
async def test_key_pages_match_full_tag_ranking_and_aggregate_top_keys() -> None:
async with _daily_activity_database(include_tag_activity=True) as database:
scope: Final = DailyActivityScope(
table=DailyActivityTable.TAG,
entity_id_field="tag",
entity_ids=None,
exclude_entity_ids=(),
api_keys=None,
start_date="2026-06-01",
end_date="2026-06-02",
model=None,
timezone_offset_minutes=None,
)
repository: Final = _repository(database)
first_page: Final = await repository.key_page(scope, offset=0, limit=3)
remaining_pages: Final = tuple(
[
await repository.key_page(scope, offset=offset, limit=3)
for offset in range(3, first_page.total_api_keys, 3)
]
)
pages: Final = (first_page, *remaining_pages)
actual_keys: Final = tuple(row.api_key for page in pages for row in page.rows)
expected_rows: Final = _TAG_RANKED_KEY_ADAPTER.validate_python(
await database.query_raw(
"""
SELECT api_key
FROM "LiteLLM_DailyTagSpend"
WHERE date >= $1 AND date <= $2 AND api_key <> $3
GROUP BY api_key
ORDER BY SUM(spend::numeric) DESC, api_key
""",
"2026-06-01",
"2026-06-02",
constants.PTU_SENTINEL_API_KEY,
)
)
expected_keys: Final = tuple(row.api_key for row in expected_rows)
independent_count: Final = _TAG_DISTINCT_KEY_COUNT_ADAPTER.validate_python(
await database.query_raw(
"""
SELECT COUNT(DISTINCT api_key)::bigint AS total_api_keys
FROM "LiteLLM_DailyTagSpend"
WHERE date >= $1 AND date <= $2 AND api_key <> $3
""",
"2026-06-01",
"2026-06-02",
constants.PTU_SENTINEL_API_KEY,
)
)[0].total_api_keys
aggregate: Final = await repository.aggregated(scope, include_entity_breakdown=False, api_key_limit=3)
aggregate_top_keys: Final = frozenset(
row.api_key for row in aggregate.grouping_rows if row.group_level == 31 and row.api_key is not None
)
empty_page: Final = await repository.key_page(scope, offset=independent_count + 3, limit=3)
assert actual_keys == expected_keys
assert len(actual_keys) == len(frozenset(actual_keys))
assert all(page.total_api_keys == independent_count for page in pages)
assert first_page.total_api_keys == independent_count
assert frozenset(row.api_key for row in first_page.rows) == aggregate_top_keys
assert empty_page.rows == ()
assert empty_page.total_api_keys == independent_count
@pytest.mark.asyncio
async def test_top_api_key_rank_is_order_independent_for_float_ties() -> None:
async with _daily_activity_database(include_tag_float_tie_activity=True) as database:
float_totals: Final = _TAG_FLOAT_SPEND_ADAPTER.validate_python(
await database.query_raw(
"""
SELECT api_key, SUM(spend)::float AS spend
FROM "LiteLLM_DailyTagSpend"
WHERE tag = $1 AND date = $2
GROUP BY api_key
""",
"tag-float-tie",
"2026-06-01",
)
)
float_spends: Final = MappingProxyType({row.api_key: row.spend for row in float_totals})
assert float_spends["key-z"] > float_spends["key-a"]
scope: Final = DailyActivityScope(
table=DailyActivityTable.TAG,
entity_id_field="tag",
entity_ids=None,
exclude_entity_ids=(),
api_keys=None,
start_date="2026-06-01",
end_date="2026-06-01",
model=None,
timezone_offset_minutes=None,
)
aggregate: Final = await _repository(database).aggregated(scope, include_entity_breakdown=True, api_key_limit=1)
key_page: Final = await _repository(database).key_page(scope, offset=0, limit=1)
top_level_keys: Final = frozenset(
row.api_key for row in aggregate.grouping_rows if row.group_level == 31 and row.api_key
)
entity_keyed_keys: Final = frozenset(
row.api_key for row in aggregate.entity_rows or () if not row.api_key_rolled and row.api_key is not None
)
expected_keys: Final = frozenset(("key-a",))
assert (top_level_keys, entity_keyed_keys) == (
expected_keys,
expected_keys,
), f"plain float SUM totals: {float_spends}"
assert frozenset(row.api_key for row in key_page.rows) == expected_keys
@pytest.mark.asyncio
@pytest.mark.parametrize(
("table", "entity_field", "entity_id", "fixture_name"),
(
(DailyActivityTable.USER, "user_id", "user-1", "daily_activity_user.json"),
(DailyActivityTable.TEAM, "team_id", "team-1", "daily_activity_team.json"),
),
)
async def test_aggregated_response_matches_base_golden(
table: DailyActivityTable, entity_field: str, entity_id: str, fixture_name: str
) -> None:
async with _daily_activity_database() as database:
result: Final = await get_daily_activity_aggregated(
_repository(database),
_scope(table, entity_field, entity_id),
entity_metadata_field=MappingProxyType({"team-1": {"team_alias": "Usage Team"}}),
include_entity_breakdown=True,
)
golden_path: Final = Path(__file__).with_name("fixtures") / fixture_name
assert result.model_dump_json() + "\n" == golden_path.read_text()
@pytest.mark.asyncio
async def test_team_entity_rollups_merge_null_and_empty_entity_ids() -> None:
async with _daily_activity_database(include_team_unassigned_activity=True) as database:
scope: Final = DailyActivityScope(
table=DailyActivityTable.TEAM,
entity_id_field="team_id",
entity_ids=None,
exclude_entity_ids=(),
api_keys=None,
start_date="2026-06-03",
end_date="2026-06-03",
model=None,
timezone_offset_minutes=None,
)
aggregate: Final = await _repository(database).aggregated(scope, include_entity_breakdown=True, api_key_limit=3)
totals: Final = tuple(row for row in aggregate.grouping_rows if row.group_level == 127)
assert len(totals) == 1
assert totals[0].spend == 23.0
assert aggregate.distinct_api_keys == 2
rolled_rows: Final = tuple(row for row in aggregate.entity_rows or () if row.api_key_rolled)
assert len(rolled_rows) == 1
assert rolled_rows[0].entity_id == ""
assert rolled_rows[0].spend == 23.0
assert rolled_rows[0].ptu_flat_cost == 13.0
assert rolled_rows[0].distinct_api_keys == 2
keyed_rows: Final = tuple(row for row in aggregate.entity_rows or () if not row.api_key_rolled)
assert {row.entity_id for row in keyed_rows} == {""}
assert {row.api_key for row in keyed_rows} == {"key-unassigned-null", "key-unassigned-empty"}

778
tests/unit/proxy/management_endpoints/test_common_daily_activity.py
View file

File diff suppressed because it is too large Load diff

29
tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -2574,19 +2574,18 @@ async def test_get_user_daily_activity_aggregated_admin_global_view(monkeypatch,
assert result is mock_response
# Verify the helper was called with the right parameters
mock_get_daily_agg.assert_called_once_with(
prisma_client=mock_prisma_client,
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None, # global view: no user_id filter
entity_metadata_field=None,
start_date="2025-02-01",
end_date="2025-02-28",
model="gpt-4",
api_key=None,
timezone_offset_minutes=480,
include_current_utc_day=include_current_utc_day,
)
mock_get_daily_agg.assert_called_once()
repository, scope = mock_get_daily_agg.call_args.args
assert repository is not None
assert scope.table.value == "litellm_dailyuserspend"
assert scope.entity_id_field == "user_id"
assert scope.entity_ids is None
assert scope.start_date == "2025-02-01"
assert scope.end_date == "2025-02-28"
assert scope.model == "gpt-4"
assert scope.api_keys is None
assert scope.timezone_offset_minutes == 480
assert scope.include_current_utc_day is include_current_utc_day
@pytest.mark.asyncio
@ -2655,7 +2654,9 @@ async def test_get_user_daily_activity_aggregated_non_admin_cannot_view_other_us
assert result is mock_response
mock_get_daily_agg.assert_called_once()
assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "regular-user-123"
repository, scope = mock_get_daily_agg.call_args.args
assert repository is not None
assert scope.entity_ids == ("regular-user-123",)
@pytest.mark.asyncio

12
tests/unit/proxy/management_endpoints/test_team_endpoints.py
View file

@ -14669,15 +14669,17 @@ async def test_get_team_daily_activity_aggregated_scopes_and_flags(mock_db_clien
)
mock_aggregated.assert_called_once()
call_kwargs = mock_aggregated.call_args[1]
assert call_kwargs["api_key"] == ["user_key_1"]
assert call_kwargs["entity_id"] == [team_id]
repository, scope = mock_aggregated.call_args.args
call_kwargs = mock_aggregated.call_args.kwargs
assert repository is not None
assert scope.api_keys == ("user_key_1",)
assert scope.entity_ids == (team_id,)
assert call_kwargs["entity_metadata_field"] == {
team_id: {"team_alias": "Test Team"}
}
assert call_kwargs["include_entity_breakdown"] is True
assert call_kwargs["timezone_offset_minutes"] == 480
assert call_kwargs["table_name"] == "litellm_dailyteamspend"
assert scope.timezone_offset_minutes == 480
assert scope.table.value == "litellm_dailyteamspend"
@pytest.mark.asyncio

11
tests/unit/proxy/proxy_server/test_lifecycle.py
View file

@ -17,19 +17,16 @@ Pins covered:
from __future__ import annotations
import asyncio
import inspect
import json
import logging
import os
import subprocess
from collections.abc import Awaitable, Callable
from typing import List, Optional, Union
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from fastapi import FastAPI
from pydantic import BaseModel
from typing_extensions import TypedDict
@ -682,16 +679,16 @@ class _SampleTD(TypedDict):
def test_resolve_typed_dict_type_finds_class_in_optional():
typ = Optional[_SampleTD]
typ = _SampleTD | None
result = _resolve_typed_dict_type(typ)
observed = {
"input_repr": "Optional[_SampleTD]",
"input_repr": "_SampleTD | None",
"result_is_sample_td": result is _SampleTD,
"result_is_class": isinstance(result, type),
}
assert normalize(observed) == {
"input_repr": "Optional[_SampleTD]",
"input_repr": "_SampleTD | None",
"result_is_sample_td": True,
"result_is_class": True,
}
@ -717,7 +714,7 @@ class _SampleModelB(BaseModel):
def test_resolve_pydantic_type_extracts_non_none_args_from_union():
typ = Union[_SampleModelA, _SampleModelB, None]
typ = _SampleModelA | _SampleModelB | None
result = _resolve_pydantic_type(typ)
observed = {

549
tests/unit/repositories/test_daily_activity_repository.py Normal file
View file

@ -0,0 +1,549 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final
import pytest
from pydantic import ValidationError
from litellm import constants
from litellm.repositories.daily_activity_repository import DailyActivityRepository
from litellm.repositories.daily_activity_sql import (
ExportCursor,
build_cache_leakage_keys_sql,
build_entity_rollup_sql,
build_export_sql,
build_key_page_sql,
build_key_search_sql,
build_model_top_keys_sql,
)
from litellm.types.repositories.daily_activity import (
DailyActivityScope,
DailyActivityTable,
ExportType,
KeyMetadataRow,
KeyPage,
KeySpendRow,
SpendLogsWindow,
)
@dataclass(frozen=True, slots=True)
class _FakeVerificationToken:
token: str
key_alias: str | None
team_id: str | None
user_id: str | None
metadata: object | None
@dataclass(frozen=True, slots=True)
class _FakeDeletedVerificationToken(_FakeVerificationToken):
deleted_at: datetime
def _scope(
*,
table: DailyActivityTable = DailyActivityTable.USER,
entity_ids: tuple[str, ...] | None = ("user-1",),
api_keys: tuple[str, ...] | None = None,
exclude_entity_ids: tuple[str, ...] = (),
model: str | None = None,
) -> DailyActivityScope:
entity_field: Final = {
DailyActivityTable.USER: "user_id",
DailyActivityTable.TEAM: "team_id",
DailyActivityTable.TAG: "tag",
DailyActivityTable.ORGANIZATION: "organization_id",
DailyActivityTable.CUSTOMER: "end_user_id",
DailyActivityTable.AGENT: "agent_id",
}[table]
return DailyActivityScope(
table=table,
entity_id_field=entity_field,
entity_ids=entity_ids,
exclude_entity_ids=exclude_entity_ids,
api_keys=api_keys,
start_date="2026-01-01",
end_date="2026-01-31",
model=model,
timezone_offset_minutes=None,
)
def _key_spend_row(api_key: str) -> dict[str, object]:
return {
"api_key": api_key,
"spend": 1.0,
"prompt_tokens": 10,
"completion_tokens": 2,
"total_tokens": 12,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
"cache_read_input_tokens": 3,
"cache_creation_input_tokens": 1,
}
def _export_row(api_key: str | None) -> dict[str, object]:
return {
"date": "2026-01-01",
"entity_id": "user-1",
"entity_alias": None,
"api_key": api_key,
"key_alias": None,
"user_id": None,
"user_email": None,
"model": None,
"spend": 1.0,
"flat_cost": 0.0,
"prompt_tokens": 10,
"completion_tokens": 2,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
"cache_read_input_tokens": 3,
"cache_creation_input_tokens": 1,
}
class _FakeTable:
def __init__(self, rows: Sequence[object] = ()) -> None:
self.rows: Final = tuple(rows)
self.find_many_calls: list[Mapping[str, object]] = []
self.count_calls: list[Mapping[str, object]] = []
self.pagination_calls: list[tuple[int | None, int | None, tuple[Mapping[str, str], ...] | None]] = []
async def find_many(
self,
*,
where: Mapping[str, object],
skip: int | None = None,
take: int | None = None,
order: tuple[Mapping[str, str], ...] | None = None,
) -> tuple[object, ...]:
self.find_many_calls.append(where)
self.pagination_calls.append((skip, take, order))
if "token" not in where:
return self.rows
token_filter: Final = where["token"]
if not isinstance(token_filter, Mapping):
return ()
token_values: Final = token_filter.get("in")
if not isinstance(token_values, list):
return ()
return tuple(row for row in self.rows if isinstance(row, _FakeVerificationToken) and row.token in token_values)
async def count(self, *, where: Mapping[str, object]) -> int:
self.count_calls.append(where)
return len(self.rows)
class _FailingTable(_FakeTable):
def __init__(self, failure: str) -> None:
super().__init__()
self.failure: Final = failure
async def find_many(
self,
*,
where: Mapping[str, object],
skip: int | None = None,
take: int | None = None,
order: tuple[Mapping[str, str], ...] | None = None,
) -> tuple[object, ...]:
raise RuntimeError(f"{self.failure}: {where!r} {skip!r} {take!r} {order!r}")
class _FakeDatabase:
def __init__(self, responses: Sequence[Sequence[Mapping[str, object]] | None] = ()) -> None:
self.responses = tuple(responses)
self.query_calls: list[tuple[str, tuple[object, ...]]] = []
self.litellm_verificationtoken = _FakeTable()
self.litellm_deletedverificationtoken = _FakeTable()
self.litellm_dailyuserspend = _FakeTable()
self.litellm_dailyteamspend = _FakeTable()
self.litellm_dailytagspend = _FakeTable()
self.litellm_dailyorganizationspend = _FakeTable()
self.litellm_dailyenduserspend = _FakeTable()
self.litellm_dailyagentspend = _FakeTable()
async def query_raw(self, query: str, *params: object) -> Sequence[Mapping[str, object]] | None:
self.query_calls.append((query, params))
response_index: Final = len(self.query_calls) - 1
if response_index >= len(self.responses):
return ()
return self.responses[response_index]
class _FakePrismaClient:
def __init__(self, database: _FakeDatabase) -> None:
self.db: Final = database
class _ProxyReads:
def __init__(self) -> None:
self.recovery_calls: list[tuple[Mapping[str, KeyMetadataRow], frozenset[str], SpendLogsWindow | None]] = []
async def recover_key_metadata(
self,
resolved: Mapping[str, KeyMetadataRow],
api_keys: frozenset[str],
window: SpendLogsWindow | None,
) -> Mapping[str, KeyMetadataRow]:
self.recovery_calls.append((resolved, api_keys, window))
return resolved
def _repository(
database: _FakeDatabase, proxy_reads: _ProxyReads | None = None
) -> tuple[DailyActivityRepository, _ProxyReads]:
reads: Final = proxy_reads if proxy_reads is not None else _ProxyReads()
return DailyActivityRepository(_FakePrismaClient(database), proxy_reads=reads), reads
@pytest.mark.asyncio
async def test_key_methods_send_builder_queries_with_caller_limits() -> None:
database = _FakeDatabase(((_key_spend_row("key-a"),), (_key_spend_row("key-b"),), (_key_spend_row("key-c"),)))
repository, _ = _repository(database)
scope = _scope()
assert await repository.search_keys(scope, search="key", limit=2) == ("key-a",)
model_keys: Final = await repository.model_top_keys(scope, model_group="model-a", by_model_group=True, limit=2)
leakage_keys: Final = await repository.cache_leakage_keys(scope, limit=2)
assert tuple(row.api_key for row in model_keys) == ("key-b",)
assert tuple(row.api_key for row in leakage_keys) == ("key-c",)
assert model_keys[0].spend == 1.0
assert leakage_keys[0].prompt_tokens - leakage_keys[0].cache_read_input_tokens == 7
assert database.query_calls == [
(
build_key_search_sql(scope, search="key", limit=2).sql,
build_key_search_sql(scope, search="key", limit=2).params,
),
(
build_model_top_keys_sql(scope, model_group="model-a", by_model_group=True, limit=2).sql,
build_model_top_keys_sql(scope, model_group="model-a", by_model_group=True, limit=2).params,
),
(
build_cache_leakage_keys_sql(scope, limit=2).sql,
build_cache_leakage_keys_sql(scope, limit=2).params,
),
]
@pytest.mark.asyncio
async def test_key_page_maps_rows_and_keeps_total_for_an_empty_page() -> None:
database = _FakeDatabase(
(
({"total_api_keys": 2, **_key_spend_row("key-a")},),
({"total_api_keys": 2, "api_key": None},),
)
)
repository, _ = _repository(database)
scope = _scope()
first_page: Final = await repository.key_page(scope, offset=0, limit=1)
empty_page: Final = await repository.key_page(scope, offset=2, limit=1)
assert first_page == KeyPage(
rows=(
KeySpendRow(
api_key="key-a",
spend=1.0,
prompt_tokens=10,
completion_tokens=2,
total_tokens=12,
api_requests=1,
successful_requests=1,
failed_requests=0,
cache_read_input_tokens=3,
cache_creation_input_tokens=1,
),
),
total_api_keys=2,
)
assert empty_page == KeyPage(rows=(), total_api_keys=2)
assert database.query_calls == [
(
build_key_page_sql(scope, offset=0, limit=1).sql,
build_key_page_sql(scope, offset=0, limit=1).params,
),
(
build_key_page_sql(scope, offset=2, limit=1).sql,
build_key_page_sql(scope, offset=2, limit=1).params,
),
]
@pytest.mark.asyncio
async def test_key_methods_reject_limits_outside_bounds() -> None:
database = _FakeDatabase()
repository, _ = _repository(database)
with pytest.raises(ValueError, match="limit"):
await repository.search_keys(_scope(), search="key", limit=0)
with pytest.raises(ValueError, match="limit"):
await repository.model_top_keys(_scope(), model_group="model-a", by_model_group=False, limit=0)
with pytest.raises(ValueError, match="limit"):
await repository.cache_leakage_keys(_scope(), limit=0)
with pytest.raises(ValueError, match="limit"):
await repository.search_keys(_scope(), search="key", limit=constants.USAGE_KEY_SEARCH_MAX + 1)
with pytest.raises(ValueError, match="limit"):
await repository.model_top_keys(
_scope(), model_group="model-a", by_model_group=False, limit=constants.USAGE_MODEL_TOP_KEYS_MAX + 1
)
with pytest.raises(ValueError, match="limit"):
await repository.cache_leakage_keys(_scope(), limit=constants.USAGE_CACHE_LEAKAGE_KEYS_MAX + 1)
assert database.query_calls == []
@pytest.mark.asyncio
async def test_key_spend_validation_rejects_malformed_rows() -> None:
repository, _ = _repository(_FakeDatabase((({"api_key": "missing-metrics"},),)))
with pytest.raises(ValidationError):
await repository.search_keys(_scope(), search="key", limit=1)
@pytest.mark.asyncio
async def test_key_metadata_prefers_active_rows_and_recovers_all_requested_keys() -> None:
database = _FakeDatabase()
active: Final = _FakeVerificationToken(
token="active",
key_alias="current",
team_id="team-active",
user_id="user-active",
metadata={"tags": ["production", "internal"]},
)
deleted_active_duplicate: Final = _FakeDeletedVerificationToken(
token="active",
key_alias="stale",
team_id="team-stale",
user_id="user-stale",
metadata={"tags": []},
deleted_at=datetime(2026, 1, 3, tzinfo=timezone.utc),
)
deleted_older: Final = _FakeDeletedVerificationToken(
token="deleted",
key_alias="older",
team_id=None,
user_id=None,
metadata={"tags": "invalid"},
deleted_at=datetime(2026, 1, 2, tzinfo=timezone.utc),
)
deleted_newer: Final = _FakeDeletedVerificationToken(
token="deleted",
key_alias="newer",
team_id=None,
user_id=None,
metadata={"tags": ["archived"]},
deleted_at=datetime(2026, 1, 4, tzinfo=timezone.utc),
)
malformed_non_list: Final = _FakeVerificationToken(
token="malformed-non-list",
key_alias=None,
team_id=None,
user_id=None,
metadata={"tags": "invalid"},
)
malformed_list: Final = _FakeVerificationToken(
token="malformed-list",
key_alias=None,
team_id=None,
user_id=None,
metadata={"tags": [1]},
)
database.litellm_verificationtoken = _FakeTable((active, malformed_non_list, malformed_list))
database.litellm_deletedverificationtoken = _FakeTable((deleted_active_duplicate, deleted_older, deleted_newer))
proxy_reads: Final = _ProxyReads()
repository, _ = _repository(database, proxy_reads)
window: Final = (datetime(2026, 1, 1), datetime(2026, 2, 1))
requested: Final = frozenset(("active", "deleted", "malformed-non-list", "malformed-list", "unresolved"))
result = await repository.key_metadata(requested, window)
assert result["active"] == KeyMetadataRow(
api_key="active",
key_alias="current",
team_id="team-active",
user_id="user-active",
user_email=None,
key_exists=True,
tags=("production", "internal"),
)
assert result["deleted"].key_alias == "newer"
assert result["deleted"].key_exists is False
assert result["deleted"].tags == ("archived",)
assert result["malformed-non-list"].tags == ()
assert result["malformed-list"].tags == ()
assert len(database.litellm_deletedverificationtoken.find_many_calls) == 1
assert set(database.litellm_deletedverificationtoken.find_many_calls[0]["token"]["in"]) == {
"deleted",
"unresolved",
}
assert proxy_reads.recovery_calls == [
(
result,
requested,
window,
)
]
@pytest.mark.asyncio
async def test_key_metadata_continues_with_active_rows_when_deleted_lookup_fails() -> None:
database = _FakeDatabase()
active: Final = _FakeVerificationToken(
token="active",
key_alias="current",
team_id=None,
user_id=None,
metadata={"tags": []},
)
database.litellm_verificationtoken = _FakeTable((active,))
database.litellm_deletedverificationtoken = _FailingTable("deleted token query failed")
repository, proxy_reads = _repository(database)
result = await repository.key_metadata(frozenset(("active", "deleted")), None)
assert result["active"].key_alias == "current"
assert tuple(proxy_reads.recovery_calls[0][0]) == ("active",)
assert proxy_reads.recovery_calls[0][1] == frozenset(("active", "deleted"))
@pytest.mark.asyncio
async def test_key_metadata_empty_set_does_not_query_tables() -> None:
database = _FakeDatabase()
repository, proxy_reads = _repository(database)
assert await repository.key_metadata(frozenset(), None) == {}
assert database.litellm_verificationtoken.find_many_calls == []
assert proxy_reads.recovery_calls == []
@pytest.mark.asyncio
async def test_key_metadata_propagates_active_token_lookup_failures() -> None:
database = _FakeDatabase()
database.litellm_verificationtoken = _FailingTable("active token query failed")
repository, _ = _repository(database)
with pytest.raises(RuntimeError, match="active token query failed"):
await repository.key_metadata(frozenset(("active",)), None)
assert database.litellm_deletedverificationtoken.find_many_calls == []
@pytest.mark.asyncio
async def test_aggregated_normalizes_a_null_raw_query_result() -> None:
database = _FakeDatabase((None,))
repository, _ = _repository(database)
result = await repository.aggregated(
_scope(), include_entity_breakdown=False, api_key_limit=constants.USAGE_TOP_API_KEYS_DEFAULT
)
assert result.grouping_rows == ()
assert result.entity_rows is None
assert result.distinct_api_keys == 0
assert len(database.query_calls) == 1
@pytest.mark.asyncio
async def test_aggregated_passes_api_key_limit_to_entity_rollup_query() -> None:
database = _FakeDatabase(((), ()))
repository, _ = _repository(database)
scope = _scope(table=DailyActivityTable.TEAM)
result = await repository.aggregated(scope, include_entity_breakdown=True, api_key_limit=3)
assert result.entity_rows == ()
assert database.query_calls[1] == (
build_entity_rollup_sql(scope, api_key_limit=3).sql,
build_entity_rollup_sql(scope, api_key_limit=3).params,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("table", "entity_field"),
[
(DailyActivityTable.USER, "user_id"),
(DailyActivityTable.TEAM, "team_id"),
(DailyActivityTable.TAG, "tag"),
(DailyActivityTable.ORGANIZATION, "organization_id"),
(DailyActivityTable.CUSTOMER, "end_user_id"),
(DailyActivityTable.AGENT, "agent_id"),
],
)
async def test_daily_rows_selects_the_table_and_applies_filters_and_pagination(
table: DailyActivityTable, entity_field: str
) -> None:
database = _FakeDatabase()
repository, _ = _repository(database)
scope = _scope(
table=table,
entity_ids=("entity-1",),
exclude_entity_ids=("excluded-1",),
api_keys=("key-1",),
model="model-1",
)
result = await repository.daily_rows(scope, page=3, page_size=2)
expected_where: Final = {
"date": {"gte": "2026-01-01", "lte": "2026-01-31"},
entity_field: {"in": ["entity-1"], "not": {"in": ["excluded-1"]}},
"model": "model-1",
"api_key": {"in": ["key-1"]},
}
tables: Final = {
DailyActivityTable.USER: database.litellm_dailyuserspend,
DailyActivityTable.TEAM: database.litellm_dailyteamspend,
DailyActivityTable.TAG: database.litellm_dailytagspend,
DailyActivityTable.ORGANIZATION: database.litellm_dailyorganizationspend,
DailyActivityTable.CUSTOMER: database.litellm_dailyenduserspend,
DailyActivityTable.AGENT: database.litellm_dailyagentspend,
}
selected_table: Final = tables[table]
assert result.total_count == 0
assert result.rows == ()
assert selected_table.count_calls == [expected_where]
assert selected_table.find_many_calls == [expected_where]
assert selected_table.pagination_calls == [(4, 2, ({"date": "desc"}, {"id": "asc"}))]
assert sum(len(daily_table.find_many_calls) for daily_table in tables.values()) == 1
@pytest.mark.asyncio
async def test_export_is_lazy_and_uses_the_last_row_as_the_next_cursor(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(constants, "USAGE_EXPORT_BATCH_SIZE", 2)
database = _FakeDatabase(
(
(_export_row("key-1"), _export_row("key-2")),
(_export_row("key-3"), _export_row("key-4")),
(_export_row("key-5"),),
)
)
repository, _ = _repository(database)
rows = repository.export_rows(_scope(), export_type=ExportType.DAILY_WITH_KEYS)
assert database.query_calls == []
assert (await rows.__anext__()).api_key == "key-1"
assert len(database.query_calls) == 1
results = [row async for row in rows]
assert [row.api_key for row in results] == ["key-2", "key-3", "key-4", "key-5"]
assert len(database.query_calls) == 3
assert database.query_calls[1][1][-4:] == ("2026-01-01", "user-1", "key-2", 2)
assert database.query_calls[2][1][-4:] == ("2026-01-01", "user-1", "key-4", 2)
assert (
build_export_sql(
_scope(),
export_type=ExportType.DAILY_WITH_KEYS,
after=ExportCursor("2026-01-01", "user-1", "key-2"),
batch_size=2,
).params
== database.query_calls[1][1]
)

411
tests/unit/repositories/test_daily_activity_sql.py Normal file
View file

@ -0,0 +1,411 @@
from datetime import datetime, timezone
from typing import Final
import pytest
from litellm import constants
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.repositories.daily_activity_sql import (
ExportCursor,
adjust_dates_for_timezone,
build_aggregated_sql,
build_cache_leakage_keys_sql,
build_entity_rollup_sql,
build_export_sql,
build_key_page_sql,
build_key_search_sql,
build_model_top_keys_sql,
build_where_clause,
)
from litellm.types.proxy.management_endpoints.common_daily_activity import SpendMetrics
from litellm.types.repositories.daily_activity import DailyActivityScope, DailyActivityTable, ExportType
def _scope(
*,
table: DailyActivityTable = DailyActivityTable.USER,
entity_ids: tuple[str, ...] | None = ("user-1",),
exclude_entity_ids: tuple[str, ...] = (),
api_keys: tuple[str, ...] | None = None,
model: str | None = None,
timezone_offset_minutes: int | None = None,
include_current_utc_day: bool = False,
start_date: str = "2026-01-01",
end_date: str = "2026-01-31",
) -> DailyActivityScope:
entity_field = {
DailyActivityTable.USER: "user_id",
DailyActivityTable.TEAM: "team_id",
DailyActivityTable.TAG: "tag",
DailyActivityTable.ORGANIZATION: "organization_id",
DailyActivityTable.CUSTOMER: "end_user_id",
DailyActivityTable.AGENT: "agent_id",
}[table]
return DailyActivityScope(
table=table,
entity_id_field=entity_field,
entity_ids=entity_ids,
exclude_entity_ids=exclude_entity_ids,
api_keys=api_keys,
start_date=start_date,
end_date=end_date,
model=model,
timezone_offset_minutes=timezone_offset_minutes,
include_current_utc_day=include_current_utc_day,
)
def test_where_clause_binds_each_filter_as_a_single_array_parameter() -> None:
scope = _scope(
entity_ids=("user-1", "user-2"),
exclude_entity_ids=("user-3",),
api_keys=("key-1", "key-2"),
model="gpt-test",
)
sql, params = build_where_clause(scope)
assert sql == (
'date >= $1 AND date <= $2 AND "user_id" = ANY($3::text[]) '
'AND NOT ("user_id" = ANY($4::text[])) AND model = $5 AND api_key = ANY($6::text[])'
)
assert params == (
"2026-01-01",
"2026-01-31",
["user-1", "user-2"],
["user-3"],
"gpt-test",
["key-1", "key-2"],
)
@pytest.mark.parametrize(
("entity_ids", "api_keys", "expected_sql", "expected_params"),
[
(None, None, "date >= $1 AND date <= $2", ("2026-01-01", "2026-01-31")),
((), None, "date >= $1 AND date <= $2 AND FALSE", ("2026-01-01", "2026-01-31")),
(None, (), "date >= $1 AND date <= $2 AND FALSE", ("2026-01-01", "2026-01-31")),
],
)
def test_where_clause_distinguishes_no_filter_from_empty_membership(
entity_ids: tuple[str, ...] | None,
api_keys: tuple[str, ...] | None,
expected_sql: str,
expected_params: tuple[object, ...],
) -> None:
scope = _scope(entity_ids=entity_ids, api_keys=api_keys)
sql, params = build_where_clause(scope)
assert sql == expected_sql
assert params == expected_params
def test_key_page_sql_orders_exact_spend_and_binds_scope_before_page() -> None:
query = build_key_page_sql(_scope(), offset=7, limit=3)
assert query.params == (
"2026-01-01",
"2026-01-31",
["user-1"],
PTU_SENTINEL_API_KEY,
3,
7,
)
assert "SUM(spend::numeric) AS rank_spend" in query.sql
assert "ORDER BY rank_spend DESC, api_key" in query.sql
assert "(SELECT COUNT(*) FROM ranked)::bigint AS total_api_keys" in query.sql
@pytest.mark.parametrize(
("offset", "limit", "error"),
(
(0, 0, "limit must be between"),
(0, constants.USAGE_KEY_PAGE_MAX + 1, "limit must be between"),
(-1, 1, "offset must be non-negative"),
),
)
def test_key_page_sql_rejects_invalid_page_bounds(offset: int, limit: int, error: str) -> None:
with pytest.raises(ValueError, match=error):
build_key_page_sql(_scope(), offset=offset, limit=limit)
def test_scope_rejects_an_entity_field_not_allowed_for_its_table() -> None:
with pytest.raises(ValueError, match="Invalid entity_id_field"):
DailyActivityScope(
table=DailyActivityTable.USER,
entity_id_field="team_id",
entity_ids=None,
exclude_entity_ids=(),
api_keys=None,
start_date="2026-01-01",
end_date="2026-01-31",
model=None,
timezone_offset_minutes=None,
)
def test_timezone_adjustment_only_extends_an_opted_in_live_range() -> None:
now = datetime(2026, 8, 6, 4, 30, tzinfo=timezone.utc)
assert adjust_dates_for_timezone("2026-07-06", "2026-08-05", 420, include_current_utc_day=True, utc_now=now) == (
"2026-07-06",
"2026-08-06",
)
assert adjust_dates_for_timezone("2026-07-01", "2026-08-04", 420, include_current_utc_day=True, utc_now=now) == (
"2026-07-01",
"2026-08-04",
)
@pytest.mark.parametrize("offset_minutes", [None, 0, -330, -540, -60, 240, 300, 480])
def test_timezone_adjustment_preserves_daily_bucket_dates(offset_minutes: int | None) -> None:
assert adjust_dates_for_timezone("2026-05-29", "2026-05-29", offset_minutes) == (
"2026-05-29",
"2026-05-29",
)
@pytest.mark.parametrize("offset_minutes", [-330, 480])
def test_timezone_adjustment_preserves_single_day_additivity(offset_minutes: int) -> None:
days: Final = ("2026-05-29", "2026-05-30", "2026-05-31", "2026-06-01", "2026-06-02")
single_day_ranges: Final = tuple(adjust_dates_for_timezone(day, day, offset_minutes) for day in days)
multi_day_range: Final = adjust_dates_for_timezone(days[0], days[-1], offset_minutes)
assert tuple(start for start, _ in single_day_ranges) == days
assert tuple(end for _, end in single_day_ranges) == days
assert (min(start for start, _ in single_day_ranges), max(end for _, end in single_day_ranges)) == multi_day_range
def test_timezone_adjustment_live_end_handles_offset_and_opt_in_cases() -> None:
pt_evening: Final = datetime(2026, 8, 6, 4, 30, tzinfo=timezone.utc)
ist_evening: Final = datetime(2026, 8, 5, 17, 0, tzinfo=timezone.utc)
utc_noon: Final = datetime(2026, 8, 5, 12, 0, tzinfo=timezone.utc)
assert adjust_dates_for_timezone(
"2026-07-06", "2026-08-05", 420, include_current_utc_day=True, utc_now=pt_evening
) == ("2026-07-06", "2026-08-06")
assert adjust_dates_for_timezone("2026-07-06", "2026-08-05", 420, utc_now=pt_evening) == (
"2026-07-06",
"2026-08-05",
)
assert adjust_dates_for_timezone(
"2026-07-01", "2026-08-04", 420, include_current_utc_day=True, utc_now=pt_evening
) == ("2026-07-01", "2026-08-04")
assert adjust_dates_for_timezone(
"2026-07-07", "2026-08-06", -330, include_current_utc_day=True, utc_now=ist_evening
) == ("2026-07-07", "2026-08-06")
assert adjust_dates_for_timezone(
"2026-07-06", "2026-08-05", None, include_current_utc_day=True, utc_now=pt_evening
) == ("2026-07-06", "2026-08-05")
assert adjust_dates_for_timezone("2026-07-06", "2026-08-05", 0, include_current_utc_day=True, utc_now=utc_noon) == (
"2026-07-06",
"2026-08-05",
)
assert adjust_dates_for_timezone(
"2026-07-06", "2026-08-09", 420, include_current_utc_day=True, utc_now=pt_evening
) == ("2026-07-06", "2026-08-09")
@pytest.mark.parametrize("offset_minutes", [None, 0, -330, 480])
def test_aggregated_query_uses_the_caller_date_bounds(offset_minutes: int | None) -> None:
query = build_aggregated_sql(
_scope(
timezone_offset_minutes=offset_minutes,
start_date="2026-05-29",
end_date="2026-05-29",
),
api_key_limit=constants.USAGE_TOP_API_KEYS_DEFAULT,
)
assert query.params[:2] == ("2026-05-29", "2026-05-29")
assert "date >= $1" in query.sql
assert "date <= $2" in query.sql
def test_aggregate_query_sums_all_savings_drivers_and_response_time() -> None:
query = build_aggregated_sql(_scope(), api_key_limit=constants.USAGE_TOP_API_KEYS_DEFAULT)
fields: Final = tuple(field for field in SpendMetrics.model_fields if field.endswith("_savings_spend")) + (
"total_response_time_ms",
"timed_requests",
)
assert fields
assert all(f"SUM({field})" in query.sql for field in fields)
def test_aggregated_query_binds_sentinel_and_api_key_limit_after_scope_values() -> None:
scope = _scope(entity_ids=None, api_keys=("key-1",))
query = build_aggregated_sql(scope, api_key_limit=3)
assert "api_key <> $4" in query.sql
assert "LIMIT $5" in query.sql
assert 'FROM "LiteLLM_DailyUserSpend"' in query.sql
assert query.params == (
"2026-01-01",
"2026-01-31",
["key-1"],
PTU_SENTINEL_API_KEY,
3,
)
@pytest.mark.parametrize("api_key_limit", [0, constants.USAGE_TOP_API_KEYS_MAX + 1])
def test_aggregated_query_rejects_api_key_limits_outside_bounds(api_key_limit: int) -> None:
with pytest.raises(ValueError, match="api_key_limit"):
build_aggregated_sql(_scope(), api_key_limit=api_key_limit)
def test_entity_rollup_bounds_keys_and_reuses_scope_filters() -> None:
query = build_entity_rollup_sql(
_scope(table=DailyActivityTable.TEAM, entity_ids=None, api_keys=("key-1", "key-2")),
api_key_limit=3,
)
assert query.sql.count("COALESCE(\"team_id\", '') AS entity_id") == 3
assert query.sql.count("GROUP BY date, COALESCE(\"team_id\", '')") == 2
assert '"team_id" AS entity_id' not in query.sql
assert "JOIN top_api_keys USING (api_key)" in query.sql
assert "api_key = ANY($3::text[])" in query.sql
assert query.sql.count("api_key = ANY($3::text[])") == 4
assert query.sql.count("api_key <> $4") == 2
assert query.sql.count("ORDER BY SUM(spend::numeric) DESC, api_key") == 1
assert "k.entity_id = e.entity_id" in query.sql
assert "LIMIT $5" in query.sql
assert query.params == ("2026-01-01", "2026-01-31", ["key-1", "key-2"], PTU_SENTINEL_API_KEY, 3)
@pytest.mark.parametrize("api_key_limit", [0, constants.USAGE_TOP_API_KEYS_MAX + 1])
def test_entity_rollup_rejects_api_key_limits_outside_bounds(api_key_limit: int) -> None:
with pytest.raises(ValueError, match="api_key_limit"):
build_entity_rollup_sql(_scope(), api_key_limit=api_key_limit)
def test_search_query_escapes_pattern_metacharacters_and_binds_limit() -> None:
query = build_key_search_sql(_scope(entity_ids=None), search=r"foo%_\bar", limit=4)
assert "OR api_key IN (" in query.sql
assert 'SELECT v.token FROM "LiteLLM_VerificationToken" v' in query.sql
assert 'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = v.user_id' in query.sql
assert 'SELECT d.token FROM "LiteLLM_DeletedVerificationToken" d' in query.sql
assert 'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = d.user_id' in query.sql
assert "d.key_alias ILIKE $3 ESCAPE" in query.sql
assert "d.user_id ILIKE $3 ESCAPE" in query.sql
assert "api_key ILIKE $3 ESCAPE" in query.sql
assert "v.key_alias ILIKE $3 ESCAPE" in query.sql
assert "v.user_id ILIKE $3 ESCAPE" in query.sql
assert "u.user_email ILIKE $3 ESCAPE" in query.sql
assert query.sql.count("ILIKE $3 ESCAPE") == 7
assert "api_key <> $4" in query.sql
assert "ORDER BY SUM(spend::numeric) DESC, api_key" in query.sql
assert "LIMIT $5" in query.sql
assert query.params == (
"2026-01-01",
"2026-01-31",
r"%foo\%\_\\bar%",
PTU_SENTINEL_API_KEY,
4,
)
def test_model_and_cache_key_queries_bind_filters_sentinel_and_limits() -> None:
model_query = build_model_top_keys_sql(
_scope(entity_ids=None), model_group="public-model", by_model_group=True, limit=5
)
leakage_query = build_cache_leakage_keys_sql(_scope(entity_ids=None), limit=20)
assert "COALESCE(NULLIF(model_group, ''), model) = $3" in model_query.sql
assert "api_key <> $4" in model_query.sql
assert "ORDER BY SUM(spend::numeric) DESC, api_key" in model_query.sql
assert model_query.params == ("2026-01-01", "2026-01-31", "public-model", PTU_SENTINEL_API_KEY, 5)
assert "HAVING SUM(prompt_tokens) - SUM(cache_read_input_tokens) > 0" in leakage_query.sql
assert "ORDER BY SUM(prompt_tokens) - SUM(cache_read_input_tokens) DESC, api_key" in leakage_query.sql
assert leakage_query.params == ("2026-01-01", "2026-01-31", PTU_SENTINEL_API_KEY, 20)
@pytest.mark.parametrize(
"builder",
[
lambda: build_key_search_sql(_scope(), search="x", limit=0),
lambda: build_model_top_keys_sql(_scope(), model_group="x", by_model_group=False, limit=0),
lambda: build_cache_leakage_keys_sql(_scope(), limit=0),
lambda: build_export_sql(_scope(), export_type=ExportType.DAILY, after=None, batch_size=0),
],
)
def test_query_builders_reject_nonpositive_limits(builder) -> None:
with pytest.raises(ValueError, match="limit must be at least 1"):
builder()
@pytest.mark.parametrize(
("export_type", "group_key", "key_filter", "joins"),
[
(ExportType.DAILY, "''", "", ""),
(ExportType.DAILY_WITH_KEYS, "scoped.api_key", "api_key <> $3", 'LEFT JOIN "LiteLLM_VerificationToken"'),
(ExportType.DAILY_WITH_MODELS, "COALESCE(scoped.model, '')", "api_key <> $3", ""),
(
ExportType.DAILY_WITH_USERS,
"COALESCE(vt.user_id, dvt.user_id, '')",
"api_key <> $3",
'LEFT JOIN "LiteLLM_VerificationToken"',
),
],
)
def test_export_groups_by_requested_key_and_binds_cursor_after_scope(
export_type: ExportType, group_key: str, key_filter: str, joins: str
) -> None:
query = build_export_sql(
_scope(entity_ids=None),
export_type=export_type,
after=ExportCursor(date="2026-01-12", entity_id="user-2", group_key="group-3"),
batch_size=2,
)
assert group_key in query.sql
assert key_filter in query.sql
assert joins in query.sql
assert "(scoped.date, COALESCE(scoped.\"user_id\", '')," in query.sql
order_keys: Final = (
"scoped.date, COALESCE(scoped.\"user_id\", '')",
*((group_key,) if export_type is not ExportType.DAILY else ()),
)
assert f"ORDER BY {', '.join(order_keys)}" in query.sql
expected_limit_index: Final = "$6" if export_type is ExportType.DAILY else "$7"
assert f"LIMIT {expected_limit_index}" in query.sql
assert query.params == (
"2026-01-01",
"2026-01-31",
*((PTU_SENTINEL_API_KEY,) if export_type is not ExportType.DAILY else ()),
"2026-01-12",
"user-2",
"group-3",
2,
)
@pytest.mark.parametrize("export_type", [ExportType.DAILY_WITH_KEYS, ExportType.DAILY_WITH_USERS])
def test_export_uses_latest_deleted_key_metadata(export_type: ExportType) -> None:
query = build_export_sql(_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2)
assert 'FROM "LiteLLM_DeletedVerificationToken"' in query.sql
assert "ORDER BY deleted_at DESC" in query.sql
assert "COALESCE(vt.user_id, dvt.user_id)" in query.sql
@pytest.mark.parametrize("export_type", tuple(ExportType))
def test_export_without_cursor_omits_cursor_predicate_and_parameters(export_type: ExportType) -> None:
query = build_export_sql(
_scope(entity_ids=None),
export_type=export_type,
after=None,
batch_size=2,
)
assert "WHERE TRUE AND (scoped.date" not in query.sql
assert query.params == (
"2026-01-01",
"2026-01-31",
*((PTU_SENTINEL_API_KEY,) if export_type is not ExportType.DAILY else ()),
2,
)

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

@ -30182,6 +30182,18 @@ export interface components {
};
/** DailySpendMetadata */
DailySpendMetadata: {
/**
* Api Key Limit
* @description When set, api_keys and every api_key_breakdown list at most this many keys, ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.
*/
api_key_limit?: number | null;
/**
* Entity Total Api Keys
* @description Distinct API keys per entity over the requested range, set when the entity breakdown is included. When an entity's count exceeds api_key_limit, its api_key_breakdown lists only its keys among the top api_key_limit keys overall.
*/
entity_total_api_keys?: {
[key: string]: number;
} | null;
/**
* Has More
* @default false
@ -30192,6 +30204,11 @@ export interface components {
* @default 1
*/
page: number;
/**
* Total Api Keys
* @description Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key lists are truncated to the highest-spend keys.
*/
total_api_keys?: number | null;
/**
* Total Api Requests
* @default 0