From a6c45b0882b92097f5e39c333f996605aa499532 Mon Sep 17 00:00:00 2001 From: Ryan Westfall Date: Fri, 31 Jul 2026 06:21:54 -0500 Subject: [PATCH] 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. --- .env.example | 8 +- llm_be/chat_backend/consumers.py | 96 ++++++- llm_be/chat_backend/consumers_graph.py | 97 ++++++- llm_be/chat_backend/oauth.py | 3 + llm_be/chat_backend/serializers.py | 101 ++++++- llm_be/chat_backend/tests/test_serializers.py | 5 +- llm_be/chat_backend/views.py | 3 + llm_be/chat_backend/views_oauth.py | 4 +- llm_be/finance/admin.py | 48 +++- llm_be/finance/apps.py | 7 + .../0002_subscription_plans_quotas.py | 86 ++++++ llm_be/finance/models.py | 170 ++++++++++++ llm_be/finance/serializers.py | 33 ++- llm_be/finance/services/plans.py | 258 +++++++++++++++++ llm_be/finance/services/quotas.py | 257 +++++++++++++++++ llm_be/finance/services/stripe_service.py | 56 +++- llm_be/finance/services/webhooks.py | 10 + llm_be/finance/signals.py | 8 + llm_be/finance/tests/test_plans_quotas.py | 261 ++++++++++++++++++ llm_be/finance/urls.py | 12 + llm_be/finance/views.py | 85 +++++- llm_be/llm_be/settings.py | 6 +- llm_be/llm_be/test_runner.py | 2 + 23 files changed, 1580 insertions(+), 36 deletions(-) create mode 100644 llm_be/finance/migrations/0002_subscription_plans_quotas.py create mode 100644 llm_be/finance/services/plans.py create mode 100644 llm_be/finance/services/quotas.py create mode 100644 llm_be/finance/signals.py create mode 100644 llm_be/finance/tests/test_plans_quotas.py diff --git a/.env.example b/.env.example index c434bfc..7670d01 100644 --- a/.env.example +++ b/.env.example @@ -44,13 +44,15 @@ OAUTH_CALLBACK_BASE_URL=http://127.0.0.1:8001 STRIPE_SECRET_KEY= STRIPE_PUBLISHABLE_KEY= STRIPE_WEBHOOK_SECRET= -# Optional: pre-created Stripe Price ID. When empty, Checkout uses -# SUBSCRIPTION_PRICE_* from settings.py ($10 USD / month by default). +# Optional: pre-created Stripe Price ID for Founders. When empty, Checkout uses +# SubscriptionPlan.price_cents / SUBSCRIPTION_PRICE_* ($10 USD / month Founders). STRIPE_PRICE_ID= # SUBSCRIPTION_PRICE_AMOUNT_CENTS=1000 # SUBSCRIPTION_PRICE_CURRENCY=usd # SUBSCRIPTION_PRICE_INTERVAL=month -# SUBSCRIPTION_PRODUCT_NAME=Chat Subscription +# SUBSCRIPTION_PRODUCT_NAME=Founders +# Enforce plan feature + prompt/token quotas on chat turns (default true). +# ENFORCE_SUBSCRIPTION_GATES=true FRONTEND_BASE_URL=http://localhost:3000 # STRIPE_CHECKOUT_SUCCESS_URL=http://localhost:3000/billing/success?session_id={CHECKOUT_SESSION_ID} # STRIPE_CHECKOUT_CANCEL_URL=http://localhost:3000/billing/cancel diff --git a/llm_be/chat_backend/consumers.py b/llm_be/chat_backend/consumers.py index c1bd1c4..48ae5ab 100644 --- a/llm_be/chat_backend/consumers.py +++ b/llm_be/chat_backend/consumers.py @@ -27,7 +27,17 @@ from .services.title_generator import title_generator from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.data_analysis_service import AsyncDataAnalysisService -from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message +from .utils import ( + extract_token_usage, + has_usable_user_prompt, + is_heartbeat_payload, + normalize_user_message, +) +from finance.services.quotas import ( + FeatureNotAllowed, + QuotaExceeded, + check_generation_allowed, +) logger = logging.getLogger(__name__) @@ -47,6 +57,35 @@ def create_conversation(prompt, email, title): return conversation.id +@database_sync_to_async +def resolve_chat_user(email=None, conversation_id=None): + if email: + user = CustomUser.objects.filter(email__iexact=email).first() + if user: + return user + if conversation_id: + conversation = ( + Conversation.objects.select_related("user") + .filter(id=conversation_id) + .first() + ) + if conversation and conversation.user_id: + return conversation.user + return None + + +@database_sync_to_async +def enforce_generation_gates(user, feature="text_generation"): + return check_generation_allowed(user, feature=feature) + + +@database_sync_to_async +def enforce_feature_gate(user, feature): + from finance.services.quotas import assert_feature_allowed + + assert_feature_allowed(user, feature) + + @database_sync_to_async def get_workspace(conversation_id): conversation = Conversation.objects.get(id=conversation_id) @@ -262,6 +301,36 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): ) return + chat_user = await resolve_chat_user( + email=email, conversation_id=conversation_id + ) + if chat_user is None: + await self.send_json_message( + json.dumps( + { + "type": "error", + "code": "user_not_found", + "content": "Unable to resolve user for this chat session.", + } + ) + ) + return + + try: + await enforce_generation_gates(chat_user, feature="text_generation") + except (QuotaExceeded, FeatureNotAllowed) as exc: + await self.send_json_message( + json.dumps( + { + "type": "error", + "code": exc.code, + "content": exc.message, + "details": getattr(exc, "details", {}), + } + ) + ) + return + if not conversation_id: # we need to create a new conversation # we will generate a name for it too @@ -334,11 +403,20 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): decoded_file = input_dict.get("decoded_file") file_type = input_dict.get("file_type") - # Feature Flag: Image Generation + # Feature Flag + plan gate: Image Generation if prompt_type == PromptType.IMAGE_GENERATION: if not getattr(settings, "ALLOW_IMAGE_GENERATION", False): return {"type": "text", "content": "Image Generation is disabled."} - # If enabled, proceed (assuming implementation exists, but user said "have it set to false for now") + try: + await enforce_feature_gate( + chat_user, "image_generation" + ) + except FeatureNotAllowed as exc: + return { + "type": "error", + "code": exc.code, + "content": exc.message, + } return {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."} if prompt_type == PromptType.SEARCH: @@ -446,7 +524,17 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): await self.send("END_OF_THE_STREAM_ENDER_GAME_42") await save_generated_message(conversation_id, full_response) - await finish_prompt_metric(prompt_metric, len(full_response)) + tokens_in, tokens_out = extract_token_usage( + response_generator_or_dict + if isinstance(response_generator_or_dict, dict) + else None + ) + await finish_prompt_metric( + prompt_metric, + len(full_response), + tokens_in=tokens_in, + tokens_out=tokens_out, + ) if bytes_data: logger.info("we have byte data") diff --git a/llm_be/chat_backend/consumers_graph.py b/llm_be/chat_backend/consumers_graph.py index 4b44bb8..8e3cc0f 100644 --- a/llm_be/chat_backend/consumers_graph.py +++ b/llm_be/chat_backend/consumers_graph.py @@ -22,7 +22,13 @@ from .services.title_generator import title_generator from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.data_analysis_service import AsyncDataAnalysisService -from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message +from .utils import ( + extract_token_usage, + has_usable_user_prompt, + is_heartbeat_payload, + normalize_user_message, +) +from finance.services.quotas import FeatureNotAllowed, QuotaExceeded, check_generation_allowed logger = logging.getLogger(__name__) @@ -40,6 +46,35 @@ def create_conversation(prompt, email, title): conversation.save() return conversation.id + +@database_sync_to_async +def resolve_chat_user(email=None, conversation_id=None): + if email: + user = CustomUser.objects.filter(email__iexact=email).first() + if user: + return user + if conversation_id: + conversation = ( + Conversation.objects.select_related("user") + .filter(id=conversation_id) + .first() + ) + if conversation and conversation.user_id: + return conversation.user + return None + + +@database_sync_to_async +def enforce_generation_gates(user, feature="text_generation"): + return check_generation_allowed(user, feature=feature) + + +@database_sync_to_async +def enforce_feature_gate(user, feature): + from finance.services.quotas import assert_feature_allowed + + assert_feature_allowed(user, feature) + @database_sync_to_async def get_workspace(conversation_id): conversation = Conversation.objects.get(id=conversation_id) @@ -172,6 +207,7 @@ class ChatState(TypedDict): response_generator: Any # AsyncGenerator or dict error: Union[str, None] model_name: str + chat_user: Any # --- LangGraph Nodes --- @@ -210,10 +246,22 @@ async def generation_node(state: ChatState) -> ChatState: decoded_file = state.get("decoded_file") file_type = state.get("file_type") - # Feature Flag: Image Generation + # Feature Flag + plan gate: Image Generation if prompt_type == PromptType.IMAGE_GENERATION: if not getattr(settings, "ALLOW_IMAGE_GENERATION", False): return {"response_generator": {"type": "text", "content": "Image Generation is disabled."}} + chat_user = state.get("chat_user") + if chat_user is not None: + try: + await enforce_feature_gate(chat_user, "image_generation") + except FeatureNotAllowed as exc: + return { + "response_generator": { + "type": "error", + "code": exc.code, + "content": exc.message, + } + } return {"response_generator": {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."}} # Feature Flag: Internet Access @@ -312,6 +360,36 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): ) return + chat_user = await resolve_chat_user( + email=email, conversation_id=conversation_id + ) + if chat_user is None: + await self.send_json_message( + json.dumps( + { + "type": "error", + "code": "user_not_found", + "content": "Unable to resolve user for this chat session.", + } + ) + ) + return + + try: + await enforce_generation_gates(chat_user, feature="text_generation") + except (QuotaExceeded, FeatureNotAllowed) as exc: + await self.send_json_message( + json.dumps( + { + "type": "error", + "code": exc.code, + "content": exc.message, + "details": getattr(exc, "details", {}), + } + ) + ) + return + if not conversation_id: title = await title_generator.generate_async(message) conversation_id = await create_conversation(message, email, title) @@ -357,7 +435,8 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): "prompt_type": None, "response_generator": None, "error": None, - "model_name": model + "model_name": model, + "chat_user": chat_user, } print("Initial State: ", initial_state) @@ -387,4 +466,14 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): await self.send("END_OF_THE_STREAM_ENDER_GAME_42") await save_generated_message(conversation_id, full_response) - await finish_prompt_metric(prompt_metric, len(full_response)) + tokens_in, tokens_out = extract_token_usage( + response_generator_or_dict + if isinstance(response_generator_or_dict, dict) + else None + ) + await finish_prompt_metric( + prompt_metric, + len(full_response), + tokens_in=tokens_in, + tokens_out=tokens_out, + ) diff --git a/llm_be/chat_backend/oauth.py b/llm_be/chat_backend/oauth.py index 9339275..0186010 100644 --- a/llm_be/chat_backend/oauth.py +++ b/llm_be/chat_backend/oauth.py @@ -307,6 +307,8 @@ def upsert_identity(user: CustomUser, profile: ProviderProfile) -> OAuthIdentity def _create_sso_user(profile: ProviderProfile) -> CustomUser: + from finance.services.plans import try_redeem_backer_email + company = Company.objects.create( name=f"{profile.email}'s workspace", state="NA", @@ -323,6 +325,7 @@ def _create_sso_user(profile: ProviderProfile) -> CustomUser: ) user.set_unusable_password() user.save() + try_redeem_backer_email(user) return user diff --git a/llm_be/chat_backend/serializers.py b/llm_be/chat_backend/serializers.py index 646db4e..216cb7b 100644 --- a/llm_be/chat_backend/serializers.py +++ b/llm_be/chat_backend/serializers.py @@ -1,11 +1,14 @@ from rest_framework_simplejwt.serializers import TokenObtainPairSerializer from rest_framework import serializers +from django.db.models import Count, Q, Sum + from .models import ( CustomUser, Announcement, Company, Conversation, Prompt, + PromptMetric, Feedback, FEEDBACK_CATEGORIES, DocumentWorkspace, @@ -48,12 +51,33 @@ class CustomUserSerializer(serializers.ModelSerializer): password = serializers.CharField(min_length=8, write_only=True) company = CompanySerializer() has_usable_password = serializers.BooleanField() + subscription = serializers.SerializerMethodField() class Meta: model = CustomUser fields = "__all__" extra_kwargs = {"password": {"write_only": True}} + def get_subscription(self, obj): + from finance.services.plans import needs_checkout, plan_to_dict + from finance.models import UserSubscription + + try: + sub = obj.subscription + except UserSubscription.DoesNotExist: + return { + "plan": None, + "status": UserSubscription.Status.NONE, + "source": UserSubscription.Source.NONE, + "needs_checkout": True, + } + return { + "plan": plan_to_dict(sub.plan) if sub.plan_id else None, + "status": sub.status, + "source": sub.source, + "needs_checkout": needs_checkout(obj), + } + class SelfServeRegistrationSerializer(serializers.Serializer): """Minimal payload for public self-serve sign-up (gated by settings).""" @@ -79,6 +103,8 @@ class SelfServeRegistrationSerializer(serializers.Serializer): return email def create(self, validated_data): + from finance.services.plans import try_redeem_backer_email + email = validated_data["email"] password = validated_data["password"] first_name = (validated_data.get("first_name") or "").strip() @@ -102,16 +128,73 @@ class SelfServeRegistrationSerializer(serializers.Serializer): company=company, is_company_manager=True, ) + try_redeem_backer_email(user) return user +def _conversation_token_totals(conversation_id: int): + """ + Sum PromptMetric tokens for a conversation. + + Returns nulls when the provider never reported usage (never fabricate 0). + """ + agg = PromptMetric.objects.filter(conversation_id=conversation_id).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)), + ) + return ( + agg["tin"] if agg["with_in"] else None, + agg["tout"] if agg["with_out"] else None, + ) + + +def _prompt_token_pair(prompt_id: int): + metric = ( + PromptMetric.objects.filter(prompt_id=prompt_id) + .order_by("-created") + .only("tokens_in", "tokens_out") + .first() + ) + if metric is None: + return None, None + return metric.tokens_in, metric.tokens_out + + class ConversationSerializer(serializers.ModelSerializer): + tokens_in = serializers.SerializerMethodField() + tokens_out = serializers.SerializerMethodField() + class Meta: model = Conversation - fields = ("title", "created", "last_modified", "id") + fields = ( + "title", + "created", + "last_modified", + "id", + "tokens_in", + "tokens_out", + ) + + def _token_pair(self, obj): + cache = self.context.setdefault("_conversation_token_cache", {}) + if obj.id not in cache: + cache[obj.id] = _conversation_token_totals(obj.id) + return cache[obj.id] + + def get_tokens_in(self, obj): + tin, _ = self._token_pair(obj) + return tin + + def get_tokens_out(self, obj): + _, tout = self._token_pair(obj) + return tout class PromptSerializer(serializers.ModelSerializer): + tokens_in = serializers.SerializerMethodField() + tokens_out = serializers.SerializerMethodField() class Meta: model = Prompt @@ -120,8 +203,24 @@ class PromptSerializer(serializers.ModelSerializer): "user_created", "created", "id", + "tokens_in", + "tokens_out", ) + def _token_pair(self, obj): + cache = self.context.setdefault("_prompt_token_cache", {}) + if obj.id not in cache: + cache[obj.id] = _prompt_token_pair(obj.id) + return cache[obj.id] + + def get_tokens_in(self, obj): + tin, _ = self._token_pair(obj) + return tin + + def get_tokens_out(self, obj): + _, tout = self._token_pair(obj) + return tout + def validate_message(self, value: str) -> str: if value is None or not str(value).strip(): raise serializers.ValidationError("Message text cannot be empty.") diff --git a/llm_be/chat_backend/tests/test_serializers.py b/llm_be/chat_backend/tests/test_serializers.py index 0a92eca..b75a0f1 100644 --- a/llm_be/chat_backend/tests/test_serializers.py +++ b/llm_be/chat_backend/tests/test_serializers.py @@ -38,7 +38,10 @@ class ConversationSerializerTestCase(TestCase): data = ConversationSerializer(conversation).data - self.assertEqual(set(data.keys()), {"title", "created", "last_modified", "id"}) + self.assertEqual( + set(data.keys()), + {"title", "created", "last_modified", "id", "tokens_in", "tokens_out"}, + ) self.assertEqual(data["title"], "Weather Inquiry") diff --git a/llm_be/chat_backend/views.py b/llm_be/chat_backend/views.py index 3813dd8..8d022f0 100644 --- a/llm_be/chat_backend/views.py +++ b/llm_be/chat_backend/views.py @@ -136,6 +136,8 @@ class CustomUserCreate(APIView): user = serializer.save() refresh = RefreshToken.for_user(user) + from finance.services.plans import needs_checkout + return Response( { "email": user.email, @@ -143,6 +145,7 @@ class CustomUserCreate(APIView): "last_name": user.last_name, "access": str(refresh.access_token), "refresh": str(refresh), + "needs_checkout": needs_checkout(user), }, status=status.HTTP_201_CREATED, ) diff --git a/llm_be/chat_backend/views_oauth.py b/llm_be/chat_backend/views_oauth.py index 781d13b..11c87c8 100644 --- a/llm_be/chat_backend/views_oauth.py +++ b/llm_be/chat_backend/views_oauth.py @@ -139,7 +139,9 @@ class OAuthCallbackView(APIView): return _redirect_error("server_error", "Unexpected OAuth error.") refresh = RefreshToken.for_user(user) - needs_checkout = "1" if created else "0" + from finance.services.plans import needs_checkout as user_needs_checkout + + needs_checkout = "1" if (created and user_needs_checkout(user)) else "0" return HttpResponseRedirect( _frontend_callback_url( access=str(refresh.access_token), diff --git a/llm_be/finance/admin.py b/llm_be/finance/admin.py index 5e86753..eea703a 100644 --- a/llm_be/finance/admin.py +++ b/llm_be/finance/admin.py @@ -1,6 +1,52 @@ from django.contrib import admin -from finance.models import Invoice, Payment +from finance.models import BackerEmail, Invoice, Payment, SubscriptionPlan, UserSubscription + + +@admin.register(SubscriptionPlan) +class SubscriptionPlanAdmin(admin.ModelAdmin): + list_display = ( + "slug", + "name", + "price_cents", + "is_public", + "is_selectable", + "allows_image_generation", + "allows_all_future_features", + "prompt_quota_per_window", + "prompt_window_hours", + "monthly_token_quota", + "sort_order", + ) + list_filter = ("is_public", "is_selectable", "allows_image_generation") + search_fields = ("slug", "name", "stripe_price_id") + readonly_fields = ("created", "last_modified") + prepopulated_fields = {"slug": ("name",)} + + +@admin.register(BackerEmail) +class BackerEmailAdmin(admin.ModelAdmin): + list_display = ("email", "note", "redeemed_at", "redeemed_user", "created") + search_fields = ("email", "note", "redeemed_user__email") + raw_id_fields = ("redeemed_user",) + readonly_fields = ("created", "last_modified", "redeemed_at", "redeemed_user") + + +@admin.register(UserSubscription) +class UserSubscriptionAdmin(admin.ModelAdmin): + list_display = ( + "user", + "plan", + "status", + "source", + "stripe_subscription_id", + "monthly_token_quota_override", + "created", + ) + list_filter = ("status", "source", "plan") + search_fields = ("user__email", "user__username", "stripe_subscription_id") + raw_id_fields = ("user", "plan") + readonly_fields = ("created", "last_modified") class PaymentInline(admin.TabularInline): diff --git a/llm_be/finance/apps.py b/llm_be/finance/apps.py index e773d3c..ece0e44 100644 --- a/llm_be/finance/apps.py +++ b/llm_be/finance/apps.py @@ -5,3 +5,10 @@ class FinanceConfig(AppConfig): default_auto_field = "django.db.models.BigAutoField" name = "finance" verbose_name = "Finance" + + def ready(self): + from django.db.models.signals import post_migrate + + from finance.signals import seed_plans_on_migrate + + post_migrate.connect(seed_plans_on_migrate, sender=self) diff --git a/llm_be/finance/migrations/0002_subscription_plans_quotas.py b/llm_be/finance/migrations/0002_subscription_plans_quotas.py new file mode 100644 index 0000000..36229e8 --- /dev/null +++ b/llm_be/finance/migrations/0002_subscription_plans_quotas.py @@ -0,0 +1,86 @@ +# Generated by Django 6.0 on 2026-07-31 11:14 + +import django.db.models.deletion +import django.utils.timezone +from django.conf import settings +from django.db import migrations, models + + +def seed_plans(apps, schema_editor): + from finance.services.plans import seed_subscription_plans + + seed_subscription_plans(update_existing=True) + + +class Migration(migrations.Migration): + + dependencies = [ + ('finance', '0001_initial'), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.CreateModel( + name='SubscriptionPlan', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('created', models.DateTimeField(default=django.utils.timezone.now)), + ('last_modified', models.DateTimeField(default=django.utils.timezone.now)), + ('slug', models.SlugField(max_length=64, unique=True)), + ('name', models.CharField(max_length=128)), + ('description', models.TextField(blank=True, default='')), + ('price_cents', models.PositiveIntegerField(default=0, help_text='List price in cents (0 for complimentary tiers).')), + ('currency', models.CharField(default='usd', max_length=8)), + ('interval', models.CharField(default='month', max_length=16)), + ('stripe_price_id', models.CharField(blank=True, default='', help_text='Optional Stripe Price id; empty uses price_data at Checkout.', max_length=255)), + ('is_public', models.BooleanField(default=False, help_text='Shown in public pricing / plan list APIs.')), + ('is_selectable', models.BooleanField(default=False, help_text='Selectable at Checkout. Backer is never selectable.')), + ('allows_text_generation', models.BooleanField(default=True)), + ('allows_image_generation', models.BooleanField(default=False)), + ('allows_all_future_features', models.BooleanField(default=False, help_text='Founders/Backer: unlock new capabilities as they ship.')), + ('prompt_quota_per_window', models.PositiveIntegerField(help_text='Max prompts allowed in each rolling window.')), + ('prompt_window_hours', models.PositiveIntegerField(default=6, help_text='Length of the rolling prompt quota window in hours.')), + ('monthly_token_quota', models.PositiveIntegerField(blank=True, help_text='Optional billing-period token cap (tokens_in + tokens_out). Null = no token-period limit (prompt window still applies).', null=True)), + ('sort_order', models.PositiveIntegerField(default=0)), + ], + options={ + 'ordering': ['sort_order', 'price_cents', 'name'], + }, + ), + migrations.CreateModel( + name='BackerEmail', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('created', models.DateTimeField(default=django.utils.timezone.now)), + ('last_modified', models.DateTimeField(default=django.utils.timezone.now)), + ('email', models.EmailField(db_index=True, max_length=254, unique=True)), + ('note', models.CharField(blank=True, default='', max_length=512)), + ('redeemed_at', models.DateTimeField(blank=True, null=True)), + ('redeemed_user', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='backer_email_entries', to=settings.AUTH_USER_MODEL)), + ], + options={ + 'verbose_name': 'Backer email', + 'verbose_name_plural': 'Backer emails', + 'ordering': ['email'], + }, + ), + migrations.CreateModel( + name='UserSubscription', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('created', models.DateTimeField(default=django.utils.timezone.now)), + ('last_modified', models.DateTimeField(default=django.utils.timezone.now)), + ('status', models.CharField(choices=[('none', 'None'), ('active', 'Active'), ('past_due', 'Past due'), ('canceled', 'Canceled')], db_index=True, default='none', max_length=32)), + ('source', models.CharField(choices=[('none', 'None'), ('stripe', 'Stripe'), ('backer', 'Backer'), ('admin', 'Admin')], default='none', max_length=32)), + ('stripe_subscription_id', models.CharField(blank=True, db_index=True, default='', max_length=255)), + ('monthly_token_quota_override', models.PositiveIntegerField(blank=True, help_text='Optional per-user override of plan monthly_token_quota.', null=True)), + ('plan', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.PROTECT, related_name='subscriptions', to='finance.subscriptionplan')), + ('user', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, related_name='subscription', to=settings.AUTH_USER_MODEL)), + ], + options={ + 'verbose_name': 'User subscription', + 'verbose_name_plural': 'User subscriptions', + }, + ), + migrations.RunPython(seed_plans, migrations.RunPython.noop), + ] diff --git a/llm_be/finance/models.py b/llm_be/finance/models.py index c7cb223..2367260 100644 --- a/llm_be/finance/models.py +++ b/llm_be/finance/models.py @@ -5,6 +5,176 @@ from django.utils import timezone from chat_backend.models import Company, TimeInfoBase +class SubscriptionPlan(TimeInfoBase): + """Catalog row for a billable (or complimentary) subscription tier.""" + + class Slug(models.TextChoices): + FOUNDERS = "founders", "Founders" + STANDARD = "standard", "Standard" + PRO = "pro", "Pro / Creator" + BUSINESS = "business", "Business Team" + BACKER = "backer", "Backer" + + slug = models.SlugField(max_length=64, unique=True, db_index=True) + name = models.CharField(max_length=128) + description = models.TextField(blank=True, default="") + price_cents = models.PositiveIntegerField( + default=0, + help_text="List price in cents (0 for complimentary tiers).", + ) + currency = models.CharField(max_length=8, default="usd") + interval = models.CharField(max_length=16, default="month") + stripe_price_id = models.CharField( + max_length=255, + blank=True, + default="", + help_text="Optional Stripe Price id; empty uses price_data at Checkout.", + ) + is_public = models.BooleanField( + default=False, + help_text="Shown in public pricing / plan list APIs.", + ) + is_selectable = models.BooleanField( + default=False, + help_text="Selectable at Checkout. Backer is never selectable.", + ) + allows_text_generation = models.BooleanField(default=True) + allows_image_generation = models.BooleanField(default=False) + allows_all_future_features = models.BooleanField( + default=False, + help_text="Founders/Backer: unlock new capabilities as they ship.", + ) + prompt_quota_per_window = models.PositiveIntegerField( + help_text="Max prompts allowed in each rolling window.", + ) + prompt_window_hours = models.PositiveIntegerField( + default=6, + help_text="Length of the rolling prompt quota window in hours.", + ) + monthly_token_quota = models.PositiveIntegerField( + null=True, + blank=True, + help_text=( + "Optional billing-period token cap (tokens_in + tokens_out). " + "Null = no token-period limit (prompt window still applies)." + ), + ) + sort_order = models.PositiveIntegerField(default=0) + + class Meta: + ordering = ["sort_order", "price_cents", "name"] + + def __str__(self) -> str: + return f"{self.name} ({self.slug})" + + def allows_feature(self, feature: str) -> bool: + if self.allows_all_future_features: + return True + if feature in ("text", "text_generation"): + return self.allows_text_generation + if feature in ("image", "image_generation"): + return self.allows_image_generation + return False + + +class BackerEmail(TimeInfoBase): + """Pre-registered emails that receive complimentary Backer (Founders-level) access.""" + + email = models.EmailField(unique=True, db_index=True) + note = models.CharField(max_length=512, blank=True, default="") + redeemed_at = models.DateTimeField(null=True, blank=True) + redeemed_user = models.ForeignKey( + settings.AUTH_USER_MODEL, + on_delete=models.SET_NULL, + null=True, + blank=True, + related_name="backer_email_entries", + ) + + class Meta: + ordering = ["email"] + verbose_name = "Backer email" + verbose_name_plural = "Backer emails" + + def __str__(self) -> str: + return self.email + + def save(self, *args, **kwargs): + if self.email: + self.email = self.email.strip().lower() + super().save(*args, **kwargs) + + +class UserSubscription(TimeInfoBase): + """Per-user plan assignment (Stripe, Backer whitelist, or admin).""" + + class Status(models.TextChoices): + NONE = "none", "None" + ACTIVE = "active", "Active" + PAST_DUE = "past_due", "Past due" + CANCELED = "canceled", "Canceled" + + class Source(models.TextChoices): + NONE = "none", "None" + STRIPE = "stripe", "Stripe" + BACKER = "backer", "Backer" + ADMIN = "admin", "Admin" + + user = models.OneToOneField( + settings.AUTH_USER_MODEL, + on_delete=models.CASCADE, + related_name="subscription", + ) + plan = models.ForeignKey( + SubscriptionPlan, + on_delete=models.PROTECT, + related_name="subscriptions", + null=True, + blank=True, + ) + status = models.CharField( + max_length=32, + choices=Status.choices, + default=Status.NONE, + db_index=True, + ) + source = models.CharField( + max_length=32, + choices=Source.choices, + default=Source.NONE, + ) + stripe_subscription_id = models.CharField( + max_length=255, + blank=True, + default="", + db_index=True, + ) + monthly_token_quota_override = models.PositiveIntegerField( + null=True, + blank=True, + help_text="Optional per-user override of plan monthly_token_quota.", + ) + + class Meta: + verbose_name = "User subscription" + verbose_name_plural = "User subscriptions" + + def __str__(self) -> str: + plan = self.plan.slug if self.plan_id else "none" + return f"UserSubscription user={self.user_id} plan={plan} ({self.status})" + + @property + def is_active(self) -> bool: + return self.status == self.Status.ACTIVE and self.plan_id is not None + + def effective_monthly_token_quota(self): + if self.monthly_token_quota_override is not None: + return self.monthly_token_quota_override + if self.plan_id: + return self.plan.monthly_token_quota + return None + + class Invoice(TimeInfoBase): """Local ledger row for a billed period / Stripe invoice or checkout session.""" diff --git a/llm_be/finance/serializers.py b/llm_be/finance/serializers.py index 6f9ea59..fe710cf 100644 --- a/llm_be/finance/serializers.py +++ b/llm_be/finance/serializers.py @@ -1,6 +1,6 @@ from rest_framework import serializers -from finance.models import Invoice, Payment +from finance.models import Invoice, Payment, SubscriptionPlan class InvoiceSerializer(serializers.ModelSerializer): @@ -49,7 +49,38 @@ class PaymentSerializer(serializers.ModelSerializer): class CheckoutSessionSerializer(serializers.Serializer): success_url = serializers.URLField(required=False, allow_blank=False) cancel_url = serializers.URLField(required=False, allow_blank=False) + plan_slug = serializers.SlugField(required=False, allow_blank=False) class PortalSessionSerializer(serializers.Serializer): return_url = serializers.URLField(required=False, allow_blank=False) + + +class SubscriptionPlanSerializer(serializers.ModelSerializer): + features = serializers.SerializerMethodField() + + class Meta: + model = SubscriptionPlan + fields = [ + "slug", + "name", + "description", + "price_cents", + "currency", + "interval", + "is_public", + "is_selectable", + "features", + "prompt_quota_per_window", + "prompt_window_hours", + "monthly_token_quota", + "sort_order", + ] + read_only_fields = fields + + def get_features(self, obj): + return { + "text_generation": obj.allows_feature("text_generation"), + "image_generation": obj.allows_feature("image_generation"), + "all_future_features": obj.allows_all_future_features, + } diff --git a/llm_be/finance/services/plans.py b/llm_be/finance/services/plans.py new file mode 100644 index 0000000..21bf833 --- /dev/null +++ b/llm_be/finance/services/plans.py @@ -0,0 +1,258 @@ +"""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, + } diff --git a/llm_be/finance/services/quotas.py b/llm_be/finance/services/quotas.py new file mode 100644 index 0000000..3c8d27b --- /dev/null +++ b/llm_be/finance/services/quotas.py @@ -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) diff --git a/llm_be/finance/services/stripe_service.py b/llm_be/finance/services/stripe_service.py index 211110d..f6170ab 100644 --- a/llm_be/finance/services/stripe_service.py +++ b/llm_be/finance/services/stripe_service.py @@ -7,7 +7,8 @@ from typing import Any import stripe from django.conf import settings -from finance.models import Invoice +from finance.models import Invoice, SubscriptionPlan +from finance.services.plans import get_plan, seed_subscription_plans class StripeNotConfiguredError(RuntimeError): @@ -24,21 +25,47 @@ def configure_stripe() -> str: return secret -def subscription_line_items() -> list[dict[str, Any]]: - """Build Checkout line_items from settings-backed subscription pricing.""" - price_id = settings.STRIPE_PRICE_ID +def resolve_checkout_plan(plan_slug: str | None = None) -> SubscriptionPlan: + """Return the plan for Checkout (defaults to public Founders).""" + seed_subscription_plans(update_existing=False) + slug = (plan_slug or SubscriptionPlan.Slug.FOUNDERS).strip().lower() + plan = get_plan(slug) + if plan is None: + raise ValueError(f"Unknown plan: {slug}") + if not plan.is_selectable: + raise ValueError(f"Plan '{plan.slug}' is not available for checkout.") + return plan + + +def subscription_line_items(plan: SubscriptionPlan) -> list[dict[str, Any]]: + """Build Checkout line_items from a SubscriptionPlan (or legacy settings).""" + price_id = (plan.stripe_price_id or "").strip() or ( + settings.STRIPE_PRICE_ID if plan.slug == SubscriptionPlan.Slug.FOUNDERS else "" + ) if price_id: return [{"price": price_id, "quantity": 1}] + # Founders without a plan stripe_price_id may still use legacy env amount. + unit_amount = plan.price_cents + product_name = plan.name + currency = plan.currency or settings.SUBSCRIPTION_PRICE_CURRENCY + interval = plan.interval or settings.SUBSCRIPTION_PRICE_INTERVAL + if plan.slug == SubscriptionPlan.Slug.FOUNDERS and not plan.stripe_price_id: + # Keep env overrides working for the live Founders price. + unit_amount = int( + getattr(settings, "SUBSCRIPTION_PRICE_AMOUNT_CENTS", None) or unit_amount + ) + product_name = ( + getattr(settings, "SUBSCRIPTION_PRODUCT_NAME", None) or product_name + ) + return [ { "price_data": { - "currency": settings.SUBSCRIPTION_PRICE_CURRENCY, - "unit_amount": settings.SUBSCRIPTION_PRICE_AMOUNT_CENTS, - "recurring": {"interval": settings.SUBSCRIPTION_PRICE_INTERVAL}, - "product_data": { - "name": settings.SUBSCRIPTION_PRODUCT_NAME, - }, + "currency": currency, + "unit_amount": unit_amount, + "recurring": {"interval": interval}, + "product_data": {"name": product_name}, }, "quantity": 1, } @@ -50,19 +77,22 @@ def create_checkout_session( user, success_url: str | None = None, cancel_url: str | None = None, + plan_slug: str | None = None, ): - """Create a Stripe Checkout Session for the subscription plan.""" + """Create a Stripe Checkout Session for a selectable subscription plan.""" configure_stripe() + plan = resolve_checkout_plan(plan_slug) metadata = { "user_id": str(user.pk), "company_id": str(user.company_id) if user.company_id else "", + "plan_slug": plan.slug, } customer_email = getattr(user, "email", None) or None session = stripe.checkout.Session.create( mode="subscription", - line_items=subscription_line_items(), + line_items=subscription_line_items(plan), success_url=success_url or settings.STRIPE_CHECKOUT_SUCCESS_URL, cancel_url=cancel_url or settings.STRIPE_CHECKOUT_CANCEL_URL, customer_email=customer_email, @@ -70,7 +100,7 @@ def create_checkout_session( metadata=metadata, subscription_data={"metadata": metadata}, ) - return session + return session, plan def resolve_stripe_customer_id(*, user) -> str | None: diff --git a/llm_be/finance/services/webhooks.py b/llm_be/finance/services/webhooks.py index 0d1bf28..7bec904 100644 --- a/llm_be/finance/services/webhooks.py +++ b/llm_be/finance/services/webhooks.py @@ -11,6 +11,7 @@ from django.db import transaction from django.utils import timezone from finance.models import Invoice, Payment +from finance.services.plans import assign_founders_from_stripe logger = logging.getLogger(__name__) User = get_user_model() @@ -211,6 +212,11 @@ def handle_checkout_session_completed(session: dict[str, Any]) -> Invoice | None ), paid_at=timezone.now(), ) + if session.get("payment_status") == "paid" or session.get("subscription"): + assign_founders_from_stripe( + user, + stripe_subscription_id=session.get("subscription") or "", + ) return invoice @@ -270,6 +276,10 @@ def handle_invoice_paid(stripe_invoice: dict[str, Any]) -> Invoice | None: stripe_charge_id=charge if isinstance(charge, str) else None, paid_at=paid_at, ) + assign_founders_from_stripe( + user, + stripe_subscription_id=stripe_invoice.get("subscription") or "", + ) return invoice diff --git a/llm_be/finance/signals.py b/llm_be/finance/signals.py new file mode 100644 index 0000000..1bb2132 --- /dev/null +++ b/llm_be/finance/signals.py @@ -0,0 +1,8 @@ +"""Finance app signals.""" + + +def seed_plans_on_migrate(sender, **kwargs): + """Ensure the subscription catalog exists after migrate.""" + from finance.services.plans import seed_subscription_plans + + seed_subscription_plans(update_existing=False) diff --git a/llm_be/finance/tests/test_plans_quotas.py b/llm_be/finance/tests/test_plans_quotas.py new file mode 100644 index 0000000..6cd1770 --- /dev/null +++ b/llm_be/finance/tests/test_plans_quotas.py @@ -0,0 +1,261 @@ +"""Tests for multi-plan catalog, Backer whitelist, quotas, and subscription API.""" + +from datetime import timedelta +from unittest.mock import MagicMock, patch + +from django.test import TestCase, override_settings +from django.urls import reverse +from django.utils import timezone +from rest_framework import status +from rest_framework.test import APITestCase + +from chat_backend.models import PromptMetric +from chat_backend.tests.factories import make_company, make_conversation, make_user +from finance.models import BackerEmail, SubscriptionPlan, UserSubscription +from finance.services.plans import ( + assign_plan, + needs_checkout, + seed_subscription_plans, + try_redeem_backer_email, +) +from finance.services.quotas import ( + FeatureNotAllowed, + QuotaExceeded, + assert_feature_allowed, + assert_within_quotas, + get_usage_snapshot, +) + + +class PlanCatalogTestCase(TestCase): + def test_seed_creates_expected_plans(self): + plans = {p.slug: p for p in seed_subscription_plans()} + self.assertEqual( + set(plans), + {"founders", "standard", "pro", "business", "backer"}, + ) + self.assertTrue(plans["founders"].is_public) + self.assertTrue(plans["founders"].is_selectable) + self.assertEqual(plans["founders"].price_cents, 1000) + self.assertEqual(plans["founders"].prompt_quota_per_window, 300) + + self.assertFalse(plans["standard"].is_public) + self.assertEqual(plans["standard"].price_cents, 1500) + self.assertEqual(plans["standard"].prompt_quota_per_window, 100) + self.assertFalse(plans["standard"].allows_image_generation) + + self.assertEqual(plans["pro"].price_cents, 4000) + self.assertEqual(plans["pro"].prompt_quota_per_window, 200) + self.assertTrue(plans["pro"].allows_image_generation) + + self.assertEqual(plans["business"].price_cents, 9900) + self.assertEqual(plans["business"].prompt_quota_per_window, 300) + + self.assertEqual(plans["backer"].price_cents, 0) + self.assertFalse(plans["backer"].is_selectable) + self.assertTrue(plans["backer"].allows_all_future_features) + + +class BackerRedeemTestCase(TestCase): + def setUp(self): + seed_subscription_plans() + self.company = make_company() + + def test_redeem_assigns_backer_and_skips_checkout(self): + BackerEmail.objects.create(email="backer@example.com") + user = make_user(email="backer@example.com", company=self.company) + sub = try_redeem_backer_email(user) + self.assertIsNotNone(sub) + self.assertEqual(sub.plan.slug, "backer") + self.assertEqual(sub.source, UserSubscription.Source.BACKER) + self.assertFalse(needs_checkout(user)) + entry = BackerEmail.objects.get(email="backer@example.com") + self.assertIsNotNone(entry.redeemed_at) + self.assertEqual(entry.redeemed_user_id, user.pk) + + def test_redeem_is_one_shot(self): + BackerEmail.objects.create(email="once@example.com") + user = make_user(email="once@example.com", company=self.company) + self.assertIsNotNone(try_redeem_backer_email(user)) + self.assertIsNone(try_redeem_backer_email(user)) + + +@override_settings(ENFORCE_SUBSCRIPTION_GATES=True) +class QuotaGateTestCase(TestCase): + def setUp(self): + seed_subscription_plans() + self.company = make_company() + self.user = make_user(company=self.company) + self.plan = SubscriptionPlan.objects.get(slug="standard") + assign_plan( + self.user, + plan=self.plan, + source=UserSubscription.Source.ADMIN, + ) + self.conversation = make_conversation(user=self.user) + + def _add_metrics(self, count, *, tokens_in=None, tokens_out=None): + now = timezone.now() + for i in range(count): + PromptMetric.objects.create( + prompt_id=1000 + i, + conversation_id=self.conversation.id, + start_time=now, + prompt_length=10, + tokens_in=tokens_in, + tokens_out=tokens_out, + has_file=False, + model_name="test", + ) + + def test_prompt_quota_blocks(self): + self._add_metrics(100) + with self.assertRaises(QuotaExceeded) as ctx: + assert_within_quotas(self.user) + self.assertEqual(ctx.exception.code, "prompt_quota_exceeded") + + def test_feature_gate_blocks_image_on_standard(self): + with self.assertRaises(FeatureNotAllowed) as ctx: + assert_feature_allowed(self.user, "image_generation") + self.assertEqual(ctx.exception.code, "feature_not_allowed") + + def test_pro_allows_image(self): + pro = SubscriptionPlan.objects.get(slug="pro") + assign_plan( + self.user, plan=pro, source=UserSubscription.Source.ADMIN + ) + assert_feature_allowed(self.user, "image_generation") + + def test_token_quota_blocks_when_reported(self): + self.plan.monthly_token_quota = 50 + self.plan.prompt_quota_per_window = 1000 + self.plan.save() + self._add_metrics(1, tokens_in=30, tokens_out=30) + with self.assertRaises(QuotaExceeded) as ctx: + assert_within_quotas(self.user) + self.assertEqual(ctx.exception.code, "token_quota_exceeded") + + def test_null_tokens_do_not_fabricate_zero_usage(self): + self._add_metrics(3, tokens_in=None, tokens_out=None) + usage = get_usage_snapshot(self.user) + self.assertIsNone(usage.tokens_in_period) + self.assertIsNone(usage.tokens_out_period) + self.assertIsNone(usage.tokens_total_period) + self.assertGreaterEqual(usage.turns_missing_token_usage, 3) + + +class PlanListAndSubscriptionApiTestCase(APITestCase): + def setUp(self): + seed_subscription_plans() + self.company = make_company() + self.user = make_user(company=self.company) + + def test_public_plans_only_founders(self): + url = reverse("finance_plans") + response = self.client.get(url) + self.assertEqual(response.status_code, status.HTTP_200_OK) + slugs = [row["slug"] for row in response.data] + self.assertEqual(slugs, ["founders"]) + + def test_subscription_me_includes_usage_nulls(self): + self.client.force_authenticate(user=self.user) + founders = SubscriptionPlan.objects.get(slug="founders") + assign_plan( + self.user, + plan=founders, + source=UserSubscription.Source.STRIPE, + ) + url = reverse("finance_subscription_me") + response = self.client.get(url) + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.data["plan"]["slug"], "founders") + self.assertFalse(response.data["needs_checkout"]) + self.assertIsNone(response.data["usage"]["tokens_in_period"]) + self.assertEqual(response.data["usage"]["prompt_quota"], 300) + + +class CheckoutUsesFoundersPlanTestCase(APITestCase): + def setUp(self): + seed_subscription_plans() + self.company = make_company() + self.user = make_user(company=self.company) + self.client.force_authenticate(user=self.user) + self.url = reverse("finance_checkout") + + @override_settings( + STRIPE_SECRET_KEY="sk_test_fake", + SUBSCRIPTION_PRICE_AMOUNT_CENTS=1000, + SUBSCRIPTION_PRICE_CURRENCY="usd", + SUBSCRIPTION_PRICE_INTERVAL="month", + SUBSCRIPTION_PRODUCT_NAME="Founders", + STRIPE_PRICE_ID="", + STRIPE_CHECKOUT_SUCCESS_URL="http://localhost:3000/ok", + STRIPE_CHECKOUT_CANCEL_URL="http://localhost:3000/cancel", + ) + @patch("finance.services.stripe_service.stripe.checkout.Session.create") + def test_checkout_defaults_to_founders(self, mock_create): + mock_session = MagicMock() + mock_session.id = "cs_test_founders" + mock_session.url = "https://checkout.stripe.com/c/pay/cs_test_founders" + mock_session.customer = None + mock_create.return_value = mock_session + + response = self.client.post(self.url, {}, format="json") + self.assertEqual(response.status_code, status.HTTP_201_CREATED) + self.assertEqual(response.data["plan_slug"], "founders") + kwargs = mock_create.call_args.kwargs + self.assertEqual(kwargs["metadata"]["plan_slug"], "founders") + self.assertEqual( + kwargs["line_items"][0]["price_data"]["unit_amount"], 1000 + ) + + def test_backer_cannot_checkout(self): + backer = SubscriptionPlan.objects.get(slug="backer") + assign_plan( + self.user, plan=backer, source=UserSubscription.Source.BACKER + ) + response = self.client.post(self.url, {}, format="json") + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertFalse(response.data["needs_checkout"]) + + +class TokenSerializerApiTestCase(APITestCase): + def setUp(self): + self.company = make_company() + self.user = make_user(company=self.company) + self.client.force_authenticate(user=self.user) + self.conversation = make_conversation(user=self.user, title="Tok") + + def test_conversation_tokens_null_when_unreported(self): + PromptMetric.objects.create( + prompt_id=1, + conversation_id=self.conversation.id, + start_time=timezone.now(), + prompt_length=5, + tokens_in=None, + tokens_out=None, + has_file=False, + model_name="t", + ) + response = self.client.get(reverse("conversations")) + self.assertEqual(response.status_code, status.HTTP_200_OK) + row = next(r for r in response.data if r["id"] == self.conversation.id) + self.assertIsNone(row["tokens_in"]) + self.assertIsNone(row["tokens_out"]) + + def test_conversation_tokens_sum_when_reported(self): + for tin, tout, pid in ((10, 20, 1), (5, None, 2), (None, 7, 3)): + PromptMetric.objects.create( + prompt_id=pid, + conversation_id=self.conversation.id, + start_time=timezone.now(), + prompt_length=5, + tokens_in=tin, + tokens_out=tout, + has_file=False, + model_name="t", + ) + response = self.client.get(reverse("conversations")) + row = next(r for r in response.data if r["id"] == self.conversation.id) + self.assertEqual(row["tokens_in"], 15) + self.assertEqual(row["tokens_out"], 27) diff --git a/llm_be/finance/urls.py b/llm_be/finance/urls.py index 2ca88d1..f475bcf 100644 --- a/llm_be/finance/urls.py +++ b/llm_be/finance/urls.py @@ -5,7 +5,9 @@ from finance.views import ( CreateCheckoutSessionView, InvoiceListView, PaymentListView, + PlanListView, StripeWebhookView, + SubscriptionMeView, ) urlpatterns = [ @@ -29,6 +31,16 @@ urlpatterns = [ PaymentListView.as_view(), name="finance_payments", ), + path( + "plans/", + PlanListView.as_view(), + name="finance_plans", + ), + path( + "subscription/", + SubscriptionMeView.as_view(), + name="finance_subscription_me", + ), path( "webhooks/stripe/", StripeWebhookView.as_view(), diff --git a/llm_be/finance/views.py b/llm_be/finance/views.py index 3cfcc60..5cb6c36 100644 --- a/llm_be/finance/views.py +++ b/llm_be/finance/views.py @@ -6,13 +6,20 @@ from rest_framework import permissions, status from rest_framework.response import Response from rest_framework.views import APIView -from finance.models import Invoice, Payment +from finance.models import Invoice, Payment, SubscriptionPlan, UserSubscription from finance.serializers import ( CheckoutSessionSerializer, InvoiceSerializer, PaymentSerializer, PortalSessionSerializer, + SubscriptionPlanSerializer, ) +from finance.services.plans import ( + needs_checkout, + plan_to_dict, + seed_subscription_plans, +) +from finance.services.quotas import get_usage_snapshot from finance.services.stripe_service import ( StripeNotConfiguredError, create_billing_portal_session, @@ -32,11 +39,34 @@ class CreateCheckoutSessionView(APIView): serializer.is_valid(raise_exception=True) try: - session = create_checkout_session( + sub = request.user.subscription + except UserSubscription.DoesNotExist: + sub = None + if ( + sub is not None + and sub.is_active + and sub.source == UserSubscription.Source.BACKER + ): + return Response( + { + "detail": ( + "This account has complimentary Backer access and " + "does not require payment." + ), + "needs_checkout": False, + }, + status=status.HTTP_400_BAD_REQUEST, + ) + + try: + session, plan = create_checkout_session( user=request.user, success_url=serializer.validated_data.get("success_url"), cancel_url=serializer.validated_data.get("cancel_url"), + plan_slug=serializer.validated_data.get("plan_slug"), ) + except ValueError as exc: + return Response({"detail": str(exc)}, status=status.HTTP_400_BAD_REQUEST) except StripeNotConfiguredError as exc: return Response( {"detail": str(exc)}, @@ -58,11 +88,11 @@ class CreateCheckoutSessionView(APIView): "company": request.user.company, "provider": Invoice.Provider.STRIPE, "status": Invoice.Status.OPEN, - "currency": settings.SUBSCRIPTION_PRICE_CURRENCY, - "amount_due": settings.SUBSCRIPTION_PRICE_AMOUNT_CENTS, + "currency": plan.currency or settings.SUBSCRIPTION_PRICE_CURRENCY, + "amount_due": plan.price_cents, "amount_paid": 0, "stripe_customer_id": getattr(session, "customer", None) or "", - "description": settings.SUBSCRIPTION_PRODUCT_NAME, + "description": plan.name, }, ) @@ -70,6 +100,7 @@ class CreateCheckoutSessionView(APIView): { "checkout_url": session.url, "session_id": session.id, + "plan_slug": plan.slug, }, status=status.HTTP_201_CREATED, ) @@ -129,6 +160,50 @@ class PaymentListView(APIView): return Response(PaymentSerializer(payments, many=True).data) +class PlanListView(APIView): + """Public-facing plan catalog (only `is_public` rows by default).""" + + permission_classes = (permissions.AllowAny,) + authentication_classes = () + + def get(self, request): + seed_subscription_plans(update_existing=False) + include_all = ( + request.user + and request.user.is_authenticated + and request.user.is_staff + and request.query_params.get("all") == "1" + ) + qs = SubscriptionPlan.objects.all() + if not include_all: + qs = qs.filter(is_public=True) + return Response(SubscriptionPlanSerializer(qs, many=True).data) + + +class SubscriptionMeView(APIView): + """Current user's plan, checkout need, and usage snapshot (#16/#17/#36).""" + + def get(self, request): + seed_subscription_plans(update_existing=False) + try: + sub = request.user.subscription + except UserSubscription.DoesNotExist: + sub = None + + usage = get_usage_snapshot(request.user) + payload = { + "plan": plan_to_dict(sub.plan) if sub and sub.plan_id else None, + "status": sub.status if sub else UserSubscription.Status.NONE, + "source": sub.source if sub else UserSubscription.Source.NONE, + "needs_checkout": needs_checkout(request.user), + "stripe_subscription_id": ( + sub.stripe_subscription_id if sub else "" + ), + "usage": usage.to_dict(), + } + return Response(payload) + + class StripeWebhookView(APIView): """Verify Stripe signatures and upsert local invoice/payment rows.""" diff --git a/llm_be/llm_be/settings.py b/llm_be/llm_be/settings.py index ba363ed..dfbacde 100644 --- a/llm_be/llm_be/settings.py +++ b/llm_be/llm_be/settings.py @@ -172,7 +172,7 @@ INSTALLED_APPS = [ "whitenoise.runserver_nostatic", "django.contrib.staticfiles", "chat_backend", - "finance", + "finance.apps.FinanceConfig", "rest_framework", "corsheaders", "rest_framework_simplejwt.token_blacklist", @@ -307,6 +307,8 @@ os.makedirs(directory_path, exist_ok=True) ALLOW_IMAGE_GENERATION = env_bool("ALLOW_IMAGE_GENERATION", False) ALLOW_INTERNET_ACCESS = env_bool("ALLOW_INTERNET_ACCESS", True) +# When True, chat turns require an active plan and respect prompt/token quotas. +ENFORCE_SUBSCRIPTION_GATES = env_bool("ENFORCE_SUBSCRIPTION_GATES", True) # Self-serve account registration (sign-up page). Default off — enable via # control-node secret (chat_backend_.env) when ready for public sign-up. @@ -346,7 +348,7 @@ SUBSCRIPTION_PRICE_INTERVAL = ( env("SUBSCRIPTION_PRICE_INTERVAL", "month") or "month" ).lower() SUBSCRIPTION_PRODUCT_NAME = ( - env("SUBSCRIPTION_PRODUCT_NAME", "Chat Subscription") or "Chat Subscription" + env("SUBSCRIPTION_PRODUCT_NAME", "Founders") or "Founders" ) FRONTEND_BASE_URL = ( diff --git a/llm_be/llm_be/test_runner.py b/llm_be/llm_be/test_runner.py index 2224620..95d6747 100644 --- a/llm_be/llm_be/test_runner.py +++ b/llm_be/llm_be/test_runner.py @@ -18,3 +18,5 @@ class ChatBackendTestRunner(DiscoverRunner): os.environ.setdefault("SKIP_RAG_INIT", "1") super().setup_test_environment(**kwargs) settings.PASSWORD_HASHERS = ["django.contrib.auth.hashers.MD5PasswordHasher"] + # Opt-in for plan/quota tests via @override_settings. + settings.ENFORCE_SUBSCRIPTION_GATES = False