Update app.py

app.py

@@ -14,11 +14,11 @@ import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 from litellm import completion
 
-
 class CodeEnvironment:
-    """Enhanced environment for executing code with both static and interactive visualization capabilities"""
+    """Safe environment for executing code with data analysis capabilities"""
     
     def __init__(self):
+        # Initialize libraries in globals
         self.globals = {
             'pd': pd,
             'np': np,
@@ -37,10 +37,14 @@ class CodeEnvironment:
             
         # Capture output
         output_buffer = io.StringIO()
+        # Redirect stdout to capture print statements
+        import sys
+        sys.stdout = output_buffer
+        
         result = {
-            'output': '',
-            'figures': [],      # For base64 static images
-            'interactive': [],  # For Plotly HTML
+            'output': '', 
+            'figures': [],      # For matplotlib figures
+            'plotly_html': [],  # For Plotly figures
             'error': None
         }
         
@@ -48,7 +52,7 @@ class CodeEnvironment:
             # Execute code
             exec(code, self.globals, self.locals)
             
-            # Capture matplotlib figures (static)
+            # Capture matplotlib figures
             for i in plt.get_fignums():
                 fig = plt.figure(i)
                 buf = io.BytesIO()
@@ -58,18 +62,16 @@ class CodeEnvironment:
                 result['figures'].append(f"data:image/png;base64,{img_str}")
                 plt.close(fig)
             
-            # Capture Plotly figures (interactive)
-            for var in list(self.locals.values()):
-                if isinstance(var, (go.Figure, px.Figure)):
-                    html = var.to_html(
+            # Capture Plotly figures
+            if 'fig' in self.locals:
+                if isinstance(self.locals['fig'], (go.Figure, px.Figure)):
+                    # Convert Plotly figure to HTML
+                    html = self.locals['fig'].to_html(
                         include_plotlyjs=True,
                         full_html=False,
-                        config={
-                            'displayModeBar': True,
-                            'responsive': True
-                        }
+                        config={'displayModeBar': True}
                     )
-                    result['interactive'].append(html)
+                    result['plotly_html'].append(html)
             
             # Get printed output
             result['output'] = output_buffer.getvalue()
@@ -78,11 +80,12 @@ class CodeEnvironment:
             result['error'] = str(e)
             
         finally:
+            # Reset stdout
+            sys.stdout = sys.__stdout__
             output_buffer.close()
             
         return result
 
-
 @dataclass
 class Tool:
     """Tool for data analysis"""
@@ -90,9 +93,8 @@ class Tool:
     description: str
     func: Callable
 
-
 class AnalysisAgent:
-    """Enhanced agent with interactive visualization capabilities"""
+    """Agent that can analyze data and execute code"""
     
     def __init__(
         self,
@@ -104,12 +106,8 @@ class AnalysisAgent:
         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, df: pd.DataFrame = None) -> str:
-        """Run analysis with enhanced visualization support"""
+        """Run analysis with code execution"""
         messages = [
             {"role": "system", "content": self._get_system_prompt()},
             {"role": "user", "content": prompt}
@@ -138,69 +136,58 @@ class AnalysisAgent:
                     if result['output']:
                         results.append(result['output'])
                     
-                    # Add interactive plots
-                    for plot in result['interactive']:
-                        results.append(f"<div class='plot-container'>{plot}</div>")
+                    # Add Plotly interactive visualizations
+                    for html in result['plotly_html']:
+                        results.append(f'<div class="plot-container">{html}</div>')
                     
-                    # Add static figures as fallback
+                    # Add static matplotlib figures as fallback
                     for fig in result['figures']:
-                        results.append(f"![Figure]({fig})")
-                        
+                        results.append(f'<img src="{fig}" style="max-width: 100%; height: auto;">')
+            
             # Combine analysis and results
-            return analysis + "\n\n" + "\n".join(results)
+            return f'<div class="analysis-text">{analysis}</div>' + "\n\n" + "\n".join(results)
             
         except Exception as e:
             return f"Error: {str(e)}"
     
     def _get_system_prompt(self) -> str:
-        """Get enhanced system prompt with interactive visualization capabilities"""
+        """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 data analysis assistant with interactive visualization capabilities.
-        
-Available tools:
-{tools_desc}
-
-Capabilities:
-- Data analysis (pandas, numpy)
-- Interactive visualization (plotly)
-- Static visualization (matplotlib, seaborn)
-- Statistical analysis (scipy)
-- Machine learning (sklearn)
+        return """You are a data analysis assistant with interactive visualization capabilities.
 
-When writing code:
-- Prefer Plotly for interactive visualizations
-- Use matplotlib/seaborn for static plots when appropriate
-- Create clear visualizations with proper labels
-- Include explanatory text
-- Handle errors gracefully
+When analyzing data, use Plotly for interactive visualizations. Here are examples:
 
-Example Plotly usage:
 ```python
 # Create interactive scatter plot
-fig = px.scatter(df, x='column1', y='column2', 
-                color='category',
-                title='Interactive Analysis')
-fig.update_layout(height=600)
+import plotly.express as px
+fig = px.scatter(df, x='Date', y='Salary', color='Title')
+fig.show()  # This will be captured and displayed
+
+# Create interactive box plot
+fig = px.box(df, x='Title', y='Salary')
 fig.show()
 
 # Create interactive time series
-fig = px.line(df, x='date', y='value',
-              color='category',
-              title='Time Series Analysis')
-fig.update_layout(height=600)
+fig = px.line(df, x='Date', y='Salary', color='Title')
 fig.show()
 ```
 
-Example Matplotlib usage:
+Remember to:
+1. Always store Plotly figures in a variable named 'fig'
+2. Use fig.show() to display the plot
+3. Create clear labels and titles
+4. Include hover information
+5. Use colors effectively
+
+For static visualizations, you can still use matplotlib:
 ```python
-# Create static plot
+import matplotlib.pyplot as plt
 plt.figure(figsize=(10, 6))
-sns.boxplot(data=df, x='category', y='value')
-plt.title('Distribution Analysis')
+plt.plot(df['Date'], df['Salary'])
 plt.show()
 ```
 """