use org.subscription instead of user.subscription

orgs/models.py

@@ -98,6 +98,15 @@ def migrate_from_appuser(self, user: "AppUser") -> Org:
             is_paying=user.is_paying,
         )
 
+    def get_dollars_spent_this_month(self) -> float:
+        today = timezone.now()
+        cents_spent = self.transactions.filter(
+            created_at__month=today.month,
+            created_at__year=today.year,
+            amount__gt=0,
+        ).aggregate(total=Sum("charged_amount"))["total"]
+        return (cents_spent or 0) / 100
+
 
 class Org(SafeDeleteModel):
     _safedelete_policy = SOFT_DELETE_CASCADE

payments/auto_recharge.py

@@ -3,12 +3,10 @@
 import sentry_sdk
 from loguru import logger
 
-from app_users.models import AppUser, PaymentProvider
+from app_users.models import PaymentProvider
 from daras_ai_v2.redis_cache import redis_lock
-from payments.tasks import (
-    send_monthly_budget_reached_email,
-    send_auto_recharge_failed_email,
-)
+from orgs.models import Org
+from payments.tasks import send_monthly_budget_reached_email
 
 
 class AutoRechargeException(Exception):
@@ -30,69 +28,68 @@ class AutoRechargeCooldownException(AutoRechargeException):
     pass
 
 
-def should_attempt_auto_recharge(user: AppUser):
+def should_attempt_auto_recharge(org: Org):
     return (
-        user.subscription
-        and user.subscription.auto_recharge_enabled
-        and user.subscription.payment_provider
-        and user.balance < user.subscription.auto_recharge_balance_threshold
+        org.subscription
+        and org.subscription.auto_recharge_enabled
+        and org.subscription.payment_provider
+        and org.balance < org.subscription.auto_recharge_balance_threshold
     )
 
 
-def run_auto_recharge_gracefully(user: AppUser):
+def run_auto_recharge_gracefully(org: Org):
     """
-    Wrapper over _auto_recharge_user, that handles exceptions so that it can:
+    Wrapper over _auto_recharge_org, that handles exceptions so that it can:
     - log exceptions
     - send emails when auto-recharge fails
     - not retry if this is run as a background task
 
     Meant to be used in conjunction with should_attempt_auto_recharge
     """
     try:
-        with redis_lock(f"gooey/auto_recharge_user/v1/{user.uid}"):
-            _auto_recharge_user(user)
+        with redis_lock(f"gooey/auto_recharge_user/v1/{org.org_id}"):
+            _auto_recharge_org(org)
     except AutoRechargeCooldownException as e:
         logger.info(
-            f"Rejected auto-recharge because auto-recharge is in cooldown period for user"
-            f"{user=}, {e=}"
+            f"Rejected auto-recharge because auto-recharge is in cooldown period for org"
+            f"{org=}, {e=}"
         )
     except MonthlyBudgetReachedException as e:
-        send_monthly_budget_reached_email(user)
+        send_monthly_budget_reached_email(org)
         logger.info(
             f"Rejected auto-recharge because user has reached monthly budget"
-            f"{user=}, spending=${e.spending}, budget=${e.budget}"
+            f"{org=}, spending=${e.spending}, budget=${e.budget}"
         )
     except Exception as e:
         traceback.print_exc()
         sentry_sdk.capture_exception(e)
-        send_auto_recharge_failed_email(user)
 
 
-def _auto_recharge_user(user: AppUser):
+def _auto_recharge_org(org: Org):
     """
     Returns whether a charge was attempted
     """
     from payments.webhooks import StripeWebhookHandler
 
     assert (
-        user.subscription.payment_provider == PaymentProvider.STRIPE
+        org.subscription.payment_provider == PaymentProvider.STRIPE
     ), "Auto recharge is only supported with Stripe"
 
     # check for monthly budget
-    dollars_spent = user.get_dollars_spent_this_month()
+    dollars_spent = org.get_dollars_spent_this_month()
     if (
-        dollars_spent + user.subscription.auto_recharge_topup_amount
-        > user.subscription.monthly_spending_budget
+        dollars_spent + org.subscription.auto_recharge_topup_amount
+        > org.subscription.monthly_spending_budget
     ):
         raise MonthlyBudgetReachedException(
             "Performing this top-up would exceed your monthly recharge budget",
-            budget=user.subscription.monthly_spending_budget,
+            budget=org.subscription.monthly_spending_budget,
             spending=dollars_spent,
         )
 
     try:
-        invoice = user.subscription.stripe_get_or_create_auto_invoice(
-            amount_in_dollars=user.subscription.auto_recharge_topup_amount,
+        invoice = org.subscription.stripe_get_or_create_auto_invoice(
+            amount_in_dollars=org.subscription.auto_recharge_topup_amount,
             metadata_key="auto_recharge",
         )
     except Exception as e:
@@ -106,9 +103,9 @@ def _auto_recharge_user(user: AppUser):
 
     # get default payment method and attempt payment
     assert invoice.status == "open"  # sanity check
-    pm = user.subscription.stripe_get_default_payment_method()
+    pm = org.subscription.stripe_get_default_payment_method()
     if not pm:
-        logger.warning(f"{user} has no default payment method, cannot auto-recharge")
+        logger.warning(f"{org} has no default payment method, cannot auto-recharge")
         return
 
     try:
@@ -119,4 +116,6 @@ def _auto_recharge_user(user: AppUser):
         ) from e
     else:
         assert invoice_data.paid
-        StripeWebhookHandler.handle_invoice_paid(uid=user.uid, invoice=invoice_data)
+        StripeWebhookHandler.handle_invoice_paid(
+            org_id=org.org_id, invoice=invoice_data
+        )

payments/tasks.py

@@ -29,6 +29,7 @@ def send_monthly_spending_notification_email(id: int):
                 "monthly_spending_notification_threshold_email.html"
             ).render(
                 user=owner.user,
+                org=org,
                 account_url=get_app_route_url(account_route),
             ),
         )
@@ -40,43 +41,27 @@ def send_monthly_spending_notification_email(id: int):
         org.subscription.save(update_fields=["monthly_spending_notification_sent_at"])
 
 
-def send_monthly_budget_reached_email(user: AppUser):
+def send_monthly_budget_reached_email(org: Org):
     from routers.account import account_route
 
-    if not user.email:
-        return
+    for owner in org.get_owners():
+        if not owner.user.email:
+            continue
 
-    email_body = templates.get_template("monthly_budget_reached_email.html").render(
-        user=user,
-        account_url=get_app_route_url(account_route),
-    )
-    send_email_via_postmark(
-        from_address=settings.SUPPORT_EMAIL,
-        to_address=user.email,
-        subject="[Gooey.AI] Monthly Budget Reached",
-        html_body=email_body,
-    )
+        email_body = templates.get_template("monthly_budget_reached_email.html").render(
+            user=owner.user,
+            org=org,
+            account_url=get_app_route_url(account_route),
+        )
+        send_email_via_postmark(
+            from_address=settings.SUPPORT_EMAIL,
+            to_address=owner.user.email,
+            subject="[Gooey.AI] Monthly Budget Reached",
+            html_body=email_body,
+        )
 
     # IMPORTANT: always use update_fields=... when updating subscription
     # info. We don't want to overwrite other changes made to subscription
     # during the same time
-    user.subscription.monthly_budget_email_sent_at = timezone.now()
-    user.subscription.save(update_fields=["monthly_budget_email_sent_at"])
-
-
-def send_auto_recharge_failed_email(user: AppUser):
-    from routers.account import account_route
-
-    if not user.email:
-        return
-
-    email_body = templates.get_template("auto_recharge_failed_email.html").render(
-        user=user,
-        account_url=get_app_route_url(account_route),
-    )
-    send_email_via_postmark(
-        from_address=settings.SUPPORT_EMAIL,
-        to_address=user.email,
-        subject="[Gooey.AI] Auto-Recharge failed",
-        html_body=email_body,
-    )
+    org.subscription.monthly_budget_email_sent_at = timezone.now()
+    org.subscription.save(update_fields=["monthly_budget_email_sent_at"])

routers/paypal.py

@@ -126,15 +126,16 @@ def create_subscription(request: Request, payload: dict = fastapi_request_json):
     if plan.deprecated:
         return JSONResponse({"error": "Deprecated plan"}, status_code=400)
 
-    if request.user.subscription and request.user.subscription.is_paid():
+    org, _ = request.user.get_or_create_personal_org()
+    if org.subscription and org.subscription.is_paid():
         return JSONResponse(
             {"error": "User already has an active subscription"}, status_code=400
         )
 
     paypal_plan_info = plan.get_paypal_plan()
     pp_subscription = paypal.Subscription.create(
         plan_id=paypal_plan_info["plan_id"],
-        custom_id=request.user.uid,
+        custom_id=org.org_id,
         plan=paypal_plan_info.get("plan", {}),
         application_context={
             "brand_name": "Gooey.AI",

templates/monthly_budget_reached_email.html

@@ -1,6 +1,6 @@
-{% set dollars_spent = user.get_dollars_spent_this_month() %}
-{% set monthly_budget = user.subscription.monthly_spending_budget %}
-{% set threshold = user.subscription.auto_recharge_balance_threshold %}
+{% set dollars_spent = org.get_dollars_spent_this_month() %}
+{% set monthly_budget = org.subscription.monthly_spending_budget %}
+{% set threshold = org.subscription.auto_recharge_balance_threshold %}
 
 <p>
   Hey, {{ user.first_name() }}!
@@ -18,7 +18,7 @@
 </p>
 
 <ul>
-  <li>Credit Balance: {{ user.balance }} credits</li>
+  <li>Credit Balance: {{ org.balance }} credits</li>
   <li>Monthly Budget: ${{ monthly_budget }}</li>
   <li>Spending this month: ${{ dollars_spent }}</li>
 </ul>

templates/monthly_spending_notification_threshold_email.html

@@ -1,16 +1,16 @@
-{% set dollars_spent = user.get_dollars_spent_this_month() %}
+{% set dollars_spent = org.get_dollars_spent_this_month() %}
 
 <p>
     Hi, {{ user.first_name() }}!
 </p>
 
 <p>
     Your spend on Gooey.AI so far this month is ${{ dollars_spent }}, exceeding your notification threshold
-    of ${{ user.subscription.monthly_spending_notification_threshold }}. 
+    of ${{ org.subscription.monthly_spending_notification_threshold }}. 
 </p>
 
 <p>
-    Your monthly budget is ${{ user.subscription.monthly_spending_budget }}, after which auto-recharge will be
+    Your monthly budget is ${{ org.subscription.monthly_spending_budget }}, after which auto-recharge will be
     paused and all runs / API calls will be rejected.
 </p>