Merge branch 'litellm_alerts_reconcile_02_final' of /Users/mfkhalil/Code/litellm-daily-alert-detector into litellm_alerts_reconcile_03

This commit is contained in:
moe-berri 2026-11-06 14:43:27 +08:00
commit 41f180fbe6
4 changed files with 49 additions and 16 deletions

View file

@ -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 = (

View file

@ -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.*

View file

@ -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

View file

@ -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)