60 lines
2.4 KiB
Python
60 lines
2.4 KiB
Python
"""Persistent-boundary-compatible budget guard; provider invocation happens after this check."""
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from threading import Lock
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BudgetLimits:
|
|
max_run_cost_usd: float = 0.54
|
|
max_daily_cost_usd: float = 1.00
|
|
max_monthly_cost_usd: float = 10.00
|
|
max_requests_per_day: int = 100
|
|
max_new_products_per_run: int = 750
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BudgetDecision:
|
|
allowed: bool
|
|
reason: str
|
|
reserved_cost_usd: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UsageEntry:
|
|
run_id: str
|
|
tenant: str
|
|
cost_usd: float
|
|
requests: int
|
|
new_products: int
|
|
timestamp: str
|
|
|
|
|
|
class BudgetLedger:
|
|
"""Atomic in-process ledger with an explicit persistence boundary for Hub storage."""
|
|
|
|
def __init__(self, limits: BudgetLimits):
|
|
self.limits = limits
|
|
self.entries: list[UsageEntry] = []
|
|
self._lock = Lock()
|
|
|
|
def check_and_reserve(self, run_id: str, tenant: str, cost_usd: float, requests: int, new_products: int) -> BudgetDecision:
|
|
with self._lock:
|
|
run_cost = sum(e.cost_usd for e in self.entries if e.run_id == run_id)
|
|
day = datetime.now(timezone.utc).date()
|
|
day_entries = [e for e in self.entries if e.timestamp[:10] == day.isoformat()]
|
|
month_entries = [e for e in self.entries if e.timestamp[:7] == day.strftime("%Y-%m")]
|
|
if run_cost + cost_usd > self.limits.max_run_cost_usd:
|
|
return BudgetDecision(False, "AI_BUDGET_EXHAUSTED", 0.0)
|
|
if sum(e.cost_usd for e in day_entries) + cost_usd > self.limits.max_daily_cost_usd:
|
|
return BudgetDecision(False, "AI_BUDGET_EXHAUSTED", 0.0)
|
|
if sum(e.cost_usd for e in month_entries) + cost_usd > self.limits.max_monthly_cost_usd:
|
|
return BudgetDecision(False, "AI_BUDGET_EXHAUSTED", 0.0)
|
|
if sum(e.requests for e in day_entries) + requests > self.limits.max_requests_per_day:
|
|
return BudgetDecision(False, "AI_BUDGET_EXHAUSTED", 0.0)
|
|
if sum(e.new_products for e in self.entries if e.run_id == run_id) + new_products > self.limits.max_new_products_per_run:
|
|
return BudgetDecision(False, "AI_BUDGET_EXHAUSTED", 0.0)
|
|
self.entries.append(UsageEntry(run_id, tenant, cost_usd, requests, new_products, datetime.now(timezone.utc).isoformat()))
|
|
return BudgetDecision(True, "RESERVED", cost_usd)
|