"""Prompt-window and token-period quota checks (shared by chat + finance APIs).""" from __future__ import annotations from dataclasses import dataclass from datetime import timedelta from typing import Any from django.conf import settings from django.db.models import Count, Q, Sum from django.utils import timezone from chat_backend.models import PromptMetric from finance.models import UserSubscription from finance.services.plans import get_or_create_user_subscription, seed_subscription_plans class QuotaExceeded(Exception): """Raised when a generation turn is blocked by quota.""" def __init__(self, code: str, message: str, *, details: dict | None = None): super().__init__(message) self.code = code self.message = message self.details = details or {} class FeatureNotAllowed(Exception): """Raised when the user's plan cannot use a feature.""" def __init__(self, code: str, message: str, *, details: dict | None = None): super().__init__(message) self.code = code self.message = message self.details = details or {} @dataclass class UsageSnapshot: prompts_in_window: int prompt_quota: int | None prompts_remaining: int | None window_hours: int tokens_in_period: int | None tokens_out_period: int | None tokens_total_period: int | None turns_missing_token_usage: int monthly_token_quota: int | None tokens_remaining: int | None period_start: Any period_end: Any def to_dict(self) -> dict[str, Any]: return { "prompts_in_window": self.prompts_in_window, "prompt_quota": self.prompt_quota, "prompts_remaining": self.prompts_remaining, "window_hours": self.window_hours, "tokens_in_period": self.tokens_in_period, "tokens_out_period": self.tokens_out_period, "tokens_total_period": self.tokens_total_period, "turns_missing_token_usage": self.turns_missing_token_usage, "monthly_token_quota": self.monthly_token_quota, "tokens_remaining": self.tokens_remaining, "period_start": self.period_start.isoformat() if self.period_start else None, "period_end": self.period_end.isoformat() if self.period_end else None, } def _user_conversation_ids(user) -> list[int]: from chat_backend.models import Conversation return list( Conversation.objects.filter(user=user, deleted=False).values_list("id", flat=True) ) def _billing_period_bounds(): """Calendar-month UTC window for token-period aggregation (#17).""" now = timezone.now() start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) if start.month == 12: end = start.replace(year=start.year + 1, month=1) else: end = start.replace(month=start.month + 1) return start, end def _sum_tokens(qs) -> tuple[int | None, int | None]: """ Sum tokens_in / tokens_out. Returns (None, None) when *no* rows reported usage — never fabricate 0. When some rows reported usage, sum only those (nulls ignored by Sum). """ agg = qs.aggregate( tin=Sum("tokens_in"), tout=Sum("tokens_out"), with_in=Count("id", filter=Q(tokens_in__isnull=False)), with_out=Count("id", filter=Q(tokens_out__isnull=False)), ) tokens_in = agg["tin"] if agg["with_in"] else None tokens_out = agg["tout"] if agg["with_out"] else None return tokens_in, tokens_out def get_usage_snapshot(user) -> UsageSnapshot: seed_subscription_plans(update_existing=False) sub = ( UserSubscription.objects.select_related("plan") .filter(user_id=user.pk) .first() ) if sub is None: sub = get_or_create_user_subscription(user) plan = sub.plan if sub.is_active else None window_hours = plan.prompt_window_hours if plan else 6 prompt_quota = plan.prompt_quota_per_window if plan else None monthly_token_quota = sub.effective_monthly_token_quota() if sub.is_active else None conversation_ids = _user_conversation_ids(user) now = timezone.now() window_start = now - timedelta(hours=window_hours) period_start, period_end = _billing_period_bounds() base = PromptMetric.objects.filter(conversation_id__in=conversation_ids) prompts_in_window = base.filter(created__gte=window_start).count() period_qs = base.filter(created__gte=period_start, created__lt=period_end) tokens_in, tokens_out = _sum_tokens(period_qs) missing = period_qs.filter( Q(tokens_in__isnull=True) | Q(tokens_out__isnull=True) ).count() if tokens_in is None and tokens_out is None: tokens_total = None else: tokens_total = (tokens_in or 0) + (tokens_out or 0) prompts_remaining = None if prompt_quota is not None: prompts_remaining = max(prompt_quota - prompts_in_window, 0) tokens_remaining = None if monthly_token_quota is not None and tokens_total is not None: tokens_remaining = max(monthly_token_quota - tokens_total, 0) elif monthly_token_quota is not None and tokens_total is None: # No provider usage yet — do not treat as 0 consumed. tokens_remaining = monthly_token_quota return UsageSnapshot( prompts_in_window=prompts_in_window, prompt_quota=prompt_quota, prompts_remaining=prompts_remaining, window_hours=window_hours, tokens_in_period=tokens_in, tokens_out_period=tokens_out, tokens_total_period=tokens_total, turns_missing_token_usage=missing, monthly_token_quota=monthly_token_quota, tokens_remaining=tokens_remaining, period_start=period_start, period_end=period_end, ) def assert_feature_allowed(user, feature: str) -> None: if not getattr(settings, "ENFORCE_SUBSCRIPTION_GATES", True): return seed_subscription_plans(update_existing=False) sub = ( UserSubscription.objects.select_related("plan") .filter(user_id=user.pk) .first() ) if sub is None or not sub.is_active or sub.plan is None: raise FeatureNotAllowed( "subscription_required", "An active subscription is required to use this feature.", details={"feature": feature}, ) if not sub.plan.allows_feature(feature): raise FeatureNotAllowed( "feature_not_allowed", f"Your plan ({sub.plan.name}) does not include {feature.replace('_', ' ')}.", details={ "feature": feature, "plan": sub.plan.slug, }, ) def assert_within_quotas(user) -> UsageSnapshot: """ Enforce prompt-window (#36) and token-period (#17) limits. Precedence: either limit may block. Missing provider token usage does not silently under-count toward a token cap — turns with null tokens are tracked in `turns_missing_token_usage` and token-cap enforcement only uses reported sums; if quota is set and usage is entirely unknown, we allow the turn but surface the gap (callers/admin can tighten later). """ seed_subscription_plans(update_existing=False) sub = ( UserSubscription.objects.select_related("plan") .filter(user_id=user.pk) .first() ) if sub is None or not sub.is_active or sub.plan is None: raise QuotaExceeded( "subscription_required", "An active subscription is required before sending prompts.", ) usage = get_usage_snapshot(user) if usage.prompt_quota is not None and usage.prompts_in_window >= usage.prompt_quota: raise QuotaExceeded( "prompt_quota_exceeded", ( f"Prompt limit reached ({usage.prompt_quota} per " f"{usage.window_hours} hours). Try again later." ), details=usage.to_dict(), ) if ( usage.monthly_token_quota is not None and usage.tokens_total_period is not None and usage.tokens_total_period >= usage.monthly_token_quota ): raise QuotaExceeded( "token_quota_exceeded", ( f"Monthly token limit reached ({usage.monthly_token_quota}). " "Upgrade or wait for the next billing period." ), details=usage.to_dict(), ) return usage def check_generation_allowed(user, *, feature: str = "text_generation") -> UsageSnapshot: """Combined feature + quota gate for a chat turn.""" if not getattr(settings, "ENFORCE_SUBSCRIPTION_GATES", True): return get_usage_snapshot(user) assert_feature_allowed(user, feature) return assert_within_quotas(user)