Update app.py

app.py

@@ -1,104 +1,163 @@
+import base64
+import io
 import os
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional
 
 import gradio as gr
+import matplotlib.pyplot as plt
+import numpy as np
 import pandas as pd
-import torch
+import seaborn as sns
 from litellm import completion
 
-# Agent Implementation
+
+# Code Execution Environment
+class CodeEnvironment:
+    """Safe environment for executing code with data analysis capabilities"""
+    
+    def __init__(self):
+        self.globals = {
+            'pd': pd,
+            'np': np,
+            'plt': plt,
+            'sns': sns,
+        }
+        self.locals = {}
+        
+    def execute(self, code: str, df: pd.DataFrame = None) -> Dict[str, Any]:
+        """Execute code and capture outputs"""
+        if df is not None:
+            self.globals['df'] = df
+            
+        # Capture output
+        output_buffer = io.StringIO()
+        result = {'output': '', 'figures': [], 'error': None}
+        
+        try:
+            # Execute code
+            exec(code, self.globals, self.locals)
+            
+            # Capture figures
+            for i in plt.get_fignums():
+                fig = plt.figure(i)
+                buf = io.BytesIO()
+                fig.savefig(buf, format='png')
+                buf.seek(0)
+                img_str = base64.b64encode(buf.read()).decode()
+                result['figures'].append(f"data:image/png;base64,{img_str}")
+                plt.close(fig)
+            
+            # Get printed output
+            result['output'] = output_buffer.getvalue()
+            
+        except Exception as e:
+            result['error'] = str(e)
+            
+        finally:
+            output_buffer.close()
+            
+        return result
+
 @dataclass
 class Tool:
-    """Simplified tool class"""
+    """Tool for data analysis"""
     name: str
     description: str
     func: Callable
 
-class MinimalAgent:
-    """Minimal agent implementation for demo purposes"""
+class AnalysisAgent:
+    """Agent that can analyze data and execute code"""
     
     def __init__(
         self,
         model_id: str = "gpt-4o-mini",
         temperature: float = 0.7,
-        max_steps: int = 5
     ):
         self.model_id = model_id
         self.temperature = temperature
-        self.max_steps = max_steps
         self.tools: List[Tool] = []
+        self.code_env = CodeEnvironment()
         
     def add_tool(self, name: str, description: str, func: Callable) -> None:
         """Add a tool to the agent"""
         self.tools.append(Tool(name=name, description=description, func=func))
         
-    def run(self, prompt: str, **kwargs) -> str:
-        """Run the agent with a prompt"""
+    def run(self, prompt: str, df: pd.DataFrame = None) -> str:
+        """Run analysis with code execution"""
         messages = [
             {"role": "system", "content": self._get_system_prompt()},
             {"role": "user", "content": prompt}
         ]
         
         try:
+            # Get response from model
             response = completion(
                 model=self.model_id,
                 messages=messages,
                 temperature=self.temperature,
             )
-            return response.choices[0].message.content
+            analysis = response.choices[0].message.content
+            
+            # Extract code blocks
+            code_blocks = self._extract_code(analysis)
+            
+            # Execute code and capture results
+            results = []
+            for code in code_blocks:
+                result = self.code_env.execute(code, df)
+                if result['error']:
+                    results.append(f"Error executing code: {result['error']}")
+                else:
+                    # Add output and figures
+                    if result['output']:
+                        results.append(result['output'])
+                    for fig in result['figures']:
+                        results.append(f"![Figure]({fig})")
+                        
+            # Combine analysis and results
+            return analysis + "\n\n" + "\n".join(results)
+            
         except Exception as e:
             return f"Error: {str(e)}"
-            
+    
     def _get_system_prompt(self) -> str:
-        """Get the system prompt including available tools"""
+        """Get system prompt with tools and capabilities"""
         tools_desc = "\n".join([
             f"- {tool.name}: {tool.description}"
             for tool in self.tools
         ])
         
-        return f"""You are a helpful AI agent that can analyze data and write code.
+        return f"""You are a data analysis assistant.
         
 Available tools:
 {tools_desc}
 
-Additional capabilities:
-- Data analysis with pandas, numpy
-- Visualization with matplotlib, seaborn
-- Machine learning with sklearn
-- Statistical analysis with scipy
-
-Provide clear explanations and code examples."""
+Capabilities:
+- Data analysis (pandas, numpy)
+- Visualization (matplotlib, seaborn)
+- Statistical analysis (scipy)
+- Machine learning (sklearn)
 
-# Analysis Functions
-def analyze_dataframe(df: pd.DataFrame, analysis_type: str) -> str:
-    """Basic DataFrame analysis"""
-    if analysis_type == "summary":
-        return str(df.describe())
-    elif analysis_type == "info":
-        buffer = []
-        df.info(buf=buffer)
-        return "\n".join(buffer)
-    return "Unknown analysis type"
-
-def plot_data(df: pd.DataFrame, plot_type: str) -> None:
-    """Basic plotting function"""
-    import matplotlib.pyplot as plt
-    import seaborn as sns
+When writing code:
+- Use markdown code blocks
+- Create clear visualizations
+- Include explanations
+- Handle errors gracefully
+"""
     
-    if plot_type == "correlation":
-        plt.figure(figsize=(10, 8))
-        sns.heatmap(df.corr(), annot=True)
-        plt.title("Correlation Heatmap")
-    elif plot_type == "distribution":
-        df.hist(figsize=(15, 10))
-        plt.tight_layout()
+    @staticmethod
+    def _extract_code(text: str) -> List[str]:
+        """Extract Python code blocks from markdown"""
+        import re
+        pattern = r'```python\n(.*?)```'
+        return re.findall(pattern, text, re.DOTALL)
 
 def process_file(file: gr.File) -> Optional[pd.DataFrame]:
     """Process uploaded file into DataFrame"""
     if not file:
         return None
-    
+        
     try:
         if file.name.endswith('.csv'):
             return pd.read_csv(file.name)
@@ -118,32 +177,20 @@ def analyze_data(
     
     if not api_key:
         return "Error: Please provide an API key."
-    
+        
     if not file:
         return "Error: Please upload a file."
-    
+        
     try:
         # Set up environment
         os.environ["OPENAI_API_KEY"] = api_key
         
         # Create agent
-        agent = MinimalAgent(
+        agent = AnalysisAgent(
             model_id="gpt-4o-mini",
             temperature=temperature
         )
         
-        # Add tools
-        agent.add_tool(
-            "analyze_dataframe",
-            "Analyze DataFrame with various metrics",
-            analyze_dataframe
-        )
-        agent.add_tool(
-            "plot_data",
-            "Create various plots from DataFrame",
-            plot_data
-        )
-        
         # Process file
         df = process_file(file)
         if df is None:
@@ -168,13 +215,12 @@ def analyze_data(
         User request: {query}
         
         Please analyze the data and provide:
-        1. A clear explanation of your approach
-        2. Code for the analysis
-        3. Visualizations where relevant
-        4. Key insights and findings
+        1. Clear explanation of approach
+        2. Code with visualizations
+        3. Key insights and findings
         """
         
-        return agent.run(prompt)
+        return agent.run(prompt, df=df)
         
     except Exception as e:
         return f"Error occurred: {str(e)}"
@@ -186,14 +232,14 @@ def create_interface():
         gr.Markdown("""
         # AI Data Analysis Assistant
         
-        Upload your data file and ask questions in natural language.
+        Upload your data file and get AI-powered analysis with visualizations.
         
         **Features:**
         - Data analysis and visualization
         - Statistical analysis
         - Machine learning capabilities
         
-        **Note**: Requires your own GPT-4 API key.
+        **Note**: Requires your own OpenAi API key.
         """)
         
         with gr.Row():
@@ -232,10 +278,10 @@ def create_interface():
         
         gr.Examples(
             examples=[
-                [None, "Show key statistics and create visualizations for numeric columns"],
-                [None, "Find correlations and patterns in the data"],
-                [None, "Identify outliers and unusual patterns"],
-                [None, "Create summary visualizations of the main variables"],
+                [None, "Show the distribution of values and key statistics"],
+                [None, "Create a correlation analysis with heatmap"],
+                [None, "Identify and visualize any outliers in the data"],
+                [None, "Generate summary plots for the main variables"],
             ],
             inputs=[file, query]
         )