From a84933b224a73fb5647c3292f0628b3bf38d1774 Mon Sep 17 00:00:00 2001
From: clr-li <111320104+clr-li@users.noreply.github.com>
Date: Wed, 14 Feb 2024 21:59:54 -0800
Subject: [PATCH] Correct credit count

---
 celeryapp/tasks.py               | 33 +++++++++++++++++++++++++++++++-
 daras_ai_v2/base.py              | 25 +-----------------------
 daras_ai_v2/send_email.py        |  6 +++---
 templates/low_balance_email.html |  2 +-
 4 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/celeryapp/tasks.py b/celeryapp/tasks.py
index 5e1690b31..8a4f6a260 100644
--- a/celeryapp/tasks.py
+++ b/celeryapp/tasks.py
@@ -1,12 +1,13 @@
 import traceback
 import typing
+import datetime
 from time import time
 from types import SimpleNamespace
 
 import sentry_sdk
 
 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
@@ -16,6 +17,9 @@
 from daras_ai_v2.settings import templates
 from gooey_ui.pubsub import realtime_push
 from gooey_ui.state import set_query_params
+from daras_ai_v2.send_email import send_low_balance_email
+from django.db.models import Sum
+from django.utils import timezone
 
 
 @app.task
@@ -101,10 +105,37 @@ def save(done=False):
                 save()
     finally:
         save(done=True)
+        low_balance_email(sr)
         if not is_api_call:
             send_email_on_completion(page, sr)
 
 
+def low_balance_email(sr: SavedRun):
+    user = AppUser.objects.get(uid=sr.uid)
+    if (
+        user.is_paying
+        and user.balance < 500
+        and (
+            user.low_balance_email_sent_at == None
+            or user.low_balance_email_sent_at
+            < timezone.now() - datetime.timedelta(days=7)
+        )
+    ):
+        total_credits_consumed = (
+            -1
+            * AppUserTransaction.objects.filter(
+                user=user,
+                created_at__gte=timezone.now() - datetime.timedelta(days=7),
+            ).aggregate(Sum("amount"))["amount__sum"]
+        )
+        send_low_balance_email(
+            user=user,
+            total_credits_consumed=total_credits_consumed,
+        )
+        user.low_balance_email_sent_at = timezone.now()
+        user.save()
+
+
 def send_email_on_completion(page: BasePage, sr: SavedRun):
     run_time_sec = sr.run_time.total_seconds()
     if (
diff --git a/daras_ai_v2/base.py b/daras_ai_v2/base.py
index ae1150b63..d6243aeb7 100644
--- a/daras_ai_v2/base.py
+++ b/daras_ai_v2/base.py
@@ -17,7 +17,6 @@
 import sentry_sdk
 from django.utils import timezone
 from django.utils.text import slugify
-from django.db.models import Sum
 from fastapi import HTTPException
 from firebase_admin import auth
 from furl import furl
@@ -58,7 +57,7 @@
 from daras_ai_v2.query_params_util import (
     extract_query_params,
 )
-from daras_ai_v2.send_email import send_reported_run_email, send_low_balance_email
+from daras_ai_v2.send_email import send_reported_run_email
 from daras_ai_v2.tabs_widget import MenuTabs
 from daras_ai_v2.user_date_widgets import (
     render_js_dynamic_dates,
@@ -1388,28 +1387,6 @@ def estimate_run_duration(self) -> int | None:
 
     def on_submit(self):
         example_id, run_id, uid = self.create_new_run()
-        if (
-            self.request.user.is_paying
-            and self.request.user.balance < 500
-            # and (
-            #     self.request.user.low_balance_email_sent_at == None
-            #     or self.request.user.low_balance_email_sent_at
-            #     < timezone.now() - datetime.timedelta(days=7)
-            # )
-        ):
-            credits_consumed = (
-                -1
-                * AppUserTransaction.objects.filter(
-                    user=self.request.user,
-                    created_at__gte=timezone.now() - datetime.timedelta(days=7),
-                ).aggregate(Sum("amount"))["amount__sum"]
-            )
-            send_low_balance_email(
-                user=self.request.user, credits_consumed=credits_consumed
-            )
-            app_user = AppUser.objects.get(uid=uid)
-            app_user.low_balance_email_sent_at = timezone.now()
-            app_user.save()
         if settings.CREDITS_TO_DEDUCT_PER_RUN and not self.check_credits():
             st.session_state[StateKeys.run_status] = None
             st.session_state[StateKeys.error_msg] = self.generate_credit_error_message(
diff --git a/daras_ai_v2/send_email.py b/daras_ai_v2/send_email.py
index fea83a855..fc2445f41 100644
--- a/daras_ai_v2/send_email.py
+++ b/daras_ai_v2/send_email.py
@@ -46,19 +46,19 @@ def send_reported_run_email(
 def send_low_balance_email(
     *,
     user: AppUser,
-    credits_consumed: int,
+    total_credits_consumed: int,
 ):
     recipeints = "support@gooey.ai, devs@gooey.ai"
     html_body = templates.get_template("low_balance_email.html").render(
         user=user,
         url="https://gooey.ai/account",
-        credits_consumed=credits_consumed,
+        total_credits_consumed=total_credits_consumed,
     )
     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 ",
+        subject="Your Gooey.AI credit balance is low",
         html_body=html_body,
     )
 
diff --git a/templates/low_balance_email.html b/templates/low_balance_email.html
index b7dfa7850..d3215c861 100644
--- a/templates/low_balance_email.html
+++ b/templates/low_balance_email.html
@@ -3,7 +3,7 @@
 </p>
 
 <p>
-    This is a friendly reminder that your Gooey.AI balance is now just {{ user.balance }}. Your account has consumed {{ credits_consumed }} credits in the last 7 days. 
+    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 7 days. 
     <br><br>
     To buy more credits, please visit <a href="{{ url }}">https://gooey.ai/account</a>.
     <br><br>