scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(alerts): bound daily metric identifiers before materialization

Browse source
This commit is contained in:
moe-berri 2026-11-06 15:51:08 +08:00
parent 25936a6cc6
commit 8f997824db
4 changed files with 175 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/repositories/provider_latency.py
View file

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

23
litellm/repositories/usage_alert_metrics.py
View file

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

107
tests/integration/database/test_performance_issues.py
View file

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

45
tests/integration/database/test_provider_latency.py
View file

@ -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()}"