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 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 re 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 def clean_code(self) -> str: """Remove trailing newlines while preserving internal structure""" return self.code.rstrip('\n') class ChartSpecification(BaseModel): """Details about requested charts""" image_description: str code: Optional[str] = None def clean_description(self) -> str: """Replace newlines in description with spaces (preserves readability)""" return self.image_description.replace('\n', ' ').strip() class AnalysisOperation(BaseModel): """Container for a single analysis operation with its code and result""" code: CodeResponse result_var: str class CsvChatResult(BaseModel): """Structured response for CSV-related AI interactions""" response_type: Literal["casual", "data_analysis", "visualization", "mixed"] casual_response: str analysis_operations: List[AnalysisOperation] charts: Optional[List[ChartSpecification]] = None def clean_casual_response(self) -> str: """Clean casual response by replacing newlines with spaces when appropriate""" # Preserve intentional line breaks (markdown-style) if '\n\n' in self.casual_response: return self.casual_response return self.casual_response.replace('\n', ' ') class PythonExecutor: """Handles execution of Python code with comprehensive data analysis libraries""" def __init__(self, df: pd.DataFrame, charts_folder: str = "generated_charts"): self.df = df self.charts_folder = Path(charts_folder) self.charts_folder.mkdir(exist_ok=True) self.exec_locals = {} def execute_code(self, code: str) -> Dict[str, Any]: output = "" error = None plots = [] stdout = io.StringIO() original_show = plt.show def custom_show(): for i, fig in enumerate(plt.get_fignums()): figure = plt.figure(fig) buf = io.BytesIO() figure.savefig(buf, format='png', bbox_inches='tight') buf.seek(0) plots.append(buf.read()) plt.close('all') try: exec_globals = { 'pd': pd, 'np': np, 'df': self.df, 'plt': plt, 'sns': sns, 'tabulate': tabulate, 'stats': stats, 'datetime': datetime, 'timedelta': timedelta, 'time': time, 'json': json, '__builtins__': __builtins__, } plt.show = custom_show with contextlib.redirect_stdout(stdout): exec(code, exec_globals, self.exec_locals) output = stdout.getvalue() except Exception as e: error = {"message": str(e), "traceback": traceback.format_exc()} finally: plt.show = original_show return { 'output': output, 'error': error, 'plots': plots, 'locals': self.exec_locals } async def save_plot_to_supabase(self, plot_data: bytes, description: str, chat_id: str) -> str: filename = f"chart_{uuid.uuid4().hex}.png" filepath = self.charts_folder / filename with open(filepath, 'wb') as f: f.write(plot_data) try: public_url = await upload_file_to_supabase( file_path=str(filepath), file_name=filename, chat_id=chat_id ) os.remove(filepath) return public_url except Exception as e: if os.path.exists(filepath): os.remove(filepath) raise Exception(f"Failed to upload plot to Supabase: {e}") def _format_result(self, result: Any) -> str: """Format result with safe newline handling""" if isinstance(result, (pd.DataFrame, pd.Series)): return result.to_string() elif isinstance(result, (dict, list)): return json.dumps(result, indent=2) # Clean string representation while preserving essential newlines str_result = str(result) if '\n' in str_result and not any(x in str_result for x in ['```', 'def ', 'class ']): return str_result.replace('\n', ' ') return str_result async def process_response(self, response: CsvChatResult, chat_id: str) -> str: """Process response with intelligent newline handling""" output_parts = [response.clean_casual_response()] # Process analysis operations for operation in response.analysis_operations: execution_result = self.execute_code(operation.code.clean_code()) result = self.exec_locals.get(operation.result_var) if execution_result['error']: output_parts.append(f"\n❌ Error in operation '{operation.result_var}':") output_parts.append(f"```python\n{execution_result['error']['message']}\n```") elif result is not None: if result is None or (hasattr(result, '__len__') and len(result) == 0): output_parts.append(f"\n⚠️ Values are missing - Operation '{operation.result_var}' returned no data") else: output_parts.append(f"\n🔹 Result for '{operation.result_var}':") output_parts.append(f"```python\n{self._format_result(result)}\n```") else: output_str = execution_result['output'].strip() if output_str: output_parts.append(f"\nOutput for '{operation.result_var}':") output_parts.append(f"```\n{output_str}\n```") # Process charts if response.charts: output_parts.append("\n📊 Visualizations:") for chart in response.charts: if chart.code: chart_result = self.execute_code(chart.code) if chart_result['plots']: for plot_data in chart_result['plots']: try: public_url = await self.save_plot_to_supabase( plot_data=plot_data, description=chart.clean_description(), chat_id=chat_id ) output_parts.append(f"\n🖼️ {chart.clean_description()}") output_parts.append(f"![{chart.clean_description()}]({public_url})") except Exception as e: output_parts.append(f"\n⚠️ Error uploading chart: {str(e)}") elif chart_result['error']: output_parts.append(f"```python\nError generating chart: {chart_result['error']['message']}\n```") else: output_parts.append(f"\n⚠️ No chart generated for '{chart.clean_description()}'") return "\n".join(output_parts)