346 lines
12 KiB
Python
346 lines
12 KiB
Python
"""预算强制执行服务
|
||
|
||
实现：
|
||
- AI 调用前预算检查
|
||
- 超支硬限制执行
|
||
- 预算使用量累计
|
||
- 80% 阈值预警与超支告警
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from datetime import datetime, timedelta
|
||
from typing import Any, Optional
|
||
|
||
from sqlalchemy import and_, desc, select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from ..models import Budget, BudgetAlert
|
||
|
||
|
||
class BudgetEnforcementService:
|
||
"""预算强制执行服务"""
|
||
|
||
def __init__(self, db: AsyncSession):
|
||
self.db = db
|
||
|
||
async def check_budget_before_ai_call(
|
||
self,
|
||
tenant_id: int,
|
||
estimated_cost_usd: float,
|
||
project_id: Optional[int] = None,
|
||
) -> dict[str, Any]:
|
||
"""AI 调用前预算检查"""
|
||
budget = await self._get_active_budget(tenant_id=tenant_id, project_id=project_id)
|
||
if not budget:
|
||
return {
|
||
"allowed": True,
|
||
"reason": "no_budget_configured",
|
||
"budget_checked": False,
|
||
}
|
||
|
||
now = datetime.utcnow()
|
||
await self._ensure_current_period(budget, now)
|
||
|
||
current_usage = float(budget.current_usage_usd or 0.0)
|
||
budget_amount = float(budget.budget_amount_usd or 0.0)
|
||
if budget_amount <= 0:
|
||
return {
|
||
"allowed": True,
|
||
"reason": "invalid_budget_amount",
|
||
"budget_checked": True,
|
||
}
|
||
|
||
if budget.status == "suspended":
|
||
return {
|
||
"allowed": False,
|
||
"reason": "budget_suspended",
|
||
"budget_checked": True,
|
||
"budget_id": budget.id,
|
||
"current_usage_usd": round(current_usage, 4),
|
||
"budget_amount_usd": round(budget_amount, 4),
|
||
}
|
||
|
||
projected_usage = current_usage + max(estimated_cost_usd, 0.0)
|
||
usage_percent = (projected_usage / budget_amount) * 100
|
||
|
||
# 阈值预警（80% 默认）
|
||
if usage_percent >= float(budget.alert_threshold_percent or 80.0):
|
||
await self.generate_budget_alert(
|
||
budget=budget,
|
||
alert_type="threshold_reached",
|
||
usage_usd=projected_usage,
|
||
usage_percent=usage_percent,
|
||
message=(
|
||
f"预算使用率达到 {usage_percent:.2f}% "
|
||
f"(阈值 {float(budget.alert_threshold_percent or 80.0):.2f}%)"
|
||
),
|
||
action_taken=None,
|
||
)
|
||
|
||
# 超支硬拦截
|
||
if (
|
||
projected_usage > budget_amount
|
||
and bool(budget.hard_limit_enabled)
|
||
and budget.status != "overridden"
|
||
):
|
||
await self.enforce_spending_limit(
|
||
budget=budget,
|
||
usage_usd=projected_usage,
|
||
usage_percent=usage_percent,
|
||
)
|
||
return {
|
||
"allowed": False,
|
||
"reason": "budget_exceeded_hard_limit",
|
||
"budget_checked": True,
|
||
"budget_id": budget.id,
|
||
"current_usage_usd": round(current_usage, 4),
|
||
"projected_usage_usd": round(projected_usage, 4),
|
||
"budget_amount_usd": round(budget_amount, 4),
|
||
"usage_percent": round(usage_percent, 2),
|
||
}
|
||
|
||
return {
|
||
"allowed": True,
|
||
"reason": "within_budget",
|
||
"budget_checked": True,
|
||
"budget_id": budget.id,
|
||
"current_usage_usd": round(current_usage, 4),
|
||
"projected_usage_usd": round(projected_usage, 4),
|
||
"budget_amount_usd": round(budget_amount, 4),
|
||
"usage_percent": round(usage_percent, 2),
|
||
}
|
||
|
||
async def update_budget_usage_after_ai_call(
|
||
self,
|
||
tenant_id: int,
|
||
actual_cost_usd: float,
|
||
project_id: Optional[int] = None,
|
||
) -> dict[str, Any]:
|
||
"""AI 调用成功后累计预算使用量"""
|
||
budget = await self._get_active_budget(tenant_id=tenant_id, project_id=project_id)
|
||
if not budget:
|
||
return {"updated": False, "reason": "no_budget_configured"}
|
||
|
||
now = datetime.utcnow()
|
||
await self._ensure_current_period(budget, now)
|
||
|
||
cost = max(float(actual_cost_usd or 0.0), 0.0)
|
||
budget.current_usage_usd = float(budget.current_usage_usd or 0.0) + cost
|
||
|
||
budget_amount = float(budget.budget_amount_usd or 0.0)
|
||
usage_percent = (float(budget.current_usage_usd) / budget_amount * 100) if budget_amount > 0 else 0.0
|
||
|
||
if usage_percent >= float(budget.alert_threshold_percent or 80.0):
|
||
await self.generate_budget_alert(
|
||
budget=budget,
|
||
alert_type="threshold_reached",
|
||
usage_usd=float(budget.current_usage_usd),
|
||
usage_percent=usage_percent,
|
||
message=(
|
||
f"预算使用率达到 {usage_percent:.2f}% "
|
||
f"(阈值 {float(budget.alert_threshold_percent or 80.0):.2f}%)"
|
||
),
|
||
action_taken=None,
|
||
)
|
||
|
||
if budget_amount > 0 and float(budget.current_usage_usd) > budget_amount and bool(budget.hard_limit_enabled):
|
||
await self.enforce_spending_limit(
|
||
budget=budget,
|
||
usage_usd=float(budget.current_usage_usd),
|
||
usage_percent=usage_percent,
|
||
)
|
||
|
||
await self.db.commit()
|
||
await self.db.refresh(budget)
|
||
|
||
return {
|
||
"updated": True,
|
||
"budget_id": budget.id,
|
||
"current_usage_usd": float(budget.current_usage_usd),
|
||
"usage_percent": round(usage_percent, 2),
|
||
"status": budget.status,
|
||
}
|
||
|
||
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]:
|
||
"""BOSS/ADMIN 超支解锁"""
|
||
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 = datetime.utcnow()
|
||
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),
|
||
}
|
||
|
||
async def _get_active_budget(self, tenant_id: int, project_id: Optional[int]) -> Optional[Budget]:
|
||
"""优先读取项目预算，否则回退租户预算"""
|
||
if project_id is not None:
|
||
result = await self.db.execute(
|
||
select(Budget).where(
|
||
and_(
|
||
Budget.tenant_id == tenant_id,
|
||
Budget.budget_type == "project",
|
||
Budget.project_id == project_id,
|
||
)
|
||
).order_by(desc(Budget.created_at))
|
||
)
|
||
project_budget = result.scalar_one_or_none()
|
||
if project_budget:
|
||
return project_budget
|
||
|
||
result = await self.db.execute(
|
||
select(Budget).where(
|
||
and_(
|
||
Budget.tenant_id == tenant_id,
|
||
Budget.budget_type == "tenant",
|
||
)
|
||
).order_by(desc(Budget.created_at))
|
||
)
|
||
return result.scalar_one_or_none()
|
||
|
||
async def _ensure_current_period(self, budget: Budget, now: datetime) -> None:
|
||
"""若预算周期已过，自动切换到新周期并清零使用量"""
|
||
start = budget.current_period_start
|
||
end = budget.current_period_end
|
||
if start is None or end is None or now >= end:
|
||
next_start, next_end = self._compute_period_window(now, budget.period)
|
||
budget.current_period_start = next_start
|
||
budget.current_period_end = next_end
|
||
budget.current_usage_usd = 0.0
|
||
# 周期切换后恢复 active，允许重新消费
|
||
budget.status = "active"
|
||
await self.db.commit()
|
||
|
||
def _compute_period_window(self, now: datetime, period: str) -> tuple[datetime, datetime]:
|
||
period = (period or "monthly").lower()
|
||
if period == "yearly":
|
||
start = datetime(now.year, 1, 1)
|
||
end = datetime(now.year + 1, 1, 1)
|
||
return start, end
|
||
if period == "quarterly":
|
||
quarter = (now.month - 1) // 3
|
||
start_month = quarter * 3 + 1
|
||
start = datetime(now.year, start_month, 1)
|
||
if start_month == 10:
|
||
end = datetime(now.year + 1, 1, 1)
|
||
else:
|
||
end = datetime(now.year, start_month + 3, 1)
|
||
return start, end
|
||
|
||
# default monthly
|
||
start = datetime(now.year, now.month, 1)
|
||
if now.month == 12:
|
||
end = datetime(now.year + 1, 1, 1)
|
||
else:
|
||
end = datetime(now.year, now.month + 1, 1)
|
||
return start, end
|
||
|