westfarn
29c69b91f7
Rename finance → monetization (keep finance_* tables via app label), add RevenueCat webhook + ledger upserts so store IAP syncs subscriptions and billing history like Stripe. Companion to chat_web_app#100 / #68.
255 lines
8.6 KiB
Python
255 lines
8.6 KiB
Python
"""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 monetization.models import UserSubscription
|
|
from monetization.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)
|