Update main.py

main.py

@@ -312,6 +312,12 @@ def generate_story_endpoint():
             "story_type": story_type
         }
         story_ref.set(story_record)
+          # --- Subtract 5 Credits from the User ---
+        user_ref = db.reference(f"users/{uid}")
+        user_data = user_ref.get() or {}
+        current_credits = user_data.get("credits", 0)
+        new_credits = max(0, current_credits - 5)
+        user_ref.update({"credits": new_credits})
 
         preview = sections[0] if sections else {}
 
@@ -320,7 +326,8 @@ def generate_story_endpoint():
             "full_story": full_story,
             "preview": preview,
             "sections": sections,
-            "generation_times": story_record["generation_times"]
+            "generation_times": story_record["generation_times"],
+            "new_credits": new_credits
         })
 
     except Exception as e:
@@ -392,6 +399,13 @@ def generate_video_endpoint():
         # Optionally, update the story record with the video URL.
         story_ref.update({"video_url": video_url})
 
+          # --- Subtract 5 Credits from the User ---
+        user_ref = db.reference(f"users/{uid}")
+        user_data = user_ref.get() or {}
+        current_credits = user_data.get("credits", 0)
+        new_credits = max(0, current_credits - 5)
+        user_ref.update({"credits": new_credits})
+
         # Clean up temporary local files.
         for f in image_files:
             os.remove(f)
@@ -399,7 +413,7 @@ def generate_video_endpoint():
             os.remove(f)
         os.remove(video_file)
 
-        return jsonify({"video_url": video_url})
+        return jsonify({"video_url": video_url,"new_credits": new_credits})
     except Exception as e:
         return jsonify({'error': str(e)}), 500