"""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.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: from django.conf import settings 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.""" from django.conf import settings if not getattr(settings, "ENFORCE_SUBSCRIPTION_GATES", True): return get_usage_snapshot(user) assert_feature_allowed(user, feature) return assert_within_quotas(user)