Update main.py

main.py

@@ -354,6 +354,8 @@ 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 - Generating report")
@@ -365,22 +367,17 @@ def generate_sozo_report(project_id):
             logger.warning(f"Unauthorized access attempt to generate report for project {project_id}")
             return jsonify({'error': 'Unauthorized'}), 401
         
-        logger.info(f"User {uid} generating report for project {project_id}")
-        
         project_ref = db.reference(f'sozo_projects/{project_id}')
         project_data = project_ref.get()
         
         if not project_data or project_data.get('uid') != uid:
-            logger.warning(f"Project {project_id} not found or unauthorized for user {uid}")
             return jsonify({'error': 'Project not found or unauthorized'}), 404
         
-        logger.info(f"Project {project_id} validated for user {uid}")
-        
-        blob_path = "/".join(project_data['originalDataUrl'].split('/')[4:])
+        blob_path = f"sozo_projects/{uid}/{project_id}/data{Path(project_data['originalFilename']).suffix}"
         blob = bucket.blob(blob_path)
         file_bytes = blob.download_as_bytes()
         
-        # This function is the source of the inconsistent response
+        # This function now returns data_context, which we need to capture
         draft_data = generate_report_draft(
             io.BytesIO(file_bytes),
             project_data['originalFilename'],
@@ -390,67 +387,17 @@ def generate_sozo_report(project_id):
             bucket
         )
         
-        # Initialize variables from the draft data, providing safe defaults
-        final_raw_md = draft_data.get('raw_md', '')
-        final_chart_urls = draft_data.get('chartUrls', {})
-
-        # CORE LOGIC: If raw_md is empty BUT we received chart_specs, it means the
-        # primary markdown generation failed. We must reconstruct it to ensure the
-        # user sees their report.
-        if not final_raw_md.strip() and draft_data.get('report_content', {}).get('chart_specs'):
-            logger.warning(f"raw_md was empty for project {project_id}. Reconstructing from chart_specs.")
-            
-            # 1. Reconstruct the Markdown String
-            final_raw_md = "## Executive Summary\n\nThis report provides a visual analysis of the key metrics from the provided dataset. Each section below highlights a specific insight through a detailed chart.\n\n"
-            chart_specs_list = draft_data['report_content']['chart_specs']
-            
-            for spec_data in chart_specs_list:
-                title = spec_data.get('title', 'Chart')
-                final_raw_md += f"### {title}\n\n"
-                final_raw_md += f"<generate_chart: \"{title}\">\n\n"
-
-            # 2. Reconstruct the Chart URLs (as they will also be missing)
-            logger.warning(f"Reconstructing missing chart URLs for project {project_id}.")
-            df = load_dataframe_safely(io.BytesIO(file_bytes), project_data['originalFilename'])
-            llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", google_api_key=os.getenv("GOOGLE_API_KEY"), temperature=0.1)
-            
-            for spec_data in chart_specs_list:
-                try:
-                    spec_title = spec_data.get('title')
-                    # Create a ChartSpecification instance from the dictionary data
-                    chart_spec = ChartSpecification(
-                        chart_type=spec_data.get('chart_type'),
-                        title=spec_data.get('title'),
-                        x_col=spec_data.get('x_col'),
-                        y_col=spec_data.get('y_col'),
-                        agg_method=spec_data.get('agg_method'),
-                        top_n=spec_data.get('top_n')
-                    )
-                    
-                    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
-                        img_path = Path(temp_file.name)
-                        if execute_chart_spec(chart_spec, df, img_path):
-                            blob_name = f"sozo_projects/{uid}/{project_id}/charts/{uuid.uuid4().hex}.png"
-                            chart_blob = bucket.blob(blob_name)
-                            chart_blob.upload_from_filename(str(img_path))
-                            final_chart_urls[spec_title] = chart_blob.public_url
-                            logger.info(f"Re-generated and uploaded chart '{spec_title}' to {chart_blob.public_url}")
-                except Exception as chart_gen_error:
-                    logger.error(f"Error re-generating chart '{spec_data.get('title')}': {chart_gen_error}")
-                finally:
-                    if 'img_path' in locals() and os.path.exists(img_path):
-                        os.unlink(img_path)
-
-        # Prepare the final, corrected data for Firebase
+        # **CORRECTION HERE: Save the data_context to Firebase**
         update_data = {
             'status': 'draft',
-            'rawMarkdown': final_raw_md,
-            'chartUrls': final_chart_urls
+            'rawMarkdown': draft_data.get('raw_md'),
+            'chartUrls': draft_data.get('chartUrls'),
+            'dataContext': draft_data.get('data_context') # Save the context for the video generator
         }
         
         project_ref.update(update_data)
         
-        logger.info(f"Project {project_id} successfully updated with draft data. Markdown length: {len(final_raw_md)}")
+        logger.info(f"Project {project_id} successfully updated with draft data.")
         
         return jsonify({
             'success': True,
@@ -458,15 +405,10 @@ def generate_sozo_report(project_id):
         }), 200
         
     except Exception as e:
-        logger.error(f"CRITICAL error generating report for project {project_id}: {str(e)}")
-        logger.error(f"Traceback: {traceback.format_exc()}")
-        db.reference(f'sozo_projects/{project_id}').update({
-            'status': 'failed',
-            'error': str(e)
-        })
+        logger.error(f"CRITICAL error generating report for project {project_id}: {traceback.format_exc()}")
+        db.reference(f'sozo_projects/{project_id}').update({'status': 'failed', 'error': str(e)})
         return jsonify({'error': str(e)}), 500
 
-# In app.py
 
 @app.route('/api/sozo/projects/<string:project_id>/generate-video', methods=['POST'])
 def generate_sozo_video(project_id):
@@ -481,23 +423,23 @@ def generate_sozo_video(project_id):
         
         data = request.get_json()
         voice_model = data.get('voice_model', 'aura-2-andromeda-en')
-        if voice_model not in ['aura-2-andromeda-en', 'aura-2-orpheus-en']: return jsonify({'error': 'Invalid voice model specified'}), 400
         
         project_ref.update({'status': 'generating_video'})
         
-        blob_path = "/".join(project_data['originalDataUrl'].split('/')[4:])
+        blob_path = f"sozo_projects/{uid}/{project_id}/data{Path(project_data['originalFilename']).suffix}"
         blob = bucket.blob(blob_path)
         file_bytes = blob.download_as_bytes()
         df = load_dataframe_safely(io.BytesIO(file_bytes), project_data['originalFilename'])
         
-        # --- THIS IS THE CORRECTED FUNCTION CALL ---
+        # **CORRECTION HERE: Pass the required 'data_context' argument**
         video_url = generate_video_from_project(
             df,
-            project_data.get('rawMarkdown', ''), # Use the correct key for the markdown
-            uid,                                # ADDED missing uid
-            project_id,                         # ADDED missing project_id
-            voice_model,                        # Now in the correct position
-            bucket                              # ADDED missing bucket
+            project_data.get('rawMarkdown', ''),
+            project_data.get('dataContext', {}), # Load the context from the project
+            uid,
+            project_id,
+            voice_model,
+            bucket                              # The missing argument from your error message
         )
         
         if not video_url: raise Exception("Video generation failed in core function.")