Update app.py

app.py

@@ -23,12 +23,53 @@ from wordpress_xmlrpc import Client
 from wordpress_xmlrpc.compat import xmlrpc_client
 from wordpress_xmlrpc.methods import media
 
-# Retrieve user_id from query parameters
-user_params = st.experimental_get_query_params()
-user_id = user_params.get('user_id', [None])[0]
+def get_user_id():
+    url = 'https://songlabai.com/wp-json/custom-api/v1/current-user-id-session'
+    try:
+        response = requests.get(url)
+        if response.status_code == 200:
+            data = response.json()
+            if 'user_id' in data:
+                return data['user_id']
+            else:
+                return None
+        else:
+            return None
+    except Exception as e:
+        print(f"Error getting user ID: {e}")
+        return None
 
-# Save user_id in session_state
-st.session_state['user_id'] = user_id
+def get_subscription_info(user_id):
+    url = f'https://songlabai.com/wp-json/custom-api/subscription/{user_id}'
+    try:
+        response = requests.get(url)
+        if response.status_code == 200:
+            data = response.json()
+            return data
+        else:
+            return None
+    except Exception as e:
+        print(f"Error getting subscription info: {e}")
+        return None
+
+def update_user_limit(user_id):
+    url = f'https://songlabai.com/wp-json/custom-api/update-limit/{user_id}'
+    try:
+        response = requests.post(url)
+        if response.status_code == 200:
+            data = response.json()
+            if data.get('success') == True:
+                print("User limit updated successfully.")
+                return True
+            else:
+                print(f"Failed to update user limit: {data}")
+                return False
+        else:
+            print(f"Failed to update user limit, status code: {response.status_code}")
+            return False
+    except Exception as e:
+        print(f"Error updating user limit: {e}")
+        return False
 
 
 # Try to get API_URL from environment variables, if not found set to a default value
@@ -273,128 +314,78 @@ def time_post_request(api_url, headers=None, payload=None):
 
 def generate_audio(genre, energy_level, tempo, description, duration):
     count_down = 300
-    # Check if user_id is present
-    user_id = st.session_state.get('user_id', None)
-    if not user_id:
-        st.error("You need to be logged in to generate audio.")
-        return
-
-    # Proceed to check the subscription
-    subscription_api_url = f"https://songlabai.com/wp-json/custom-api/subscription/{user_id}"
-
-    try:
-        response = requests.get(subscription_api_url)
-        if response.status_code == 200:
-            subscription_info = response.json()
-            subscription_plan_id = subscription_info.get('subscription_plan_id')
-            has_exceeded_daily_limit = subscription_info.get('has_exceeded_daily_limit', False)
-            
-            # Enforce limits based on subscription_plan_id
-            if subscription_plan_id == '576':
-                if has_exceeded_daily_limit in [True, 'true', 'True']:
-                    st.error("You have exceeded your daily limit. Please upgrade your plan.")
-                    return
-                else:
-                    # Proceed to generate audio
-                    pass
-            elif subscription_plan_id in ['304', '305', '306']:
-                # User has unlimited generation
-                pass
-            else:
-                st.error("Invalid subscription plan. Please contact support.")
-                return
-        else:
-            st.error("Failed to retrieve your subscription information. Please try again later.")
-            return
-    except Exception as e:
-        st.error(f"An error occurred: {e}")
-        return
     prompt = f"Genre: {genre}, Energy Level: {energy_level}, Tempo: {tempo}, Description: {description},"
     payload = {"inputs": {"prompt": prompt, "duration": duration}}
     with st.spinner("Generating audio ..."):
         response = time_post_request(API_URL, headers, payload)
     placeholder1 = st.empty()
-    if response.status_code != 200:
-        # Temp Jay Start
+
+    attempt = 1
+    max_attempts = 3
+    while response.status_code != 200 and attempt < max_attempts:
+        # Log the error
+        print(f"Attempt {attempt}: Failed to generate audio.")
         print(f"Request Headers: {response.request.headers}")
         print(f"Request Body: {response.request.body}")
         print(f"Response Code: {response.status_code}")
         print(f"Response Body: {response.text}")
-        # Temp Jay End
 
-        
-        print(str(response.content))
+        # Wait for count_down seconds
         count_down_placeholder = st.empty()
         for seconds in range(count_down):
             count_down_placeholder.info(
-                f"The server is currently loading. Retrying in ⏳ {count_down - seconds} seconds."
+                f"Attempt {attempt}: The server is currently loading. Retrying in ⏳ {count_down - seconds} seconds."
             )
             time.sleep(1)
         count_down_placeholder.empty()
+
+        # Try again
         with st.spinner("Generating audio ..."):
             response = time_post_request(API_URL, headers, payload)
-        if response.status_code != 200:
-            count_down_placeholder = st.empty()
-            for seconds in range(count_down):
-                count_down_placeholder.info(
-                    f"2nd attempt failed. trying again one more time in ⏳ {count_down - seconds} seconds."
-                )
-                time.sleep(1)
-            count_down_placeholder.empty()
-            with st.spinner("Generating audio ..."):
-                response = time_post_request(API_URL, headers, payload)
-            if response.status_code != 200:
-                placeholder1.error(
-                    "Failed to generate audio after multiple attempts. Please try again later."
-                )
-            else:
-                placeholder1.success(f"✔️ Audio Generated, Loading...")
-                load_and_play_generated_audio(response)
-                placeholder1.empty()
-                if subscription_plan_id == '576':
-                    update_limit_api_url = f"https://songlabai.com/wp-json/custom-api/update-limit/{user_id}"
-                    try:
-                        update_response = requests.post(update_limit_api_url)
-                        if update_response.status_code == 200:
-                            st.success("Your daily limit has been updated.")
-                        else:
-                            st.error("Failed to update your usage limit. Please try again later.")
-                    except Exception as e:
-                        st.error(f"An error occurred while updating your limit: {e}")
-                    except Exception as e:
-                        st.error(f"An error occurred while updating your limit: {e}")
+        attempt += 1
 
+    if response.status_code != 200:
+        placeholder1.error(
+            "Failed to generate audio after multiple attempts. Please try again later."
+        )
+        return
+
+    # After successful response (after retries if necessary), perform API checks
+
+    # Get user ID
+    user_id = get_user_id()
+    if not user_id:
+        st.error("Please log in to generate audio.")
+        return
+
+    # Check subscription via API 2
+    subscription_info = get_subscription_info(user_id)
+    if not subscription_info:
+        st.error("Could not retrieve subscription information.")
+        return
+
+    subscription_plan_id = subscription_info['subscription_plan_id']
+    has_exceeded_daily_limit = subscription_info['has_exceeded_daily_limit']
+
+    if subscription_plan_id == '576':
+        if has_exceeded_daily_limit == True:
+            st.error("You have reached your daily limit for audio generation.")
+            return
         else:
+            # Allow the user to proceed
             placeholder1.success(f"✔️ Audio Generated, Loading...")
             load_and_play_generated_audio(response)
             placeholder1.empty()
-            if subscription_plan_id == '576':
-                update_limit_api_url = f"https://songlabai.com/wp-json/custom-api/update-limit/{user_id}"
-                try:
-                    update_response = requests.post(update_limit_api_url)
-                    if update_response.status_code == 200:
-                        st.success("Your daily limit has been updated.")
-                    else:
-                        st.error("Failed to update your usage limit. Please try again later.")
-                except Exception as e:
-                    st.error(f"An error occurred while updating your limit: {e}")
-    else:
+            # After generating audio, update the limit via API 3
+            update_user_limit(user_id)
+    elif subscription_plan_id in ['304', '305', '306']:
+        # Unlimited generation
         placeholder1.success(f"✔️ Audio Generated, Loading...")
         load_and_play_generated_audio(response)
         placeholder1.empty()
-        if subscription_plan_id == '576':
-            update_limit_api_url = f"https://songlabai.com/wp-json/custom-api/update-limit/{user_id}"
-            try:
-                update_response = requests.post(update_limit_api_url)
-                if update_response.status_code == 200:
-                    st.success("Your daily limit has been updated.")
-                else:
-                    st.error("Failed to update your usage limit. Please try again later.")
-            except Exception as e:
-                st.error(f"An error occurred while updating your limit: {e}")
-
-        
-
+    else:
+        st.error("Your subscription does not allow audio generation.")
+        return
 
 def load_and_play_generated_audio(response):