diff --git a/litellm/constants.py b/litellm/constants.py index bbeb4846e27..edfcf67f212 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2084,16 +2084,12 @@ 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_LIMIT: Final[int] = int(os.getenv("USAGE_TOP_API_KEYS_LIMIT", "100")) # 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 # Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide # expiry cannot produce an alert too large for the channel delivering it. PTU_LAPSED_ALERT_LIMIT: Final[int] = 10 -DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID: Final[str] = "daily_global_spend_reconcile_job" -DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS: Final[int] = 3600 -DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM: Final[str] = "daily_global_spend_reconciled_through" # Slack allowed when deciding a sentinel row is stale. The row's updated_at and the # run's cutoff are stamped by different hosts, so clock skew between them must not let # one run delete a charge another just wrote. A stale row is hours old and a concurrent diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 06e157498aa..688e95a6acd 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3050,18 +3050,6 @@ }, "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" - }, "has_more": { "default": false, "title": "Has More", @@ -3072,18 +3060,6 @@ "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", diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 5a19d743105..b22d57170a4 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -9,9 +9,8 @@ from fastapi import HTTPException, status from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT +from litellm.constants import PTU_SENTINEL_API_KEY from litellm.proxy._types import CommonProxyErrors -from litellm.proxy.spend_tracking.daily_global_spend_rollup import GLOBAL_SPEND_TABLE_NAME, reconciled_through from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_emails, recover_double_hashed_key_metadata, @@ -149,9 +148,15 @@ class _AggregatedSpendData(TypedDict): totals: SpendMetrics -class _RollupMetricsRow(SimpleNamespace): +class _GroupingSetsRow(SimpleNamespace): date: str api_key: str | None + model: str | None + model_group: str | None + custom_llm_provider: str | None + mcp_namespaced_tool_name: str | None + endpoint: str | None + group_level: int spend: float | None prompt_tokens: int | None completion_tokens: int | None @@ -169,46 +174,12 @@ class _RollupMetricsRow(SimpleNamespace): timed_requests: int | None -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 - - -class _EntityRollupRow(_RollupMetricsRow): +class _EntityRollupRow(_GroupingSetsRow): entity_id: str | None api_key_rolled: int -class _AggregatedQueryKwargs(TypedDict): - table_name: ReadOnly[str] - entity_id_field: ReadOnly[str] - entity_id: ReadOnly[str | list[str] | None] - start_date: ReadOnly[str] - end_date: ReadOnly[str] - model: ReadOnly[str | None] - api_key: ReadOnly[str | list[str] | None] - exclude_entity_ids: ReadOnly[list[str] | None] - timezone_offset_minutes: ReadOnly[int | None] - include_current_utc_day: ReadOnly[bool] - - -_SqlQuery = tuple[str, list[str]] - - -async def _query_raw_optional( - prisma_client: PrismaClient, query: _SqlQuery | None -) -> list[dict[str, object]] | None: # mutable-ok: prisma query_raw return shape - if query is None: - return None - return await prisma_client.db.query_raw(query[0], *query[1]) - - -def _reported_flat_cost(record: DailySpendRecord | _RollupMetricsRow) -> float: +def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float: """Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled. Both read paths funnel through here: the paginated path reads the ``ptu_flat_cost`` @@ -731,8 +702,71 @@ def _ptu_flat_cost_select(table_name: str) -> str: return "0::float AS ptu_flat_cost" -def _rollup_metric_select(table_name: str) -> str: - return f""" +def _build_aggregated_sql_query( + *, + table_name: str, + entity_id_field: str, + entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + start_date: str, + end_date: str, + model: str | None, + api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path + timezone_offset_minutes: int | None = None, + include_current_utc_day: bool = False, +) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params + """Build a parameterized SQL GROUP BY query for aggregated daily activity. + + Groups by (date, api_key, model, model_group, custom_llm_provider, + mcp_namespaced_tool_name, endpoint) with SUMs on all metric columns. + The entity_id column is intentionally omitted from GROUP BY to collapse + rows across entities — this is where the biggest row reduction comes from. + + Returns: + Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw(). + """ + pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name) + if pg_table is None: + raise ValueError(f"Unknown table name: {table_name}") + + adjusted_start, adjusted_end = _adjust_dates_for_timezone( + start_date, end_date, timezone_offset_minutes, include_current_utc_day + ) + + where_clause, sql_params = _build_aggregated_where_clause( + entity_id_field=entity_id_field, + entity_id=entity_id, + adjusted_start=adjusted_start, + adjusted_end=adjusted_end, + model=model, + api_key=api_key, + exclude_entity_ids=exclude_entity_ids, + ) + + # Postgres computes every rollup level the response needs — per-date + # totals, per-(date, model), per-(date, model, api_key), per-provider, + # etc. — in a single pass via GROUPING SETS. The GROUPING() bitmask + # encodes which level a row belongs to so Python can dispatch rows + # straight into their buckets without re-summing. The leaf grouping + # is omitted on purpose: nothing in the response shape needs it once + # all the rollups are present. + # + # TODO: drop the successful_requests/failed_requests aggregates (and the + # total_successful_requests metadata they feed) once the admin UI reads SGR + # only from LiteLLM_DailyGatewayRequests. The remaining spend, token and + # api_requests rollups are still served from here. + sql_query: Final = f""" + SELECT + date, + api_key, + model, + COALESCE(NULLIF(model_group, ''), model) AS model_group, + custom_llm_provider, + mcp_namespaced_tool_name, + endpoint, + GROUPING(date, api_key, model, COALESCE(NULLIF(model_group, ''), model), + custom_llm_provider, mcp_namespaced_tool_name, + endpoint) AS group_level, SUM(spend)::float AS spend, {_ptu_flat_cost_select(table_name)}, SUM(prompt_tokens)::bigint AS prompt_tokens, @@ -748,175 +782,27 @@ def _rollup_metric_select(table_name: str) -> str: 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""" - - -_MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" - - -_KEY_FREE_SOURCE_COLUMNS: Final = ( - "date", - "model", - "model_group", - "custom_llm_provider", - "mcp_namespaced_tool_name", - "endpoint", - "spend", - "prompt_tokens", - "completion_tokens", - "cache_read_input_tokens", - "cache_creation_input_tokens", - "compression_saved_tokens", - "compression_savings_spend", - "prompt_caching_savings_spend", - "gateway_injected_caching_savings_spend", - "autorouter_savings_spend", - "api_requests", - "successful_requests", - "failed_requests", - "total_response_time_ms", - "timed_requests", -) - - -async def global_rollup_reconciled_through(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: - """The last day ``LiteLLM_DailyGlobalSpend`` can answer the key-free arm for, or None to - read it all from the per-key table. - - Only an unfiltered read of the user table sums to the same rows as the global table. The - marker read is served from the config cache, so this is not a database round trip per request. - """ - if query["table_name"] != "litellm_dailyuserspend": - return None - if query["entity_id"] is not None or query["api_key"] is not None or query["exclude_entity_ids"]: - return None - try: - return await reconciled_through(prisma_client) - except Exception as exc: # noqa: BLE001 # the per-key table is always a correct answer, so never fail the read - verbose_proxy_logger.warning("Could not read the daily global spend marker, using the per-key table: %s", exc) - return None - - -def _key_free_source(pg_table: str, where_clause: str, marker_param: str | None) -> str: - """The relation the key-free arm aggregates: the per-key table alone, or the global rollup - for days through the marker plus the per-key table for the days still open after it.""" - if marker_param is None: - return f'"{pg_table}"\n WHERE {where_clause}' - columns: Final = ", ".join(_KEY_FREE_SOURCE_COLUMNS) - return f"""( - SELECT {columns} - FROM "{GLOBAL_SPEND_TABLE_NAME}" - WHERE {where_clause} AND date <= {marker_param} - UNION ALL - SELECT {columns} - FROM "{pg_table}" - WHERE {where_clause} AND date > {marker_param} - ) AS key_free_source""" - - -def _build_aggregated_sql_query( - *, - table_name: str, - entity_id_field: str, - entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path - start_date: str, - end_date: str, - model: str | None, - api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path - exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path - timezone_offset_minutes: int | None = None, - include_current_utc_day: bool = False, - global_rollup_through: str | None = None, -) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params - """Build the GROUPING SETS query for aggregated daily activity. - - Returns: - Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw(). - """ - pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name) - if pg_table is None: - raise ValueError(f"Unknown table name: {table_name}") - - adjusted_start, adjusted_end = _adjust_dates_for_timezone( - start_date, end_date, timezone_offset_minutes, include_current_utc_day - ) - - where_clause, where_params = _build_aggregated_where_clause( - entity_id_field=entity_id_field, - entity_id=entity_id, - adjusted_start=adjusted_start, - adjusted_end=adjusted_end, - model=model, - api_key=api_key, - exclude_entity_ids=exclude_entity_ids, - ) - sentinel_param: Final = f"${len(where_params) + 1}" - marker_param: Final = None if global_rollup_through is None else f"${len(where_params) + 2}" - metric_select: Final = _rollup_metric_select(table_name) - - # TODO: drop the successful_requests/failed_requests aggregates (and the - # total_successful_requests metadata they feed) once the admin UI reads SGR - # only from LiteLLM_DailyGatewayRequests. The remaining spend, token and - # api_requests rollups are still served from here. - sql_query: 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 {_key_free_source(pg_table, where_clause, marker_param)} - 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 ( - 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) DESC, api_key - LIMIT {USAGE_TOP_API_KEYS_LIMIT} - ) - 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) + SUM(timed_requests)::bigint AS timed_requests + FROM "{pg_table}" WHERE {where_clause} GROUP BY GROUPING SETS ( + (date), (date, api_key), + (date, model), (date, model, api_key), - (date, {_MODEL_GROUP_EXPR}, api_key), + (date, COALESCE(NULLIF(model_group, ''), model)), + (date, COALESCE(NULLIF(model_group, ''), model), api_key), + (date, custom_llm_provider), (date, custom_llm_provider, api_key), + (date, mcp_namespaced_tool_name), (date, mcp_namespaced_tool_name, api_key), - (date, endpoint, api_key) - )) + (date, endpoint), + (date, endpoint, api_key), + () + ) """ - marker_params: Final = () if global_rollup_through is None else (global_rollup_through,) - return sql_query, [*where_params, PTU_SENTINEL_API_KEY, *marker_params] + return sql_query, sql_params def _build_entity_rollup_sql_query( @@ -961,7 +847,23 @@ def _build_entity_rollup_sql_query( "{entity_id_field}" AS entity_id, date, api_key, - GROUPING(api_key) AS api_key_rolled,{_rollup_metric_select(table_name)} + GROUPING(api_key) AS api_key_rolled, + SUM(spend)::float AS spend, + {_ptu_flat_cost_select(table_name)}, + SUM(prompt_tokens)::bigint AS prompt_tokens, + SUM(completion_tokens)::bigint AS completion_tokens, + SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens, + SUM(cache_creation_input_tokens)::bigint AS cache_creation_input_tokens, + SUM(compression_saved_tokens)::bigint AS compression_saved_tokens, + SUM(compression_savings_spend)::float AS compression_savings_spend, + SUM(prompt_caching_savings_spend)::float AS prompt_caching_savings_spend, + SUM(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 FROM "{pg_table}" WHERE {where_clause} GROUP BY GROUPING SETS ( @@ -1063,7 +965,6 @@ async def _aggregate_spend_records( # current grouping set's key), 0 when the column is part of the key. _GROUP_GRAND_TOTAL: Final = 127 # 0b1111111 — all rolled up _GROUP_DATE: Final = 63 # 0b0111111 — only date kept -_API_KEY_ROLLED_UP_BIT: Final = 32 # 0b0100000 _GROUP_DATE_API_KEY: Final = 31 # 0b0011111 _GROUP_DATE_MODEL: Final = 47 # 0b0101111 _GROUP_DATE_MODEL_API_KEY: Final = 15 # 0b0001111 @@ -1077,7 +978,7 @@ _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110 _GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110 -def _record_to_spend_metrics(record: _RollupMetricsRow) -> SpendMetrics: +def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics: """Build a SpendMetrics directly from one already-aggregated rollup row. SUM() over zero rows is SQL NULL, so rollup rows (notably the grand-total @@ -1431,6 +1332,10 @@ async def get_daily_activity_aggregated( ) -> SpendAnalyticsPaginatedResponse: """Aggregated variant that returns the full result set (no pagination). + Uses SQL GROUP BY to aggregate rows in the database rather than fetching + all individual rows into Python. This collapses rows across entities + (users/teams/orgs), reducing ~150k rows to ~2-3k grouped rows. + include_entity_breakdown runs a small companion rollup query and folds `breakdown.entities` onto the response, as entity-scoped views like Team Usage need. @@ -1449,7 +1354,7 @@ async def get_daily_activity_aggregated( ) try: - query_kwargs: Final = _AggregatedQueryKwargs( + sql_query, sql_params = _build_aggregated_sql_query( table_name=table_name, entity_id_field=entity_id_field, entity_id=entity_id, @@ -1461,19 +1366,36 @@ async def get_daily_activity_aggregated( timezone_offset_minutes=timezone_offset_minutes, include_current_utc_day=include_current_utc_day, ) - sql_query, sql_params = _build_aggregated_sql_query( - **query_kwargs, - global_rollup_through=await global_rollup_reconciled_through(prisma_client, query_kwargs), - ) - entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None - raw_rows, raw_entity_rows = await asyncio.gather( - prisma_client.db.query_raw(sql_query, *sql_params), - _query_raw_optional(prisma_client, entity_query), + entity_query: Final = ( + _build_entity_rollup_sql_query( + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + start_date=start_date, + end_date=end_date, + model=model, + api_key=api_key, + exclude_entity_ids=exclude_entity_ids, + timezone_offset_minutes=timezone_offset_minutes, + include_current_utc_day=include_current_utc_day, + ) + if include_entity_breakdown + else None ) - records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or ())] - total_api_keys: Final = next((r.distinct_api_keys for r in records if r.distinct_api_keys is not None), 0) + # Execute the GROUPING SETS query (one row per rollup level), alongside + # the per-entity companion rollup when the caller wants entities. + raw_rows, raw_entity_rows = ( + await asyncio.gather( + prisma_client.db.query_raw(sql_query, *sql_params), + prisma_client.db.query_raw(entity_query[0], *entity_query[1]), + ) + if entity_query is not None + else (await prisma_client.db.query_raw(sql_query, *sql_params), None) + ) + + records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or [])] # The grouping-sets dispatcher places each row directly in its bucket # using the row's GROUPING() bitmask. No Python-side summing needed. @@ -1527,8 +1449,6 @@ async def get_daily_activity_aggregated( page=1, total_pages=1, has_more=False, - api_key_limit=USAGE_TOP_API_KEYS_LIMIT, - total_api_keys=total_api_keys, ), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index af25d418a63..613580fc401 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -265,7 +265,6 @@ from litellm.constants import ( APSCHEDULER_MISFIRE_GRACE_TIME, APSCHEDULER_REPLACE_EXISTING, CLI_SSO_SESSION_TTL_SECONDS, - DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, @@ -696,9 +695,6 @@ from litellm.proxy.route_priority import hot_routes_first from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start -from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( - run_scheduled_daily_global_spend_reconcile, -) from litellm.proxy.spend_tracking.spend_counter_batch import ( PendingSpendIncrement, active_spend_counter_batch, @@ -10229,12 +10225,6 @@ class ProxyStartupEvent: await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) - cls._initialize_daily_global_spend_reconcile_job( - scheduler=scheduler, - proxy_logging_obj=proxy_logging_obj, - prisma_client=prisma_client, - ) - ### PTU DAILY ROLLUP ### from litellm.proxy.spend_tracking.ptu_feature_flag import ( is_ptu_cost_attribution_enabled, @@ -10576,39 +10566,6 @@ class ProxyStartupEvent: "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)" ) - @classmethod - def _initialize_daily_global_spend_reconcile_job( - cls, - scheduler: AsyncIOScheduler, - proxy_logging_obj: ProxyLogging, - prisma_client: PrismaClient, - ) -> None: - async def alert(message: str) -> None: - await proxy_logging_obj.alerting_handler( - message=message, - level="High", - alert_type=AlertType.failed_tracking_spend, - ) - - async def reconcile() -> None: - await run_scheduled_daily_global_spend_reconcile( - prisma_client, - pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, - alert=alert, - ) - - scheduler.add_job( - reconcile, - "cron", - hour=0, - minute=30, - timezone="UTC", - id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - next_run_time=datetime.now(timezone.utc) + timedelta(minutes=2), - ) - @classmethod async def _initialize_slack_alerting_jobs( cls, diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py deleted file mode 100644 index 376b113ed02..00000000000 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ /dev/null @@ -1,293 +0,0 @@ -"""Roll closed UTC days of ``LiteLLM_DailyUserSpend`` up into ``LiteLLM_DailyGlobalSpend``. - -Only days that are over get rolled up, so a pod still flushing per-key spend for the current -day can never leave the global table short; usage reads serve days through the recorded -marker from the global table and later days live from the per-key table. Per-key rows are -dated by request start, so spend can land on a day that was already rolled up (a flush -straddling midnight, a retry after an outage). Each run therefore also rewrites every closed -day that has rows touched since the previous run's scan, whatever the date. The marker lives -in ``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on -a large deployment the first backfill is minutes of work. -""" - -from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from datetime import date, timedelta -from typing import TYPE_CHECKING, Final - -from pydantic import BaseModel, ConfigDict, ValidationError - -from litellm._logging import verbose_proxy_logger -from litellm.constants import ( - DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, - DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS, - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, -) - -if TYPE_CHECKING: - from litellm.caching.redis_cache import RedisCache - from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager - from litellm.proxy.utils import PrismaClient - -GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend" -# The unique constraint, in constraint order. NULL never matches itself in a unique index, so -# every column is normalized to '' or the same group would be inserted again on every run. -_KEY_COLUMNS: Final = ("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") -_METRIC_COLUMNS: Final = ( - "prompt_tokens", - "completion_tokens", - "cache_read_input_tokens", - "cache_creation_input_tokens", - "compression_saved_tokens", - "api_requests", - "successful_requests", - "failed_requests", - "total_response_time_ms", - "timed_requests", - "compression_savings_spend", - "prompt_caching_savings_spend", - "gateway_injected_caching_savings_spend", - "autorouter_savings_spend", - "spend", -) - - -def _quoted(columns: tuple[str, ...]) -> str: - return ", ".join(f'"{column}"' for column in columns) - - -def _reconcile_day_sql() -> str: - normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in _KEY_COLUMNS) - sums: Final = ", ".join(f'SUM("{column}")' for column in _METRIC_COLUMNS) - overwrite: Final = ", ".join(f'"{column}" = EXCLUDED."{column}"' for column in _METRIC_COLUMNS) - return ( - f'INSERT INTO "{GLOBAL_SPEND_TABLE_NAME}" ("id", {_quoted(_KEY_COLUMNS)}, {_quoted(_METRIC_COLUMNS)}, ' - '"updated_at")\n' - f"SELECT gen_random_uuid()::text, {normalized_keys}, {sums}, (NOW() AT TIME ZONE 'UTC')\n" - 'FROM "LiteLLM_DailyUserSpend" WHERE "date" = $1\n' - f"GROUP BY {normalized_keys}\n" - f"ON CONFLICT ({_quoted(_KEY_COLUMNS)}) DO UPDATE SET {overwrite}, " - "\"updated_at\" = (NOW() AT TIME ZONE 'UTC')" - ) - - -RECONCILE_DAY_SQL: Final = _reconcile_day_sql() -_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now, (NOW() AT TIME ZONE 'UTC')::date::text AS today" -_ALL_CLOSED_DAYS_SQL: Final = 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ORDER BY "date"' -# Pod clocks drift from the database clock and from each other, so rows are picked up from a -# little before the previous scan; rewriting a day twice is idempotent. -_PENDING_DAYS_SQL: Final = ( - 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ' - 'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') ' - 'ORDER BY "date"' -) -# Runs can overlap (Redis unreachable, lock expired on a long backfill), so the database keeps the -# later of the stored and the incoming day and scan time in one statement; GREATEST skips NULL. -_ADVANCE_MARKER_SQL: Final = ( - 'INSERT INTO "LiteLLM_Config" ("param_name", "param_value") ' - "VALUES ($1, jsonb_build_object('reconciled_through', $2::text, 'scanned_at', $3::text)) " - 'ON CONFLICT ("param_name") DO UPDATE SET "param_value" = jsonb_build_object(' - "'reconciled_through', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'reconciled_through', " - "EXCLUDED.\"param_value\" ->> 'reconciled_through'), " - "'scanned_at', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'scanned_at', " - "EXCLUDED.\"param_value\" ->> 'scanned_at'))" -) - - -class ReconciledThrough(BaseModel): - """``reconciled_through`` is the last closed UTC day the global table covers. ``scanned_at`` is - the database clock when the scan behind the last fully successful run started: every per-key - row written before it, on any day through the marker, is in the global table.""" - - model_config = ConfigDict(frozen=True, extra="ignore") - - reconciled_through: str - scanned_at: str | None = None - - -class _MarkerRow(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) - - param_value: object = None - - -class _DateRow(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore") - - date: str - - -class _NowRow(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore") - - now: str - today: str - - -@dataclass(frozen=True, slots=True) -class ReconcileResult: - days_reconciled: tuple[str, ...] - reconciled_through: str | None - failed_day: str | None = None - - -@dataclass(frozen=True, slots=True) -class _PendingScan: - marker: ReconciledThrough | None - scanned_at: str - days: tuple[str, ...] - - -def _marker_from_param_value(value: object) -> ReconciledThrough | None: - try: - return ( - ReconciledThrough.model_validate_json(value) - if isinstance(value, str) - else ReconciledThrough.model_validate(value) - ) - except ValidationError: - return None - - -async def read_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: - from litellm.proxy.utils import get_config_param - - row: Final = await get_config_param(prisma_client, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - return None if row is None else _marker_from_param_value(_MarkerRow.model_validate(row).param_value) - - -async def reconciled_through(prisma_client: "PrismaClient") -> str | None: - """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" - marker: Final = await read_marker(prisma_client) - return None if marker is None else marker.reconciled_through - - -async def _advance_marker(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None: - """Move the stored marker to the last of ``days`` and to ``scanned_at`` where those are later - than what is stored, so a slower overlapping run can only add to a faster run's marker.""" - from litellm.proxy.utils import invalidate_config_param - - await prisma_client.db.execute_raw( - _ADVANCE_MARKER_SQL, - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, - max(days) if days else None, - scanned_at, - ) - await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - -async def _db_now(prisma_client: "PrismaClient") -> _NowRow: - rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) - return _NowRow.model_validate(rows[0]) - - -async def _scan_pending(prisma_client: "PrismaClient") -> _PendingScan: - """Every closed UTC day (strictly before the database's today) still to roll up, oldest first: - days past the marker, plus any day with per-key rows written since the scan behind the marker. - Before a run has fully succeeded there is no such scan, so every closed day is rolled up.""" - marker: Final = await read_marker(prisma_client) - db_now: Final = await _db_now(prisma_client) - last_closed_day: Final = (date.fromisoformat(db_now.today) - timedelta(days=1)).isoformat() - rows: Final = ( - await prisma_client.db.query_raw(_ALL_CLOSED_DAYS_SQL, last_closed_day) - if marker is None or marker.scanned_at is None - else await prisma_client.db.query_raw( - _PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at - ) - ) - return _PendingScan(marker, db_now.now, tuple(_DateRow.model_validate(row).date for row in rows)) - - -async def pending_days(prisma_client: "PrismaClient") -> tuple[str, ...]: - return (await _scan_pending(prisma_client)).days - - -async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: - """Rewrite one day of the global table from the per-key sums. Idempotent: a rerun - overwrites every group with the same totals.""" - await prisma_client.db.execute_raw(RECONCILE_DAY_SQL, day) - - -async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> ReconcileResult: - """Roll up every pending day, advancing the marker after each; a failing day stops the run - with the marker on the last good day so the next run resumes there. The scan time is only - recorded once every pending day is done, so late rows a failed run saw are found again.""" - scan: Final = await _scan_pending(prisma_client) - done: Final = await _reconcile_until_failure(prisma_client, scan) - if len(done) < len(scan.days): - marker: Final = await reconciled_through(prisma_client) - return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) - if scan.marker is not None or done: - await _advance_marker(prisma_client, done, scanned_at=scan.scanned_at) - return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) - - -async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]: - for index, day in enumerate(scan.days): - if not await _reconcile_and_record(prisma_client, scan.days[: index + 1]): - return scan.days[:index] - return scan.days - - -async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: tuple[str, ...]) -> bool: - day: Final = done_with_this[-1] - try: - await reconcile_day(prisma_client, day) - await _advance_marker(prisma_client, done_with_this, scanned_at=None) - except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done - verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) - return False - return True - - -async def run_scheduled_daily_global_spend_reconcile( - prisma_client: "PrismaClient", - pod_lock_manager: "PodLockManager | None" = None, - alert: Callable[[str], Awaitable[None]] | None = None, -) -> ReconcileResult | None: - """Run the reconcile under a cross-pod lock so one proxy does the work; the lock only saves - effort (each day is an idempotent rewrite), so an unreachable Redis runs unguarded rather than skipping.""" - redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache - if pod_lock_manager is None or redis_cache is None: - return await _run_and_alert(prisma_client, alert=alert) - - acquired: Final = await pod_lock_manager.acquire_lock( - cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, ttl=DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS - ) - if not acquired and await _lock_is_held(pod_lock_manager, redis_cache): - verbose_proxy_logger.info("Daily global spend reconcile: another pod holds the lock, skipping this run") - return None - try: - return await _run_and_alert(prisma_client, alert=alert) - finally: - if acquired: - await pod_lock_manager.release_lock(cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) - - -async def _lock_is_held(pod_lock_manager: "PodLockManager", redis_cache: "RedisCache") -> bool: - try: - lock_key: Final = pod_lock_manager.get_redis_lock_key(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) - return bool(await redis_cache.async_get_cache(lock_key)) - except Exception as exc: # noqa: BLE001 # an unreadable lock must not skip the run - verbose_proxy_logger.warning("Daily global spend reconcile: could not read the lock: %s", exc) - return False - - -async def _run_and_alert( - prisma_client: "PrismaClient", - *, - alert: Callable[[str], Awaitable[None]] | None, -) -> ReconcileResult: - result: Final = await run_daily_global_spend_reconcile(prisma_client) - if result.days_reconciled: - verbose_proxy_logger.info( - "Daily global spend reconcile: rolled up %d day(s), reconciled through %s", - len(result.days_reconciled), - result.reconciled_through, - ) - if result.failed_day is not None and alert is not None: - await alert( - f"Daily global spend reconcile stopped at {result.failed_day}; usage totals keep reading the per-key " - f"table for ranges past {result.reconciled_through or 'the beginning'} until the next run succeeds." - ) - return result diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 2a4f6b2944a..28488ba7de6 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -101,16 +101,6 @@ 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.", - ) class SpendAnalyticsPaginatedResponse(BaseModel): diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index baaf3f4ba2f..1884011411d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1,21 +1,13 @@ -import pathlib -import re -from collections.abc import Sequence from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock -import psycopg import pytest -from psycopg.rows import dict_row -from pytest_postgresql import factories -from litellm.constants import ( - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, - PTU_SENTINEL_API_KEY, - USAGE_TOP_API_KEYS_LIMIT, -) +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + from litellm.proxy.management_endpoints.common_daily_activity import ( _adjust_dates_for_timezone, _build_aggregated_sql_query, @@ -25,12 +17,8 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, - global_rollup_reconciled_through, update_metrics, ) -from litellm.proxy.spend_tracking.daily_global_spend_rollup import RECONCILE_DAY_SQL -from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR -from litellm.proxy.utils import evict_config_param from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, SpendMetrics, @@ -181,7 +169,6 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): "endpoint": "/v1/chat/completions", "api_key": None, "group_level": 62, - "distinct_api_keys": None, "spend": 15.0, "prompt_tokens": 150, "completion_tokens": 75, @@ -194,7 +181,31 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): "endpoint": "/v1/embeddings", "api_key": None, "group_level": 62, - "distinct_api_keys": None, + "spend": 3.0, + "prompt_tokens": 30, + "completion_tokens": 0, + "api_requests": 1, + "successful_requests": 1, + }, + # (date, endpoint, api_key) — populates the per-key sub-bucket + { + **base, + "date": "2024-01-01", + "endpoint": "/v1/chat/completions", + "api_key": "key-1", + "group_level": 30, + "spend": 15.0, + "prompt_tokens": 150, + "completion_tokens": 75, + "api_requests": 2, + "successful_requests": 2, + }, + { + **base, + "date": "2024-01-01", + "endpoint": "/v1/embeddings", + "api_key": "key-2", + "group_level": 30, "spend": 3.0, "prompt_tokens": 30, "completion_tokens": 0, @@ -208,7 +219,6 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): "endpoint": None, "api_key": None, "group_level": 63, - "distinct_api_keys": None, "spend": 18.0, "prompt_tokens": 180, "completion_tokens": 75, @@ -222,40 +232,12 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown(): "endpoint": None, "api_key": None, "group_level": 127, - "distinct_api_keys": None, "spend": 18.0, "prompt_tokens": 180, "completion_tokens": 75, "api_requests": 3, "successful_requests": 3, }, - # (date, endpoint, api_key) — populates the per-key sub-bucket - { - **base, - "date": "2024-01-01", - "endpoint": "/v1/chat/completions", - "api_key": "key-1", - "group_level": 30, - "distinct_api_keys": 2, - "spend": 15.0, - "prompt_tokens": 150, - "completion_tokens": 75, - "api_requests": 2, - "successful_requests": 2, - }, - { - **base, - "date": "2024-01-01", - "endpoint": "/v1/embeddings", - "api_key": "key-2", - "group_level": 30, - "distinct_api_keys": 2, - "spend": 3.0, - "prompt_tokens": 30, - "completion_tokens": 0, - "api_requests": 1, - "successful_requests": 1, - }, ] mock_prisma.db.query_raw = AsyncMock(return_value=mock_rows) @@ -492,7 +474,9 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash( return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")] ) mock_prisma.db.query_raw = AsyncMock( - return_value=[{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}] + return_value=[ + {"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"} + ] ) result = await get_api_key_metadata( @@ -869,7 +853,6 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): "endpoint": "/v1/chat/completions", "api_key": None, "group_level": 62, - "distinct_api_keys": None, "spend": 10.0, "prompt_tokens": 100, "completion_tokens": 50, @@ -882,7 +865,6 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): "endpoint": "/v1/chat/completions", "api_key": "deleted-key-hash", "group_level": 30, - "distinct_api_keys": 1, "spend": 10.0, "prompt_tokens": 100, "completion_tokens": 50, @@ -1327,11 +1309,42 @@ class TestBuildAggregatedSqlQuery: "user-1", "bedrock/global.anthropic.claude-opus-4-8", "sk-test", - PTU_SENTINEL_API_KEY, ] assert "model = $4" in sql assert "api_key = $5" in sql + def test_model_group_rollups_fall_back_to_model_name(self): + """Aggregated model_groups rollups must fall back to model for group-less rows. + + The (date, model_group) grouping level cannot recover the model column + after the fact (it is rolled up), so the fallback has to happen in SQL; + without it, group-less rows silently vanish from the model_groups + breakdown that the usage UI now renders by default. Group-less rows are + stored as empty strings, not NULL (spend_tracking_utils defaults + model_group to ""), so a plain COALESCE is not enough: the fallback must + be NULLIF-wrapped to catch both + """ + sql, _ = _build_aggregated_sql_query( + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + start_date="2026-07-01", + end_date="2026-07-01", + model=None, + api_key=None, + ) + + normalized = " ".join(sql.split()) + fallback = "COALESCE(NULLIF(model_group, ''), model)" + assert f"{fallback} AS model_group" in normalized + assert ( + f"GROUPING(date, api_key, model, {fallback}, " + "custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level" in normalized + ) + assert f"(date, {fallback}), (date, {fallback}, api_key)," in normalized + assert "(date, model_group)" not in normalized + assert "COALESCE(model_group, model)" not in normalized + class TestAggregatedEmptyEntityFilter: _BUILDERS: Final = (_build_aggregated_sql_query, _build_entity_rollup_sql_query) @@ -1351,8 +1364,7 @@ class TestAggregatedEmptyEntityFilter: normalized = " ".join(sql.split()) assert "IN ()" not in normalized assert '"team_id" IN' not in normalized - sentinel_params = [PTU_SENTINEL_API_KEY] if build is _build_aggregated_sql_query else [] - assert params == ["2026-08-01", "2026-08-19", *sentinel_params] + assert params == ["2026-08-01", "2026-08-19"] @pytest.mark.parametrize("build", _BUILDERS) def test_empty_entity_list_matches_nothing_rather_than_everything(self, build): @@ -1383,8 +1395,7 @@ class TestAggregatedEmptyEntityFilter: normalized = " ".join(sql.split()) assert '"team_id" IN ($3, $4)' in normalized assert "FALSE" not in normalized - sentinel_params = [PTU_SENTINEL_API_KEY] if build is _build_aggregated_sql_query else [] - assert params == ["2026-08-01", "2026-08-19", "team-alpha", "team-beta", *sentinel_params] + assert params == ["2026-08-01", "2026-08-19", "team-alpha", "team-beta"] @pytest.mark.asyncio @@ -1409,7 +1420,6 @@ async def test_get_daily_activity_aggregated_empty_result_set(): "mcp_namespaced_tool_name": None, "endpoint": None, "group_level": 127, - "distinct_api_keys": None, "spend": None, "prompt_tokens": None, "completion_tokens": None, @@ -1454,484 +1464,6 @@ async def test_get_daily_activity_aggregated_empty_result_set(): assert result.metadata.total_compression_saved_tokens == 0 -_aggregated_postgresql_proc: Final = factories.postgresql_proc() -_aggregated_postgresql: Final = factories.postgresql("_aggregated_postgresql_proc") - -_DAILY_USER_SPEND_DDL: Final = """ - CREATE TABLE "LiteLLM_DailyUserSpend" ( - id TEXT PRIMARY KEY, - user_id TEXT, - date TEXT NOT NULL, - api_key TEXT NOT NULL, - model TEXT, - model_group TEXT, - custom_llm_provider TEXT, - mcp_namespaced_tool_name TEXT, - endpoint TEXT, - prompt_tokens BIGINT DEFAULT 0, - completion_tokens BIGINT DEFAULT 0, - cache_read_input_tokens BIGINT DEFAULT 0, - cache_creation_input_tokens BIGINT DEFAULT 0, - compression_saved_tokens BIGINT DEFAULT 0, - compression_savings_spend DOUBLE PRECISION DEFAULT 0, - prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, - spend DOUBLE PRECISION DEFAULT 0, - api_requests BIGINT DEFAULT 0, - successful_requests BIGINT DEFAULT 0, - failed_requests BIGINT DEFAULT 0, - total_response_time_ms BIGINT DEFAULT 0, - timed_requests BIGINT DEFAULT 0 - ) -""" - - -def _seed_daily_user_spend(conn: psycopg.Connection, rows: Sequence[tuple[object, ...]]) -> None: - with conn.cursor() as cur: - cur.execute(_DAILY_USER_SPEND_DDL) - cur.executemany( - """ - INSERT INTO "LiteLLM_DailyUserSpend" - (id, user_id, date, api_key, model, model_group, custom_llm_provider, - endpoint, prompt_tokens, spend, api_requests, successful_requests) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) - """, - rows, - ) - conn.commit() - - -def _psycopg_query_raw(conn: psycopg.Connection, row_counts: list[int]): - """Run the proxy's $N-parameterized SQL through psycopg, recording each result size.""" - - async def query_raw(sql: str, *params: str) -> list[dict[str, object]]: - converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) - with conn.cursor(row_factory=dict_row) as cur: - cur.execute( - converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query - {f"p{i}": v for i, v in enumerate(params, start=1)}, - ) - rows: Final = cur.fetchall() - row_counts.append(len(rows)) - return rows - - return query_raw - - -@pytest.mark.asyncio -async def test_get_daily_activity_aggregated_bounds_api_key_rollups( - _aggregated_postgresql: psycopg.Connection, -): - """Run the GROUPING SETS statement against real Postgres with more keys than the cap. - - key-004 and key-005 tie on spend exactly at the USAGE_TOP_API_KEYS_LIMIT - cutoff; the api_key tiebreaker must keep key-004 and drop key-005. The PTU - sentinel outspends every key but must not take a slot. Excluded keys and the - sentinel still count toward the totals and the model rollup, which come from - the key-free arm. - """ - n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 5 - key_rows: Final = [ - ( - f"row-{i:03d}", - f"user-{i:03d}", - "2026-06-01", - f"key-{i:03d}", - "gpt-5", - "", - "openai", - "/v1/chat/completions", - 10, - 6.0 if i == 4 else float(i + 1), - 1, - 1, - ) - for i in range(n_keys) - ] - sentinel_row: Final = ( - "row-ptu", - None, - "2026-06-01", - PTU_SENTINEL_API_KEY, - "gpt-5", - "", - "azure", - None, - 0, - 1000.0, - 0, - 0, - ) - _seed_daily_user_spend(_aggregated_postgresql, [*key_rows, sentinel_row]) - key_spend: Final = sum(6.0 if i == 4 else float(i + 1) for i in range(n_keys)) - - row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts) - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - - result = await get_daily_activity_aggregated( - prisma_client=mock_prisma, - table_name="litellm_dailyuserspend", - entity_id_field="user_id", - entity_id=None, - entity_metadata_field=None, - start_date="2026-06-01", - end_date="2026-06-01", - model=None, - api_key=None, - ) - - # Key-free arm: (), (date), (date, model), (date, model_group), two providers, - # one mcp NULL bucket, endpoint plus its NULL bucket = 9 rows regardless of key count. - # Per-key arm: six per-key grouping sets, each capped at the limit. - assert row_counts == [9 + 6 * USAGE_TOP_API_KEYS_LIMIT] - - assert result.metadata.total_spend == pytest.approx(key_spend + 1000.0) - assert result.metadata.total_api_requests == n_keys - assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT - assert result.metadata.total_api_keys == n_keys - - expected_top: Final = {f"key-{i:03d}" for i in range(6, n_keys)} | {"key-004"} - day: Final = result.results[0] - assert day.metrics.spend == pytest.approx(key_spend + 1000.0) - assert set(day.breakdown.api_keys) == expected_top - assert day.breakdown.api_keys["key-004"].metrics.spend == 6.0 - assert "key-005" not in day.breakdown.api_keys - assert PTU_SENTINEL_API_KEY not in day.breakdown.api_keys - - assert day.breakdown.models["gpt-5"].metrics.spend == pytest.approx(key_spend + 1000.0) - assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == expected_top - assert day.breakdown.providers["openai"].metrics.spend == pytest.approx(key_spend) - assert set(day.breakdown.providers["openai"].api_key_breakdown) == expected_top - assert day.breakdown.endpoints["/v1/chat/completions"].metrics.api_requests == n_keys - - -@pytest.mark.asyncio -async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both_arms( - _aggregated_postgresql: psycopg.Connection, -): - """An explicit api_key filter must scope the key-free totals and the per-key - rollups to that key alone, so the two arms never disagree.""" - rows: Final = [ - ( - f"row-{i}", - f"user-{i}", - "2026-06-01", - f"key-{i}", - "gpt-5", - "", - "openai", - "/v1/chat/completions", - 10, - float(i + 1), - 1, - 1, - ) - for i in range(3) - ] - _seed_daily_user_spend(_aggregated_postgresql, rows) - - row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts) - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - - result = await get_daily_activity_aggregated( - prisma_client=mock_prisma, - table_name="litellm_dailyuserspend", - entity_id_field="user_id", - entity_id=None, - entity_metadata_field=None, - start_date="2026-06-01", - end_date="2026-06-01", - model=None, - api_key="key-1", - ) - - assert result.metadata.total_spend == 2.0 - assert result.metadata.total_api_keys == 1 - day: Final = result.results[0] - assert set(day.breakdown.api_keys) == {"key-1"} - assert day.breakdown.api_keys["key-1"].metrics.spend == 2.0 - assert day.breakdown.models["gpt-5"].metrics.spend == 2.0 - assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == {"key-1"} - - -def _prisma_with_marker(marker: str | None) -> MagicMock: - prisma = MagicMock() - prisma.db = MagicMock() - prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - row = ( - None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') - ) - prisma.get_generic_data = AsyncMock(return_value=row) - return prisma - - -def _unfiltered_user_query(**overrides): - return { - "table_name": "litellm_dailyuserspend", - "entity_id_field": "user_id", - "entity_id": None, - "start_date": "2026-06-01", - "end_date": "2026-06-02", - "model": None, - "api_key": None, - "exclude_entity_ids": None, - "timezone_offset_minutes": None, - "include_current_utc_day": False, - **overrides, - } - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("marker", "overrides", "expected"), - [ - ("2026-06-02", {}, "2026-06-02"), - ("2026-06-02", {"model": "gpt-5"}, "2026-06-02"), - ("2026-05-01", {}, "2026-05-01"), - (None, {}, None), - ("2026-06-02", {"api_key": "sk-1"}, None), - ("2026-06-02", {"api_key": []}, None), - ("2026-06-02", {"entity_id": "u-1"}, None), - ("2026-06-02", {"exclude_entity_ids": ["u-1"]}, None), - ("2026-06-02", {"table_name": "litellm_dailyteamspend", "entity_id_field": "team_id"}, None), - ], -) -async def test_global_rollup_marker_is_used_only_for_unfiltered_user_reads(marker, overrides, expected): - """Anything that filters by key or entity has no counterpart in the global table; the - SQL splits the range at the marker itself, so the marker passes through unchanged.""" - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - prisma = _prisma_with_marker(marker) - - assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query(**overrides)) == expected - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - -@pytest.mark.asyncio -async def test_global_rollup_marker_read_failure_falls_back_to_the_per_key_table(): - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - prisma = _prisma_with_marker(None) - prisma.get_generic_data = AsyncMock(side_effect=RuntimeError("db down")) - - assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query()) is None - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - -_GLOBAL_SPEND_MIGRATION: Final = ( - pathlib.Path(__file__).resolve().parents[4] - / "litellm-proxy-extras" - / "litellm_proxy_extras" - / "migrations" - / "20260915000000_add_daily_global_spend" - / "migration.sql" -) - - -@pytest.mark.asyncio -async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_table_and_open_days_live( - _aggregated_postgresql: psycopg.Connection, -): - """Day 1 is rolled up and day 2 is still open (never rolled up), so a marker of day 1 must - give the same response as reading everything per-key: day 1 from the global table, day 2 - live, one grand total across both. Per-key rows that land after the rollup then tell the - two sources apart: a late day 1 row is invisible to totals until the next reconcile while a - late day 2 row shows up at once, and both keys rank in the key breakdown, which stays - per-key throughout.""" - n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3 - rows: Final = [ - ( - f"row-{day}-{i:03d}", - f"user-{i % 7}", - day, - f"key-{i:03d}", - "gpt-5" if i % 2 else "claude", - "" if i % 3 else "gpt-5", - "openai" if i % 2 else None, - "/v1/chat/completions" if i % 5 else None, - 10, - float(i + 1), - 1, - 1, - ) - for day in ("2026-06-01", "2026-06-02") - for i in range(n_keys) - ] - _seed_daily_user_spend(_aggregated_postgresql, rows) - with _aggregated_postgresql.cursor() as cur: - cur.execute( - 'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms = prompt_tokens * 25, ' - "timed_requests = api_requests" - ) - cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - cur.execute( - re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg - {"p1": "2026-06-01"}, - ) - _aggregated_postgresql.commit() - - async def read(marker: str | None): - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - prisma = _prisma_with_marker(marker) - prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, []) - return await get_daily_activity_aggregated( - prisma_client=prisma, - entity_metadata_field=None, - **_unfiltered_user_query(), - ) - - from_per_key = await read(None) - from_global = await read("2026-06-01") - - assert from_global.model_dump() == from_per_key.model_dump() - seeded_spend: Final = 2 * sum(float(i + 1) for i in range(n_keys)) - assert from_global.metadata.total_spend == pytest.approx(seeded_spend) - assert from_global.metadata.total_response_time_ms == 2 * n_keys * 10 * 25 - assert from_global.metadata.total_timed_requests == 2 * n_keys - assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"} - assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT - assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} - - with _aggregated_postgresql.cursor() as cur: - cur.executemany( - """ - INSERT INTO "LiteLLM_DailyUserSpend" - (id, user_id, date, api_key, model, model_group, custom_llm_provider, - endpoint, prompt_tokens, spend, api_requests, successful_requests) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) - """, - [ - ("late-1", "user-late", "2026-06-01", "key-late-1", "gpt-5", "", "openai", None, 10, 1000.0, 1, 1), - ("late-2", "user-late", "2026-06-02", "key-late-2", "gpt-5", "", "openai", None, 10, 500.0, 1, 1), - ], - ) - _aggregated_postgresql.commit() - - late_per_key = await read(None) - late_global = await read("2026-06-01") - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - assert late_per_key.metadata.total_spend == pytest.approx(seeded_spend + 1000.0 + 500.0) - assert late_global.metadata.total_spend == pytest.approx(seeded_spend + 500.0) - by_day: Final = {day.date.isoformat(): day for day in late_global.results} - assert by_day["2026-06-01"].metrics.spend == pytest.approx(seeded_spend / 2) - assert by_day["2026-06-02"].metrics.spend == pytest.approx(seeded_spend / 2 + 500.0) - assert by_day["2026-06-01"].breakdown.api_keys["key-late-1"].metrics.spend == pytest.approx(1000.0) - assert by_day["2026-06-02"].breakdown.api_keys["key-late-2"].metrics.spend == pytest.approx(500.0) - assert late_global.metadata.total_api_keys == n_keys + 2 - - -@pytest.mark.asyncio -async def test_get_daily_activity_aggregated_reports_exact_limit_key_count_as_complete( - _aggregated_postgresql: psycopg.Connection, -): - """With exactly USAGE_TOP_API_KEYS_LIMIT keys nothing is dropped, and the - response must say so: total_api_keys equals the limit rather than exceeding it.""" - rows: Final = [ - ( - f"row-{i:03d}", - f"user-{i:03d}", - "2026-06-01", - f"key-{i:03d}", - "gpt-5", - "", - "openai", - "/v1/chat/completions", - 10, - float(i + 1), - 1, - 1, - ) - for i in range(USAGE_TOP_API_KEYS_LIMIT) - ] - _seed_daily_user_spend(_aggregated_postgresql, rows) - - row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts) - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - - result = await get_daily_activity_aggregated( - prisma_client=mock_prisma, - table_name="litellm_dailyuserspend", - entity_id_field="user_id", - entity_id=None, - entity_metadata_field=None, - start_date="2026-06-01", - end_date="2026-06-01", - model=None, - api_key=None, - ) - - assert result.metadata.total_api_keys == USAGE_TOP_API_KEYS_LIMIT - assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT - assert len(result.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT - - -@pytest.mark.asyncio -async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_model_name( - _aggregated_postgresql: psycopg.Connection, -): - """Rows stored with an empty or NULL model_group must land in the model_groups - breakdown under their model name instead of vanishing from the usage UI.""" - rows: Final = [ - ( - "row-0", - "user-0", - "2026-06-01", - "key-0", - "gpt-5", - "gpt-5-eu", - "openai", - "/v1/chat/completions", - 10, - 7.0, - 1, - 1, - ), - ("row-1", "user-1", "2026-06-01", "key-1", "gpt-5", "", "openai", "/v1/chat/completions", 10, 3.0, 1, 1), - ("row-2", "user-2", "2026-06-01", "key-2", "claude-x", None, "anthropic", "/v1/messages", 10, 2.0, 1, 1), - ] - _seed_daily_user_spend(_aggregated_postgresql, rows) - - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, []) - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - - result = await get_daily_activity_aggregated( - prisma_client=mock_prisma, - table_name="litellm_dailyuserspend", - entity_id_field="user_id", - entity_id=None, - entity_metadata_field=None, - start_date="2026-06-01", - end_date="2026-06-01", - model=None, - api_key=None, - ) - - breakdown: Final = result.results[0].breakdown - assert set(breakdown.model_groups) == {"gpt-5-eu", "gpt-5", "claude-x"} - assert breakdown.model_groups["gpt-5-eu"].metrics.spend == 7.0 - assert breakdown.model_groups["gpt-5"].metrics.spend == 3.0 - assert breakdown.model_groups["claude-x"].metrics.spend == 2.0 - assert set(breakdown.model_groups["gpt-5"].api_key_breakdown) == {"key-1"} - assert set(breakdown.models) == {"gpt-5", "claude-x"} - assert breakdown.models["gpt-5"].metrics.spend == 10.0 - - def _no_spend_record(): """A rollup row for a key with no spend, where SUM() returns NULL (None).""" return SimpleNamespace( @@ -2717,7 +2249,7 @@ def test_entity_rollup_sql_query_and_api_key_list_filter(): api_key=[], ) assert "FALSE" in empty_sql - assert empty_params == ["2024-01-01", "2024-01-31", PTU_SENTINEL_API_KEY] + assert empty_params == ["2024-01-01", "2024-01-31"] @pytest.mark.asyncio @@ -2751,10 +2283,10 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown(): "successful_requests": 0, } main_rows = [ - {**base, "date": None, "group_level": 127, "distinct_api_keys": None, "spend": 18.0}, - {**base, "date": "2024-01-01", "group_level": 63, "distinct_api_keys": None, "spend": 18.0}, - {**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 47, "distinct_api_keys": None, "spend": 18.0}, - {**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 31, "distinct_api_keys": 1, "spend": 12.0}, + {**base, "date": None, "group_level": 127, "spend": 18.0}, + {**base, "date": "2024-01-01", "group_level": 63, "spend": 18.0}, + {**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 47, "spend": 18.0}, + {**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 31, "spend": 12.0}, ] entity_base = { key: value diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index 6121608b658..deb7289d2d1 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -28,7 +28,6 @@ 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 @@ -1043,57 +1042,6 @@ async def test_spend_report_locks_are_never_released(): proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited() -def _init_daily_global_spend_reconcile_job() -> tuple[AsyncIOScheduler, MagicMock, MagicMock]: - scheduler = AsyncIOScheduler() - proxy_logging_obj = MagicMock() - proxy_logging_obj.alerting_handler = AsyncMock() - prisma_client = MagicMock() - ProxyStartupEvent._initialize_daily_global_spend_reconcile_job( - scheduler=scheduler, - proxy_logging_obj=proxy_logging_obj, - prisma_client=prisma_client, - ) - return scheduler, proxy_logging_obj, prisma_client - - -def test_daily_global_spend_reconcile_job_is_scheduled_nightly_with_an_immediate_catch_up_run(): - """Startup schedules the LiteLLM_DailyGlobalSpend backfill a couple of minutes out, so a - fresh deploy switches usage reads to the global table without waiting for the nightly - run, and after that it fires once a day at 00:30 UTC, when the previous UTC day is closed.""" - from datetime import datetime, timedelta, timezone - - from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID - - scheduler, _, _ = _init_daily_global_spend_reconcile_job() - job = scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) - assert job is not None - - assert timedelta(0) < job.next_run_time - datetime.now(timezone.utc) <= timedelta(minutes=2) - after_catch_up = datetime(2026, 9, 16, 12, 0, tzinfo=timezone.utc) - assert job.trigger.get_next_fire_time(None, after_catch_up) == datetime(2026, 9, 17, 0, 30, tzinfo=timezone.utc) - just_after_a_run = datetime(2026, 9, 17, 0, 30, 1, tzinfo=timezone.utc) - assert job.trigger.get_next_fire_time(None, just_after_a_run) == datetime(2026, 9, 18, 0, 30, tzinfo=timezone.utc) - - -@pytest.mark.asyncio -async def test_daily_global_spend_reconcile_job_runs_under_the_pod_lock_and_alerts_through_the_proxy(monkeypatch): - from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID - - scheduler, proxy_logging_obj, prisma_client = _init_daily_global_spend_reconcile_job() - run = AsyncMock() - monkeypatch.setattr(ps, "run_scheduled_daily_global_spend_reconcile", run) - - await scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID).func() - - run.assert_awaited_once() - assert run.await_args.args == (prisma_client,) - assert run.await_args.kwargs["pod_lock_manager"] is proxy_logging_obj.db_spend_update_writer.pod_lock_manager - await run.await_args.kwargs["alert"]("day 2026-09-01 failed") - proxy_logging_obj.alerting_handler.assert_awaited_once() - assert proxy_logging_obj.alerting_handler.await_args.kwargs["message"] == "day 2026-09-01 failed" - assert proxy_logging_obj.alerting_handler.await_args.kwargs["level"] == "High" - - @pytest.mark.asyncio async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch): """The boot-time send goes through the same gate, so a losing pod sends nothing at all: diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py deleted file mode 100644 index 3da587435ad..00000000000 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ /dev/null @@ -1,532 +0,0 @@ -"""Tests for the LiteLLM_DailyGlobalSpend reconcile job (LIT-7818).""" - -import json -import pathlib -import re -from datetime import date -from typing import Final -from unittest.mock import AsyncMock, MagicMock - -import psycopg -import pytest -from psycopg.rows import dict_row -from pytest_postgresql import factories - -from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM -from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key -from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( - _ADVANCE_MARKER_SQL, - RECONCILE_DAY_SQL, - read_marker, - reconciled_through, - run_daily_global_spend_reconcile, - run_scheduled_daily_global_spend_reconcile, -) -from litellm.proxy.utils import evict_config_param - -USER_TABLE: Final = DAILY_SPEND_TABLES["user"] -TODAY: Final = date(2026, 9, 15) - - -class _FakeConfigRow: - def __init__(self, param_name: str, param_value: object) -> None: - self.param_name = param_name - self.param_value = param_value - - -class _FakeConfigTable: - def __init__(self) -> None: - self.rows: dict[str, object] = {} - - def advance(self, param_name: str, through: str | None, scanned_at: str | None) -> None: - """What ``_ADVANCE_MARKER_SQL`` does in Postgres: keep the later of stored and incoming per field.""" - stored = self.rows.get(param_name) - current: dict[str, str | None] = json.loads(stored) if isinstance(stored, str) else {} - self.rows[param_name] = json.dumps( - { - "reconciled_through": _greatest(current.get("reconciled_through"), through), - "scanned_at": _greatest(current.get("scanned_at"), scanned_at), - } - ) - - -def _greatest(stored: str | None, incoming: str | None) -> str | None: - present = [value for value in (stored, incoming) if value is not None] - return max(present) if present else None - - -class _FakeDb: - """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, - so "rows written since the last scan" behaves like Postgres would. The database's own - date decides which day is still open, never the pod's clock.""" - - def __init__(self, prisma: "_FakePrisma") -> None: - self._prisma = prisma - self.litellm_config = _FakeConfigTable() - - async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]: - if sql.startswith("SELECT (NOW()"): - self._prisma.clock += 1 - return [{"now": f"clock-{self._prisma.clock:04d}", "today": self._prisma.today.isoformat()}] - rows = self._prisma.user_rows - if len(params) == 1: - (last,) = params - return [{"date": d} for d in sorted(rows) if d <= last] - last, marker, scanned_at = params - return [ - {"date": d} for d, written in sorted(rows.items()) if d <= last and (d > marker or written >= scanned_at) - ] - - async def execute_raw(self, sql: str, *params: str | None) -> int: - if sql == _ADVANCE_MARKER_SQL: - param_name, through, scanned_at = params - assert param_name is not None - self.litellm_config.advance(param_name, through, scanned_at) - return 1 - (day,) = params - if day is None or day in self._prisma.failing_days: - raise RuntimeError(f"day {day} exploded") - self._prisma.reconciled.append(day) - landing = self._prisma.marker_landing_on_day.get(day) - if landing is not None: - self.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = landing - return 1 - - -class _FakePrisma: - """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw. - ``marker_landing_on_day`` stores another pod's marker the moment this run rewrites that day.""" - - def __init__( - self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset(), today: date = TODAY - ) -> None: - self.clock = 0 - self.today = today - self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} - self.failing_days = failing_days - self.marker_landing_on_day: dict[str, str] = {} - self.reconciled: list[str] = [] - self.db = _FakeDb(self) - - def write_late_row(self, day: str) -> None: - """A per-key row for ``day`` lands now, after whatever scans already happened.""" - self.clock += 1 - self.user_rows[day] = f"clock-{self.clock:04d}" - - async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: - stored = self.db.litellm_config.rows.get(value) - return None if stored is None else _FakeConfigRow(value, stored) - - -@pytest.fixture(autouse=True) -async def _fresh_marker_cache(): - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - yield - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - - -@pytest.mark.asyncio -async def test_first_run_rolls_up_every_closed_day_and_never_the_database_s_today(): - """Before any marker exists every closed day with per-key rows is rolled up. Today is left - out: pods are still flushing it, so it is served live from the per-key table until it closes. - The database clock says which day that is; a pod booting with its clock a day ahead must not - roll the open day up and mark it reconciled.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15")) - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14") - assert result.failed_day is None - assert result.reconciled_through == "2026-09-14" - assert await reconciled_through(prisma) == "2026-09-14" - assert "2026-09-15" not in prisma.reconciled - - -@pytest.mark.asyncio -async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed(): - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14"), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.reconciled.clear() - prisma.today = TODAY - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-14",) - assert await reconciled_through(prisma) == "2026-09-14" - - -@pytest.mark.asyncio -async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_run(): - """Per-key rows carry the request start date, so a delayed flush or retry can add spend to a - day far behind the marker. That day is rewritten, and the marker never moves back for it.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13"), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.reconciled.clear() - prisma.today = TODAY - prisma.write_late_row("2026-09-01") - prisma.write_late_row("2026-09-03") - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-03") - assert "2026-09-05" not in prisma.reconciled - assert await reconciled_through(prisma) == "2026-09-13" - - -@pytest.mark.asyncio -async def test_a_late_row_seen_by_a_failed_run_is_seen_again_by_the_next_one(): - """The scan time only advances when every pending day was rewritten, otherwise a late row - found by the failed run would be counted as handled.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13"), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.today = TODAY - prisma.write_late_row("2026-09-01") - prisma.failing_days = frozenset({"2026-09-01"}) - failed = await run_daily_global_spend_reconcile(prisma) - prisma.failing_days = frozenset() - prisma.reconciled.clear() - - result = await run_daily_global_spend_reconcile(prisma) - - assert failed.failed_day == "2026-09-01" - assert failed.reconciled_through == "2026-09-13" - assert result.days_reconciled == ("2026-09-01",) - assert result.failed_day is None - - -@pytest.mark.asyncio -async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again(): - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) - prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-13"}' - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-13") - marker = await read_marker(prisma) - assert marker is not None and marker.reconciled_through == "2026-09-13" and marker.scanned_at is not None - - -@pytest.mark.asyncio -async def test_a_run_with_no_new_closed_days_keeps_the_marker(): - prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.reconciled.clear() - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == () - assert result.reconciled_through == "2026-09-13" - - -@pytest.mark.asyncio -async def test_a_failing_day_stops_the_run_and_leaves_the_marker_on_the_last_good_day(): - """The marker may never claim a day that was not rewritten: reads past it would then trust - a global table missing that day's spend.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01",) - assert result.failed_day == "2026-09-02" - assert result.reconciled_through == "2026-09-01" - assert prisma.reconciled == ["2026-09-01"] - assert await reconciled_through(prisma) == "2026-09-01" - - -@pytest.mark.asyncio -async def test_the_next_run_resumes_from_the_failed_day(): - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) - await run_daily_global_spend_reconcile(prisma) - prisma.failing_days = frozenset() - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03") - assert await reconciled_through(prisma) == "2026-09-03" - - -@pytest.mark.asyncio -async def test_a_slower_overlapping_run_never_rewinds_the_marker_a_faster_run_stored(): - """Two pods can reconcile at once (Redis unreachable, or the lock expired on a long backfill). - When the faster one has already stored a later marker, the slower one may only add to it. Putting - its own older prefix back, or dropping the scan time, would send usage reads for every day in - between back to the per-key table until the next run.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-03"})) - prisma.marker_landing_on_day = { - "2026-09-02": '{"reconciled_through": "2026-09-14", "scanned_at": "clock-0009"}', - } - - result = await run_daily_global_spend_reconcile(prisma) - - assert result.days_reconciled == ("2026-09-01", "2026-09-02") - assert result.reconciled_through == "2026-09-14" - marker = await read_marker(prisma) - assert marker is not None and (marker.reconciled_through, marker.scanned_at) == ("2026-09-14", "clock-0009") - - -@pytest.mark.asyncio -async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): - """When the rewrite of a late day fails the marker must stay put and the operator must hear about it.""" - prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14)) - await run_daily_global_spend_reconcile(prisma) - prisma.today = TODAY - prisma.write_late_row("2026-09-12") - prisma.failing_days = frozenset({"2026-09-12"}) - alert = AsyncMock() - - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert) - - assert result is not None - assert result.days_reconciled == () - assert result.failed_day == "2026-09-12" - assert result.reconciled_through == "2026-09-13" - alert.assert_awaited_once() - assert "2026-09-12" in alert.await_args.args[0] - - -@pytest.mark.asyncio -async def test_a_clean_run_does_not_alert(): - prisma = _FakePrisma(user_days=("2026-09-13",)) - alert = AsyncMock() - - await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert) - - alert.assert_not_awaited() - - -def _pod_lock(acquired: bool) -> MagicMock: - lock = MagicMock() - lock.redis_cache = MagicMock() - lock.redis_cache.async_get_cache = AsyncMock(return_value="other-pod") - lock.get_redis_lock_key = MagicMock(return_value="lock-key") - lock.acquire_lock = AsyncMock(return_value=acquired) - lock.release_lock = AsyncMock() - return lock - - -@pytest.mark.asyncio -async def test_scheduled_run_skips_when_another_pod_holds_the_lock(): - prisma = _FakePrisma(user_days=("2026-09-13",)) - lock = _pod_lock(acquired=False) - - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) - - assert result is None - assert prisma.reconciled == [] - lock.release_lock.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins(): - prisma = _FakePrisma(user_days=("2026-09-13",)) - lock = _pod_lock(acquired=True) - - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) - - assert result is not None and result.days_reconciled == ("2026-09-13",) - lock.release_lock.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read(): - """A Redis outage must not stall the backfill: the day rewrite is idempotent, so running - twice is only wasted effort while skipping forever leaves usage on the slow path.""" - prisma = _FakePrisma(user_days=("2026-09-13",)) - lock = _pod_lock(acquired=False) - lock.redis_cache.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) - - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) - - assert result is not None and result.days_reconciled == ("2026-09-13",) - lock.release_lock.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_marker_is_read_back_from_the_json_string_the_config_table_stores(): - prisma = _FakePrisma(user_days=()) - prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-10"}' - - assert await reconciled_through(prisma) == "2026-09-10" - - -@pytest.mark.asyncio -async def test_an_unparseable_marker_reads_as_never_reconciled(): - prisma = _FakePrisma(user_days=()) - prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"something_else": 1}' - - assert await reconciled_through(prisma) is None - - -_rollup_postgresql_proc: Final = factories.postgresql_proc() -_rollup_postgresql: Final = factories.postgresql("_rollup_postgresql_proc") - -_MIGRATIONS_DIR: Final = ( - pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" -) -_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql" - -_DAILY_USER_SPEND_DDL: Final = """ - CREATE TABLE "LiteLLM_DailyUserSpend" ( - id TEXT PRIMARY KEY, - user_id TEXT, - date TEXT NOT NULL, - api_key TEXT NOT NULL, - model TEXT, - model_group TEXT, - custom_llm_provider TEXT, - mcp_namespaced_tool_name TEXT, - endpoint TEXT, - prompt_tokens BIGINT DEFAULT 0, - completion_tokens BIGINT DEFAULT 0, - cache_read_input_tokens BIGINT DEFAULT 0, - cache_creation_input_tokens BIGINT DEFAULT 0, - compression_saved_tokens BIGINT DEFAULT 0, - compression_savings_spend DOUBLE PRECISION DEFAULT 0, - prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, - spend DOUBLE PRECISION DEFAULT 0, - api_requests BIGINT DEFAULT 0, - successful_requests BIGINT DEFAULT 0, - failed_requests BIGINT DEFAULT 0, - total_response_time_ms BIGINT DEFAULT 0, - timed_requests BIGINT DEFAULT 0, - created_at TIMESTAMP DEFAULT now(), - updated_at TIMESTAMP, - UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) - ) -""" - -_PER_KEY_SUMS_SQL: Final = """ - SELECT COALESCE(model, '') AS model, COALESCE(model_group, '') AS model_group, - COALESCE(custom_llm_provider, '') AS custom_llm_provider, - SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, SUM(api_requests) AS api_requests, - SUM(total_response_time_ms) AS total_response_time_ms, SUM(timed_requests) AS timed_requests - FROM "LiteLLM_DailyUserSpend" WHERE date = %s - GROUP BY 1, 2, 3 ORDER BY 1, 2, 3 -""" -_GLOBAL_ROWS_SQL: Final = """ - SELECT model, model_group, custom_llm_provider, spend, prompt_tokens, api_requests, - total_response_time_ms, timed_requests - FROM "LiteLLM_DailyGlobalSpend" WHERE date = %s ORDER BY 1, 2, 3 -""" - - -def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None: - converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) - conn.execute( - converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query - {f"p{i}": v for i, v in enumerate(params, start=1)}, - ) - conn.commit() - - -def _user_txn(**overrides): - return { - "user_id": "u-1", - "date": "2026-09-14", - "api_key": "sk-1", - "model": "gpt-5", - "model_group": "gpt-5", - "custom_llm_provider": "openai", - "mcp_namespaced_tool_name": "", - "endpoint": "/chat/completions", - "prompt_tokens": 10, - "completion_tokens": 20, - "spend": 1.0, - "api_requests": 1, - "successful_requests": 1, - "failed_requests": 0, - "total_response_time_ms": 800, - "timed_requests": 1, - **overrides, - } - - -def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]: - return [ - ( - r["model"], - r["model_group"], - r["custom_llm_provider"], - float(r["spend"]), - int(r["prompt_tokens"]), - int(r["api_requests"]), - int(r["total_response_time_ms"]), - int(r["timed_requests"]), - ) # pyright: ignore[reportArgumentType] # dict_row values are untyped - for r in rows - ] - - -def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_postgresql: psycopg.Connection): - """Against real Postgres and the shipped migration: writer-shaped rows and legacy rows - (NULL and '' dimension spellings) fold into one global day, running the day twice changes - nothing, and other days are left alone.""" - conn: Final = _rollup_postgresql - conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal - conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - conn.commit() - - written_batch = merge_by_conflict_key( - USER_TABLE, - (_user_txn(api_key="sk-1", spend=1.0), _user_txn(api_key="sk-2", user_id="u-2", spend=2.0, prompt_tokens=20)), - ) - _execute_dollar_sql(conn, *build_bulk_upsert(USER_TABLE, written_batch)) - - conn.execute( - """ - INSERT INTO "LiteLLM_DailyUserSpend" - (id, user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, - endpoint, prompt_tokens, spend, api_requests) - VALUES - ('legacy-1', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', NULL, 'openai', NULL, NULL, 5, 4.0, 1), - ('legacy-2', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', '', 'openai', '', '', 5, 8.0, 1), - ('legacy-3', 'u-9', '2026-09-13', 'sk-9', 'claude', '', 'anthropic', '', '', 7, 16.0, 1) - """ - ) - conn.commit() - - _execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",)) - _execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",)) - - with conn.cursor(row_factory=dict_row) as cur: - global_rows = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-14",)).fetchall() - per_key = cur.execute(_PER_KEY_SUMS_SQL, ("2026-09-14",)).fetchall() - untouched = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-13",)).fetchall() - - assert _normalized(global_rows) == _normalized(per_key) - assert sum(float(r["spend"]) for r in global_rows) == pytest.approx(15.0) # pyright: ignore[reportArgumentType] # dict_row values are untyped - assert sum(int(r["total_response_time_ms"]) for r in global_rows) == 1600 # pyright: ignore[reportArgumentType] # dict_row values are untyped - assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")] - assert untouched == [] - - -_CONFIG_DDL: Final = 'CREATE TABLE "LiteLLM_Config" (param_name TEXT PRIMARY KEY, param_value JSONB)' -_MARKER_SQL: Final = 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s' - - -def test_advance_marker_sql_only_ever_moves_the_stored_marker_forward(_rollup_postgresql: psycopg.Connection): - """Against real Postgres: the statement a slower overlapping run issues after the faster run - already stored a later marker leaves that marker alone, whether it carries an older scan time or - none at all, while a run that is further along moves both fields on.""" - conn: Final = _rollup_postgresql - conn.execute(_CONFIG_DDL) # pyright: ignore[reportArgumentType] # DDL literal - conn.commit() - param: Final = DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM - - def stored() -> object: - with conn.cursor(row_factory=dict_row) as cur: - row = cur.execute(_MARKER_SQL, (param,)).fetchone() - return None if row is None else row["param_value"] - - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-01", None)) - assert stored() == {"reconciled_through": "2026-09-01", "scanned_at": None} - - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-14", "2026-09-15 00:30:02.5")) - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-02", None)) - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-03", "2026-09-15 00:30:01.25")) - assert stored() == {"reconciled_through": "2026-09-14", "scanned_at": "2026-09-15 00:30:02.5"} - - _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-15", "2026-09-16 00:30:00.75")) - assert stored() == {"reconciled_through": "2026-09-15", "scanned_at": "2026-09-16 00:30:00.75"} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx index 1c39c6eb36d..a97f23b3334 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx @@ -178,28 +178,4 @@ describe("CacheLeakageCard", () => { screen.queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."), ).not.toBeInTheDocument(); }); - - it("says which keys are missing from the key ranking when the proxy capped the per-key lists", () => { - const day = dayWithKeys("2026-07-12", { - "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), - }); - renderWith([day], { apiKeyTruncation: { limit: 100, total: 3000 } }); - - expect(screen.getByRole("note")).toHaveTextContent( - "Only the 100 highest-spend keys of 3,000 are loaded, so a lower-spend key that leaks more is not listed here.", - ); - - fireEvent.click(screen.getByRole("tab", { name: "By model" })); - - expect(screen.queryByRole("note")).not.toBeInTheDocument(); - }); - - it("keeps the key ranking note off when every key was loaded", () => { - const day = dayWithKeys("2026-07-12", { - "hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }), - }); - renderWith([day]); - - expect(screen.queryByRole("note")).not.toBeInTheDocument(); - }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx index a0877b04648..3f27449ebe1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx @@ -81,7 +81,7 @@ const SortableHead = ({ }; const CacheLeakageCard: React.FC = ({ activity }) => { - const { dateValue, onDateChange, results, loading, isFetchingMore, apiKeyTruncation } = activity; + const { dateValue, onDateChange, results, loading, isFetchingMore } = activity; const [dimension, setDimension] = useState("key"); const [sort, setSort] = useState({ column: "potentialSavings", dir: "desc" }); const leakage = useMemo(() => computeCacheLeakage(results, dimension), [results, dimension]); @@ -123,13 +123,6 @@ const CacheLeakageCard: React.FC = ({ activity }) => { - {dimension === "key" && apiKeyTruncation !== undefined && ( -

- Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "} - {apiKeyTruncation.total.toLocaleString()} are loaded, so a lower-spend key that leaks more is not listed - here. Raise USAGE_TOP_API_KEYS_LIMIT on the proxy to load more keys. -

- )} {rows.length > 0 && isFetchingMore && (

Data is still loading; rows and totals will update as the rest of the range arrives. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx index 4059303d5a5..00902aa9fdd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx @@ -4,13 +4,12 @@ import { describe, expect, it, vi } from "vitest"; const mockUsePaginatedDailyActivity = vi.fn(); const mockCancel = vi.fn(); -let mockMetadata: Record = {}; vi.mock("@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity", () => ({ usePaginatedDailyActivity: (args: unknown) => { mockUsePaginatedDailyActivity(args); return { - data: { results: [], metadata: mockMetadata }, + data: { results: [] }, loading: false, isFetchingMore: false, progress: { currentPage: 4, totalPages: 9 }, @@ -81,18 +80,4 @@ describe("useDailyActivityRange", () => { expect(mockUsePaginatedDailyActivity).toHaveBeenLastCalledWith(expect.objectContaining({ enabled: false })); }); - - it("reports how many keys the proxy left out of the per-key lists", () => { - mockMetadata = { api_key_limit: 100, total_api_keys: 3000 }; - const { result } = renderHook(() => useDailyActivityRange("test-token", "u1", "proxy_admin")); - - expect(result.current.apiKeyTruncation).toEqual({ limit: 100, total: 3000 }); - }); - - it("reports no key truncation when every key fit under the proxy limit", () => { - mockMetadata = { api_key_limit: 100, total_api_keys: 100 }; - const { result } = renderHook(() => useDailyActivityRange("test-token", "u1", "proxy_admin")); - - expect(result.current.apiKeyTruncation).toBeUndefined(); - }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts index 92dd24b8d6d..9f793a68bf5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts @@ -1,7 +1,6 @@ import { useMemo, useState } from "react"; import { userDailyActivityAggregatedCall, userDailyActivityCall } from "@/components/networking"; -import { ApiKeyTruncation, getApiKeyTruncation } from "@/components/EntityUsageExport/exportBlockedReason"; import { DailyData } from "@/components/UsagePage/types"; import { spendScopeUserId } from "@/utils/roles"; import { usePaginatedDailyActivity } from "@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity"; @@ -23,7 +22,6 @@ export interface DailyActivityRange { cancelled: boolean; failed: boolean; cancel: () => void; - apiKeyTruncation?: ApiKeyTruncation; } /** @@ -80,7 +78,6 @@ export const useScopedDailyActivityRange = ( cancelled, failed, cancel, - apiKeyTruncation: getApiKeyTruncation(data.metadata?.api_key_limit, data.metadata?.total_api_keys), }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx index 6bd16095351..8f5bb9455d0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx @@ -650,23 +650,6 @@ describe("EntityUsage", () => { expect(screen.getAllByText("Activity Metrics")[1]).toBeInTheDocument(); }); - it("tells the team view how many keys the proxy left out of the per-key lists", async () => { - mockTeamDailyActivityAggregatedCall.mockResolvedValue({ - ...mockSpendData, - metadata: { ...mockSpendData.metadata, api_key_limit: 100, total_api_keys: 3000 }, - }); - render(); - - await waitFor(() => { - expect(mockTeamDailyActivityAggregatedCall).toHaveBeenCalled(); - }); - act(() => { - fireEvent.click(screen.getByText("Key Activity")); - }); - - expect(await screen.findByRole("note")).toHaveTextContent("Only the 100 highest-spend keys of 3,000 are loaded"); - }); - // An inactive tab panel is marked aria-selected="false" by one tab library and hidden by the // other, so treat either as "not on screen" and the assertion holds whichever one is rendering. const isShowing = (element: HTMLElement): boolean => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 27460b21108..6687bd4df03 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -25,7 +25,7 @@ import TeamMultiSelect from "@/components/common_components/team_multi_select"; import UserDropdown from "@/components/common_components/UserDropdown"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import { UsageExportHeader } from "@/components/EntityUsageExport"; -import { getApiKeyTruncation, getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; +import { getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; import type { EntityType } from "@/components/EntityUsageExport/types"; import { agentDailyActivityCall, @@ -71,8 +71,6 @@ interface EntitySpendData { total_successful_requests: number; total_failed_requests: number; total_tokens: number; - api_key_limit?: number | null; - total_api_keys?: number | null; }; } @@ -162,7 +160,6 @@ const EntityUsage: React.FC = ({ }); const spendData = spendDataRaw as unknown as EntitySpendData; - const apiKeyTruncation = getApiKeyTruncation(spendData.metadata?.api_key_limit, spendData.metadata?.total_api_keys); const { data: agentSpendDataRaw, @@ -662,18 +659,12 @@ const EntityUsage: React.FC = ({ { key: "keys", label: "Key Activity", - content: ( - - ), + content: , }, { key: "endpoints", label: "Endpoint Activity", content: }, ]; - const spendFetchState = { coversRange, cancelled, failed, apiKeyTruncation }; + const spendFetchState = { coversRange, cancelled, failed }; return (

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index 228d8acf146..018dfc84740 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -30,7 +30,7 @@ import { ActivityMetrics, processActivityData } from "@/components/activity_metr import CloudZeroExportModal from "@/components/cloudzero_export_modal"; import UserDropdown from "@/components/common_components/UserDropdown"; import EntityUsageExportModal from "@/components/EntityUsageExport"; -import { getApiKeyTruncation, getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; +import { getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel"; import { Team } from "@/components/key_team_helpers/key_list"; import { @@ -256,10 +256,6 @@ const UsagePage: React.FC = ({ teams, organizations }) => { coversRange: activeAggregated !== null || paginatedResult.coversRange, cancelled: paginatedResult.cancelled, failed: paginatedResult.failed, - apiKeyTruncation: getApiKeyTruncation( - userSpendData.metadata?.api_key_limit, - userSpendData.metadata?.total_api_keys, - ), }; const exportBlockedReason = getExportBlockedReason(spendFetchState); @@ -865,7 +861,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { - + diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts index 8491b31f5f9..e39b01a5dea 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts @@ -1,12 +1,11 @@ import { describe, expect, it } from "vitest"; -import { getApiKeyTruncation, getExportBlockedReason, type UsageFetchState } from "./exportBlockedReason"; +import { getExportBlockedReason, type UsageFetchState } from "./exportBlockedReason"; const state = (overrides: Partial = {}): UsageFetchState => ({ coversRange: true, cancelled: false, failed: false, - apiKeyTruncation: undefined, ...overrides, }); @@ -32,27 +31,4 @@ describe("getExportBlockedReason", () => { expect(reason).toMatch(/failed to load/i); expect(reason).not.toMatch(/stopped/i); }); - - it("blocks when the aggregated endpoint dropped keys, since a per-team CSV would miss them", () => { - const reason = getExportBlockedReason(state({ apiKeyTruncation: { limit: 100, total: 3000 } })); - - expect(reason).toMatch(/100 highest-spend keys of 3000/); - expect(reason).toMatch(/USAGE_TOP_API_KEYS_LIMIT/); - }); -}); - -describe("getApiKeyTruncation", () => { - it("reports truncation once the proxy saw more keys than it returned", () => { - expect(getApiKeyTruncation(100, 101)).toEqual({ limit: 100, total: 101 }); - }); - - it("stays quiet when exactly the cap exists, since every key is on screen", () => { - expect(getApiKeyTruncation(100, 100)).toBeUndefined(); - expect(getApiKeyTruncation(100, 7)).toBeUndefined(); - }); - - it("stays quiet when the response carries no cap, as the paginated fallback does", () => { - expect(getApiKeyTruncation(undefined, undefined)).toBeUndefined(); - expect(getApiKeyTruncation(100, null)).toBeUndefined(); - }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts index 6c5a5f83231..71408ba8f3f 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts @@ -1,31 +1,13 @@ -export interface ApiKeyTruncation { - limit: number; - total: number; -} - export interface UsageFetchState { coversRange: boolean; cancelled: boolean; failed: boolean; - apiKeyTruncation: ApiKeyTruncation | undefined; } -export const getApiKeyTruncation = (apiKeyLimit: unknown, totalApiKeys: unknown): ApiKeyTruncation | undefined => { - if (typeof apiKeyLimit !== "number" || typeof totalApiKeys !== "number") return undefined; - return totalApiKeys > apiKeyLimit ? { limit: apiKeyLimit, total: totalApiKeys } : undefined; -}; - -export const getExportBlockedReason = ({ - coversRange, - cancelled, - failed, - apiKeyTruncation, -}: UsageFetchState): string | undefined => { +export const getExportBlockedReason = ({ coversRange, cancelled, failed }: UsageFetchState): string | undefined => { if (failed) return "Some spend data failed to load, so an export would under-report. Reload the page to try again."; if (cancelled) return "Loading was stopped before the whole range arrived, so an export would under-report. Reload the page to load it all."; if (!coversRange) return "Spend data is still loading, so an export would under-report. Wait for it to finish."; - if (apiKeyTruncation !== undefined) - return `Only the ${apiKeyTruncation.limit} highest-spend keys of ${apiKeyTruncation.total} were loaded, so a per-team export would under-report. Raise USAGE_TOP_API_KEYS_LIMIT on the proxy to load more keys.`; return undefined; }; diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx index 830139143e9..693ac20a360 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx @@ -68,14 +68,4 @@ describe("KeyActivityPanel", () => { expect(screen.getByLabelText("Search keys")).toHaveValue(""); expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alicehash-bob"); }); - - it("says how many keys the proxy left out when only the top spenders were loaded", () => { - render(); - expect(screen.getByRole("note")).toHaveTextContent("Only the 2 highest-spend keys of 3,000 are loaded"); - }); - - it("shows no truncation note when every key is loaded", () => { - render(); - expect(screen.queryByRole("note")).not.toBeInTheDocument(); - }); }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx index 8b2141f8528..8287a04d0c7 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx @@ -2,7 +2,6 @@ import { Search, X } from "lucide-react"; import React, { useMemo, useState } from "react"; import { ActivityMetrics } from "@/components/activity_metrics"; -import type { ApiKeyTruncation } from "@/components/EntityUsageExport/exportBlockedReason"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { filterKeyActivity } from "../keyActivityFilter"; @@ -11,14 +10,9 @@ import type { ModelActivityData } from "../types"; interface KeyActivityPanelProps { keyMetrics: Record; hidePromptCachingMetrics?: boolean; - apiKeyTruncation?: ApiKeyTruncation; } -const KeyActivityPanel: React.FC = ({ - keyMetrics, - hidePromptCachingMetrics = false, - apiKeyTruncation, -}) => { +const KeyActivityPanel: React.FC = ({ keyMetrics, hidePromptCachingMetrics = false }) => { const [query, setQuery] = useState(""); const filtered = useMemo(() => filterKeyActivity(keyMetrics, query), [keyMetrics, query]); const totalKeys = Object.keys(keyMetrics).length; @@ -49,12 +43,6 @@ const KeyActivityPanel: React.FC = ({ Showing {shownKeys.toLocaleString()} of {totalKeys.toLocaleString()} keys - {apiKeyTruncation !== undefined && ( - - Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "} - {apiKeyTruncation.total.toLocaleString()} are loaded - - )}
{isFiltering && totalKeys > 0 && shownKeys === 0 ? (

diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 57ff2fe9d87..503636c9f71 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27961,11 +27961,6 @@ 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; /** * Has More * @default false @@ -27976,11 +27971,6 @@ 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