diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index 0b285261658..e2e29b9b3ba 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -216,8 +216,10 @@ class DailyActivityRepository: rows: Final = _KEY_SPEND_ADAPTER.validate_python(await self._query(query)) return tuple(row.api_key for row in rows) - async def key_page(self, scope: DailyActivityScope, *, offset: int, limit: int) -> KeyPage: - query: Final = build_key_page_sql(scope, offset=offset, limit=limit) + async def key_page( + self, scope: DailyActivityScope, *, offset: int, limit: int, bounded_identifiers: bool = False + ) -> KeyPage: + query: Final = build_key_page_sql(scope, offset=offset, limit=limit, bounded_identifiers=bounded_identifiers) result: Final = await self._query(query) total_api_keys_value: Final = result[0].get("total_api_keys") if result else 0 total_api_keys: Final = ( @@ -226,8 +228,10 @@ class DailyActivityRepository: rows: Final = _KEY_SPEND_ADAPTER.validate_python(tuple(row for row in result if row.get("api_key") is not None)) return KeyPage(rows=rows, total_api_keys=total_api_keys) - async def user_page(self, scope: DailyActivityScope, *, offset: int, limit: int) -> UserPage: - query: Final = build_user_page_sql(scope, offset=offset, limit=limit) + async def user_page( + self, scope: DailyActivityScope, *, offset: int, limit: int, bounded_identifiers: bool = False + ) -> UserPage: + query: Final = build_user_page_sql(scope, offset=offset, limit=limit, bounded_identifiers=bounded_identifiers) result: Final = await self._query(query) total_users_value: Final = result[0].get("total_users") if result else 0 total_users: Final = ( diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index da58389c09e..c00ce989a83 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -411,7 +411,9 @@ def _key_spend_select() -> str: COALESCE(SUM(timed_requests), 0)::bigint AS timed_requests""" -def build_key_page_sql(scope: DailyActivityScope, *, offset: int, limit: int) -> SqlQuery: +def build_key_page_sql( + scope: DailyActivityScope, *, offset: int, limit: int, bounded_identifiers: bool = False +) -> SqlQuery: if not 1 <= limit <= constants.USAGE_KEY_PAGE_MAX: raise ValueError(f"limit must be between 1 and {constants.USAGE_KEY_PAGE_MAX}") if offset < 0: @@ -420,11 +422,14 @@ def build_key_page_sql(scope: DailyActivityScope, *, offset: int, limit: int) -> sentinel_param: Final = len(where_params) + 1 limit_param: Final = sentinel_param + 1 offset_param: Final = limit_param + 1 + identity_filter: Final = ( + "AND length(api_key)<=CASE WHEN api_key LIKE 'raw:%' THEN 1020 ELSE 1024 END" if bounded_identifiers else "" + ) sql: Final = f""" WITH ranked AS ( SELECT api_key,{_key_spend_select()}, SUM(spend::numeric) AS rank_spend FROM "{PRISMA_TO_PG_TABLE[scope.table]}" - WHERE {where_clause} AND api_key <> ${sentinel_param} + WHERE {where_clause} AND api_key <> ${sentinel_param} {identity_filter} GROUP BY api_key ) SELECT (SELECT COUNT(*) FROM ranked)::bigint AS total_api_keys, page.* @@ -451,7 +456,9 @@ def _user_spend_select() -> str: COALESCE(SUM(timed_requests), 0)::bigint AS timed_requests""" -def build_user_page_sql(scope: DailyActivityScope, *, offset: int, limit: int) -> SqlQuery: +def build_user_page_sql( + scope: DailyActivityScope, *, offset: int, limit: int, bounded_identifiers: bool = False +) -> SqlQuery: if not 1 <= limit <= constants.USAGE_USER_PAGE_MAX: raise ValueError(f"limit must be between 1 and {constants.USAGE_USER_PAGE_MAX}") if offset < 0: @@ -460,11 +467,12 @@ def build_user_page_sql(scope: DailyActivityScope, *, offset: int, limit: int) - limit_param: Final = len(where_params) + 1 offset_param: Final = limit_param + 1 group_expr: Final = f"NULLIF(\"{scope.entity_id_field}\", '')" + identity_filter: Final = f"AND ({group_expr} IS NULL OR length({group_expr})<=1024)" if bounded_identifiers else "" sql: Final = f""" WITH ranked AS ( SELECT {group_expr} AS user_id,{_user_spend_select()}, SUM(spend::numeric) AS rank_spend FROM "{PRISMA_TO_PG_TABLE[scope.table]}" - WHERE {where_clause} + WHERE {where_clause} {identity_filter} GROUP BY {group_expr} ) SELECT (SELECT COUNT(*) FROM ranked)::bigint AS total_users, page.* diff --git a/litellm/repositories/issue_activity.py b/litellm/repositories/issue_activity.py index a6b9bf96fe7..c06cb51e9f6 100644 --- a/litellm/repositories/issue_activity.py +++ b/litellm/repositories/issue_activity.py @@ -94,7 +94,9 @@ class IssueActivityRepository: ) -> DailyActivityKeyPageResponse: if not 1 <= limit <= 100 or offset < 0: raise ValueError("Read between one and 100 keys with a nonnegative offset") - page: Final = await self.activity.key_page(issue_activity_scope(cohort, start, end), limit=limit, offset=offset) + page: Final = await self.activity.key_page( + issue_activity_scope(cohort, start, end), limit=limit, offset=offset, bounded_identifiers=True + ) metadata: Final = await self._key_labels(tuple(row.api_key for row in page.rows if row.api_key)) return DailyActivityKeyPageResponse( api_keys=[ @@ -118,7 +120,7 @@ class IssueActivityRepository: if cohort.scope_type == "team": raise ValueError("Historical team aggregates do not retain user membership") page: Final = await self.activity.user_page( - issue_activity_scope(cohort, start, end), limit=limit, offset=offset + issue_activity_scope(cohort, start, end), limit=limit, offset=offset, bounded_identifiers=True ) metadata: Final = await self._user_labels(tuple(row.user_id for row in page.rows if row.user_id)) return DailyActivityUserPageResponse( @@ -149,7 +151,10 @@ class IssueActivityRepository: await self.database.query_raw( """SELECT w.api_key,jsonb_build_object( 'key_alias',LEFT(COALESCE(k.key_alias,d.key_alias),256), - 'team_id',COALESCE(k.team_id,d.team_id),'user_id',COALESCE(k.user_id,d.user_id), + 'team_id',CASE WHEN length(COALESCE(k.team_id,d.team_id))<=1024 + THEN COALESCE(k.team_id,d.team_id) END, + 'user_id',CASE WHEN length(COALESCE(k.user_id,d.user_id))<=1024 + THEN COALESCE(k.user_id,d.user_id) END, 'user_email',CASE WHEN length(u.user_email)<=320 THEN u.user_email END, 'key_exists',CASE WHEN k.token IS NOT NULL THEN true WHEN d.id IS NOT NULL THEN false END ) AS metadata diff --git a/tests/integration/database/test_issues.py b/tests/integration/database/test_issues.py index 128787f1524..2409ab21a0f 100644 --- a/tests/integration/database/test_issues.py +++ b/tests/integration/database/test_issues.py @@ -17,6 +17,7 @@ from prisma.types import DatasourceOverride from psycopg import sql from pydantic import BaseModel, TypeAdapter, ValidationError +from litellm.proxy.management_endpoints.common_daily_activity import daily_activity_repository from litellm.repositories.daily_activity_repository import DailyActivityQueryDatabase from litellm.repositories.daily_activity_sql import build_bounded_daily_activity_sql from litellm.repositories.issue_activity import IssueActivityRepository @@ -912,13 +913,13 @@ async def test_issue_resource_metadata_bounds_labels_and_deleted_history( store: Final = IssuesRepository.from_prisma(db.prisma) activity: Final = IssueActivityRepository.from_prisma(db.prisma) measured: Final = _MeasuredActivityDatabase(activity.database) - repository: Final = IssueActivityRepository(activity.activity, measured) + repository: Final = IssueActivityRepository(daily_activity_repository(db.prisma, query_database=measured), measured) identity: Final = uuid4().hex reporting_key: Final = sha256(identity.encode()).hexdigest() day: Final = date.today() - timedelta(days=3) with gateway.scenario() as scenario: owner: Final = scenario.user(user_role="internal_user") - issue: Final = _issue(IssueCohort(scope_type="user", scope_id=owner)).model_copy( + issue: Final = _issue(IssueCohort(scope_type="provider", scope_id=identity)).model_copy( update={"range_start": day, "range_end": day} ) try: @@ -942,8 +943,10 @@ async def test_issue_resource_metadata_bounds_labels_and_deleted_history( await _query( db, """INSERT INTO "LiteLLM_DailyUserSpend" - (id,user_id,date,api_key,model,api_requests,successful_requests,updated_at) - VALUES($1,$2,$3,$4,'example-model',10,10,CURRENT_TIMESTAMP)""", + (id,user_id,date,api_key,model,custom_llm_provider,api_requests,successful_requests,updated_at) + VALUES($1,$2,$3,$4,'example-model',$1,10,10,CURRENT_TIMESTAMP), + ($1||'-wide',repeat(chr(128640),65536),$3,repeat(chr(128640),65536), + 'example-model',$1,20,20,CURRENT_TIMESTAMP)""", identity, owner, day.isoformat(), @@ -963,13 +966,26 @@ async def test_issue_resource_metadata_bounds_labels_and_deleted_history( assert users.users[0].user_id == owner and users.users[0].user_alias == "u" * 256 assert users.users[0].user_email is None assert measured.outputs[-1][0] == 1 and measured.outputs[-1][1] < 2048 + assert all(size < 2048 for _, size in measured.outputs) + totals: Final = await repository.daily_activity(issue.cohort, day, day) + assert totals.metadata.total_api_requests == 30 for endpoint in ("keys", "users"): response: Final = gateway.request("GET", f"/issues/{issue.id}/{endpoint}", params={"limit": "1"}) assert response.status_code == 200 and len(response.content) < 2048 assert "padding" not in response.text + await _query( + db, + """UPDATE "LiteLLM_DeletedVerificationToken" + SET user_id=repeat(chr(128640),65536),team_id=repeat(chr(128640),65536) + WHERE id=$1||'-400'""", + identity, + ) + legacy: Final = await repository.keys(issue.cohort, day, day, 1, 0) + assert legacy.api_keys[0].metadata.user_id is None and legacy.api_keys[0].metadata.team_id is None + assert all(size < 2048 for _, size in measured.outputs[-2:]) finally: await _remove_issues(db, (issue.id,)) - await _query(db, 'DELETE FROM "LiteLLM_DailyUserSpend" WHERE id=$1', identity) + await _query(db, "DELETE FROM \"LiteLLM_DailyUserSpend\" WHERE id IN ($1,$1||'-wide')", identity) await _query(db, 'DELETE FROM "LiteLLM_DeletedVerificationToken" WHERE token=$1', reporting_key)