From 8f997824dbb06cbd95f43de472b91db5105a0100 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 6 Nov 2026 15:51:08 +0800 Subject: [PATCH] fix(alerts): bound daily metric identifiers before materialization --- litellm/repositories/provider_latency.py | 4 + litellm/repositories/usage_alert_metrics.py | 23 +++- .../database/test_performance_issues.py | 107 +++++++++++++++++- .../database/test_provider_latency.py | 45 ++++++++ 4 files changed, 175 insertions(+), 4 deletions(-) diff --git a/litellm/repositories/provider_latency.py b/litellm/repositories/provider_latency.py index eef64462aab..ebfc52e9668 100644 --- a/litellm/repositories/provider_latency.py +++ b/litellm/repositories/provider_latency.py @@ -53,6 +53,8 @@ _MEMBERS: Final = TypeAdapter(tuple[_Member, ...]) _CANDIDATES: Final = """ SELECT DISTINCT custom_llm_provider AS provider, model FROM "LiteLLM_DailyUserSpend" WHERE date=$1 AND timed_requests>0 AND custom_llm_provider<>'' AND model<>'' + AND CASE WHEN octet_length(custom_llm_provider)<=1024 THEN length(custom_llm_provider)<=256 ELSE false END + AND CASE WHEN octet_length(model)<=4096 THEN length(model)<=1024 ELSE false END AND ($2::text IS NULL OR (custom_llm_provider,model)>($2,$3)) ORDER BY provider,model LIMIT $4 """ @@ -201,6 +203,8 @@ class ProviderLatencyRepository: await self._database.query_raw( """SELECT DISTINCT custom_llm_provider AS provider,model FROM "LiteLLM_DailyUserSpend" WHERE date=$1 AND custom_llm_provider=$2 AND model IS NOT NULL AND model<>'' AND timed_requests>0 + AND CASE WHEN octet_length(custom_llm_provider)<=1024 THEN length(custom_llm_provider)<=256 ELSE false END + AND CASE WHEN octet_length(model)<=4096 THEN length(model)<=1024 ELSE false END ORDER BY model LIMIT 101""", day.isoformat(), provider, diff --git a/litellm/repositories/usage_alert_metrics.py b/litellm/repositories/usage_alert_metrics.py index d2a9af41913..2f766df702d 100644 --- a/litellm/repositories/usage_alert_metrics.py +++ b/litellm/repositories/usage_alert_metrics.py @@ -64,6 +64,16 @@ class IssueSourceProvider: provider: str +def _bounded_identity(scope: Scope, entity: str) -> str: + # Check the stored byte length before decoding text; escaped key IDs add four characters. + maximum: Final = ( + f"CASE WHEN {entity}='unattributed' OR {entity} LIKE 'raw:%' THEN 1020 ELSE 1024 END" + if scope == "key" + else "1024" + ) + return f"CASE WHEN octet_length({entity})<=4096 THEN length({entity})<={maximum} ELSE false END" + + def _window_payload(scope: Scope, windows: tuple[IssueMetricWindow, ...]) -> str: if len(windows) > 500 or any(not 1 <= (window.day - window.start).days <= 30 for window in windows): raise ValueError("Read at most 500 issue windows with 1–30 reference days per batch") @@ -154,6 +164,7 @@ def _key_attribution_query(scope: Scope, metric: Literal["spend", "requests"]) - FROM population p JOIN ranked r USING(reference_id) WHERE r.rank=1 ), windows AS MATERIALIZED ( SELECT DISTINCT api_key,entity_id,start_day,target_day FROM candidate WHERE api_key<>'' + AND {_bounded_identity("key", "api_key")} ), key_population AS ( SELECT w.*, {team_days} AS team_days, @@ -183,10 +194,11 @@ _ROWS: Final = TypeAdapter(tuple[DailyAlertMetrics, ...]) _BATCH_SIZE: Final = 10 _DEPRECATION_USAGE: Final = TypeAdapter(tuple[ModelDeprecationUsage, ...]) -_DEPRECATION_USAGE_QUERY: Final = """ +_DEPRECATION_USAGE_QUERY: Final = f""" WITH activity AS MATERIALIZED ( SELECT model, date, api_key, user_id, api_requests FROM "LiteLLM_DailyUserSpend" WHERE date >= $1 AND date <= $2 AND model = ANY($3::text[]) AND api_requests > 0 + AND {_bounded_identity("model", "model")} ), daily AS ( SELECT model, date, SUM(api_requests)::bigint AS requests FROM activity GROUP BY model, date ), latest AS ( @@ -201,6 +213,7 @@ WITH activity AS MATERIALIZED ( SELECT model, COUNT(DISTINCT NULLIF(NULLIF(team_id, ''), 'litellm-dashboard'))::int AS teams FROM "LiteLLM_DailyTeamSpend" WHERE date >= $1 AND date <= $2 AND model = ANY($3::text[]) AND api_requests > 0 + AND {_bounded_identity("model", "model")} GROUP BY model ) SELECT totals.*, latest.last_used_date, latest.last_day_requests, COALESCE(teams.teams, 0) AS teams @@ -227,6 +240,7 @@ def _query(scope: Scope) -> str: WITH entities AS MATERIALIZED ( SELECT DISTINCT {entity} AS entity_id FROM {table} WHERE date >= $1 AND date <= $2 AND {entity} >= COALESCE($3::text, '') AND ($3::text IS NULL OR {entity} <> $3) + AND {_bounded_identity(scope, entity)} ORDER BY entity_id LIMIT $4 ) SELECT entities.entity_id, activity.date AS day, @@ -327,7 +341,8 @@ class UsageAlertMetricsRepository: ), ranked AS ( SELECT *,SUM(timed) OVER(PARTITION BY reference_id) AS total FROM grouped ) SELECT reference_id,provider FROM ranked - WHERE provider<>'' AND timed>0 AND timed>=0.8*total""", + WHERE provider<>'' AND timed>0 AND timed>=0.8*total + AND CASE WHEN octet_length(provider)<=1024 THEN length(provider)<=256 ELSE false END""", _window_payload(scope, windows), ) ) @@ -368,7 +383,9 @@ class UsageAlertMetricsRepository: ) exact_query: Final = ( query[: query.index("WITH entities")] - + "WITH entities AS MATERIALIZED (SELECT unnest($3::text[]) AS entity_id)" + + f"""WITH entities AS MATERIALIZED ( + SELECT entity_id FROM unnest($3::text[]) AS ids(entity_id) + WHERE {_bounded_identity(scope, "entity_id")})""" + query[query.index("\nSELECT entities") :] ) async with db_span("daily_activity_query", sql_relation(query)): diff --git a/tests/integration/database/test_performance_issues.py b/tests/integration/database/test_performance_issues.py index 47c2a4386a5..3c2539ac071 100644 --- a/tests/integration/database/test_performance_issues.py +++ b/tests/integration/database/test_performance_issues.py @@ -14,8 +14,9 @@ from integration._support.prisma_clients import connected_proxy_prisma from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.proxy.spend_tracking.usage_alert_job import UsageAlertJob from litellm.proxy.utils import PrismaClient +from litellm.repositories.daily_activity_buckets import api_key_bucket from litellm.repositories.issues import IssuesRepository -from litellm.repositories.usage_alert_metrics import IssueMetricWindow +from litellm.repositories.usage_alert_metrics import IssueMetricWindow, UsageAlertMetricsRepository from litellm.repositories.usage_alerts import UsageAlertsRepository from litellm.types.proxy.usage_alerts import UsageAlertObservation, UsageAlertRule, UsageAlertRuleInput, UsageAlertScope @@ -42,6 +43,110 @@ _TABLES: Final = ( ) +@pytest.mark.asyncio(loop_scope="session") +@pytest.mark.parametrize("scope", ("user", "team", "key", "model")) +async def test_metric_pages_skip_oversized_ids_without_losing_valid_later_entities( + db: ManagementDatabase, scope: UsageAlertScope +) -> None: + identities = ( + "a" * 1025, + "\U0001f600" * 1025, + "raw:" + "x" * 1016, + "raw:" + "x" * 1017, + "oversized-" + "x" * 65536, + "oversized-" + "\U0001f600" * 8192, + *(f"valid-{index:02}" for index in range(10)), + "\U0001f600" * 1024, + ) + expected = tuple( + sorted( + api_key_bucket(value) if scope == "key" else value + for value in identities + if len(api_key_bucket(value) if scope == "key" else value) <= 1024 + ) + ) + async with _isolated(db) as prisma: + for table, entity in (("LiteLLM_DailyUserSpend", "user_id"), ("LiteLLM_DailyTeamSpend", "team_id")): + await prisma.writer_db.query_raw( + f"""INSERT INTO "{table}" (id,{entity},date,api_key,model,spend,api_requests,updated_at) + SELECT $1||n,identity,$2,identity,identity,2,3,CURRENT_TIMESTAMP + FROM unnest($3::text[]) WITH ORDINALITY AS items(identity,n) RETURNING id""", + table, + _DAY.isoformat(), + identities, + ) + metrics = UsageAlertMetricsRepository(prisma.writer_db) + rows = [] + cursor = None + for _ in range(3): + page = await metrics.page(scope, _DAY, 1, cursor=cursor) + rows.extend(page.rows) + if page.complete: + break + cursor = page.cursor + else: + pytest.fail("bounded metric pagination did not finish") + assert tuple(sorted(row.entity_id for row in rows)) == expected + assert all((row.spend, row.requests) == (2, 3) for row in rows) + assert "\U0001f600" * 1024 in expected + exact = tuple(api_key_bucket(value) if scope == "key" else value for value in identities[:6]) + ("valid-00",) + chosen = await metrics.page(scope, _DAY, 1, entity_ids=exact) + assert {row.entity_id for row in chosen.rows} == set(expected).intersection(exact) + deployment = await metrics.page("deployment", _DAY, 1) + assert len(deployment.rows) == 1 + assert deployment.rows[0].requests == 3 * len(identities) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_attribution_bounds_keep_invalid_identity_traffic_in_population_totals(db: ManagementDatabase) -> None: + async with _isolated(db) as prisma: + await prisma.writer_db.query_raw( + """INSERT INTO "LiteLLM_DailyUserSpend" + (id,user_id,date,api_key,model,custom_llm_provider,spend,api_requests,timed_requests,updated_at) + SELECT owner||d||k,owner,($1::date-d)::text, + CASE WHEN k='bad' THEN owner||repeat('x',1025) ELSE owner||'-valid' END, + 'example-model',CASE WHEN k='bad' THEN repeat('p',257) ELSE 'valid-provider' END, + CASE WHEN d>0 THEN 10 WHEN (owner='invalid-dominant')=(k='bad') THEN 100 ELSE 20 END, + 1,CASE WHEN (owner='invalid-dominant')=(k='bad') THEN 90 ELSE 10 END,CURRENT_TIMESTAMP + FROM generate_series(0,3) d CROSS JOIN (VALUES('invalid-dominant'),('valid-dominant')) owners(owner) + CROSS JOIN (VALUES('bad'),('good')) keys(k) RETURNING id""", + _DAY.isoformat(), + ) + windows = tuple( + IssueMetricWindow(owner, owner, _DAY - timedelta(days=3), _DAY) + for owner in ("invalid-dominant", "valid-dominant") + ) + metrics = UsageAlertMetricsRepository(prisma.writer_db) + attributions = await metrics.issue_key_attributions("user", windows) + assert len(attributions) == 1 + attribution = attributions[0] + assert (attribution.reference_id, attribution.key_id, attribution.equivalent) == ( + "valid-dominant", + "valid-dominant-valid", + False, + ) + assert attribution.positive_change_share == pytest.approx(0.9) + assert (attribution.baseline_value, attribution.observed_value) == (20, 120) + providers = await metrics.issue_source_providers("user", windows) + assert tuple((row.reference_id, row.provider) for row in providers) == (("valid-dominant", "valid-provider"),) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_retirement_usage_omits_oversized_models_before_materialization(db: ManagementDatabase) -> None: + models = ("a" * 65536, "valid-model", "\U0001f600" * 1024) + async with _isolated(db) as prisma: + await prisma.writer_db.query_raw( + """INSERT INTO "LiteLLM_DailyUserSpend" (id,user_id,date,api_key,model,api_requests,updated_at) + SELECT 'retirement-'||n,'owner',$1,'key',model,7,CURRENT_TIMESTAMP + FROM unnest($2::text[]) WITH ORDINALITY AS items(model,n) RETURNING id""", + _DAY.isoformat(), + models, + ) + rows = await UsageAlertMetricsRepository(prisma.writer_db).deprecation_usage(models, _DAY, _DAY) + assert frozenset(row.model for row in rows) == frozenset(models[1:]) + assert tuple(row.requests for row in rows) == (7, 7) + + @asynccontextmanager async def _isolated(db: ManagementDatabase) -> AsyncIterator[PrismaClient]: schema: Final = "issue_test_" + uuid4().hex diff --git a/tests/integration/database/test_provider_latency.py b/tests/integration/database/test_provider_latency.py index 8175d8525bd..ec4986b7c29 100644 --- a/tests/integration/database/test_provider_latency.py +++ b/tests/integration/database/test_provider_latency.py @@ -9,6 +9,51 @@ from litellm.types.proxy.issues import IssueModelPair from litellm.types.proxy.usage_alerts import UsageAlertRuleInput +@pytest.mark.asyncio(loop_scope="session") +async def test_provider_pair_pages_skip_oversized_names_and_continue(db: ManagementDatabase) -> None: + prefix = f"provider-width-{uuid4()}" + day = date(2038, 1, 12) + provider = prefix + "-valid" + wide_provider = "\U0001f600" * 256 + wide_model = "\U0001f600" * 1024 + pairs = ( + ("a" * 257, "first"), + (provider, "a" * 1025), + (provider, "valid-model"), + (provider, wide_model), + (wide_provider, wide_model), + ("\U0001f600" * 257, "last"), + ("\U0001f600" * 8192, "last-large"), + (provider, "\U0001f600" * 8192), + ) + repository = ProviderLatencyRepository(db.prisma.writer_db) + try: + await db.prisma.writer_db.query_raw( + """INSERT INTO "LiteLLM_DailyUserSpend" + (id,user_id,date,api_key,custom_llm_provider,model,api_requests,timed_requests,updated_at) + SELECT $1||n,$1,$2,$1,provider,model,1,1,CURRENT_TIMESTAMP + FROM unnest($3::text[],$4::text[]) WITH ORDINALITY AS items(provider,model,n) RETURNING id""", + prefix, + day.isoformat(), + tuple(pair[0] for pair in pairs), + tuple(pair[1] for pair in pairs), + ) + expected = tuple(IssueModelPair(provider=p, model=m) for p, m in pairs[2:5]) + first = await repository.candidate_pairs(day, limit=1) + second = await repository.candidate_pairs(day, first[0], limit=1) + third = await repository.candidate_pairs(day, second[0], limit=1) + assert len(first) == len(second) == len(third) == 1 + assert frozenset((*first, *second, *third)) == frozenset(expected) + assert await repository.candidate_pairs(day, third[0], limit=1) == () + assert frozenset(await repository.provider_pairs(provider, day)) == frozenset(expected[:2]) + assert await repository.provider_pairs(wide_provider, day) == expected[2:] + assert await repository.provider_pairs("a" * 257, day) == () + finally: + await db.prisma.writer_db.query_raw( + 'DELETE FROM "LiteLLM_DailyUserSpend" WHERE id LIKE $1 RETURNING id', prefix + "%" + ) + + @pytest.mark.asyncio(loop_scope="session") async def test_provider_latency_requires_established_within_model_regressions(db: ManagementDatabase) -> None: prefix = f"provider-latency-{uuid4()}"