diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index 5cdf165a822..809f78e378c 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -1,11 +1,12 @@ import asyncio from collections.abc import AsyncIterator, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from itertools import groupby from types import MappingProxyType from typing import Final, Protocol -from pydantic import StrictStr, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, StrictStr, TypeAdapter, ValidationError from typing_extensions import assert_never from litellm import constants @@ -18,7 +19,7 @@ from litellm.repositories.daily_activity_sql import ( adjust_dates_for_timezone, build_aggregated_sql, build_cache_leakage_keys_sql, - build_customer_filtered_keys_sql, + build_customer_daily_rows_sql, build_customer_team_rollup_sql, build_entity_rollup_sql, build_export_sql, @@ -106,7 +107,6 @@ class DailyActivityDatabase(Protocol): _GROUPING_ADAPTER: Final = TypeAdapter(tuple[GroupingSetsRow, ...]) _CUSTOMER_TEAM_ADAPTER: Final = TypeAdapter(tuple[CustomerTeamRollupRow, ...]) -_FILTERED_KEYS_ADAPTER: Final = TypeAdapter(tuple[dict[str, StrictStr], ...]) _ENTITY_ADAPTER: Final = TypeAdapter(tuple[EntityRollupRow, ...]) _GUARDRAIL_SPEND_ADAPTER: Final = TypeAdapter(tuple[GuardrailSpendRow, ...]) _KEY_SPEND_ADAPTER: Final = TypeAdapter(tuple[KeySpendRow, ...]) @@ -116,6 +116,44 @@ _EXPORT_ADAPTER: Final = TypeAdapter(tuple[ExportRow, ...]) _METADATA_TAGS_ADAPTER: Final = TypeAdapter(list[StrictStr]) +class _CustomerDailyRow(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + date: str + api_key: str + end_user_id: str | None + prompt_tokens: int + completion_tokens: int + cache_read_input_tokens: int + cache_creation_input_tokens: int + compression_saved_tokens: int + compression_savings_spend: float + prompt_caching_savings_spend: float + gateway_injected_caching_savings_spend: float + autorouter_savings_spend: float + spend: float + api_requests: int + successful_requests: int + failed_requests: int + total_response_time_ms: int + timed_requests: int + model: str | None = None + model_group: str | None = None + custom_llm_provider: str | None = None + mcp_namespaced_tool_name: str | None = None + endpoint: str | None = None + + +@dataclass(frozen=True, slots=True) +class _CustomerDailyRowsResult: + total_count: int + rows: tuple[_CustomerDailyRow, ...] + + +_CUSTOMER_DAILY_ROWS_ADAPTER: Final = TypeAdapter(tuple[_CustomerDailyRowsResult]) + + def _metadata_tags(value: object) -> tuple[str, ...]: stable_value: Final = value if not isinstance(value, list): @@ -329,6 +367,11 @@ class DailyActivityRepository: ) async def daily_rows(self, scope: DailyActivityScope, *, page: int, page_size: int) -> DailyRowsPage: + if scope.table is DailyActivityTable.CUSTOMER and scope.team_ids is not None: + result: Final = _CUSTOMER_DAILY_ROWS_ADAPTER.validate_python( + await self._query(build_customer_daily_rows_sql(scope, page=page, page_size=page_size)) + )[0] + return DailyRowsPage(total_count=result.total_count, rows=result.rows) table: Final = _daily_rows_table(self._prisma_client, scope.table) adjusted_start, adjusted_end = adjust_dates_for_timezone( scope.start_date, @@ -346,16 +389,6 @@ class DailyActivityRepository: if scope.exclude_entity_ids else {} ) - filtered_keys: Final = ( - tuple( - row["api_key"] - for row in _FILTERED_KEYS_ADAPTER.validate_python( - await self._query(build_customer_filtered_keys_sql(scope)) - ) - ) - if scope.table is DailyActivityTable.CUSTOMER and scope.team_ids is not None - else scope.api_keys - ) conditions: Final = { "date": {"gte": adjusted_start, "lte": adjusted_end}, **({scope.entity_id_field: {"in": list(scope.entity_ids)}} if scope.entity_ids is not None else {}), @@ -367,7 +400,7 @@ class DailyActivityRepository: if scope.model else {} ), - **({"api_key": {"in": list(filtered_keys)}} if filtered_keys is not None else {}), + **({"api_key": {"in": list(scope.api_keys)}} if scope.api_keys is not None else {}), } count, rows = await asyncio.gather( table.count(where=conditions), diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index e198385a622..5d1f85658ea 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -189,11 +189,18 @@ def build_customer_team_rollup_sql(scope: DailyActivityScope, *, row_ids: tuple[ ) -def build_customer_filtered_keys_sql(scope: DailyActivityScope) -> SqlQuery: +def build_customer_daily_rows_sql(scope: DailyActivityScope, *, page: int, page_size: int) -> SqlQuery: where_clause, params = build_where_clause(scope) return SqlQuery( - sql=f'SELECT DISTINCT api_key FROM "{PRISMA_TO_PG_TABLE[scope.table]}" WHERE {where_clause}', - params=params, + sql=f'''WITH filtered AS MATERIALIZED ( + SELECT * FROM "{PRISMA_TO_PG_TABLE[scope.table]}" WHERE {where_clause} + ) + SELECT (SELECT COUNT(*) FROM filtered)::bigint AS total_count, + COALESCE((SELECT jsonb_agg(page_rows) FROM ( + SELECT * FROM filtered ORDER BY date DESC, id ASC + OFFSET ${len(params) + 1} LIMIT ${len(params) + 2} + ) page_rows), '[]'::jsonb) AS rows''', + params=(*params, (page - 1) * page_size, page_size), ) diff --git a/tests/integration/spend/test_customer_spend_by_team.py b/tests/integration/spend/test_customer_spend_by_team.py index 0eeca325b49..4c78af44e25 100644 --- a/tests/integration/spend/test_customer_spend_by_team.py +++ b/tests/integration/spend/test_customer_spend_by_team.py @@ -1,6 +1,7 @@ import uuid from hashlib import sha256 from datetime import datetime, timedelta, timezone +from time import perf_counter from typing import Final import pytest @@ -119,3 +120,66 @@ def test_customer_spend_by_current_team_and_team_filter(gateway: Gateway, route: legacy_teams: Final = legacy_body.results[0].breakdown.teams or {} assert set(legacy_teams) == {team} assert legacy_teams[team].metrics.spend == legacy_body.metadata.total_spend + + +@pytest.mark.parametrize("route", ("/customer/daily/activity", "/customer/daily/activity/aggregated")) +def test_large_customer_team_reports_keep_filters_and_pagination(gateway: Gateway, route: str) -> None: + with gateway.scenario() as scenario: + team: Final = scenario.team(team_alias="large-customer-team") + customer: Final = f"large-customer-{uuid.uuid4().hex}" + scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_VerificationToken" WHERE team_id=%s', (team,)) + scenario.cleanups.callback( + write_rows, 'DELETE FROM "LiteLLM_DailyEndUserSpend" WHERE end_user_id=%s', (customer,) + ) + write_rows( + 'INSERT INTO "LiteLLM_VerificationToken" (token,models,team_id) ' + "SELECT encode(sha256(convert_to(%s||i,'UTF8')),'hex'),ARRAY[]::text[],%s " + "FROM generate_series(1,33000) i", + (customer, team), + ) + write_rows( + 'INSERT INTO "LiteLLM_DailyEndUserSpend" ' + "(id,end_user_id,date,api_key,model,spend,api_requests,updated_at) " + "SELECT %s||lpad(i::text,5,'0'),%s,'2026-02-03'," + "encode(sha256(convert_to(%s||i,'UTF8')),'hex')," + "CASE WHEN i%%2=1 THEN 'large-model-a' ELSE 'large-model-b' END," + "CASE WHEN i%%2=1 THEN 0.25 ELSE 0.5 END,1,now() FROM generate_series(1,33000) i", + (customer, customer, customer), + ) + write_rows('ANALYZE "LiteLLM_VerificationToken"', ()) + write_rows('ANALYZE "LiteLLM_DeletedVerificationToken"', ()) + write_rows('ANALYZE "LiteLLM_DailyEndUserSpend"', ()) + key: Final = sha256(f"{customer}1".encode()).hexdigest() + base: Final = ( + f"{route}?end_user_ids={customer}&team_ids={team}&start_date=2026-02-03&end_date=2026-02-03&page_size=10" + ) + aggregated: Final = route.endswith("aggregated") + for suffix, spend, pages in ( + ("", 12375.0 if aggregated else 3.75, 3300), + ("&model=large-model-a", 4125.0 if aggregated else 2.5, 1650), + (f"&api_key={key}", 0.25, 1), + (f"&api_key={key}&model=large-model-a", 0.25, 1), + (f"&api_key={key}&model=large-model-b", 0.0, 0), + ): + started: Final = perf_counter() + response: Final = gateway.request("GET", base + suffix) + print(f"large-team {route} {suffix!r}: {response.status_code}, {perf_counter() - started:.3f}s") + assert response.status_code == 200, response.text + body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(response.content) + assert body.metadata.total_spend == spend + assert ( + sum( + sum(bucket.metrics.spend for bucket in (day.breakdown.teams or {}).values()) for day in body.results + ) + == spend + ) + if not aggregated: + assert body.metadata.total_pages == pages + assert body.metadata.has_more == (pages > 1) + if not aggregated: + second: Final = gateway.request("GET", base + "&page=2") + assert second.status_code == 200, second.text + second_body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(second.content) + assert second_body.metadata.total_spend == 3.75 + expected_keys: Final = {sha256(f"{customer}{i}".encode()).hexdigest() for i in range(11, 21)} + assert set(second_body.results[0].breakdown.api_keys) == expected_keys diff --git a/tests/unit/repositories/test_daily_activity_repository.py b/tests/unit/repositories/test_daily_activity_repository.py index d93a75476a2..b68bb948b7b 100644 --- a/tests/unit/repositories/test_daily_activity_repository.py +++ b/tests/unit/repositories/test_daily_activity_repository.py @@ -1,5 +1,5 @@ from collections.abc import Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timezone from typing import Final @@ -662,6 +662,36 @@ async def test_daily_rows_selects_the_table_and_applies_filters_and_pagination( assert sum(len(daily_table.find_many_calls) for daily_table in tables.values()) == 1 +@pytest.mark.asyncio +@pytest.mark.parametrize("empty", (False, True)) +async def test_customer_team_page_returns_loaded_metrics_and_retains_count_on_empty_pages(empty: bool) -> None: + stored: Final = { + **_key_spend_row("key-a"), + "id": "row-a", + "end_user_id": "customer-a", + "date": "2026-01-01", + "compression_saved_tokens": 2, + "compression_savings_spend": 0.5, + "prompt_caching_savings_spend": 0.25, + "gateway_injected_caching_savings_spend": 0.125, + "autorouter_savings_spend": 0.0, + "total_response_time_ms": 10, + "timed_requests": 1, + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + } + database: Final = _FakeDatabase((({"total_count": 33000, "rows": [] if empty else [stored]},),)) + repository, _ = _repository(database) + scope: Final = replace(_scope(table=DailyActivityTable.CUSTOMER), team_ids=("team-a",)) + + result: Final = await repository.daily_rows(scope, page=3301 if empty else 1, page_size=10) + + assert result.total_count == 33000 + assert tuple((row.id, row.api_key, row.spend, row.prompt_tokens) for row in result.rows) == ( + () if empty else (("row-a", "key-a", 1.0, 10),) + ) + + @pytest.mark.asyncio async def test_daily_rows_exclusion_without_entity_filter_keeps_null_entity_rows() -> None: database = _FakeDatabase()