Update main.py

main.py

@@ -356,43 +356,114 @@ def delete_sozo_project(project_id):
 
 @app.route('/api/sozo/projects/<string:project_id>/generate-report', methods=['POST'])
 def generate_sozo_report(project_id):
-    logger.info(f"Endpoint /generate-report POST for project {project_id}")
+    logger.info(f"POST /api/sozo/projects/{project_id}/generate-report - Generating report")
+    
     try:
         token = request.headers.get('Authorization', '').split(' ')[1]
         uid = verify_token(token)
-        if not uid: return jsonify({'error': 'Unauthorized'}), 401
-        logger.info(f"Token verified for user {uid} for report generation.")
+        if not uid:
+            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"User {uid} failed to generate report: Project {project_id} not found or not owned.")
+            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"User {uid}: Ownership verified for project {project_id}. Starting report generation.")
-        project_ref.update({'status': 'generating_report', 'updatedAt': datetime.utcnow().isoformat()})
+        logger.info(f"Project {project_id} validated for user {uid}")
         
-        blob_path = f"sozo_projects/{uid}/{project_id}/data{Path(project_data['originalFilename']).suffix}"
-        logger.info(f"User {uid}: Downloading data from {blob_path}")
+        blob_path = "/".join(project_data['originalDataUrl'].split('/')[4:])
         blob = bucket.blob(blob_path)
         file_bytes = blob.download_as_bytes()
-        df = load_dataframe_safely(io.BytesIO(file_bytes), project_data['originalFilename'])
         
-        logger.info(f"User {uid}: Calling core logic to generate report draft for project {project_id}.")
-        draft_data = generate_report_draft(df, project_data['userContext'])
+        # This function is the source of the inconsistent response
+        draft_data = generate_report_draft(
+            io.BytesIO(file_bytes),
+            project_data['originalFilename'],
+            project_data['userContext'],
+            uid,
+            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
+        update_data = {
+            'status': 'draft',
+            'rawMarkdown': final_raw_md,
+            'chartUrls': final_chart_urls
+        }
         
-        logger.info(f"User {uid}: Saving report content specification to database for project {project_id}.")
-        project_ref.child('report_content').set(draft_data)
+        project_ref.update(update_data)
         
-        project_ref.update({'status': 'draft_complete', 'updatedAt': datetime.utcnow().isoformat()})
+        logger.info(f"Project {project_id} successfully updated with draft data. Markdown length: {len(final_raw_md)}")
         
-        full_project_data = project_ref.get()
-        logger.info(f"User {uid}: Report generation for project {project_id} complete.")
-        return jsonify({'success': True, 'project': full_project_data}), 200
+        return jsonify({
+            'success': True,
+            'project': {**project_data, **update_data}
+        }), 200
         
     except Exception as e:
-        db.reference(f'sozo_projects/{project_id}').update({'status': 'failed', 'error': str(e), 'updatedAt': datetime.utcnow().isoformat()})
-        logger.error(f"CRITICAL ERROR generating report for {project_id}: {traceback.format_exc()}")
+        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)
+        })
         return jsonify({'error': str(e)}), 500
 
 @app.route('/api/sozo/projects/<string:project_id>/generate-video', methods=['POST'])