import pandas as pd import matplotlib.pyplot as plt import seaborn as sns import datetime as dt import os import uuid import json import logging import sys import traceback import shutil import ast # <--- Added for syntax checking from io import StringIO from typing import List, Dict, Any, Optional from pydantic import BaseModel from starlette.concurrency import run_in_threadpool import numpy as np # --- Internal Imports --- from supabase_service import upload_file_to_supabase # Configure Logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger("Report_Generator") # Configure Matplotlib backend for thread safety (Headless) os.environ['MPLBACKEND'] = 'agg' # Disable plt.show globally to prevent blocking plt.show = lambda: None # ============================================================================== # REQUEST MODELS # ============================================================================== class ReportRequest(BaseModel): code: str csv_url: str chat_id: str # ============================================================================== # RESPONSE MODELS # ============================================================================== class FileProps(BaseModel): fileName: str filePath: str fileType: str # 'csv' | 'image' class Files(BaseModel): csv_files: List[FileProps] image_files: List[FileProps] class FileBoxProps(BaseModel): files: Files # ============================================================================== # EXECUTION ENGINE # ============================================================================== class PythonREPL: """Secure Python REPL with file generation tracking""" def __init__(self, df: pd.DataFrame): self.df = df # Create a unique directory for this execution to avoid thread collisions self.output_dir = os.path.abspath(f'generated_outputs/{uuid.uuid4()}') os.makedirs(self.output_dir, exist_ok=True) self.local_env = { "pd": pd, "df": self.df.copy(), "plt": plt, "os": os, "uuid": uuid, "sns": sns, "json": json, "dt": dt, "np": np, "output_dir": self.output_dir # INJECT output_dir so code can use it } def execute(self, code: str) -> Dict[str, Any]: logger.info(f'Executing code in: {self.output_dir}') old_stdout = sys.stdout sys.stdout = mystdout = StringIO() file_tracker = { 'csv_files': set(), 'image_files': set() } # Wrap code to ensure non-interactive plotting inside the exec scope wrapped_code = f""" import matplotlib.pyplot as plt plt.switch_backend('agg') {code} plt.close('all') """ error = False error_msg = None try: # --- IMPROVEMENT: Syntax Validation --- # This checks if the generated code is valid Python before running it. # It catches the "missing parenthesis" error gracefully. try: ast.parse(wrapped_code) except SyntaxError as e: raise SyntaxError(f"Generated code is incomplete or invalid: {e}") # Execute the code exec(wrapped_code, self.local_env) # Check for generated files in the specific directory if os.path.exists(self.output_dir): for fname in os.listdir(self.output_dir): if fname.endswith('.csv'): file_tracker['csv_files'].add(fname) elif fname.lower().endswith(('.png', '.jpg', '.jpeg')): file_tracker['image_files'].add(fname) except Exception as e: error = True error_msg = traceback.format_exc() # --- IMPROVEMENT: Better Logging --- # Log the code causing the error so you can debug the LLM prompt logger.error("============= CODE EXECUTION FAILED =============") logger.error(f"Error Message: {str(e)}") logger.error("---------------- FAILED CODE ----------------") for i, line in enumerate(wrapped_code.split('\n')): logger.error(f"{i+1}: {line}") logger.error("=================================================") finally: sys.stdout = old_stdout return { "output": mystdout.getvalue(), "error": error, "error_message": error_msg, "output_dir": self.output_dir, "files": { "csv": [os.path.join(self.output_dir, f) for f in file_tracker['csv_files']], "images": [os.path.join(self.output_dir, f) for f in file_tracker['image_files']] } } def cleanup(self): """Remove the temporary directory""" if os.path.exists(self.output_dir): try: shutil.rmtree(self.output_dir) except Exception as e: logger.warning(f"Cleanup failed: {e}") # ============================================================================== # MAIN LOGIC (Synchronous Worker) # ============================================================================== def _generate_report_sync(code: str, csv_url: str, chat_id: str) -> FileBoxProps: """ Blocking worker function. 1. Reads CSV. 2. Runs Code. 3. Uploads files synchronously. """ repl = None csv_props = [] image_props = [] try: # 1. Load Data df = pd.read_csv(csv_url) # 2. Initialize REPL repl = PythonREPL(df) # 3. Execute Code result = repl.execute(code) if result['error']: # Log is already handled in repl.execute return FileBoxProps(files=Files(csv_files=[], image_files=[])) # 4. Process & Upload CSVs for csv_path in result['files']['csv']: if os.path.exists(csv_path): file_name = os.path.basename(csv_path) unique_name = f"{uuid.uuid4()}_{file_name}" try: # Sync call - NO AWAIT public_url = upload_file_to_supabase( file_path=csv_path, file_name=unique_name, chat_id=chat_id ) csv_props.append(FileProps( fileName=file_name, filePath=public_url, fileType="csv" )) except Exception as e: logger.error(f"Failed upload CSV {file_name}: {e}") # 5. Process & Upload Images for img_path in result['files']['images']: if os.path.exists(img_path): file_name = os.path.basename(img_path) unique_name = f"{uuid.uuid4()}_{file_name}" try: # Sync call - NO AWAIT public_url = upload_file_to_supabase( file_path=img_path, file_name=unique_name, chat_id=chat_id ) image_props.append(FileProps( fileName=file_name, filePath=public_url, fileType="image" )) except Exception as e: logger.error(f"Failed upload Image {file_name}: {e}") return FileBoxProps( files=Files( csv_files=csv_props, image_files=image_props ) ) except Exception as e: logger.error(f"System Error in Report Generator: {e}") return FileBoxProps(files=Files(csv_files=[], image_files=[])) finally: # 6. Cleanup local files if repl: repl.cleanup() # ============================================================================== # ASYNC WRAPPER # ============================================================================== async def execute_report_generation(code: str, csv_url: str, chat_id: str) -> FileBoxProps: """ Async entry point that runs the blocking logic in a separate thread. This enables concurrency. """ return await run_in_threadpool( _generate_report_sync, code=code, csv_url=csv_url, chat_id=chat_id )