"""Author RAG Chatbot SaaS — Token budget helpers. Single source of truth for budget math and post-chat usage updates. """ from app.models.client_access import ClientAccess BUDGET_EXHAUSTED_PREFIX = "budget_exhausted:" TOKENS_REDIS_PREFIX = "tokens:" def total_budget(access: ClientAccess) -> int: """Total available tokens including bonus grants.""" return (access.token_budget or 0) + (access.bonus_tokens or 0) def tokens_remaining(access: ClientAccess) -> int: """Tokens left in the current billing period.""" return max(0, total_budget(access) - (access.tokens_used or 0)) def percent_used(access: ClientAccess) -> float: """Percentage of total budget consumed.""" budget = total_budget(access) if budget <= 0: return 0.0 return round((access.tokens_used or 0) / budget * 100, 1) def usage_summary(access: ClientAccess) -> dict: """Standard token usage dict for API responses.""" budget = total_budget(access) used = access.tokens_used or 0 return { "budget": budget, "used": used, "remaining": max(0, budget - used), "percent_used": percent_used(access), "has_active_subscription": True, "subscription_status": "active", "plan": getattr(access, "plan", None), "expires_at": getattr(access, "expires_at", None) and access.expires_at.isoformat(), } NO_SUBSCRIPTION_SUMMARY: dict = { "budget": 0, "used": 0, "remaining": 0, "percent_used": 0, "has_active_subscription": False, "subscription_status": "none", "plan": None, "expires_at": None, } def usage_summary_or_empty(access: ClientAccess | None) -> dict: """Return usage stats for an active grant, or zeros when unsubscribed.""" if not access: return dict(NO_SUBSCRIPTION_SUMMARY) return usage_summary(access) async def record_token_usage( db, redis, author_id: str, token_count: int, ) -> None: """Increment usage counters and set budget_exhausted when at 100%. Updates client_access.tokens_used (DB) and tokens:{author_id}:current (Redis). """ if token_count <= 0: return from app.repositories.access_repo import AccessRepository token_key = f"{TOKENS_REDIS_PREFIX}{author_id}:current" await redis.incrby(token_key, token_count) await redis.expire(token_key, 32 * 24 * 3600) access_repo = AccessRepository(db) access = await access_repo.get_active_for_author(author_id) if not access: return # R-014/R-052: Atomic SQL UPDATE instead of Python read-modify-write from sqlalchemy import update from app.models.client_access import ClientAccess as CA await db.execute( update(CA) .where(CA.id == access.id, CA.author_id == access.author_id) .values(tokens_used=CA.tokens_used + token_count) ) await db.flush() await db.refresh(access) budget = total_budget(access) if budget > 0 and access.tokens_used >= budget: await redis.set(f"{BUDGET_EXHAUSTED_PREFIX}{author_id}", "1", ex=32 * 24 * 3600) async def clear_budget_exhausted(redis, author_id: str) -> None: """Clear the budget exhausted flag (e.g. after bonus tokens).""" await redis.delete(f"{BUDGET_EXHAUSTED_PREFIX}{author_id}") async def clear_token_usage_counters(redis, author_id: str) -> None: """Clear Redis usage counters after SuperAdmin reset.""" await redis.delete( f"{BUDGET_EXHAUSTED_PREFIX}{author_id}", f"{TOKENS_REDIS_PREFIX}{author_id}:current", ) # ── Budget Warning (B8 fix) ─────────────────────────────────────────────────── _BUDGET_WARN_PREFIX = "budget_warned:" async def check_and_warn_budget(redis, access: ClientAccess, author) -> None: """Send a one-time warning email when an author hits 80% token usage. B8 fix: Previously authors had zero warning before their chatbot stopped working. This function checks usage after each chat turn and sends a single email when the threshold is crossed. Redis prevents duplicate emails within the same grant period. Gracefully skips if: - Usage is below threshold - Email was already sent this grant period - Email service is not configured (SMTP not set) Args: redis: Async Redis connection. access: Active ClientAccess grant object. author: Author User object (for email address). """ from app.config import get_settings cfg = get_settings() pct = percent_used(access) if pct < cfg.TOKEN_WARNING_THRESHOLD_PCT * 100: return warn_key = f"{_BUDGET_WARN_PREFIX}{access.id}" already_warned = await redis.get(warn_key) if already_warned: return # Already sent warning this grant period if not getattr(author, "notify_token_alerts", True): return try: from datetime import datetime, timezone from app.services.email_service import EmailService budget = total_budget(access) used = access.tokens_used or 0 remaining = tokens_remaining(access) days_left = 30 if access.expires_at: days_left = max((access.expires_at - datetime.now(timezone.utc)).days, 1) daily = used / max(30 - days_left, 1) if used else 0 forecast = "soon" if daily > 0 and remaining / daily < 7 else f"in ~{days_left} days" email_svc = EmailService() email_svc.send_token_budget_warning( to=author.email, name=author.full_name or author.email, pct_used=int(round(pct, 0)), tokens_used=used, tokens_total=budget, forecast_date=forecast, ) except Exception as e: import structlog structlog.get_logger(__name__).warning( "Budget warning email failed — skipping (non-fatal)", error=str(e) ) return from datetime import datetime, timezone remaining_secs = 86400 if access.expires_at: delta = (access.expires_at - datetime.now(timezone.utc)).total_seconds() remaining_secs = max(3600, int(delta)) await redis.setex(warn_key, remaining_secs, "1") try: from app.tasks.webhook_delivery_task import deliver_webhook_event deliver_webhook_event.delay( author.id, "low_budget_warning", { "percent_used": round(pct, 1), "tokens_used": access.tokens_used or 0, "tokens_total": budget, }, ) except Exception: pass