Update main.py

main.py

@@ -404,8 +404,6 @@ def delete_sozo_project(project_id):
         logger.error(f"CRITICAL ERROR during delete project {project_id}: {traceback.format_exc()}")
         return jsonify({'error': str(e)}), 500
 
-# In main.py
-
 @app.route('/api/sozo/projects/<string:project_id>/generate-report', methods=['POST'])
 def generate_sozo_report(project_id):
     logger.info(f"POST /api/sozo/projects/{project_id}/generate-report - Received request")
@@ -416,6 +414,16 @@ def generate_sozo_report(project_id):
         if not uid:
             return jsonify({'error': 'Unauthorized'}), 401
         
+        # Check user credits first
+        user_ref = db.reference(f'users/{uid}')
+        user_data = user_ref.get()
+        if not user_data:
+            return jsonify({'error': 'User not found'}), 404
+        
+        current_credits = user_data.get('credits', 0)
+        if current_credits < 2:
+            return jsonify({'error': 'Insufficient credits. Report generation requires 2 credits.'}), 402
+        
         project_ref = db.reference(f'sozo_projects/{project_id}')
         project_data = project_ref.get()
         
@@ -454,11 +462,15 @@ def generate_sozo_report(project_id):
         
         project_ref.update(update_data)
         
-        logger.info(f"Project {project_id} successfully updated with draft data.")
+        # Deduct credits ONLY after successful generation
+        user_ref.update({'credits': current_credits - 2})
+        
+        logger.info(f"Project {project_id} successfully updated with draft data. User {uid} charged 2 credits.")
         
         return jsonify({
             'success': True,
-            'project': {**project_data, **update_data}
+            'project': {**project_data, **update_data},
+            'credits_remaining': current_credits - 2
         }), 200
         
     except Exception as e:
@@ -478,6 +490,16 @@ def generate_sozo_video(project_id):
         uid = verify_token(token)
         if not uid: return jsonify({'error': 'Unauthorized'}), 401
         
+        # Check user credits first
+        user_ref = db.reference(f'users/{uid}')
+        user_data = user_ref.get()
+        if not user_data:
+            return jsonify({'error': 'User not found'}), 404
+        
+        current_credits = user_data.get('credits', 0)
+        if current_credits < 5:
+            return jsonify({'error': 'Insufficient credits. Video generation requires 5 credits.'}), 402
+        
         project_ref = db.reference(f'sozo_projects/{project_id}')
         project_data = project_ref.get()
         if not project_data or project_data.get('uid') != uid: return jsonify({'error': 'Project not found or unauthorized'}), 404
@@ -506,13 +528,87 @@ def generate_sozo_video(project_id):
         if not video_url: raise Exception("Video generation failed in core function.")
         
         project_ref.update({'status': 'video_complete', 'videoUrl': video_url})
-        return jsonify({'success': True, 'video_url': video_url}), 200
+        
+        # Deduct credits ONLY after successful generation
+        user_ref.update({'credits': current_credits - 5})
+        
+        return jsonify({
+            'success': True, 
+            'video_url': video_url,
+            'credits_remaining': current_credits - 5
+        }), 200
         
     except Exception as e:
         db.reference(f'sozo_projects/{project_id}').update({'status': 'failed', 'error': str(e)})
         traceback.print_exc()
         return jsonify({'error': str(e)}), 500
 
+
+@app.route('/api/sozo/projects/<string:project_id>/generate-slides', methods=['POST'])
+def generate_sozo_slides(project_id):
+    logger.info(f"POST /api/sozo/projects/{project_id}/generate-slides - Generating slides")
+    try:
+        token = request.headers.get('Authorization', '').split(' ')
+        uid = verify_token(token)
+        if not uid:
+            return jsonify({'error': 'Unauthorized'}), 401
+        
+        # Check user credits first
+        user_ref = db.reference(f'users/{uid}')
+        user_data = user_ref.get()
+        if not user_data:
+            return jsonify({'error': 'User not found'}), 404
+        
+        current_credits = user_data.get('credits', 0)
+        if current_credits < 5:
+            return jsonify({'error': 'Insufficient credits. Slide generation requires 5 credits.'}), 402
+        
+        project_ref = db.reference(f'sozo_projects/{project_id}')
+        project_data = project_ref.get()
+        
+        if not project_data or project_data.get('uid') != uid:
+            return jsonify({'error': 'Project not found or unauthorized'}), 404
+        
+        raw_md = project_data.get('rawMarkdown')
+        chart_urls = project_data.get('chartUrls', {})
+        
+        if not raw_md:
+            return jsonify({'error': 'Report must be generated before slides can be created.'}), 400
+            
+        # The planner AI needs an LLM instance
+        llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=os.getenv("GOOGLE_API_KEY"), temperature=0.2)
+
+        slides_data = generate_slides_from_report(
+            raw_md,
+            chart_urls,
+            uid,
+            project_id,
+            bucket,
+            llm
+        )
+        
+        if not slides_data:
+            raise Exception("Slide generation failed in core function.")
+            
+        # Save the slides data to the project in Firebase
+        project_ref.update({'slides': slides_data})
+        
+        # Deduct credits ONLY after successful generation
+        user_ref.update({'credits': current_credits - 5})
+        
+        logger.info(f"Project {project_id} successfully updated with {len(slides_data)} slides. User {uid} charged 5 credits.")
+        
+        return jsonify({
+            'success': True,
+            'slides': slides_data,
+            'credits_remaining': current_credits - 5
+        }), 200
+        
+    except Exception as e:
+        logger.error(f"CRITICAL error generating slides for project {project_id}: {traceback.format_exc()}")
+        db.reference(f'sozo_projects/{project_id}').update({'status': 'failed_slides', 'error': str(e)})
+        return jsonify({'error': str(e)}), 500
+
 @app.route('/api/sozo/projects/<string:project_id>/charts', methods=['POST'])
 def regenerate_sozo_chart(project_id):
     logger.info(f"Endpoint /charts POST for project {project_id}")
@@ -622,58 +718,6 @@ def update_narration_audio(project_id):
         logger.error(f"CRITICAL ERROR updating narration for {project_id}: {traceback.format_exc()}")
         return jsonify({'error': str(e)}), 500
 
-# In main.py, add this new endpoint
-
-@app.route('/api/sozo/projects/<string:project_id>/generate-slides', methods=['POST'])
-def generate_sozo_slides(project_id):
-    logger.info(f"POST /api/sozo/projects/{project_id}/generate-slides - Generating slides")
-    try:
-        token = request.headers.get('Authorization', '').split(' ')
-        uid = verify_token(token)
-        if not uid:
-            return jsonify({'error': 'Unauthorized'}), 401
-        
-        project_ref = db.reference(f'sozo_projects/{project_id}')
-        project_data = project_ref.get()
-        
-        if not project_data or project_data.get('uid') != uid:
-            return jsonify({'error': 'Project not found or unauthorized'}), 404
-        
-        raw_md = project_data.get('rawMarkdown')
-        chart_urls = project_data.get('chartUrls', {})
-        
-        if not raw_md:
-            return jsonify({'error': 'Report must be generated before slides can be created.'}), 400
-            
-        # The planner AI needs an LLM instance
-        llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=os.getenv("GOOGLE_API_KEY"), temperature=0.2)
-
-        slides_data = generate_slides_from_report(
-            raw_md,
-            chart_urls,
-            uid,
-            project_id,
-            bucket,
-            llm
-        )
-        
-        if not slides_data:
-            raise Exception("Slide generation failed in core function.")
-            
-        # Save the slides data to the project in Firebase
-        project_ref.update({'slides': slides_data})
-        
-        logger.info(f"Project {project_id} successfully updated with {len(slides_data)} slides.")
-        
-        return jsonify({
-            'success': True,
-            'slides': slides_data
-        }), 200
-        
-    except Exception as e:
-        logger.error(f"CRITICAL error generating slides for project {project_id}: {traceback.format_exc()}")
-        db.reference(f'sozo_projects/{project_id}').update({'status': 'failed_slides', 'error': str(e)})
-        return jsonify({'error': str(e)}), 500
 # -----------------------------------------------------------------------------
 # 5. UNIVERSAL ENDPOINTS (Waitlist, Feedback, Credits)
 # -----------------------------------------------------------------------------