mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
refactor(alerts): propagate repository simplification
# Conflicts: # litellm/repositories/usage_alerts.py
This commit is contained in:
moe-berri
commit
d20ec80f90
6 changed files with 120 additions and 270 deletions
|
|
@ -7,7 +7,6 @@ from litellm.types.proxy.issues import (
|
|||
IssueActivityResponse,
|
||||
IssueCohort,
|
||||
IssueDeliveryBinding,
|
||||
IssueDeliveryInput,
|
||||
IssueFamily,
|
||||
IssueFamilyList,
|
||||
IssueFamilySetting,
|
||||
|
|
@ -20,7 +19,6 @@ from litellm.types.proxy.issues import (
|
|||
IssueProjectionSignal,
|
||||
IssueScope,
|
||||
IssueSignal,
|
||||
IssueSignalInput,
|
||||
IssueSignalList,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
|
|
@ -58,13 +56,6 @@ class IssueStore(Protocol):
|
|||
|
||||
async def status_scopes(self, issue_id: str) -> tuple[IssueCohort, ...]: ...
|
||||
|
||||
async def upsert_issue(
|
||||
self,
|
||||
data: IssueInput,
|
||||
reopen: bool = False,
|
||||
status: UsageAlertStatus | None = None,
|
||||
) -> Issue: ...
|
||||
|
||||
async def upsert_issues(
|
||||
self,
|
||||
inputs: tuple[IssueInput, ...],
|
||||
|
|
@ -72,12 +63,8 @@ class IssueStore(Protocol):
|
|||
status: UsageAlertStatus | None = None,
|
||||
) -> tuple[Issue, ...]: ...
|
||||
|
||||
async def find_active_issue(self, cohort_key: str) -> Issue | None: ...
|
||||
|
||||
async def find_active_issues(self, cohort_keys: tuple[str, ...]) -> tuple[Issue, ...]: ...
|
||||
|
||||
async def add_signals(self, issue_id: str, signals: tuple[IssueSignalInput, ...]) -> None: ...
|
||||
|
||||
async def add_signal_batch(self, signals: tuple[IssueSignal, ...]) -> None: ...
|
||||
|
||||
async def list_signals(
|
||||
|
|
@ -115,14 +102,6 @@ class IssueStore(Protocol):
|
|||
self, family: PerformanceIssueFamily, data: IssueFamilyUpdate
|
||||
) -> IssueFamilySetting | None: ...
|
||||
|
||||
async def queue_delivery(
|
||||
self,
|
||||
data: IssueDeliveryInput,
|
||||
rule_id: str,
|
||||
rule_revision: str,
|
||||
source_id: str,
|
||||
) -> None: ...
|
||||
|
||||
async def queue_deliveries(self, bindings: tuple[IssueDeliveryBinding, ...]) -> None: ...
|
||||
|
||||
async def pending_deliveries(self, limit: int = 100) -> tuple[IssuePendingDelivery, ...]: ...
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ from litellm.types.proxy.issues import (
|
|||
Issue,
|
||||
IssueCohort,
|
||||
IssueDeliveryBinding,
|
||||
IssueDeliveryInput,
|
||||
IssueFamily,
|
||||
IssueFamilyList,
|
||||
IssueFamilySetting,
|
||||
|
|
@ -29,7 +28,6 @@ from litellm.types.proxy.issues import (
|
|||
IssueProjectionSignal,
|
||||
IssueScope,
|
||||
IssueSignal,
|
||||
IssueSignalInput,
|
||||
IssueSignalList,
|
||||
)
|
||||
from litellm.types.proxy.usage_alert_defaults import DEFAULT_USAGE_RULE_IDS, PerformanceIssueFamily
|
||||
|
|
@ -110,12 +108,6 @@ class IssuesRepository:
|
|||
def from_prisma(cls, prisma: "PrismaClient") -> "IssuesRepository":
|
||||
return cls(cast(AlertDatabase, prisma.writer_db)) # cast-ok: generated Prisma delegate implements query_raw
|
||||
|
||||
async def upsert_issue(
|
||||
self, data: IssueInput, reopen: bool = False, status: UsageAlertStatus | None = None
|
||||
) -> Issue:
|
||||
rows: Final = await self.upsert_issues((data,), reopen=reopen, status=status)
|
||||
return rows[0]
|
||||
|
||||
async def upsert_issues(
|
||||
self,
|
||||
inputs: tuple[IssueInput, ...],
|
||||
|
|
@ -228,10 +220,6 @@ class IssuesRepository:
|
|||
for issue in issues
|
||||
)
|
||||
|
||||
async def find_active_issue(self, cohort_key: str) -> Issue | None:
|
||||
rows: Final = await self.find_active_issues((cohort_key,))
|
||||
return rows[0] if rows else None
|
||||
|
||||
async def find_active_issues(self, cohort_keys: tuple[str, ...]) -> tuple[Issue, ...]:
|
||||
if len(cohort_keys) > 500:
|
||||
raise ValueError("Read at most 500 issue cohorts per batch")
|
||||
|
|
@ -305,11 +293,6 @@ class IssuesRepository:
|
|||
)
|
||||
)
|
||||
|
||||
async def add_signals(self, issue_id: str, signals: tuple[IssueSignalInput, ...]) -> None:
|
||||
await self.add_signal_batch(
|
||||
tuple(IssueSignal.model_validate({**signal.model_dump(), "issue_id": issue_id}) for signal in signals)
|
||||
)
|
||||
|
||||
async def add_signal_batch(self, signals: tuple[IssueSignal, ...]) -> None:
|
||||
if not signals:
|
||||
return
|
||||
|
|
@ -579,17 +562,6 @@ class IssuesRepository:
|
|||
|
||||
return await update_family(self.database, family, data)
|
||||
|
||||
async def queue_delivery(
|
||||
self,
|
||||
data: IssueDeliveryInput,
|
||||
rule_id: str,
|
||||
rule_revision: str,
|
||||
source_id: str,
|
||||
) -> None:
|
||||
await self.queue_deliveries(
|
||||
(IssueDeliveryBinding(delivery=data, rule_id=rule_id, rule_revision=rule_revision, source_id=source_id),)
|
||||
)
|
||||
|
||||
async def queue_deliveries(self, bindings: tuple[IssueDeliveryBinding, ...]) -> None:
|
||||
if not bindings:
|
||||
return
|
||||
|
|
|
|||
|
|
@ -20,7 +20,6 @@ from litellm.repositories.daily_activity_buckets import api_key_from_bucket
|
|||
from litellm.types.proxy.usage_alert_defaults import DEFAULT_USAGE_RULE_IDS
|
||||
from litellm.types.proxy.usage_alerts import (
|
||||
UsageAlert,
|
||||
UsageAlertChannel,
|
||||
UsageAlertDefaultSummary,
|
||||
UsageAlertEvaluationStatus,
|
||||
UsageAlertGroup,
|
||||
|
|
@ -32,7 +31,6 @@ from litellm.types.proxy.usage_alerts import (
|
|||
UsageAlertMetric,
|
||||
UsageAlertModelList,
|
||||
UsageAlertModelOption,
|
||||
UsageAlertNotification,
|
||||
UsageAlertObservation,
|
||||
UsageAlertProgress,
|
||||
UsageAlertProjectionCursor,
|
||||
|
|
@ -107,11 +105,6 @@ class _Identifier(BaseModel, frozen=True):
|
|||
id: str
|
||||
|
||||
|
||||
class _Pending(_Record, frozen=True):
|
||||
rule: Mapping[str, object]
|
||||
notified_channels: tuple[UsageAlertChannel, ...]
|
||||
|
||||
|
||||
class _Ownership(BaseModel, frozen=True):
|
||||
user_id: str | None
|
||||
team_id: str | None
|
||||
|
|
@ -320,16 +313,6 @@ class UsageAlertsRepository:
|
|||
default_summary=count[0].default_summary,
|
||||
)
|
||||
|
||||
async def list_enabled_rules(self, limit: int = 100, offset: int = 0) -> tuple[UsageAlertRule, ...]:
|
||||
rows: Final = TypeAdapter(tuple[_Record, ...]).validate_python(
|
||||
await self.database.query_raw(
|
||||
'SELECT record FROM "LiteLLM_UsageAlertRule" WHERE enabled AND NOT deleted ORDER BY id LIMIT $1 OFFSET $2',
|
||||
limit,
|
||||
offset,
|
||||
)
|
||||
)
|
||||
return tuple(UsageAlertRule.model_validate(row.record) for row in rows)
|
||||
|
||||
async def enabled_rule_page(
|
||||
self, scope: UsageAlertScope, wildcard: bool, cursor: str | None, limit: int = 500
|
||||
) -> tuple[UsageAlertRule, ...]:
|
||||
|
|
@ -713,21 +696,6 @@ class UsageAlertsRepository:
|
|||
def _checkpoint(rule: UsageAlertRule, day: date) -> str:
|
||||
return f"{rule.id}:{rule.updated_at.isoformat()}:{day.isoformat()}"
|
||||
|
||||
async def is_day_complete(self, rule: UsageAlertRule, day: date) -> bool:
|
||||
rows: Final = TypeAdapter(tuple[_Count, ...]).validate_python(
|
||||
await self.database.query_raw(
|
||||
'SELECT 1 AS count FROM "LiteLLM_UsageAlertCheckpoint" WHERE id=$1 AND completed',
|
||||
self._checkpoint(rule, day),
|
||||
)
|
||||
)
|
||||
return bool(rows)
|
||||
|
||||
async def finish_day(self, rule: UsageAlertRule, day: date) -> None:
|
||||
await self.database.query_raw(
|
||||
'INSERT INTO "LiteLLM_UsageAlertCheckpoint" (id,completed) VALUES ($1,true) ON CONFLICT(id) DO NOTHING RETURNING id',
|
||||
self._checkpoint(rule, day),
|
||||
)
|
||||
|
||||
async def completed_rule_ids(self, rules: tuple[UsageAlertRule, ...], day: date) -> frozenset[str]:
|
||||
if len(rules) > 500:
|
||||
raise ValueError("A checkpoint batch must contain at most 500 rules")
|
||||
|
|
@ -811,56 +779,6 @@ class UsageAlertsRepository:
|
|||
record.model_dump_json(),
|
||||
)
|
||||
|
||||
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 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,
|
||||
)
|
||||
)
|
||||
return tuple(row["scope_id"] for row in rows)
|
||||
|
||||
async def record_evaluation_status(
|
||||
self, rule: UsageAlertRule, day: date, status: UsageAlertEvaluationStatus
|
||||
) -> None:
|
||||
await self.database.query_raw(
|
||||
"""UPDATE "LiteLLM_UsageAlertRule" SET record=record || jsonb_build_object('last_evaluated_day',$2::text,'evaluation_status',$3::text) WHERE id=$1 AND (record->>'updated_at')::timestamptz=$4::timestamptz AND NOT deleted RETURNING id""",
|
||||
rule.id,
|
||||
day.isoformat(),
|
||||
status,
|
||||
rule.updated_at.isoformat(),
|
||||
)
|
||||
|
||||
async def record_observation(
|
||||
self,
|
||||
rule: UsageAlertRule,
|
||||
entity_id: str | None,
|
||||
day: date,
|
||||
observed: float,
|
||||
baseline: float,
|
||||
threshold: float,
|
||||
breached: bool,
|
||||
) -> UsageAlert | None:
|
||||
observation: Final = UsageAlertObservation(
|
||||
entity_id=entity_id, day=day, observed=observed, baseline=baseline, threshold=threshold, breached=breached
|
||||
)
|
||||
changed: Final = await self.record_observations(rule, (observation,))
|
||||
if changed:
|
||||
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 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
|
||||
|
||||
async def record_observations(
|
||||
self, rule: UsageAlertRule, observations: tuple[UsageAlertObservation, ...]
|
||||
) -> tuple[UsageAlert, ...]:
|
||||
|
|
@ -916,50 +834,6 @@ class UsageAlertsRepository:
|
|||
)
|
||||
return tuple(UsageAlert.model_validate(row.record) for row in rows)
|
||||
|
||||
async def pending_notifications(self, limit: int = 100) -> tuple[UsageAlertNotification, ...]:
|
||||
rows: Final = TypeAdapter(tuple[_Pending, ...]).validate_python(
|
||||
await self.database.query_raw(
|
||||
"""SELECT i.record,r.record AS rule,i.notified_channels FROM "LiteLLM_UsageAlertIncident" i JOIN "LiteLLM_UsageAlertRule" r ON r.id=i.rule_id
|
||||
WHERE i.status='open' AND r.enabled AND NOT r.deleted
|
||||
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)
|
||||
AND COALESCE((i.record->>'next_notification_attempt')::timestamptz,'-infinity'::timestamptz)<=CURRENT_TIMESTAMP
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM jsonb_array_elements_text(r.record->'notification_channels') channel WHERE NOT (channel=ANY(i.notified_channels)))
|
||||
ORDER BY i.updated_at,i.id LIMIT $1""",
|
||||
limit,
|
||||
)
|
||||
)
|
||||
return tuple(
|
||||
UsageAlertNotification(
|
||||
incident=UsageAlert.model_validate(row.record),
|
||||
rule=UsageAlertRule.model_validate(row.rule),
|
||||
pending_channels=tuple(
|
||||
channel
|
||||
for channel in UsageAlertRule.model_validate(row.rule).notification_channels
|
||||
if channel not in row.notified_channels
|
||||
),
|
||||
)
|
||||
for row in rows
|
||||
)
|
||||
|
||||
async def mark_notified(self, incident_id: str, channel: UsageAlertChannel) -> None:
|
||||
await self.database.query_raw(
|
||||
"""UPDATE "LiteLLM_UsageAlertIncident" i SET
|
||||
notified_channels=array_append(i.notified_channels,$2::text),
|
||||
record=i.record || jsonb_build_object('notified_channels',array_append(i.notified_channels,$2::text), 'last_notification_attempt',CURRENT_TIMESTAMP) ||
|
||||
CASE WHEN NOT EXISTS (
|
||||
SELECT 1 FROM "LiteLLM_UsageAlertRule" r,
|
||||
jsonb_array_elements_text(r.record->'notification_channels') channel
|
||||
WHERE r.id=i.rule_id AND NOT (channel=ANY(array_append(i.notified_channels,$2::text)))
|
||||
) THEN jsonb_build_object('notification_error',NULL,'next_notification_attempt',NULL)
|
||||
ELSE '{}'::jsonb END
|
||||
WHERE i.id=$1 AND NOT ($2::text=ANY(i.notified_channels)) RETURNING i.id""",
|
||||
incident_id,
|
||||
channel,
|
||||
)
|
||||
|
||||
async def prune(self, retention_days: int = 90, limit: int = 500) -> None:
|
||||
await self.database.query_raw(
|
||||
"""DELETE FROM "LiteLLM_UsageAlertIncident" WHERE id IN (SELECT id FROM "LiteLLM_UsageAlertIncident" WHERE status<>'open' AND updated_at<CURRENT_TIMESTAMP-($1::int * INTERVAL '1 day') ORDER BY updated_at LIMIT $2) RETURNING id""",
|
||||
|
|
|
|||
|
|
@ -216,13 +216,6 @@ class UsageAlertModelList(BaseModel, frozen=True):
|
|||
end_date: date
|
||||
|
||||
|
||||
class UsageAlertNotification(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
incident: UsageAlert
|
||||
rule: UsageAlertRule
|
||||
pending_channels: tuple[UsageAlertChannel, ...]
|
||||
|
||||
|
||||
class UsageAlertObservation(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, allow_inf_nan=False)
|
||||
entity_id: str | None
|
||||
|
|
|
|||
|
|
@ -27,17 +27,19 @@ from litellm.types.proxy.issues import (
|
|||
IssueActivityResponse,
|
||||
IssueCohort,
|
||||
IssueComparison,
|
||||
IssueDeliveryBinding,
|
||||
IssueDeliveryInput,
|
||||
IssueFamilyPolicyUpdate,
|
||||
IssueFamilyUpdate,
|
||||
IssueInput,
|
||||
IssueList,
|
||||
IssueModelPair,
|
||||
IssueSignal,
|
||||
IssueSignalInput,
|
||||
IssueSignalList,
|
||||
)
|
||||
from litellm.types.proxy.usage_alert_defaults import default_usage_rule_id, default_usage_rules
|
||||
from litellm.types.proxy.usage_alerts import UsageAlert, UsageAlertObservation, UsageAlertRuleInput
|
||||
from litellm.types.proxy.usage_alerts import UsageAlert, UsageAlertObservation, UsageAlertRule, UsageAlertRuleInput
|
||||
|
||||
|
||||
def _issue(cohort: IssueCohort) -> IssueInput:
|
||||
|
|
@ -71,6 +73,16 @@ def _issue(cohort: IssueCohort) -> IssueInput:
|
|||
)
|
||||
|
||||
|
||||
def _signals(issue_id: str, inputs: tuple[IssueSignalInput, ...]) -> tuple[IssueSignal, ...]:
|
||||
return tuple(IssueSignal.model_validate({**signal.model_dump(), "issue_id": issue_id}) for signal in inputs)
|
||||
|
||||
|
||||
def _delivery(data: IssueDeliveryInput, rule: UsageAlertRule, source: UsageAlert) -> IssueDeliveryBinding:
|
||||
return IssueDeliveryBinding(
|
||||
delivery=data, rule_id=rule.id, rule_revision=rule.updated_at.isoformat(), source_id=source.id
|
||||
)
|
||||
|
||||
|
||||
async def _query(db: ManagementDatabase, query: str, *args: object) -> object:
|
||||
return await IssuesRepository.from_prisma(db.prisma).database.query_raw(query, *args)
|
||||
|
||||
|
|
@ -141,18 +153,20 @@ async def test_issue_http_scope_pagination_and_source_visibility(gateway: Gatewa
|
|||
"baseline_value": 9999,
|
||||
}
|
||||
)
|
||||
await store.add_signals(
|
||||
first.id,
|
||||
(
|
||||
IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),
|
||||
IssueSignalInput(
|
||||
source="usage_alert",
|
||||
source_id=broader.id,
|
||||
usage_alert=broader,
|
||||
grouping_reason="dominant_key",
|
||||
explanation="Parent contribution evidence.",
|
||||
await store.add_signal_batch(
|
||||
_signals(
|
||||
first.id,
|
||||
(
|
||||
IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),
|
||||
IssueSignalInput(
|
||||
source="usage_alert",
|
||||
source_id=broader.id,
|
||||
usage_alert=broader,
|
||||
grouping_reason="dominant_key",
|
||||
explanation="Parent contribution evidence.",
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
path: Final = f"/issues/{first.id}"
|
||||
own: Final = gateway.request("GET", path, key=owner_key)
|
||||
|
|
@ -312,7 +326,7 @@ async def test_issue_activity_bounds_models_without_losing_totals(gateway: Gatew
|
|||
model_count,
|
||||
)
|
||||
try:
|
||||
await store.upsert_issue(issue)
|
||||
await store.upsert_issues((issue,))
|
||||
scope: Final = issue_activity_scope(cohort, day, day + timedelta(days=2))
|
||||
original: Final = await get_daily_activity_aggregated(repository.activity, scope, api_key_limit=20)
|
||||
response: Final = gateway.request("GET", f"/issues/{issue.id}/activity")
|
||||
|
|
@ -425,7 +439,7 @@ async def test_issue_provider_pairs_filter_activity_pages_counts_and_exports(
|
|||
store: Final = IssuesRepository.from_prisma(db.prisma)
|
||||
issue: Final = _issue(cohort).model_copy(update={"range_start": day, "range_end": day})
|
||||
try:
|
||||
await store.upsert_issue(issue)
|
||||
await store.upsert_issues((issue,))
|
||||
activity_http: Final = gateway.request("GET", f"/issues/{issue.id}/activity")
|
||||
assert activity_http.status_code == 200 and activity_http.json()["metadata"]["total_api_requests"] == 100
|
||||
keys_http: Final = gateway.request("GET", f"/issues/{issue.id}/keys", params={"limit": "1"})
|
||||
|
|
@ -503,7 +517,7 @@ async def test_issue_delivery_guards_snapshot_freeze_and_delayed_recovery(db: Ma
|
|||
next_issue_id: Final = str(uuid4())
|
||||
try:
|
||||
source: Final = (await alerts.record_observations(rule, (observation,)))[0]
|
||||
initial: Final = await store.upsert_issue(original)
|
||||
initial: Final = (await store.upsert_issues((original,)))[0]
|
||||
assert tuple(
|
||||
item.alert.id
|
||||
for item in await store.active_signals(
|
||||
|
|
@ -512,11 +526,11 @@ async def test_issue_delivery_guards_snapshot_freeze_and_delayed_recovery(db: Ma
|
|||
) == (source.id,)
|
||||
assert await store.active_signals(original.observed_on, scope_ids=()) == ()
|
||||
assert await store.active_signals(original.observed_on, scope_ids=("other-example-user",)) == ()
|
||||
await store.add_signals(
|
||||
initial.id, (IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),)
|
||||
await store.add_signal_batch(
|
||||
_signals(initial.id, (IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),))
|
||||
)
|
||||
await store.queue_delivery(delivery, rule.id, rule.updated_at.isoformat(), source.id)
|
||||
await store.queue_delivery(invalid, rule.id, rule.updated_at.isoformat(), source.id)
|
||||
await store.queue_deliveries((_delivery(delivery, rule, source),))
|
||||
await store.queue_deliveries((_delivery(invalid, rule, source),))
|
||||
pending: Final = await store.pending_delivery(delivery.id)
|
||||
assert pending is not None and pending.alert.id == source.id and pending.issue.id == initial.id
|
||||
assert await store.pending_delivery(invalid.id) is None
|
||||
|
|
@ -526,7 +540,7 @@ async def test_issue_delivery_guards_snapshot_freeze_and_delayed_recovery(db: Ma
|
|||
assert enabled is not None and await store.pending_delivery(delivery.id) is None
|
||||
reevaluated: Final = await alerts.record_observations(enabled, (observation,))
|
||||
assert len(reevaluated) == 1 and reevaluated[0].id == source.id
|
||||
await store.queue_delivery(delivery, enabled.id, enabled.updated_at.isoformat(), source.id)
|
||||
await store.queue_deliveries((_delivery(delivery, enabled, source),))
|
||||
assert await store.pending_delivery(delivery.id) is not None
|
||||
await alerts.set_status(source.id, "dismissed")
|
||||
assert await store.pending_delivery(delivery.id) is None
|
||||
|
|
@ -541,7 +555,7 @@ async def test_issue_delivery_guards_snapshot_freeze_and_delayed_recovery(db: Ma
|
|||
"comparisons": (original.comparisons[0].model_copy(update={"observed_value": 999}),),
|
||||
}
|
||||
)
|
||||
continued: Final = await store.upsert_issue(later)
|
||||
continued: Final = (await store.upsert_issues((later,)))[0]
|
||||
assert continued.status == "dismissed" and continued.comparisons == initial.comparisons
|
||||
assert continued.title == initial.title and continued.observed_on == initial.observed_on
|
||||
assert continued.last_seen > initial.last_seen
|
||||
|
|
@ -558,7 +572,7 @@ async def test_issue_delivery_guards_snapshot_freeze_and_delayed_recovery(db: Ma
|
|||
retried is not None and retried.delivery.attempts == 1 and retried.delivery.last_error == "delivery_failed"
|
||||
)
|
||||
await store.mark_delivery_sent(delivery.id)
|
||||
await store.queue_delivery(delivery, enabled.id, enabled.updated_at.isoformat(), source.id)
|
||||
await store.queue_deliveries((_delivery(delivery, enabled, source),))
|
||||
assert await store.pending_delivery(delivery.id) is None
|
||||
healthy: Final = observation.model_copy(
|
||||
update={"day": observation.day + timedelta(days=1), "breached": False, "observed": 120}
|
||||
|
|
@ -567,24 +581,28 @@ async def test_issue_delivery_guards_snapshot_freeze_and_delayed_recovery(db: Ma
|
|||
await store.recover_issues(observation.day + timedelta(days=2))
|
||||
recovered: Final = await store.get_issue(initial.id)
|
||||
assert recovered is not None and not recovered.active and recovered.status == "resolved"
|
||||
assert await store.find_active_issue(original.cohort_key) is None
|
||||
assert await store.find_active_issues((original.cohort_key,)) == ()
|
||||
repeated: Final = (
|
||||
await alerts.record_observations(
|
||||
enabled, (observation.model_copy(update={"day": observation.day + timedelta(days=2)}),)
|
||||
)
|
||||
)[0]
|
||||
assert repeated.id != source.id
|
||||
new_episode: Final = await store.upsert_issue(
|
||||
original.model_copy(
|
||||
update={
|
||||
"id": next_issue_id,
|
||||
"observed_on": repeated.day,
|
||||
"observed_at": later.observed_at + timedelta(minutes=1),
|
||||
}
|
||||
new_episode: Final = (
|
||||
await store.upsert_issues(
|
||||
(
|
||||
original.model_copy(
|
||||
update={
|
||||
"id": next_issue_id,
|
||||
"observed_on": repeated.day,
|
||||
"observed_at": later.observed_at + timedelta(minutes=1),
|
||||
}
|
||||
),
|
||||
)
|
||||
)
|
||||
)
|
||||
)[0]
|
||||
assert new_episode.active and new_episode.status == "open" and new_episode.id != initial.id
|
||||
assert (await store.find_active_issue(original.cohort_key)) == new_episode
|
||||
assert await store.find_active_issues((original.cohort_key,)) == (new_episode,)
|
||||
finally:
|
||||
await _remove_issues(db, (original.id, next_issue_id))
|
||||
await alerts.delete_rule(rule.id)
|
||||
|
|
@ -664,14 +682,11 @@ async def test_late_default_override_stops_queued_delivery_and_recovers_on_fresh
|
|||
await store.upsert_issues(issues)
|
||||
for issue in issues:
|
||||
source: UsageAlert = next(source for source in sources if source.scope_id == issue.cohort.scope_id)
|
||||
await store.add_signals(
|
||||
issue.id, (IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),)
|
||||
await store.add_signal_batch(
|
||||
_signals(issue.id, (IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),))
|
||||
)
|
||||
await store.queue_delivery(
|
||||
IssueDeliveryInput(id=issue.id, issue_id=issue.id, channel="slack"),
|
||||
rule.id,
|
||||
rule.updated_at.isoformat(),
|
||||
source.id,
|
||||
await store.queue_deliveries(
|
||||
(_delivery(IssueDeliveryInput(id=issue.id, issue_id=issue.id, channel="slack"), rule, source),)
|
||||
)
|
||||
assert await store.pending_delivery(issue.id) is not None
|
||||
definition: Final = UsageAlertRuleInput(
|
||||
|
|
@ -796,7 +811,7 @@ async def test_family_edit_requires_a_new_evaluation(db: ManagementDatabase) ->
|
|||
await alerts.ensure_default_rules(default_usage_rules())
|
||||
rule: Final = await alerts.get_rule(default_usage_rule_id("key", "spend", "increase"))
|
||||
assert rule is not None
|
||||
await alerts.record_evaluation_status(rule, date.today() - timedelta(days=1), "complete")
|
||||
await alerts.record_evaluations(((rule, "complete"),), date.today() - timedelta(days=1), completed=True)
|
||||
result: Final = await store.update_family(
|
||||
"spend", IssueFamilyUpdate(policies=(IssueFamilyPolicyUpdate(direction="increase", multiplier=5),))
|
||||
)
|
||||
|
|
@ -826,7 +841,7 @@ async def test_deleted_key_owner_retains_historical_issue_access(gateway: Gatewa
|
|||
owner,
|
||||
)
|
||||
try:
|
||||
await store.upsert_issue(issue)
|
||||
await store.upsert_issues((issue,))
|
||||
path: Final = f"/issues/{issue.id}"
|
||||
detail: Final = gateway.request("GET", path, key=owner_key)
|
||||
assert detail.status_code == 200, "Deleted key owner must retain historical issue access"
|
||||
|
|
@ -884,23 +899,23 @@ async def test_issue_status_rechecks_changed_scope_set_and_caps_sources(db: Mana
|
|||
),
|
||||
)
|
||||
)[0]
|
||||
await store.upsert_issue(issue)
|
||||
await store.add_signals(
|
||||
issue.id, (IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),)
|
||||
await store.upsert_issues((issue,))
|
||||
await store.add_signal_batch(
|
||||
_signals(issue.id, (IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),))
|
||||
)
|
||||
authorized: Final = await store.status_scopes(issue.id)
|
||||
assert authorized == (issue.cohort,)
|
||||
assert await store.set_status(issue.id, "dismissed", "open", authorized_scopes=authorized) is not None
|
||||
context: Final = source.model_copy(update={"id": str(uuid4()), "metric": "spend"})
|
||||
same_scope: Final = (IssueSignalInput(source="usage_alert", source_id=context.id, usage_alert=context),)
|
||||
await store.add_signals(issue.id, same_scope)
|
||||
await store.add_signal_batch(_signals(issue.id, same_scope))
|
||||
unchanged: Final = await store.get_issue(issue.id)
|
||||
assert unchanged is not None and unchanged.status == "dismissed"
|
||||
await store.add_signals(issue.id, same_scope)
|
||||
await store.add_signal_batch(_signals(issue.id, same_scope))
|
||||
assert await store.get_issue(issue.id) == unchanged
|
||||
broader: Final = source.model_copy(update={"id": str(uuid4()), "scope_type": "deployment", "scope_id": None})
|
||||
await store.add_signals(
|
||||
issue.id, (IssueSignalInput(source="usage_alert", source_id=broader.id, usage_alert=broader),)
|
||||
await store.add_signal_batch(
|
||||
_signals(issue.id, (IssueSignalInput(source="usage_alert", source_id=broader.id, usage_alert=broader),))
|
||||
)
|
||||
assert await store.set_status(issue.id, "dismissed", "open", authorized_scopes=authorized) is None
|
||||
current: Final = await store.get_issue(issue.id)
|
||||
|
|
@ -908,9 +923,11 @@ async def test_issue_status_rechecks_changed_scope_set_and_caps_sources(db: Mana
|
|||
all_scopes: Final = await store.status_scopes(issue.id)
|
||||
assert await store.set_status(issue.id, "dismissed", "open", authorized_scopes=all_scopes) is not None
|
||||
broader_context: Final = broader.model_copy(update={"id": str(uuid4()), "metric": "spend"})
|
||||
await store.add_signals(
|
||||
issue.id,
|
||||
(IssueSignalInput(source="usage_alert", source_id=broader_context.id, usage_alert=broader_context),),
|
||||
await store.add_signal_batch(
|
||||
_signals(
|
||||
issue.id,
|
||||
(IssueSignalInput(source="usage_alert", source_id=broader_context.id, usage_alert=broader_context),),
|
||||
)
|
||||
)
|
||||
repeated_scope: Final = await store.get_issue(issue.id)
|
||||
assert repeated_scope is not None and repeated_scope.status == "dismissed"
|
||||
|
|
@ -943,22 +960,28 @@ async def test_issue_status_rechecks_changed_scope_set_and_caps_sources(db: Mana
|
|||
locked: Final = IssuesRepository(
|
||||
cast(AlertDatabase, transaction) # cast-ok: generated Prisma transaction query delegate
|
||||
)
|
||||
await locked.add_signals(
|
||||
issue.id, (IssueSignalInput(source="usage_alert", source_id=racing.id, usage_alert=racing),)
|
||||
await locked.add_signal_batch(
|
||||
_signals(
|
||||
issue.id, (IssueSignalInput(source="usage_alert", source_id=racing.id, usage_alert=racing),)
|
||||
)
|
||||
)
|
||||
assert update.result() is None
|
||||
raced: Final = await store.get_issue(issue.id)
|
||||
assert raced is not None and raced.status == "open"
|
||||
await store.add_signals(
|
||||
issue.id,
|
||||
tuple(
|
||||
IssueSignalInput(
|
||||
source="usage_alert",
|
||||
source_id=f"scope-{index}",
|
||||
usage_alert=source.model_copy(update={"id": f"scope-{index}", "scope_id": f"other-user-{index}"}),
|
||||
)
|
||||
for index in range(20)
|
||||
),
|
||||
await store.add_signal_batch(
|
||||
_signals(
|
||||
issue.id,
|
||||
tuple(
|
||||
IssueSignalInput(
|
||||
source="usage_alert",
|
||||
source_id=f"scope-{index}",
|
||||
usage_alert=source.model_copy(
|
||||
update={"id": f"scope-{index}", "scope_id": f"other-user-{index}"}
|
||||
),
|
||||
)
|
||||
for index in range(20)
|
||||
),
|
||||
)
|
||||
)
|
||||
limited: Final = await store.status_scopes(issue.id)
|
||||
assert len(limited) == 11
|
||||
|
|
@ -1015,12 +1038,15 @@ async def test_resource_recovery_ignores_broader_and_traffic_context(db: Managem
|
|||
if rule is not None
|
||||
]
|
||||
)
|
||||
await store.upsert_issue(issue)
|
||||
await store.add_signals(
|
||||
issue.id,
|
||||
tuple(
|
||||
IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source) for source in sources
|
||||
),
|
||||
await store.upsert_issues((issue,))
|
||||
await store.add_signal_batch(
|
||||
_signals(
|
||||
issue.id,
|
||||
tuple(
|
||||
IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source)
|
||||
for source in sources
|
||||
),
|
||||
)
|
||||
)
|
||||
primary: Final = rules[0]
|
||||
assert primary is not None
|
||||
|
|
@ -1127,12 +1153,11 @@ async def test_issue_prune_removes_expired_children_in_bounded_passes(
|
|||
),
|
||||
)
|
||||
)[0]
|
||||
await store.add_signals(
|
||||
expired.id,
|
||||
(IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),),
|
||||
await store.add_signal_batch(
|
||||
_signals(expired.id, (IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source),))
|
||||
)
|
||||
for delivery in deliveries:
|
||||
await store.queue_delivery(delivery, rule.id, rule.updated_at.isoformat(), source.id)
|
||||
await store.queue_deliveries((_delivery(delivery, rule, source),))
|
||||
assert await store.set_status(expired.id, "resolved", "open") is not None
|
||||
assert await store.set_status(recent.id, "resolved", "open") is not None
|
||||
await _query(
|
||||
|
|
@ -1231,7 +1256,7 @@ async def test_issue_resource_metadata_bounds_labels_and_deleted_history(
|
|||
day.isoformat(),
|
||||
reporting_key,
|
||||
)
|
||||
await store.upsert_issue(issue)
|
||||
await store.upsert_issues((issue,))
|
||||
keys: Final = await repository.keys(issue.cohort, day, day, 1, 0)
|
||||
assert keys.total_api_keys == 1 and len(keys.api_keys) == 1
|
||||
key: Final = keys.api_keys[0]
|
||||
|
|
@ -1293,12 +1318,15 @@ async def test_disabled_later_breach_requires_new_healthy_primary_evidence(db: M
|
|||
)
|
||||
try:
|
||||
sources: Final = tuple([(await alerts.record_observations(rule, (breach,)))[0] for rule in (first, second)])
|
||||
await store.upsert_issue(issue)
|
||||
await store.add_signals(
|
||||
issue.id,
|
||||
tuple(
|
||||
IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source) for source in sources
|
||||
),
|
||||
await store.upsert_issues((issue,))
|
||||
await store.add_signal_batch(
|
||||
_signals(
|
||||
issue.id,
|
||||
tuple(
|
||||
IssueSignalInput(source="usage_alert", source_id=source.id, usage_alert=source)
|
||||
for source in sources
|
||||
),
|
||||
)
|
||||
)
|
||||
await alerts.record_observations(
|
||||
first,
|
||||
|
|
|
|||
|
|
@ -397,7 +397,7 @@ async def test_same_day_rule_revision_requires_a_fresh_observation(db: Managemen
|
|||
await alerts.set_status(original.alert.id, "dismissed")
|
||||
revised: Final = await alerts.save_rule(definition.model_copy(update={"multiplier": 5}), None, existing=rule)
|
||||
assert revised is not None
|
||||
await alerts.finish_day(revised, _DAY)
|
||||
await alerts.record_evaluations(((revised, "complete"),), _DAY, completed=True)
|
||||
assert await issues.active_signals(_DAY) == ()
|
||||
await alerts.record_observations(revised, (observation.model_copy(update={"observed": 6, "threshold": 5}),))
|
||||
updated: Final = (await issues.active_signals(_DAY))[0]
|
||||
|
|
@ -572,7 +572,11 @@ async def test_missing_drop_coverage_does_not_block_established_spend_issue(db:
|
|||
)
|
||||
job: Final = UsageAlertJob.from_prisma(prisma, SlackAlerting(alerting=[], alert_types=[]))
|
||||
await job._run(_NOW)
|
||||
rules: Final = await UsageAlertsRepository.from_prisma(prisma).list_enabled_rules(limit=100)
|
||||
rules: Final = tuple(
|
||||
rule
|
||||
for rule in (await UsageAlertsRepository.from_prisma(prisma).list_rules(limit=100)).data
|
||||
if rule.enabled
|
||||
)
|
||||
assert any(rule.direction == "decrease" and rule.evaluation_status == "insufficient_coverage" for rule in rules)
|
||||
issues: Final = await IssuesRepository.from_prisma(prisma).list_issues()
|
||||
assert any(issue.family == "spend" and issue.cohort.scope_id == _KEY_A for issue in issues.data)
|
||||
|
|
@ -608,7 +612,7 @@ async def test_entity_page_progress_resumes_old_day_without_early_completion(db:
|
|||
partial: Final = await job._store.get_rule(rule.id)
|
||||
assert partial is not None and partial.evaluation_status == "delayed"
|
||||
assert partial.evaluation_progress is not None and partial.evaluation_progress.cursor is not None
|
||||
assert not await job._store.is_day_complete(rule, _DAY)
|
||||
assert not await job._store.completed_rule_ids((rule,), _DAY)
|
||||
assert "evaluation_progress" not in partial.model_dump()
|
||||
assert "evaluation_progress" not in UsageAlertRule.model_json_schema(mode="serialization")["properties"]
|
||||
first: Final = await job._store.list_alerts(limit=500)
|
||||
|
|
@ -629,7 +633,7 @@ async def test_entity_page_progress_resumes_old_day_without_early_completion(db:
|
|||
await job._store.finish_projection_work(next_work)
|
||||
assert (await IssuesRepository.from_prisma(prisma).list_issues()).total == 20
|
||||
await job._run(_NOW + timedelta(days=1))
|
||||
assert await job._store.is_day_complete(rule, _DAY)
|
||||
assert await job._store.completed_rule_ids((rule,), _DAY)
|
||||
final: Final = await job._store.list_alerts(limit=500)
|
||||
assert sum(1 for item in final.data if item.rule_id == rule.id) == 31
|
||||
count: Final = await prisma.writer_db.query_raw(
|
||||
|
|
@ -744,7 +748,7 @@ async def test_rule_edit_discards_old_cursor_and_rejects_stale_progress(db: Mana
|
|||
assert current is not None and current.evaluation_progress is None
|
||||
assert current.updated_at == edited.updated_at
|
||||
await store.record_evaluations(((rule, "complete"),), _DAY, completed=True)
|
||||
assert not await store.is_day_complete(edited, _DAY)
|
||||
assert not await store.completed_rule_ids((edited,), _DAY)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
|
|
@ -800,7 +804,7 @@ async def test_partial_broad_default_waits_for_its_key_family_only(db: Managemen
|
|||
assert await projector.run_pending(_NOW, work, job._store, asyncio.get_running_loop().time() + 10)
|
||||
partial: Final = await sources.list_issues()
|
||||
assert len(partial.data) == 1 and partial.data[0].cohort.scope_type == "key"
|
||||
while not await job._store.is_day_complete(spend[0], _DAY):
|
||||
while not await job._store.completed_rule_ids((spend[0],), _DAY):
|
||||
current: Final = await job._store.get_rule(spend[0].id)
|
||||
assert current is not None
|
||||
await job._scope_page((current,), _DAY, True)
|
||||
|
|
@ -820,7 +824,7 @@ async def test_projection_resume_keeps_spend_and_traffic_in_one_episode(db: Mana
|
|||
await _seed(prisma)
|
||||
job: Final = UsageAlertJob.from_prisma(prisma, SlackAlerting(alerting=[], alert_types=[]))
|
||||
await job._store.ensure_default_rules(job._defaults)
|
||||
await job._evaluate(_DAY, asyncio.get_running_loop().time() + 10)
|
||||
await job._evaluate(_DAY, asyncio.get_running_loop().time() + 60)
|
||||
work: Final = await job._store.pending_projection_work()
|
||||
projector: Final = job._performance_projection
|
||||
assert work is not None and projector is not None
|
||||
|
|
@ -851,7 +855,7 @@ async def test_projection_resume_keeps_spend_and_traffic_in_one_episode(db: Mana
|
|||
_NOW + timedelta(hours=1), latest, restarted._store, asyncio.get_running_loop().time() + 10
|
||||
)
|
||||
assert await restarted._performance_projection.run_pending(
|
||||
_NOW + timedelta(hours=2), latest, restarted._store, asyncio.get_running_loop().time() + 10
|
||||
_NOW + timedelta(hours=2), latest, restarted._store, asyncio.get_running_loop().time() + 60
|
||||
)
|
||||
after: Final = await IssuesRepository.from_prisma(prisma).list_issues()
|
||||
assert tuple(item.id for item in after.data) == tuple(item.id for item in before.data)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue