mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
fix: preserve alert identity when rule meaning changes
This commit is contained in:
moe-berri
ishaan-j-berri
parent
efaa14a947
commit
68b03b7eb8
2 changed files with 82 additions and 5 deletions
|
|
@ -3,7 +3,7 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from typing import ( # PrismaWrapper forwards generated client methods dynamically
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
Protocol,
|
||||
|
|
@ -92,7 +92,20 @@ def activity_path(
|
|||
|
||||
def _observation_payload(rule: UsageAlertRule, item: UsageAlertObservation, now: datetime) -> Mapping[str, object]:
|
||||
alert_id: Final = str(
|
||||
uuid5(NAMESPACE_URL, f"litellm:usage-alert:{rule.id}:{item.entity_id or ''}:{item.day.isoformat()}")
|
||||
uuid5(
|
||||
NAMESPACE_URL,
|
||||
json.dumps(
|
||||
[
|
||||
"litellm:usage-alert",
|
||||
rule.id,
|
||||
rule.scope_type,
|
||||
item.entity_id,
|
||||
rule.metric,
|
||||
rule.direction,
|
||||
item.day.isoformat(),
|
||||
]
|
||||
),
|
||||
)
|
||||
)
|
||||
incident: Final = UsageAlert(
|
||||
id=alert_id,
|
||||
|
|
@ -329,7 +342,7 @@ class UsageAlertsRepository:
|
|||
async def list_breached_entity_ids(self, rule_id: str, limit: int = 500, offset: int = 0) -> tuple[str | None, ...]:
|
||||
rows: Final = TypeAdapter(tuple[Mapping[str, str | None], ...]).validate_python(
|
||||
await self.database.query_raw(
|
||||
'SELECT DISTINCT scope_id FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND breached ORDER BY scope_id NULLS FIRST LIMIT $2 OFFSET $3',
|
||||
"""SELECT DISTINCT i.scope_id FROM "LiteLLM_UsageAlertIncident" i JOIN "LiteLLM_UsageAlertRule" r ON r.id=i.rule_id WHERE i.rule_id=$1 AND i.breached AND i.scope_type=r.scope_type AND i.metric=r.record->>'metric' AND i.record->>'direction'=r.record->>'direction' AND (r.scope_id IS NULL OR i.scope_id=r.scope_id) ORDER BY i.scope_id NULLS FIRST LIMIT $2 OFFSET $3""",
|
||||
rule_id,
|
||||
limit,
|
||||
offset,
|
||||
|
|
@ -366,9 +379,12 @@ class UsageAlertsRepository:
|
|||
return changed[0]
|
||||
rows: Final = TypeAdapter(tuple[_Record, ...]).validate_python(
|
||||
await self.database.query_raw(
|
||||
'SELECT record FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND scope_key=$2 ORDER BY evaluation_day DESC,id DESC LIMIT 1',
|
||||
"""SELECT record FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND scope_key=$2 AND scope_type=$3 AND metric=$4 AND record->>'direction'=$5 ORDER BY evaluation_day DESC,id DESC LIMIT 1""",
|
||||
rule.id,
|
||||
entity_id or "",
|
||||
rule.scope_type,
|
||||
rule.metric,
|
||||
rule.direction,
|
||||
)
|
||||
)
|
||||
return UsageAlert.model_validate(rows[0].record) if rows else None
|
||||
|
|
@ -389,7 +405,10 @@ class UsageAlertsRepository:
|
|||
prior AS (
|
||||
SELECT x.*,p.id AS prior_id,p.evaluation_day,p.breached AS prior_breached,p.record AS prior_record
|
||||
FROM input x LEFT JOIN LATERAL (
|
||||
SELECT * FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND scope_key=x.scope_key ORDER BY evaluation_day DESC,id DESC LIMIT 1
|
||||
SELECT * FROM "LiteLLM_UsageAlertIncident" WHERE rule_id=$1 AND scope_key=x.scope_key
|
||||
AND scope_type=x.record->>'scope_type' AND scope_id IS NOT DISTINCT FROM x.record->>'scope_id'
|
||||
AND metric=x.record->>'metric' AND record->>'direction'=x.record->>'direction'
|
||||
ORDER BY evaluation_day DESC,id DESC LIMIT 1
|
||||
) p ON true
|
||||
WHERE EXISTS (SELECT 1 FROM "LiteLLM_UsageAlertRule" WHERE id=$1 AND enabled AND NOT deleted AND (record->>'updated_at')::timestamptz=$3::timestamptz)
|
||||
), updated AS (
|
||||
|
|
|
|||
|
|
@ -1,8 +1,14 @@
|
|||
from datetime import date
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.actors import ManagementDatabase
|
||||
from integration._support.client import Gateway, string_value
|
||||
|
||||
from litellm.repositories.usage_alerts import UsageAlertsRepository
|
||||
from litellm.types.proxy.usage_alerts import UsageAlertObservation, UsageAlertRuleInput
|
||||
|
||||
|
||||
def test_scoped_rules_use_real_management_auth_and_preserve_read_only_roles(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
|
|
@ -78,3 +84,55 @@ def test_team_admin_can_manage_team_alerts_but_member_can_only_read(gateway: Gat
|
|||
assert gateway.request("DELETE", path, key=member_key).status_code == 403
|
||||
finally:
|
||||
assert gateway.request("DELETE", path, key=admin_key).status_code == 204
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
@pytest.mark.parametrize(
|
||||
"changed_field,new_value",
|
||||
[("metric", "latency"), ("direction", "decrease"), ("scope_type", "user"), ("scope_id", "semantic-other")],
|
||||
)
|
||||
async def test_rule_identity_edits_preserve_old_incident_and_filter_consistency(
|
||||
db: ManagementDatabase, changed_field: str, new_value: str
|
||||
) -> None:
|
||||
repository: Final = UsageAlertsRepository.from_prisma(db.prisma)
|
||||
definition: Final = UsageAlertRuleInput(
|
||||
name="Semantic identity",
|
||||
scope_type="team",
|
||||
scope_id="semantic-example",
|
||||
metric="spend",
|
||||
direction="increase",
|
||||
notification_channels=("slack",),
|
||||
)
|
||||
rule: Final = await repository.save_rule(definition, "qa-admin")
|
||||
assert rule is not None
|
||||
initial: Final = UsageAlertObservation(
|
||||
entity_id=rule.scope_id, day=date(2026, 11, 4), observed=40, baseline=5, threshold=15, breached=True
|
||||
)
|
||||
try:
|
||||
original: Final = (await repository.record_observations(rule, (initial,)))[0]
|
||||
edited: Final = await repository.save_rule(
|
||||
UsageAlertRuleInput.model_validate({**definition.model_dump(), changed_field: new_value}), "qa-admin", rule
|
||||
)
|
||||
assert edited is not None
|
||||
observation: Final = initial.model_copy(update={"entity_id": edited.scope_id})
|
||||
replacement: Final = (await repository.record_observations(edited, (observation,)))[0]
|
||||
assert replacement.id != original.id
|
||||
assert await repository.get_alert(original.id) == original
|
||||
before_page: Final = await repository.list_alerts(
|
||||
scope_type=original.scope_type, scope_id=original.scope_id, metric=original.metric
|
||||
)
|
||||
after_page: Final = await repository.list_alerts(
|
||||
scope_type=replacement.scope_type, scope_id=replacement.scope_id, metric=replacement.metric
|
||||
)
|
||||
assert original.id in {item.id for item in before_page.data}
|
||||
assert replacement.id in {item.id for item in after_page.data}
|
||||
assert all(
|
||||
item.scope_type == original.scope_type and item.metric == original.metric for item in before_page.data
|
||||
)
|
||||
healthy: Final = observation.model_copy(update={"day": date(2026, 11, 5), "breached": False})
|
||||
await repository.record_observations(edited, (healthy,))
|
||||
assert (await repository.get_alert(original.id)) == original
|
||||
resolved: Final = await repository.get_alert(replacement.id)
|
||||
assert resolved is not None and resolved.status == "resolved"
|
||||
finally:
|
||||
await repository.delete_rule(rule.id)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue