Files
chat_backend/llm_be/finance/services/plans.py
T
westfarn a6c45b0882
CI / test (pull_request) Successful in 10s
Unit Tests / test (pull_request) Successful in 10s
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.
2026-07-31 06:21:54 -05:00

259 lines
7.9 KiB
Python

"""Subscription plan catalog helpers, seeding, and user assignment."""
from __future__ import annotations
import logging
from typing import Any
from django.db import transaction
from django.utils import timezone
from finance.models import BackerEmail, SubscriptionPlan, UserSubscription
logger = logging.getLogger(__name__)
# Seed catalog for #36. Standard/Pro/Business stay hidden until explicitly enabled.
PLAN_SEED: list[dict[str, Any]] = [
{
"slug": SubscriptionPlan.Slug.FOUNDERS,
"name": "Founders",
"description": (
"Unlimited product access for early supporters: text plus all future "
"capabilities as they ship. $10/mo."
),
"price_cents": 1000,
"is_public": True,
"is_selectable": True,
"allows_text_generation": True,
"allows_image_generation": True,
"allows_all_future_features": True,
"prompt_quota_per_window": 300,
"prompt_window_hours": 6,
"monthly_token_quota": None,
"sort_order": 10,
},
{
"slug": SubscriptionPlan.Slug.STANDARD,
"name": "Standard",
"description": (
"Secure conversational chat and coding assistance for developers "
"and privacy-conscious individuals."
),
"price_cents": 1500,
"is_public": False,
"is_selectable": False,
"allows_text_generation": True,
"allows_image_generation": False,
"allows_all_future_features": False,
"prompt_quota_per_window": 100,
"prompt_window_hours": 6,
"monthly_token_quota": 1_000_000,
"sort_order": 20,
},
{
"slug": SubscriptionPlan.Slug.PRO,
"name": "Pro / Creator",
"description": (
"Higher message caps and multi-modal workflows for heavy users, "
"including image generation when available."
),
"price_cents": 4000,
"is_public": False,
"is_selectable": False,
"allows_text_generation": True,
"allows_image_generation": True,
"allows_all_future_features": False,
"prompt_quota_per_window": 200,
"prompt_window_hours": 6,
"monthly_token_quota": 3_000_000,
"sort_order": 30,
},
{
"slug": SubscriptionPlan.Slug.BUSINESS,
"name": "Business Team",
"description": (
"Team seats, centralized auth, priority support, and absolute data "
"privacy for local companies handling sensitive data."
),
"price_cents": 9900,
"is_public": False,
"is_selectable": False,
"allows_text_generation": True,
"allows_image_generation": True,
"allows_all_future_features": False,
"prompt_quota_per_window": 300,
"prompt_window_hours": 6,
"monthly_token_quota": 5_000_000,
"sort_order": 40,
},
{
"slug": SubscriptionPlan.Slug.BACKER,
"name": "Backer",
"description": (
"Complimentary Founders-level access for pre-approved emails. "
"Not shown at checkout."
),
"price_cents": 0,
"is_public": False,
"is_selectable": False,
"allows_text_generation": True,
"allows_image_generation": True,
"allows_all_future_features": True,
"prompt_quota_per_window": 300,
"prompt_window_hours": 6,
"monthly_token_quota": None,
"sort_order": 5,
},
]
def seed_subscription_plans(*, update_existing: bool = True) -> list[SubscriptionPlan]:
"""Idempotently create/update the canonical plan catalog."""
plans: list[SubscriptionPlan] = []
for row in PLAN_SEED:
slug = row["slug"]
defaults = {k: v for k, v in row.items() if k != "slug"}
plan, created = SubscriptionPlan.objects.get_or_create(
slug=slug,
defaults=defaults,
)
if not created and update_existing:
for key, value in defaults.items():
setattr(plan, key, value)
plan.save()
plans.append(plan)
return plans
def get_plan(slug: str) -> SubscriptionPlan | None:
return SubscriptionPlan.objects.filter(slug=slug).first()
def get_or_create_user_subscription(user) -> UserSubscription:
sub, _ = UserSubscription.objects.get_or_create(user=user)
return sub
def assign_plan(
user,
*,
plan: SubscriptionPlan,
source: str,
status: str = UserSubscription.Status.ACTIVE,
stripe_subscription_id: str = "",
) -> UserSubscription:
sub = get_or_create_user_subscription(user)
sub.plan = plan
sub.source = source
sub.status = status
if stripe_subscription_id:
sub.stripe_subscription_id = stripe_subscription_id
sub.save()
return sub
def user_has_active_plan(user) -> bool:
try:
sub = user.subscription
except UserSubscription.DoesNotExist:
return False
return sub.is_active
def needs_checkout(user) -> bool:
"""True when the user must complete paid Checkout to use the product."""
try:
sub = user.subscription
except UserSubscription.DoesNotExist:
return True
if not sub.is_active:
return True
# Complimentary / already-paid tiers skip Checkout.
if sub.source in (
UserSubscription.Source.BACKER,
UserSubscription.Source.ADMIN,
UserSubscription.Source.STRIPE,
):
return False
return True
@transaction.atomic
def try_redeem_backer_email(user) -> UserSubscription | None:
"""
If the user's email is on the Backer whitelist and unused, assign Backer plan.
Returns the UserSubscription when redeemed, else None.
"""
email = (getattr(user, "email", "") or "").strip().lower()
if not email:
return None
entry = (
BackerEmail.objects.select_for_update()
.filter(email__iexact=email, redeemed_at__isnull=True)
.first()
)
if entry is None:
return None
seed_subscription_plans(update_existing=False)
plan = get_plan(SubscriptionPlan.Slug.BACKER)
if plan is None:
logger.error("Backer plan missing from catalog; cannot redeem %s", email)
return None
sub = assign_plan(
user,
plan=plan,
source=UserSubscription.Source.BACKER,
status=UserSubscription.Status.ACTIVE,
)
entry.redeemed_at = timezone.now()
entry.redeemed_user = user
entry.save(update_fields=["redeemed_at", "redeemed_user", "last_modified"])
logger.info("Redeemed Backer email %s for user %s", email, user.pk)
return sub
def assign_founders_from_stripe(
user,
*,
stripe_subscription_id: str = "",
) -> UserSubscription:
seed_subscription_plans(update_existing=False)
plan = get_plan(SubscriptionPlan.Slug.FOUNDERS)
if plan is None:
raise RuntimeError("Founders plan missing from catalog")
return assign_plan(
user,
plan=plan,
source=UserSubscription.Source.STRIPE,
status=UserSubscription.Status.ACTIVE,
stripe_subscription_id=stripe_subscription_id or "",
)
def plan_to_dict(plan: SubscriptionPlan | None) -> dict[str, Any] | None:
if plan is None:
return None
return {
"slug": plan.slug,
"name": plan.name,
"description": plan.description,
"price_cents": plan.price_cents,
"currency": plan.currency,
"interval": plan.interval,
"is_public": plan.is_public,
"is_selectable": plan.is_selectable,
"features": {
"text_generation": plan.allows_feature("text_generation"),
"image_generation": plan.allows_feature("image_generation"),
"all_future_features": plan.allows_all_future_features,
},
"prompt_quota_per_window": plan.prompt_quota_per_window,
"prompt_window_hours": plan.prompt_window_hours,
"monthly_token_quota": plan.monthly_token_quota,
"sort_order": plan.sort_order,
}