python_code_executor_service.py

import os
from dotenv import load_dotenv
import uuid
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Dict, Any, List, Literal, Optional, Union
import pandas as pd
import numpy as np
import json
import io
import contextlib
import traceback
import time
from datetime import datetime, timedelta
import seaborn as sns
import scipy.stats as stats
from pydantic import BaseModel
from tabulate import tabulate
import asyncio

from supabase_service import upload_file_to_supabase

# Load environment variables from .env file
load_dotenv()

class CodeResponse(BaseModel):
    """Container for code-related responses"""
    language: str = "python"
    code: str


class ChartSpecification(BaseModel):
    """Details about requested charts"""
    image_description: str
    code: Optional[str] = None


class AnalysisOperation(BaseModel):
    """Container for a single analysis operation with its code and result"""
    code: CodeResponse
    result_var: Union[str, List[str]]  # Allow multiple result variables


class CsvChatResult(BaseModel):
    """Structured response for CSV-related AI interactions"""
    casual_response: str
    analysis_operations: Optional[AnalysisOperation] = None
    charts: Optional[ChartSpecification] = None


class PythonExecutor:
    """Handles execution of Python code with comprehensive data analysis libraries"""
    
    def __init__(self, df: pd.DataFrame, charts_folder: str = "generated_charts"):
        """
        Initialize the PythonExecutor with a DataFrame
        
        Args:
            df (pd.DataFrame): The DataFrame to operate on
            charts_folder (str): Folder to save charts in
        """
        self.df = df.copy()  # Use copy to avoid modifying original
        self.charts_folder = Path(charts_folder)
        self.charts_folder.mkdir(exist_ok=True, parents=True)
        self.exec_locals = {}
        self._setup_matplotlib()
        
    def _setup_matplotlib(self):
        """Configure matplotlib for non-interactive use"""
        plt.ioff()  # Turn off interactive mode
        plt.rcParams['figure.figsize'] = [10, 6]
        plt.rcParams['figure.dpi'] = 100
        plt.rcParams['savefig.bbox'] = 'tight'
        
    def execute_code(self, code: str) -> Dict[str, Any]:
        """
        Execute Python code with full data analysis context and return results
        
        Args:
            code (str): Python code to execute
            
        Returns:
            dict: Dictionary containing execution results and any generated plots
        """
        output = ""
        error = None
        plots = []
        
        # Capture stdout
        stdout = io.StringIO()
        
        # Store original plt.show
        original_show = plt.show
        
        def custom_show():
            """Custom show function that saves plots instead of displaying them"""
            nonlocal plots
            for i, fig in enumerate(plt.get_fignums()):
                figure = plt.figure(fig)
                # Save plot to bytes buffer
                buf = io.BytesIO()
                figure.savefig(buf, format='png', bbox_inches='tight', dpi=100)
                buf.seek(0)
                plots.append(buf.getvalue())
            plt.close('all')
        
        try:
            # Create comprehensive execution context with data analysis libraries
            exec_globals = {
                # Core data analysis
                'pd': pd,
                'np': np,
                'df': self.df,
                
                # Visualization
                'plt': plt,
                'sns': sns,
                'tabulate': tabulate,
                
                # Statistics
                'stats': stats,
                
                # Date/time
                'datetime': datetime,
                'timedelta': timedelta,
                'time': time,
                
                # Utilities
                'json': json,
                '__builtins__': __builtins__,
            }
            
            # Update with current locals to maintain state between executions
            exec_globals.update(self.exec_locals)
            
            # Replace plt.show with custom implementation
            plt.show = custom_show
            
            # Execute code and capture output
            with contextlib.redirect_stdout(stdout):
                compiled_code = compile(code, '<string>', 'exec')
                exec(compiled_code, exec_globals, self.exec_locals)
            
            output = stdout.getvalue()
            
        except Exception as e:
            error = {
                "message": str(e),
                "traceback": traceback.format_exc()
            }
            # Clean up any open figures on error
            plt.close('all')
        finally:
            # Always restore original plt.show
            plt.show = original_show
            # Ensure all figures are closed
            plt.close('all')
            
        return {
            'output': output,
            'error': error,
            'plots': plots,
            'locals': dict(self.exec_locals)  # Return copy to avoid mutation
        }
    
    async def save_plot_to_supabase(self, plot_data: bytes, description: str, chat_id: str) -> str:
        """
        Save plot to Supabase storage and return the public URL
        
        Args:
            plot_data (bytes): Image data in bytes
            description (str): Description of the plot
            chat_id (str): ID of the chat session
            
        Returns:
            str: Public URL of the uploaded chart
        """
        # Generate unique filename
        filename = f"chart_{uuid.uuid4().hex}.png"
        filepath = self.charts_folder / filename
        
        # Save the plot locally first
        try:
            with open(filepath, 'wb') as f:
                f.write(plot_data)
                
            # Upload to Supabase with timeout
            try:
                public_url = await asyncio.wait_for(
                    upload_file_to_supabase(
                        file_path=str(filepath),
                        file_name=filename,
                        chat_id=chat_id
                    ),
                    timeout=30.0  # 30 second timeout
                )
                
                # Remove the local file after upload
                try:
                    os.remove(filepath)
                except OSError:
                    pass  # Ignore removal errors
                
                return public_url
            except asyncio.TimeoutError:
                raise Exception("Upload timed out after 30 seconds")
            except Exception as e:
                raise Exception(f"Failed to upload plot to Supabase: {e}")
                
        except Exception as e:
            # Clean up local file if exists
            if os.path.exists(filepath):
                try:
                    os.remove(filepath)
                except OSError:
                    pass
            raise Exception(f"Failed to save plot: {e}")
    
    def _format_result(self, result: Any) -> str:
     """Format the result for display"""
     if isinstance(result, pd.DataFrame):
        return result.to_string()
     elif isinstance(result, pd.Series):
        return result.to_string()
     elif isinstance(result, (dict, list)):
        # Custom JSON encoder to handle special types
        def json_serializer(obj):
            """Handle special types that aren't JSON serializable"""
            if isinstance(obj, (pd.Timestamp, datetime)):
                return obj.isoformat()
            elif isinstance(obj, (np.integer, np.int64, np.int32)):
                return int(obj)
            elif isinstance(obj, (np.floating, np.float64, np.float32)):
                return float(obj)
            elif isinstance(obj, np.ndarray):
                return obj.tolist()
            elif isinstance(obj, pd.Series):
                return obj.to_dict()
            elif isinstance(obj, pd.DataFrame):
                return obj.to_dict('records')
            elif hasattr(obj, '__dict__'):
                return str(obj)
            else:
                return str(obj)
        
        try:
            return json.dumps(result, indent=2, default=json_serializer)
        except Exception as e:
            # Fallback to string representation if JSON serialization fails
            return f"Result (JSON serialization failed: {str(e)}):\n{str(result)}"
     elif isinstance(result, (pd.Timestamp, datetime)):
        return result.isoformat()
     elif isinstance(result, (np.integer, np.int64, np.int32)):
        return str(int(result))
     elif isinstance(result, (np.floating, np.float64, np.float32)):
        return str(float(result))
     elif isinstance(result, np.ndarray):
        return str(result)
     elif hasattr(result, '__str__'):
        return str(result)
     else:
        return repr(result)
    
    def _get_result_variables(self, result_var: Union[str, List[str]]) -> Dict[str, Any]:
        """Get result variables from execution locals"""
        results = {}
        
        if isinstance(result_var, str):
            # Handle comma-separated variable names in string
            if ',' in result_var:
                var_names = [name.strip() for name in result_var.split(',')]
            else:
                var_names = [result_var.strip()]
        else:
            var_names = result_var
            
        for var_name in var_names:
            if var_name in self.exec_locals:
                results[var_name] = self.exec_locals[var_name]
                
        return results
    
    async def process_response(self, response: CsvChatResult, chat_id: str) -> str:
        """Process the response with proper variable handling and error checking"""
        output_parts = [response.casual_response]
    
        # Process analysis operation if it exists
        if response.analysis_operations is not None:
            try:
                operation = response.analysis_operations
                if operation and operation.code and operation.code.code:
                    execution_result = self.execute_code(operation.code.code)
                
                    # Check for execution errors
                    if execution_result.get('error'):
                        output_parts.append(f"\n**Error in analysis operation:**")
                        output_parts.append("```python\n" + execution_result['error']['message'] + "\n```")
                    else:
                        # Get all result variables
                        result_vars = self._get_result_variables(operation.result_var)
                    
                        if result_vars:
                            for var_name, result in result_vars.items():
                                if result is not None:
                                    # Handle empty/None results
                                    if (hasattr(result, '__len__') and len(result) == 0):
                                        output_parts.append(f"\n**Warning:** Variable '{var_name}' contains empty data")
                                    else:
                                        output_parts.append(f"\n**{var_name}:**")
                                        formatted_result = self._format_result(result)
                                        # Add language identifier for proper syntax highlighting
                                        output_parts.append("```python\n" + formatted_result + "\n```")
                                else:
                                    output_parts.append(f"\n**Warning:** Variable '{var_name}' is None or not found")
                        else:
                            # Check if there was console output
                            output_str = execution_result.get('output', '').strip()
                            if output_str:
                                output_parts.append(f"\n**Execution output:**")
                                output_parts.append("```python\n" + output_str + "\n```")
                            else:
                                output_parts.append(f"\n**Note:** Analysis operation executed but no results found for: {operation.result_var}")
                else:
                    output_parts.append("\n**Warning:** Invalid analysis operation - missing code or result variable")
            except Exception as e:
                output_parts.append(f"\n**Error:** Error processing analysis operation: {str(e)}")
                if hasattr(operation, 'result_var'):
                    output_parts.append(f"Expected variables: {operation.result_var}")
                    
        # Process chart if it exists
        if response.charts is not None:
            chart = response.charts
            try:
                if chart and (chart.code or chart.image_description):
                    if chart.code:
                        chart_result = self.execute_code(chart.code)
                        if chart_result.get('plots'):
                            # Only add the description header once before all charts
                            if chart.image_description:
                                output_parts.append(f"\n**Chart:** {chart.image_description}")
                            
                            # Then add all chart images without repeating the description
                            for i, plot_data in enumerate(chart_result['plots']):
                                try:
                                    public_url = await self.save_plot_to_supabase(
                                        plot_data=plot_data,
                                        description=chart.image_description,
                                        chat_id=chat_id
                                    )
                                    output_parts.append(f"![{chart.image_description}]({public_url})")
                                except Exception as e:
                                    output_parts.append(f"\n**Warning:** Error uploading chart {i+1}: {str(e)}")
                        elif chart_result.get('error'):
                            output_parts.append("```python\n" + f"Error generating {chart.image_description}: {chart_result['error']['message']}" + "\n```")
                        else:
                            output_parts.append(f"\n**Warning:** No chart generated for '{chart.image_description}'")
                    else:
                        output_parts.append(f"\n**Warning:** No code provided for chart: {chart.image_description}")
                else:
                    output_parts.append("\n**Warning:** Invalid chart specification")
            except Exception as e:
                output_parts.append(f"\n**Error:** Error processing chart: {str(e)}")
    
        return "\n".join(output_parts)