142 lines
5.2 KiB
Python
142 lines
5.2 KiB
Python
"""Budget check/update flows."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from .shared import Any, Optional, now_cn_naive
|
|
|
|
|
|
class BudgetEnforcementChecksMixin:
|
|
"""Pre-call checks and post-call usage updates."""
|
|
|
|
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 = now_cn_naive()
|
|
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
|
|
|
|
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 = now_cn_naive()
|
|
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,
|
|
}
|