From bf92748b8ad94c767b76c58c8de66c2ed8baa7cf Mon Sep 17 00:00:00 2001 From: shivam Date: Wed, 14 Oct 2026 20:43:26 +0000 Subject: [PATCH 1/2] feat(usage): add a day-by-day by key and model usage export Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 1 + .../repositories/daily_activity_repository.py | 27 +++-- litellm/repositories/daily_activity_sql.py | 90 ++++++++-------- litellm/types/repositories/daily_activity.py | 1 + .../spend/test_daily_activity_repository.py | 100 +++++++++++++++++- .../spend/test_daily_activity_routes.py | 14 +++ .../test_daily_activity_repository.py | 49 ++++++++- .../repositories/test_daily_activity_sql.py | 68 +++++++++--- .../ExportTypeSelector.test.tsx | 14 +++ .../EntityUsageExport/ExportTypeSelector.tsx | 5 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 11 files changed, 294 insertions(+), 77 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index effb30fe8cb..8c100883cea 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3705,6 +3705,7 @@ "daily", "daily_with_keys", "daily_with_models", + "daily_with_keys_and_models", "daily_with_users" ], "title": "ExportType", diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index 10c603fe56f..152dbb7d30f 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -138,16 +138,23 @@ def _daily_rows_table( def _next_export_cursor(batch: tuple[ExportRow, ...], export_type: ExportType) -> ExportCursor: last: Final = batch[-1] - cursor_key: Final = ( - last.api_key - if export_type is ExportType.DAILY_WITH_KEYS - else last.model - if export_type is ExportType.DAILY_WITH_MODELS - else last.user_id - if export_type is ExportType.DAILY_WITH_USERS - else "" - ) - return ExportCursor(date=last.date, entity_id=last.entity_id, group_key=cursor_key or "") + match export_type: + case ExportType.DAILY: + return ExportCursor(date=last.date, entity_id=last.entity_id, group_key=()) + case ExportType.DAILY_WITH_KEYS: + return ExportCursor(date=last.date, entity_id=last.entity_id, group_key=(last.api_key or "",)) + case ExportType.DAILY_WITH_MODELS: + return ExportCursor(date=last.date, entity_id=last.entity_id, group_key=(last.model or "",)) + case ExportType.DAILY_WITH_KEYS_AND_MODELS: + return ExportCursor( + date=last.date, + entity_id=last.entity_id, + group_key=(last.api_key or "", last.model or ""), + ) + case ExportType.DAILY_WITH_USERS: + return ExportCursor(date=last.date, entity_id=last.entity_id, group_key=(last.user_id or "",)) + case _: + assert_never(export_type) class DailyActivityRepository: diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index 3d27e7d937e..db53cdef7f9 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -13,6 +13,21 @@ from litellm.types.repositories.daily_activity import DailyActivityScope, DailyA _API_KEY_ROLLED_UP_BIT: Final = 32 _MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" +_EXPORT_KEY_USER_FIELDS: Final = ( + "MAX(COALESCE(vt.key_alias, dvt.key_alias)) AS key_alias, " + "MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email" +) +_EXPORT_KEY_METADATA_JOINS: Final[tuple[str, ...]] = ( + 'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key', + """LEFT JOIN LATERAL ( + SELECT key_alias, user_id + FROM "LiteLLM_DeletedVerificationToken" + WHERE token = scoped.api_key + ORDER BY deleted_at DESC + LIMIT 1 + ) dvt ON vt.token IS NULL""", + 'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = COALESCE(vt.user_id, dvt.user_id)', +) @dataclass(frozen=True, slots=True) @@ -25,7 +40,7 @@ class SqlQuery: class ExportCursor: date: str entity_id: str - group_key: str + group_key: tuple[str, ...] PRISMA_TO_PG_TABLE: Final[Mapping[DailyActivityTable, str]] = MappingProxyType( @@ -423,10 +438,10 @@ def build_export_sql( ) -> SqlQuery: _bounded_limit(batch_size) where_clause, where_params = build_where_clause(scope) - group_key, output_key, user_fields, type_joins = _export_grouping(export_type) + group_keys, output_key, user_fields, type_joins = _export_grouping(export_type) grouping_keys: Final = ( f"scoped.date, COALESCE(scoped.\"{scope.entity_id_field}\", '')", - *((group_key,) if export_type is not ExportType.DAILY else ()), + *group_keys, ) entity_joins: Final = ( ('LEFT JOIN "LiteLLM_TeamTable" tt ON tt.team_id = scoped.team_id',) @@ -445,14 +460,19 @@ def build_export_sql( ) parameter_indexes: Final = count(len(where_params) + 1) sentinel_param: Final = next(parameter_indexes) if export_type is not ExportType.DAILY else None - cursor_indexes: Final = tuple(islice(parameter_indexes, 3)) if after is not None else () + cursor_indexes: Final = tuple(islice(parameter_indexes, 2 + len(group_keys))) if after is not None else () limit_param: Final = next(parameter_indexes) cursor_clause, cursor_params = _export_cursor_clause( - scope, after=after, cursor_indexes=cursor_indexes, group_key=group_key + scope, after=after, cursor_indexes=cursor_indexes, group_keys=group_keys ) sentinel_clause: Final = f" AND api_key <> ${sentinel_param}" if sentinel_param is not None else "" table: Final = PRISMA_TO_PG_TABLE[scope.table] flat_cost: Final = _ptu_flat_cost_select(scope.table, aggregate=False) + model_expression: Final = ( + "NULLIF(COALESCE(scoped.model, ''), '')" + if export_type in (ExportType.DAILY_WITH_MODELS, ExportType.DAILY_WITH_KEYS_AND_MODELS) + else "NULL::text" + ) sql: Final = f""" WITH scoped AS ( SELECT * FROM "{table}" @@ -464,7 +484,7 @@ def build_export_sql( {alias_expression} AS entity_alias, {output_key} AS api_key, {user_fields}, - {"NULLIF(COALESCE(scoped.model, ''), '')" if export_type is ExportType.DAILY_WITH_MODELS else "NULL::text"} AS model, + {model_expression} AS model, COALESCE(SUM(scoped.spend), 0)::float AS spend, {flat_cost} AS flat_cost, COALESCE(SUM(scoped.prompt_tokens), 0)::bigint AS prompt_tokens, @@ -492,56 +512,40 @@ def build_export_sql( ) -def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str, ...]]: +def _export_grouping(export_type: ExportType) -> tuple[tuple[str, ...], str, str, tuple[str, ...]]: if export_type is ExportType.DAILY: return ( - "''", + (), "NULL::text", "NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email", (), ) - if export_type is ExportType.DAILY_WITH_KEYS: + if export_type in (ExportType.DAILY_WITH_KEYS, ExportType.DAILY_WITH_KEYS_AND_MODELS): + group_keys: Final = ( + ("scoped.api_key",) + if export_type is ExportType.DAILY_WITH_KEYS + else ("scoped.api_key", "COALESCE(scoped.model, '')") + ) return ( - "scoped.api_key", + group_keys, "NULLIF(scoped.api_key, '')", - "MAX(COALESCE(vt.key_alias, dvt.key_alias)) AS key_alias, " - "MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email", - ( - 'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key', - """LEFT JOIN LATERAL ( - SELECT key_alias, user_id - FROM "LiteLLM_DeletedVerificationToken" - WHERE token = scoped.api_key - ORDER BY deleted_at DESC - LIMIT 1 - ) dvt ON vt.token IS NULL""", - 'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = COALESCE(vt.user_id, dvt.user_id)', - ), + _EXPORT_KEY_USER_FIELDS, + _EXPORT_KEY_METADATA_JOINS, ) if export_type is ExportType.DAILY_WITH_MODELS: return ( - "COALESCE(scoped.model, '')", + ("COALESCE(scoped.model, '')",), "NULL::text", "NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email", (), ) if export_type is ExportType.DAILY_WITH_USERS: return ( - "COALESCE(vt.user_id, dvt.user_id, '')", + ("COALESCE(vt.user_id, dvt.user_id, '')",), "NULL::text", "NULL::text AS key_alias, MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, " "MAX(u.user_email) AS user_email", - ( - 'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key', - """LEFT JOIN LATERAL ( - SELECT key_alias, user_id - FROM "LiteLLM_DeletedVerificationToken" - WHERE token = scoped.api_key - ORDER BY deleted_at DESC - LIMIT 1 - ) dvt ON vt.token IS NULL""", - 'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = COALESCE(vt.user_id, dvt.user_id)', - ), + _EXPORT_KEY_METADATA_JOINS, ) assert_never(export_type) raise AssertionError("unreachable") @@ -552,13 +556,15 @@ def _export_cursor_clause( *, after: ExportCursor | None, cursor_indexes: tuple[int, ...], - group_key: str, + group_keys: tuple[str, ...], ) -> tuple[str, tuple[object, ...]]: if after is None: return "", () - first_cursor_index: Final = cursor_indexes[0] - clause: Final = ( - f""" AND (scoped.date, COALESCE(scoped."{scope.entity_id_field}", ''), {group_key}) """ - f"> (${first_cursor_index}, ${cursor_indexes[1]}, ${cursor_indexes[2]})" + cursor_expressions: Final = ( + "scoped.date", + f"""COALESCE(scoped."{scope.entity_id_field}", '')""", + *group_keys, ) - return clause, (after.date, after.entity_id, after.group_key) + cursor_placeholders: Final = tuple(f"${index}" for index in cursor_indexes) + clause: Final = f" AND ({', '.join(cursor_expressions)}) > ({', '.join(cursor_placeholders)})" + return clause, (after.date, after.entity_id, *after.group_key) diff --git a/litellm/types/repositories/daily_activity.py b/litellm/types/repositories/daily_activity.py index 4111b83a891..d93e8963cc6 100644 --- a/litellm/types/repositories/daily_activity.py +++ b/litellm/types/repositories/daily_activity.py @@ -104,6 +104,7 @@ class ExportType(str, Enum): DAILY = "daily" DAILY_WITH_KEYS = "daily_with_keys" DAILY_WITH_MODELS = "daily_with_models" + DAILY_WITH_KEYS_AND_MODELS = "daily_with_keys_and_models" DAILY_WITH_USERS = "daily_with_users" diff --git a/tests/integration/spend/test_daily_activity_repository.py b/tests/integration/spend/test_daily_activity_repository.py index 1ed85d4a3bc..d007f7a1724 100644 --- a/tests/integration/spend/test_daily_activity_repository.py +++ b/tests/integration/spend/test_daily_activity_repository.py @@ -262,6 +262,102 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch: assert metadata["key-target"].tags == ("archived",) +@pytest.mark.asyncio +async def test_key_and_model_export_keeps_group_metrics_and_pages_by_both_keys( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async with _daily_activity_database() as database: + await database.execute_raw( + """ + INSERT INTO "LiteLLM_DailyTeamSpend" ( + id, team_id, date, api_key, model, model_group, custom_llm_provider, + mcp_namespaced_tool_name, endpoint, prompt_tokens, completion_tokens, + cache_read_input_tokens, cache_creation_input_tokens, spend, api_requests, + successful_requests, failed_requests, ptu_flat_cost, updated_at + ) VALUES + ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19::timestamp), + ($20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38::timestamp), + ($39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57::timestamp) + """, + "key-model-a", + "team-composite", + "2026-06-01", + "key-1", + "model-a", + "", + "provider-a", + None, + "/v1/chat/completions", + 10, + 1, + 0, + 0, + 2.0, + 2, + 2, + 0, + 0.0, + "2026-06-01 12:00:00", + "key-model-b", + "team-composite", + "2026-06-01", + "key-1", + "model-b", + "", + "provider-a", + None, + "/v1/chat/completions", + 20, + 2, + 0, + 0, + 3.0, + 3, + 3, + 0, + 0.0, + "2026-06-01 12:00:00", + "key-model-c", + "team-composite", + "2026-06-01", + "key-2", + "model-a", + "", + "provider-a", + None, + "/v1/chat/completions", + 30, + 3, + 0, + 0, + 5.0, + 5, + 5, + 0, + 0.0, + "2026-06-01 12:00:00", + ) + repository: Final = _repository(database) + team_scope: Final = _scope(DailyActivityTable.TEAM, "team_id", "team-composite") + monkeypatch.setattr(constants, "USAGE_EXPORT_BATCH_SIZE", 1) + + daily_rows: Final = tuple( + [row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY)] + ) + key_model_rows: Final = tuple( + [row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_KEYS_AND_MODELS)] + ) + + assert len(daily_rows) == 1 + assert tuple((row.api_key, row.model, row.spend, row.api_requests) for row in key_model_rows) == ( + ("key-1", "model-a", 2.0, 2), + ("key-1", "model-b", 3.0, 3), + ("key-2", "model-a", 5.0, 5), + ) + assert sum(row.spend for row in key_model_rows) == daily_rows[0].spend == 10.0 + assert sum(row.api_requests for row in key_model_rows) == daily_rows[0].api_requests == 10 + + @pytest.mark.asyncio @pytest.mark.parametrize( ("search", "api_keys", "expected_keys"), @@ -684,9 +780,7 @@ async def test_user_page_ranks_users_and_folds_null_ids() -> None: repository: Final = _repository(database) full_page: Final = await repository.user_page(scope, offset=0, limit=50) empty_page: Final = await repository.user_page(scope, offset=full_page.total_users + 1, limit=50) - scoped_page: Final = await repository.user_page( - replace(scope, entity_ids=("user-alpha",)), offset=0, limit=50 - ) + scoped_page: Final = await repository.user_page(replace(scope, entity_ids=("user-alpha",)), offset=0, limit=50) assert tuple((row.user_id, row.spend) for row in full_page.rows) == ( (None, 45.0), diff --git a/tests/integration/spend/test_daily_activity_routes.py b/tests/integration/spend/test_daily_activity_routes.py index 787a9ca7479..d149ac33580 100644 --- a/tests/integration/spend/test_daily_activity_routes.py +++ b/tests/integration/spend/test_daily_activity_routes.py @@ -210,6 +210,20 @@ def _assert_entity_activity_routes( export_rows: Final = tuple(csv.reader(io.StringIO(exported.text))) assert len(export_rows) == len(persisted_rows) + 1, exported.text + key_model_export: Final = gateway.request( + "GET", + f"{prefix}/daily/activity/export", + params={**params, "export_type": "daily_with_keys_and_models"}, + ) + assert key_model_export.status_code == 200, key_model_export.text + key_model_rows: Final = tuple(csv.reader(io.StringIO(key_model_export.text))) + api_key_index: Final = key_model_rows[0].index("api_key") + model_index: Final = key_model_rows[0].index("model") + key_model_pairs: Final = tuple((row[api_key_index], row[model_index]) for row in key_model_rows[1:]) + assert key_model_pairs + assert all(api_key and model_name for api_key, model_name in key_model_pairs), key_model_export.text + assert len(key_model_pairs) == len(frozenset(key_model_pairs)), key_model_export.text + @pytest.mark.timeout(90) def test_daily_activity_routes_cover_all_entities_and_bounded_key_search(gateway: Gateway, tmp_path: Path) -> None: diff --git a/tests/unit/repositories/test_daily_activity_repository.py b/tests/unit/repositories/test_daily_activity_repository.py index d93a75476a2..d7ed3f52311 100644 --- a/tests/unit/repositories/test_daily_activity_repository.py +++ b/tests/unit/repositories/test_daily_activity_repository.py @@ -104,7 +104,13 @@ def _user_spend_row(user_id: str | None, **overrides: object) -> dict[str, objec } -def _export_row(api_key: str | None) -> dict[str, object]: +def _export_row( + api_key: str | None, + *, + model: str | None = None, + spend: float = 1.0, + api_requests: int = 1, +) -> dict[str, object]: return { "date": "2026-01-01", "entity_id": "user-1", @@ -113,12 +119,12 @@ def _export_row(api_key: str | None) -> dict[str, object]: "key_alias": None, "user_id": None, "user_email": None, - "model": None, - "spend": 1.0, + "model": model, + "spend": spend, "flat_cost": 0.0, "prompt_tokens": 10, "completion_tokens": 2, - "api_requests": 1, + "api_requests": api_requests, "successful_requests": 1, "failed_requests": 0, "cache_read_input_tokens": 3, @@ -704,8 +710,41 @@ async def test_export_is_lazy_and_uses_the_last_row_as_the_next_cursor(monkeypat build_export_sql( _scope(), export_type=ExportType.DAILY_WITH_KEYS, - after=ExportCursor("2026-01-01", "user-1", "key-2"), + after=ExportCursor("2026-01-01", "user-1", ("key-2",)), batch_size=2, ).params == database.query_calls[1][1] ) + + +@pytest.mark.asyncio +async def test_export_cursor_tracks_model_when_a_key_has_multiple_models(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(constants, "USAGE_EXPORT_BATCH_SIZE", 1) + database = _FakeDatabase( + ( + (_export_row("key-1", model="model-a"),), + (_export_row("key-1", model="model-b"),), + (_export_row("key-2", model="model-a"),), + ) + ) + repository, _ = _repository(database) + rows = repository.export_rows(_scope(), export_type=ExportType.DAILY_WITH_KEYS_AND_MODELS) + + assert tuple([(row.api_key, row.model) async for row in rows]) == ( + ("key-1", "model-a"), + ("key-1", "model-b"), + ("key-2", "model-a"), + ) + assert len(database.query_calls) == 4 + assert database.query_calls[1][1][-5:] == ("2026-01-01", "user-1", "key-1", "model-a", 1) + assert database.query_calls[2][1][-5:] == ("2026-01-01", "user-1", "key-1", "model-b", 1) + assert database.query_calls[3][1][-5:] == ("2026-01-01", "user-1", "key-2", "model-a", 1) + assert ( + build_export_sql( + _scope(), + export_type=ExportType.DAILY_WITH_KEYS_AND_MODELS, + after=ExportCursor("2026-01-01", "user-1", ("key-1", "model-a")), + batch_size=1, + ).params + == database.query_calls[1][1] + ) diff --git a/tests/unit/repositories/test_daily_activity_sql.py b/tests/unit/repositories/test_daily_activity_sql.py index c57db08e976..4edba9dfa31 100644 --- a/tests/unit/repositories/test_daily_activity_sql.py +++ b/tests/unit/repositories/test_daily_activity_sql.py @@ -398,52 +398,88 @@ def test_query_builders_reject_nonpositive_limits(builder) -> None: @pytest.mark.parametrize( - ("export_type", "group_key", "key_filter", "joins"), + ("export_type", "group_keys", "group_key_values", "key_filter", "joins"), [ - (ExportType.DAILY, "''", "", ""), - (ExportType.DAILY_WITH_KEYS, "scoped.api_key", "api_key <> $3", 'LEFT JOIN "LiteLLM_VerificationToken"'), - (ExportType.DAILY_WITH_MODELS, "COALESCE(scoped.model, '')", "api_key <> $3", ""), + (ExportType.DAILY, (), (), "", ""), + ( + ExportType.DAILY_WITH_KEYS, + ("scoped.api_key",), + ("group-3",), + "api_key <> $3", + 'LEFT JOIN "LiteLLM_VerificationToken"', + ), + ( + ExportType.DAILY_WITH_MODELS, + ("COALESCE(scoped.model, '')",), + ("group-3",), + "api_key <> $3", + "", + ), + ( + ExportType.DAILY_WITH_KEYS_AND_MODELS, + ("scoped.api_key", "COALESCE(scoped.model, '')"), + ("group-3", "group-4"), + "api_key <> $3", + 'LEFT JOIN "LiteLLM_VerificationToken"', + ), ( ExportType.DAILY_WITH_USERS, - "COALESCE(vt.user_id, dvt.user_id, '')", + ("COALESCE(vt.user_id, dvt.user_id, '')",), + ("group-3",), "api_key <> $3", 'LEFT JOIN "LiteLLM_VerificationToken"', ), ], ) def test_export_groups_by_requested_key_and_binds_cursor_after_scope( - export_type: ExportType, group_key: str, key_filter: str, joins: str + export_type: ExportType, + group_keys: tuple[str, ...], + group_key_values: tuple[str, ...], + key_filter: str, + joins: str, ) -> None: query = build_export_sql( _scope(entity_ids=None), export_type=export_type, - after=ExportCursor(date="2026-01-12", entity_id="user-2", group_key="group-3"), + after=ExportCursor(date="2026-01-12", entity_id="user-2", group_key=group_key_values), batch_size=2, ) - assert group_key in query.sql + assert all(group_key in query.sql for group_key in group_keys) assert key_filter in query.sql assert joins in query.sql - assert "(scoped.date, COALESCE(scoped.\"user_id\", '')," in query.sql - order_keys: Final = ( - "scoped.date, COALESCE(scoped.\"user_id\", '')", - *((group_key,) if export_type is not ExportType.DAILY else ()), + cursor_columns: Final = ( + "scoped.date", + """COALESCE(scoped."user_id", '')""", + *group_keys, ) - assert f"ORDER BY {', '.join(order_keys)}" in query.sql - expected_limit_index: Final = "$6" if export_type is ExportType.DAILY else "$7" + assert f"GROUP BY {', '.join(cursor_columns)}" in query.sql + assert f"ORDER BY {', '.join(cursor_columns)}" in query.sql + cursor_start_index: Final = 3 + int(export_type is not ExportType.DAILY) + cursor_placeholders: Final = tuple( + f"${index}" for index in range(cursor_start_index, cursor_start_index + len(cursor_columns)) + ) + assert f"AND ({', '.join(cursor_columns)}) > ({', '.join(cursor_placeholders)})" in query.sql + expected_limit_index: Final = f"${5 + int(export_type is not ExportType.DAILY) + len(group_keys)}" assert f"LIMIT {expected_limit_index}" in query.sql + assert ("NULLIF(COALESCE(scoped.model, ''), '') AS model" in query.sql) is ( + export_type in (ExportType.DAILY_WITH_MODELS, ExportType.DAILY_WITH_KEYS_AND_MODELS) + ) assert query.params == ( "2026-01-01", "2026-01-31", *((PTU_SENTINEL_API_KEY,) if export_type is not ExportType.DAILY else ()), "2026-01-12", "user-2", - "group-3", + *group_key_values, 2, ) -@pytest.mark.parametrize("export_type", [ExportType.DAILY_WITH_KEYS, ExportType.DAILY_WITH_USERS]) +@pytest.mark.parametrize( + "export_type", + [ExportType.DAILY_WITH_KEYS, ExportType.DAILY_WITH_KEYS_AND_MODELS, ExportType.DAILY_WITH_USERS], +) def test_export_uses_latest_deleted_key_metadata(export_type: ExportType) -> None: query = build_export_sql(_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2) diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx index 997c048048f..19471a290fa 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx @@ -14,6 +14,7 @@ describe("ExportTypeSelector", () => { expect(screen.getByText(/Day-by-day breakdown by team$/)).toBeInTheDocument(); expect(screen.getByText(/Day-by-day breakdown by team and key/)).toBeInTheDocument(); expect(screen.getByText(/Day-by-day by team and model/)).toBeInTheDocument(); + expect(screen.getByText(/Day-by-day by team, key and model/)).toBeInTheDocument(); }); it("should display the correct entity type for different entities", () => { @@ -29,6 +30,18 @@ describe("ExportTypeSelector", () => { expect(onChange).toHaveBeenCalledWith("daily_with_keys"); }); + it("should render and select the combined key and model option for users", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders(); + + const option = screen.getByRole("radio", { name: /Day-by-day by user, key and model/i }); + expect(screen.getByText("Daily metrics for each user, split by API key and model")).toBeInTheDocument(); + await user.click(option); + + expect(onChange).toHaveBeenCalledWith("daily_with_keys_and_models"); + }); + it("should have the correct radio checked", () => { renderWithProviders(); expect(screen.getByRole("radio", { name: /Day-by-day by team and model/i })).toBeChecked(); @@ -55,6 +68,7 @@ describe("ExportTypeSelector", () => { ).toBeInTheDocument(); expect(screen.getByRole("radio", { name: /Day-by-day breakdown by user and key/i })).toBeInTheDocument(); expect(screen.getByRole("radio", { name: /Day-by-day by user and model/i })).toBeInTheDocument(); + expect(screen.getByRole("radio", { name: /Day-by-day by user, key and model/i })).toBeInTheDocument(); }); it("should offer the per-user scope for tags", () => { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx index 3162f00c056..d01c4cdcec1 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx @@ -25,6 +25,11 @@ const ExportTypeSelector: React.FC = ({ value, onChange title: `Day-by-day by ${entityType} and model`, description: "Daily metrics split by model", }, + { + value: "daily_with_keys_and_models", + title: `Day-by-day by ${entityType}, key and model`, + description: `Daily metrics for each ${entityType}, split by API key and model`, + }, { value: "daily_with_users", title: `Day-by-day breakdown by ${entityType} and user`, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 29d659ec1d9..c00039a7a8b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32015,7 +32015,7 @@ export interface components { * ExportType * @enum {string} */ - ExportType: "daily" | "daily_with_keys" | "daily_with_models" | "daily_with_users"; + ExportType: "daily" | "daily_with_keys" | "daily_with_models" | "daily_with_keys_and_models" | "daily_with_users"; /** * FacetListResponse * @description The distinct values one column takes over a filtered query. `data` holds bare values, not entity rows. From b6d941170b69ef8e6c21730657189ecd5d45b4e4 Mon Sep 17 00:00:00 2001 From: shivam Date: Wed, 14 Oct 2026 21:10:05 +0000 Subject: [PATCH 2/2] test(usage): cover daily activity export cursor variants Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_daily_activity_repository.py | 58 ++++++++++++++++++- 1 file changed, 56 insertions(+), 2 deletions(-) diff --git a/tests/unit/repositories/test_daily_activity_repository.py b/tests/unit/repositories/test_daily_activity_repository.py index d7ed3f52311..f523cd88e3d 100644 --- a/tests/unit/repositories/test_daily_activity_repository.py +++ b/tests/unit/repositories/test_daily_activity_repository.py @@ -108,16 +108,18 @@ def _export_row( api_key: str | None, *, model: str | None = None, + user_id: str | None = None, + entity_id: str = "user-1", spend: float = 1.0, api_requests: int = 1, ) -> dict[str, object]: return { "date": "2026-01-01", - "entity_id": "user-1", + "entity_id": entity_id, "entity_alias": None, "api_key": api_key, "key_alias": None, - "user_id": None, + "user_id": user_id, "user_email": None, "model": model, "spend": spend, @@ -717,6 +719,58 @@ async def test_export_is_lazy_and_uses_the_last_row_as_the_next_cursor(monkeypat ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("export_type", "scope", "rows", "expected_cursor_params"), + [ + ( + ExportType.DAILY, + _scope(), + (_export_row("key-1"),), + (("2026-01-01", "user-1", 1),), + ), + ( + ExportType.DAILY_WITH_MODELS, + _scope(), + (_export_row("key-1", model="model-a"), _export_row("key-1", model="model-b")), + ( + ("2026-01-01", "user-1", "model-a", 1), + ("2026-01-01", "user-1", "model-b", 1), + ), + ), + ( + ExportType.DAILY_WITH_USERS, + _scope(table=DailyActivityTable.TEAM), + ( + _export_row("key-1", user_id="user-a", entity_id="team-1"), + _export_row("key-2", user_id="user-b", entity_id="team-1"), + ), + ( + ("2026-01-01", "team-1", "user-a", 1), + ("2026-01-01", "team-1", "user-b", 1), + ), + ), + ], +) +async def test_export_pagination_preserves_existing_grouping_cursors( + monkeypatch: pytest.MonkeyPatch, + export_type: ExportType, + scope: DailyActivityScope, + rows: tuple[dict[str, object], ...], + expected_cursor_params: tuple[tuple[object, ...], ...], +) -> None: + monkeypatch.setattr(constants, "USAGE_EXPORT_BATCH_SIZE", 1) + database = _FakeDatabase(tuple((row,) for row in rows)) + repository, _ = _repository(database) + + exported_rows: Final = tuple([row async for row in repository.export_rows(scope, export_type=export_type)]) + + assert len(exported_rows) == len(rows) + cursor_param_count: Final = len(expected_cursor_params[0]) + next_page_params: Final = tuple(query[1][-cursor_param_count:] for query in database.query_calls[1:]) + assert next_page_params == expected_cursor_params + + @pytest.mark.asyncio async def test_export_cursor_tracks_model_when_a_key_has_multiple_models(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(constants, "USAGE_EXPORT_BATCH_SIZE", 1)