Add multi-plan subscriptions, quotas, and token usage APIs
Implements #16/#17/#36: Founders/Standard/Pro/Business/Backer catalog, Backer email whitelist, prompt-window + token-period gates, and tokens_in/out on conversation/prompt + subscription usage APIs.
This commit is contained in:
@@ -0,0 +1,257 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user