added together ai agent

python_code_executor_service.py

@@ -3,27 +3,37 @@ import matplotlib.pyplot as plt
 from pathlib import Path
 from typing import Dict, Any, List, 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 statsmodels.api as sm
+import scipy.stats as stats
 from pydantic import BaseModel
 
+
 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
     description: str
 
+
 class CsvChatResult(BaseModel):
     """Structured response for CSV-related AI interactions"""
     response_type: str  # Literal["casual", "data_analysis", "visualization", "mixed"]
@@ -31,8 +41,9 @@ class CsvChatResult(BaseModel):
     analysis_operations: List[AnalysisOperation]
     charts: Optional[List[ChartSpecification]] = None
 
+
 class PythonExecutor:
-    """Handles execution of Python code and dummy image generation for CSV analysis"""
+    """Handles execution of Python code with comprehensive data analysis libraries"""
     
     def __init__(self, df: pd.DataFrame, charts_folder: str = "generated_charts"):
         """
@@ -48,7 +59,7 @@ class PythonExecutor:
         
     def execute_code(self, code: str) -> Dict[str, Any]:
         """
-        Execute Python code and return the output and any generated plots
+        Execute Python code with full data analysis context and return results
         
         Args:
             code (str): Python code to execute
@@ -78,12 +89,28 @@ class PythonExecutor:
             plt.close('all')
         
         try:
-            # Create execution context with common libraries and the DataFrame
+            # 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,
+                
+                # Statistics
+                'sm': sm,
+                'stats': stats,
+                
+                # Date/time
+                'datetime': datetime,
+                'timedelta': timedelta,
+                'time': time,
+                
+                # Utilities
                 'json': json,
-                'df': self.df,  # Include the DataFrame in the execution context
                 '__builtins__': __builtins__,
             }