130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
"""Alert and override flows for budget enforcement."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from .shared import Any, Budget, BudgetAlert, Optional, and_, desc, now_cn_naive, select, timedelta
|
|
|
|
|
|
class BudgetEnforcementAlertsMixin:
|
|
"""Hard-limit enforcement, alerts, and override handling."""
|
|
|
|
async def enforce_spending_limit(
|
|
self,
|
|
budget: Budget,
|
|
usage_usd: float,
|
|
usage_percent: float,
|
|
) -> None:
|
|
"""执行超支硬性限制：标记预算为 suspended 并记录告警。"""
|
|
budget.status = "suspended"
|
|
|
|
await self.generate_budget_alert(
|
|
budget=budget,
|
|
alert_type="exceeded",
|
|
usage_usd=usage_usd,
|
|
usage_percent=usage_percent,
|
|
message=f"预算已超支，当前使用率 {usage_percent:.2f}%",
|
|
action_taken="mark_exceeded",
|
|
)
|
|
await self.generate_budget_alert(
|
|
budget=budget,
|
|
alert_type="suspended",
|
|
usage_usd=usage_usd,
|
|
usage_percent=usage_percent,
|
|
message="预算超支，已暂停新的 AI 调用",
|
|
action_taken="suspend_ai_calls",
|
|
)
|
|
await self.db.commit()
|
|
|
|
async def generate_budget_alert(
|
|
self,
|
|
budget: Budget,
|
|
alert_type: str,
|
|
usage_usd: float,
|
|
usage_percent: float,
|
|
message: str,
|
|
action_taken: Optional[str],
|
|
) -> BudgetAlert:
|
|
"""生成预算预警（同类 active 预警存在则复用更新）。"""
|
|
result = await self.db.execute(
|
|
select(BudgetAlert)
|
|
.where(
|
|
and_(
|
|
BudgetAlert.budget_id == budget.id,
|
|
BudgetAlert.alert_type == alert_type,
|
|
BudgetAlert.status == "active",
|
|
)
|
|
)
|
|
.order_by(desc(BudgetAlert.created_at))
|
|
)
|
|
existing = result.scalar_one_or_none()
|
|
if existing:
|
|
existing.usage_usd = usage_usd
|
|
existing.usage_percent = usage_percent
|
|
existing.message = message
|
|
existing.action_taken = action_taken
|
|
await self.db.commit()
|
|
await self.db.refresh(existing)
|
|
return existing
|
|
|
|
alert = BudgetAlert(
|
|
tenant_id=budget.tenant_id,
|
|
budget_id=budget.id,
|
|
alert_type=alert_type,
|
|
usage_usd=usage_usd,
|
|
usage_percent=usage_percent,
|
|
message=message,
|
|
is_read=False,
|
|
action_taken=action_taken,
|
|
status="active",
|
|
)
|
|
self.db.add(alert)
|
|
await self.db.commit()
|
|
await self.db.refresh(alert)
|
|
return alert
|
|
|
|
async def override_budget_limit(
|
|
self,
|
|
budget_id: int,
|
|
tenant_id: int,
|
|
operator_user_id: int,
|
|
note: str = "",
|
|
duration_hours: int = 24,
|
|
) -> dict[str, Any]:
|
|
"""OWNER/REVIEWER 超支解锁。"""
|
|
result = await self.db.execute(
|
|
select(Budget).where(
|
|
and_(
|
|
Budget.id == budget_id,
|
|
Budget.tenant_id == tenant_id,
|
|
)
|
|
)
|
|
)
|
|
budget = result.scalar_one_or_none()
|
|
if not budget:
|
|
raise ValueError(f"预算 {budget_id} 不存在")
|
|
|
|
now = now_cn_naive()
|
|
expires_at = now + timedelta(hours=max(duration_hours, 1))
|
|
history = dict(budget.override_history or {})
|
|
records = list(history.get("records", []))
|
|
records.append(
|
|
{
|
|
"operator_user_id": operator_user_id,
|
|
"note": note,
|
|
"overridden_at": now.isoformat(),
|
|
"expires_at": expires_at.isoformat(),
|
|
}
|
|
)
|
|
history["records"] = records
|
|
budget.override_history = history
|
|
budget.status = "overridden"
|
|
|
|
await self.db.commit()
|
|
await self.db.refresh(budget)
|
|
return {
|
|
"budget_id": budget.id,
|
|
"status": budget.status,
|
|
"expires_at": expires_at.isoformat(),
|
|
"override_count": len(records),
|
|
}
|