Merge pull request #277 from GooeyAI/send-low-balance-email

Send Low Balance Email

app_users/admin.py

@@ -41,6 +41,7 @@ class AppUserAdmin(admin.ModelAdmin):
         "view_transactions",
         "open_in_firebase",
         "open_in_stripe",
+        "low_balance_email_sent_at",
     ]
 
     @admin.display(description="User Runs")

app_users/migrations/0012_appuser_low_balance_email_sent_at.py

@@ -0,0 +1,18 @@
+# Generated by Django 4.2.7 on 2024-02-14 07:23
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('app_users', '0011_appusertransaction_charged_amount_and_more'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='appuser',
+            name='low_balance_email_sent_at',
+            field=models.DateTimeField(blank=True, null=True),
+        ),
+    ]

app_users/models.py

@@ -89,6 +89,8 @@ class AppUser(models.Model):
     stripe_customer_id = models.CharField(max_length=255, default="", blank=True)
     is_paying = models.BooleanField("paid", default=False)
 
+    low_balance_email_sent_at = models.DateTimeField(null=True, blank=True)
+
     created_at = models.DateTimeField(
         "created", editable=False, blank=True, default=timezone.now
     )
@@ -263,6 +265,10 @@ class AppUserTransaction(models.Model):
 
     class Meta:
         verbose_name = "Transaction"
+        indexes = [
+            models.Index(fields=["user", "amount", "-created_at"]),
+            models.Index(fields=["-created_at"]),
+        ]
 
     def __str__(self):
         return f"{self.invoice_id} ({self.amount})"

celeryapp/tasks.py

@@ -1,4 +1,4 @@
-from fastapi import HTTPException
+import datetime
 import html
 import traceback
 import typing
@@ -7,16 +7,20 @@
 
 import requests
 import sentry_sdk
+from django.db.models import Sum
+from django.utils import timezone
+from fastapi import HTTPException
 
 import gooey_ui as st
-from app_users.models import AppUser
+from app_users.models import AppUser, AppUserTransaction
 from bots.models import SavedRun
 from celeryapp.celeryconfig import app
 from daras_ai.image_input import truncate_text_words
 from daras_ai_v2 import settings
 from daras_ai_v2.base import StateKeys, BasePage
 from daras_ai_v2.exceptions import UserError
 from daras_ai_v2.send_email import send_email_via_postmark
+from daras_ai_v2.send_email import send_low_balance_email
 from daras_ai_v2.settings import templates
 from gooey_ui.pubsub import realtime_push
 from gooey_ui.state import set_query_params
@@ -126,6 +130,7 @@ def save(done=False):
         save(done=True)
         if not is_api_call:
             send_email_on_completion(page, sr)
+        run_low_balance_email_check(uid)
 
 
 def err_msg_for_exc(e: Exception):
@@ -153,6 +158,43 @@ def err_msg_for_exc(e: Exception):
         return f"{type(e).__name__}: {e}"
 
 
+def run_low_balance_email_check(uid: str):
+    # don't send email if feature is disabled
+    if not settings.LOW_BALANCE_EMAIL_ENABLED:
+        return
+    user = AppUser.objects.get(uid=uid)
+    # don't send email if user is not paying or has enough balance
+    if not user.is_paying or user.balance > settings.LOW_BALANCE_EMAIL_CREDITS:
+        return
+    last_purchase = (
+        AppUserTransaction.objects.filter(user=user, amount__gt=0)
+        .order_by("-created_at")
+        .first()
+    )
+    email_date_cutoff = timezone.now() - datetime.timedelta(
+        days=settings.LOW_BALANCE_EMAIL_DAYS
+    )
+    # send email if user has not been sent email in last X days or last purchase was after last email sent
+    if (
+        # user has not been sent any email
+        not user.low_balance_email_sent_at
+        # user was sent email before X days
+        or (user.low_balance_email_sent_at < email_date_cutoff)
+        # user has made a purchase after last email sent
+        or (last_purchase and last_purchase.created_at > user.low_balance_email_sent_at)
+    ):
+        # calculate total credits consumed in last X days
+        total_credits_consumed = abs(
+            AppUserTransaction.objects.filter(
+                user=user, amount__lt=0, created_at__gte=email_date_cutoff
+            ).aggregate(Sum("amount"))["amount__sum"]
+            or 0
+        )
+        send_low_balance_email(user=user, total_credits_consumed=total_credits_consumed)
+        user.low_balance_email_sent_at = timezone.now()
+        user.save(update_fields=["low_balance_email_sent_at"])
+
+
 def send_email_on_completion(page: BasePage, sr: SavedRun):
     run_time_sec = sr.run_time.total_seconds()
     if (

conftest.py

@@ -9,6 +9,7 @@
 from auth import auth_backend
 from celeryapp import app
 from daras_ai_v2.base import BasePage
+from daras_ai_v2.send_email import pytest_outbox
 
 
 @pytest.fixture(scope="session")
@@ -68,6 +69,11 @@ def runner(*args, **kwargs):
             t.join()
 
 
+@pytest.fixture(autouse=True)
+def clear_pytest_outbox():
+    pytest_outbox.clear()
+
+
 # class DummyDatabaseBlocker(pytest_django.plugin._DatabaseBlocker):
 #     class _dj_db_wrapper:
 #         def ensure_connection(self):

daras_ai_v2/send_email.py

@@ -1,4 +1,5 @@
 import smtplib
+import sys
 import typing
 from email.mime.application import MIMEApplication
 from email.mime.multipart import MIMEMultipart
@@ -12,6 +13,7 @@
 from daras_ai_v2 import settings
 from daras_ai_v2.settings import templates
 from gooey_ui import UploadedFile
+from routers.billing import account_url
 
 
 def send_reported_run_email(
@@ -43,6 +45,31 @@ def send_reported_run_email(
     )
 
 
+def send_low_balance_email(
+    *,
+    user: AppUser,
+    total_credits_consumed: int,
+):
+    recipeints = "[email protected], [email protected]"
+    html_body = templates.get_template("low_balance_email.html").render(
+        user=user,
+        url=account_url,
+        total_credits_consumed=total_credits_consumed,
+        settings=settings,
+    )
+    send_email_via_postmark(
+        from_address=settings.SUPPORT_EMAIL,
+        to_address=user.email or recipeints,
+        bcc=recipeints,
+        subject="Your Gooey.AI credit balance is low",
+        html_body=html_body,
+    )
+
+
+is_running_pytest = "pytest" in sys.modules
+pytest_outbox = []
+
+
 def send_email_via_postmark(
     *,
     from_address: str,
@@ -56,6 +83,21 @@ def send_email_via_postmark(
         "outbound", "gooey-ai-workflows", "announcements"
     ] = "outbound",
 ):
+    if is_running_pytest:
+        pytest_outbox.append(
+            dict(
+                from_address=from_address,
+                to_address=to_address,
+                cc=cc,
+                bcc=bcc,
+                subject=subject,
+                html_body=html_body,
+                text_body=text_body,
+                message_stream=message_stream,
+            ),
+        )
+        return
+
     r = requests.post(
         "https://api.postmarkapp.com/email",
         headers={

daras_ai_v2/settings.py

@@ -268,6 +268,10 @@
 ANON_USER_FREE_CREDITS = config("ANON_USER_FREE_CREDITS", 25, cast=int)
 LOGIN_USER_FREE_CREDITS = config("LOGIN_USER_FREE_CREDITS", 1000, cast=int)
 
+LOW_BALANCE_EMAIL_CREDITS = config("LOW_BALANCE_EMAIL_CREDITS", 200, cast=int)
+LOW_BALANCE_EMAIL_DAYS = config("LOW_BALANCE_EMAIL_DAYS", 7, cast=int)
+LOW_BALANCE_EMAIL_ENABLED = config("LOW_BALANCE_EMAIL_ENABLED", True, cast=bool)
+
 stripe.api_key = config("STRIPE_SECRET_KEY", None)
 STRIPE_ENDPOINT_SECRET = config("STRIPE_ENDPOINT_SECRET", None)
 

routers/billing.py

@@ -121,9 +121,7 @@ def account(request: Request):
         "user_credits": request.user.balance,
         "subscription": get_user_subscription(request.user),
         "is_admin": is_admin,
-        "canonical_url": str(
-            furl(settings.APP_BASE_URL) / router.url_path_for(account.__name__)
-        ),
+        "canonical_url": account_url,
     }
 
     return templates.TemplateResponse("account.html", context)

templates/low_balance_email.html

@@ -0,0 +1,17 @@
+<p>
+    Hey {{ user.display_name }}!
+</p>
+
+<p>
+    This is a friendly reminder that your Gooey.AI balance is now just {{ user.balance }}.
+    Your account has consumed {{ total_credits_consumed }} credits in the last {{ settings.LOW_BALANCE_EMAIL_DAYS }} days.
+    <br><br>
+    To buy more credits, please visit <a href="{{ url }}">https://gooey.ai/account</a>.
+    <br><br>
+    As always, email us at [email protected] if you have any questions too. 
+    <br><br>
+    Thanks again for your business,
+    <br>
+    Sean and the Gooey.AI team
+</p>
+{{  "{{{ pm:unsubscribe }}}" }}

tests/test_low_balance_email_check.py

@@ -0,0 +1,127 @@
+from django.utils import timezone
+
+from app_users.models import AppUserTransaction
+from bots.models import AppUser
+from celeryapp.tasks import run_low_balance_email_check
+from daras_ai_v2 import settings
+from daras_ai_v2.send_email import pytest_outbox
+
+
+def test_dont_send_email_if_feature_is_disabled(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user", is_paying=True, balance=0, is_anonymous=False
+    )
+    settings.LOW_BALANCE_EMAIL_ENABLED = False
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_dont_send_email_if_user_is_not_paying(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user", is_paying=False, balance=0, is_anonymous=False
+    )
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_dont_send_email_if_user_has_enough_balance(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user", is_paying=True, balance=500, is_anonymous=False
+    )
+    settings.LOW_BALANCE_EMAIL_CREDITS = 100
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_dont_send_email_if_user_has_been_emailed_recently(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user",
+        is_paying=True,
+        balance=66,
+        is_anonymous=False,
+        low_balance_email_sent_at=timezone.now(),
+    )
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    settings.LOW_BALANCE_EMAIL_DAYS = 1
+    settings.LOW_BALANCE_EMAIL_CREDITS = 100
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_send_email_if_user_has_been_email_recently_but_made_a_purchase(
+    transactional_db,
+):
+    user = AppUser.objects.create(
+        uid="test_user",
+        is_paying=True,
+        balance=22,
+        is_anonymous=False,
+        low_balance_email_sent_at=timezone.now(),
+    )
+    AppUserTransaction.objects.create(
+        invoice_id="test_invoice_1",
+        user=user,
+        amount=100,
+        created_at=timezone.now(),
+        end_balance=100,
+    )
+    AppUserTransaction.objects.create(
+        invoice_id="test_invoice_2",
+        user=user,
+        amount=-78,
+        created_at=timezone.now(),
+        end_balance=22,
+    )
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    settings.LOW_BALANCE_EMAIL_DAYS = 1
+    settings.LOW_BALANCE_EMAIL_CREDITS = 100
+    run_low_balance_email_check(user.uid)
+
+    assert len(pytest_outbox) == 1
+    assert " 22" in pytest_outbox[0]["html_body"]
+    assert " 78" in pytest_outbox[0]["html_body"]
+
+    pytest_outbox.clear()
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox
+
+
+def test_send_email(transactional_db):
+    user = AppUser.objects.create(
+        uid="test_user",
+        is_paying=True,
+        balance=66,
+        is_anonymous=False,
+    )
+    AppUserTransaction.objects.create(
+        invoice_id="test_invoice_1",
+        user=user,
+        amount=-100,
+        created_at=timezone.now() - timezone.timedelta(days=2),
+        end_balance=150 + 66,
+    )
+    AppUserTransaction.objects.create(
+        invoice_id="test_invoice_2",
+        user=user,
+        amount=-150,
+        created_at=timezone.now(),
+        end_balance=66,
+    )
+
+    settings.LOW_BALANCE_EMAIL_ENABLED = True
+    settings.LOW_BALANCE_EMAIL_DAYS = 1
+    settings.LOW_BALANCE_EMAIL_CREDITS = 100
+
+    run_low_balance_email_check(user.uid)
+    assert len(pytest_outbox) == 1
+    body = pytest_outbox[0]["html_body"]
+    assert " 66" in body
+    assert " 150" in body
+    assert " pm:unsubscribe" in body
+    assert " 100" not in body
+
+    pytest_outbox.clear()
+    run_low_balance_email_check(user.uid)
+    assert not pytest_outbox