Pricing improvements and bug fixes

app_users/migrations/0017_alter_appuser_subscription.py

@@ -0,0 +1,20 @@
+# Generated by Django 4.2.7 on 2024-06-10 09:21
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('payments', '0004_alter_subscription_auto_recharge_balance_threshold'),
+        ('app_users', '0016_appuser_disable_rate_limits'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='appuser',
+            name='subscription',
+            field=models.OneToOneField(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='user', to='payments.subscription'),
+        ),
+    ]

celeryapp/tasks.py

@@ -18,15 +18,14 @@
 from celeryapp.celeryconfig import app
 from daras_ai.image_input import truncate_text_words
 from daras_ai_v2 import settings
-from daras_ai_v2.auto_recharge import auto_recharge_user
 from daras_ai_v2.base import StateKeys, BasePage
 from daras_ai_v2.exceptions import UserError
-from daras_ai_v2.redis_cache import redis_lock
 from daras_ai_v2.send_email import send_email_via_postmark, 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
 from gooeysite.bg_db_conn import db_middleware, next_db_safe
+from payments.tasks import run_auto_recharge_gracefully
 
 
 @app.task
@@ -97,13 +96,13 @@ def save(done=False):
         page.dump_state_to_sr(st.session_state | output, sr)
 
     try:
-        gen = page.run(st.session_state)
+        gen = page.run_with_auto_recharge(st.session_state)
         save()
         while True:
             # record time
             start_time = time()
             try:
-                # advance the generator (to further progress of run())
+                # advance the generator (to further progress of run_with_auto_recharge())
                 yield_val = next_db_safe(gen)
                 # increment total time taken after every iteration
                 run_time += time() - start_time
@@ -143,7 +142,9 @@ def save(done=False):
         save(done=True)
         if not is_api_call:
             send_email_on_completion(page, sr)
+
         run_low_balance_email_check(uid)
+        run_auto_recharge_gracefully(uid)
 
 
 def err_msg_for_exc(e: Exception):
@@ -253,11 +254,3 @@ def send_integration_attempt_email(*, user_id: int, platform: Platform, run_url:
         subject=f"{user.display_name} Attempted to Connect to {platform.label}",
         html_body=html_body,
     )
-
-
-@app.task
-def auto_recharge(*, user_id: int):
-    redis_lock_key = f"gooey/auto_recharge/{user_id}"
-    with redis_lock(redis_lock_key):
-        user = AppUser.objects.get(id=user_id)
-        auto_recharge_user(user)

daras_ai_v2/base.py

@@ -11,6 +11,7 @@
 from time import sleep
 from types import SimpleNamespace
 
+from loguru import logger
 import sentry_sdk
 from django.db.models import Sum
 from django.utils import timezone
@@ -24,6 +25,7 @@
 )
 from starlette.requests import Request
 
+from daras_ai_v2.exceptions import UserError
 import gooey_ui as st
 from app_users.models import AppUser, AppUserTransaction
 from bots.models import (
@@ -37,7 +39,6 @@
 from daras_ai.text_format import format_number_with_suffix
 from daras_ai_v2 import settings, urls
 from daras_ai_v2.api_examples_widget import api_example_generator
-from daras_ai_v2.auto_recharge import user_should_auto_recharge
 from daras_ai_v2.breadcrumbs import render_breadcrumbs, get_title_breadcrumbs
 from daras_ai_v2.copy_to_clipboard_button_widget import (
     copy_to_clipboard_button,
@@ -71,6 +72,11 @@
 from gooey_ui.components.pills import pill
 from gooey_ui.pubsub import realtime_pull
 from routers.account import AccountTabs
+from payments.auto_recharge import (
+    AutoRechargeException,
+    auto_recharge_user,
+    should_attempt_auto_recharge,
+)
 from routers.root import RecipeTabs
 
 DEFAULT_META_IMG = (
@@ -1355,6 +1361,29 @@ def _render_help(self):
     def render_usage_guide(self):
         raise NotImplementedError
 
+    def run_with_auto_recharge(self, state: dict) -> typing.Iterator[str | None]:
+        if not self.check_credits() and should_attempt_auto_recharge(self.request.user):
+            yield "Low balance detected. Recharging..."
+            try:
+                auto_recharge_user(uid=self.request.user.uid)
+            except AutoRechargeException as e:
+                # raise this error only if another auto-recharge
+                # procedure didn't complete successfully
+                self.request.user.refresh_from_db()
+                if not self.check_credits():
+                    raise UserError(str(e)) from e
+            else:
+                self.request.user.refresh_from_db()
+
+        if not self.check_credits():
+            example_id, run_id, uid = extract_query_params(gooey_get_query_params())
+            error_msg = self.generate_credit_error_message(
+                example_id=example_id, run_id=run_id, uid=uid
+            )
+            raise UserError(error_msg)
+
+        yield from self.run(state)
+
     def run(self, state: dict) -> typing.Iterator[str | None]:
         # initialize request and response
         request = self.RequestModel.parse_obj(state)
@@ -1505,19 +1534,17 @@ def estimate_run_duration(self) -> int | None:
         pass
 
     def on_submit(self):
-        from celeryapp.tasks import auto_recharge
-
         try:
             example_id, run_id, uid = self.create_new_run(enable_rate_limits=True)
         except RateLimitExceeded as e:
             st.session_state[StateKeys.run_status] = None
             st.session_state[StateKeys.error_msg] = e.detail.get("error", "")
             return
 
-        if user_should_auto_recharge(self.request.user):
-            auto_recharge.delay(user_id=self.request.user.id)
-
-        if settings.CREDITS_TO_DEDUCT_PER_RUN and not self.check_credits():
+        if not self.check_credits() and not should_attempt_auto_recharge(
+            self.request.user
+        ):
+            # insufficient balance for this run and auto-recharge isn't setup
             st.session_state[StateKeys.run_status] = None
             st.session_state[StateKeys.error_msg] = self.generate_credit_error_message(
                 example_id, run_id, uid
@@ -1978,6 +2005,9 @@ def run_as_api_tab(self):
         manage_api_keys(self.request.user)
 
     def check_credits(self) -> bool:
+        if not settings.CREDITS_TO_DEDUCT_PER_RUN:
+            return True
+
         assert self.request, "request must be set to check credits"
         assert self.request.user, "request.user must be set to check credits"
         return self.request.user.balance >= self.get_price_roundoff(st.session_state)