Update app.py

app.py

@@ -12,7 +12,6 @@ import plotly.express as px
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 from litellm import completion
-
 class DataAnalyzer:
     """Handles data analysis and visualization"""
     
@@ -20,133 +19,133 @@ class DataAnalyzer:
         self.data: Optional[pd.DataFrame] = None
         self.width = 800
         self.height = 500
-        self.template = "plotly_white"
-    
-    def create_scatter(self, x_col: str, y_col: str, color_col: Optional[str] = None, 
-                      title: str = "") -> go.Figure:
-        """Create scatter plot"""
-        fig = px.scatter(
-            self.data,
-            x=x_col,
-            y=y_col,
-            color=color_col,
+        
+    def create_histogram(self, column: str, bins: int = 30, title: str = "") -> str:
+        """Create histogram with Plotly"""
+        if self.data is None:
+            raise ValueError("No data loaded")
+            
+        fig = go.Figure()
+        
+        fig.add_trace(go.Histogram(
+            x=self.data[column],
+            nbinsx=bins,
+            name=column
+        ))
+        
+        fig.update_layout(
             title=title,
-            template=self.template,
+            xaxis_title=column,
+            yaxis_title="Count",
+            width=self.width,
             height=self.height,
-            width=self.width
+            template="plotly_white"
         )
-        fig.update_layout(
-            hovermode='closest',
-            showlegend=True if color_col else False
-        )
-        return fig
+        
+        # Convert to HTML string
+        return fig.to_html(include_plotlyjs=True, full_html=False)
     
-    def create_line(self, x_col: str, y_cols: List[str], title: str = "") -> go.Figure:
-        """Create line plot"""
+    def create_scatter(self, x_col: str, y_col: str, color_col: Optional[str] = None, 
+                      title: str = "") -> str:
+        """Create scatter plot with Plotly"""
+        if self.data is None:
+            raise ValueError("No data loaded")
+            
         fig = go.Figure()
         
-        for y_col in y_cols:
-            fig.add_trace(
-                go.Scatter(
-                    x=self.data[x_col],
-                    y=self.data[y_col],
-                    name=y_col,
-                    mode='lines+markers'
-                )
-            )
+        if color_col:
+            for category in self.data[color_col].unique():
+                mask = self.data[color_col] == category
+                fig.add_trace(go.Scatter(
+                    x=self.data[mask][x_col],
+                    y=self.data[mask][y_col],
+                    mode='markers',
+                    name=str(category),
+                    text=self.data[mask][color_col]
+                ))
+        else:
+            fig.add_trace(go.Scatter(
+                x=self.data[x_col],
+                y=self.data[y_col],
+                mode='markers'
+            ))
         
         fig.update_layout(
             title=title,
-            template=self.template,
-            height=self.height,
+            xaxis_title=x_col,
+            yaxis_title=y_col,
             width=self.width,
-            hovermode='x unified',
-            showlegend=True
-        )
-        return fig
-    
-    def create_bar(self, x_col: str, y_col: str, color_col: Optional[str] = None,
-                  title: str = "") -> go.Figure:
-        """Create bar plot"""
-        fig = px.bar(
-            self.data,
-            x=x_col,
-            y=y_col,
-            color=color_col,
-            title=title,
-            template=self.template,
             height=self.height,
-            width=self.width
-        )
-        fig.update_layout(
-            hovermode='closest',
-            showlegend=True if color_col else False
+            template="plotly_white",
+            hovermode='closest'
         )
-        return fig
+        
+        return fig.to_html(include_plotlyjs=True, full_html=False)
     
-    def create_histogram(self, column: str, bins: int = 30, 
-                        title: str = "") -> go.Figure:
-        """Create histogram"""
-        fig = px.histogram(
-            self.data,
-            x=column,
-            nbins=bins,
+    def create_box(self, x_col: str, y_col: str, title: str = "") -> str:
+        """Create box plot with Plotly"""
+        if self.data is None:
+            raise ValueError("No data loaded")
+            
+        fig = go.Figure()
+        
+        # Create box plot for each category
+        for category in self.data[x_col].unique():
+            fig.add_trace(go.Box(
+                y=self.data[self.data[x_col] == category][y_col],
+                name=str(category),
+                boxpoints='all',  # show all points
+                jitter=0.3,
+                pointpos=-1.8
+            ))
+        
+        fig.update_layout(
             title=title,
-            template=self.template,
-            height=self.height,
+            yaxis_title=y_col,
+            xaxis_title=x_col,
             width=self.width,
-            marginal="box"  # Add box plot on the margin
-        )
-        fig.update_layout(
-            hovermode='closest',
+            height=self.height,
+            template="plotly_white",
             showlegend=False
         )
-        return fig
+        
+        return fig.to_html(include_plotlyjs=True, full_html=False)
     
-    def create_box(self, x_col: str, y_col: str, color_col: Optional[str] = None,
-                  title: str = "") -> go.Figure:
-        """Create box plot"""
-        fig = px.box(
-            self.data,
-            x=x_col,
-            y=y_col,
-            color=color_col,
-            title=title,
-            template=self.template,
-            height=self.height,
-            width=self.width,
-            points="all"  # Show all points
-        )
+    def create_line(self, x_col: str, y_col: str, color_col: Optional[str] = None, 
+                    title: str = "") -> str:
+        """Create line plot with Plotly"""
+        if self.data is None:
+            raise ValueError("No data loaded")
+            
+        fig = go.Figure()
+        
+        if color_col:
+            for category in self.data[color_col].unique():
+                mask = self.data[color_col] == category
+                fig.add_trace(go.Scatter(
+                    x=self.data[mask][x_col],
+                    y=self.data[mask][y_col],
+                    mode='lines+markers',
+                    name=str(category)
+                ))
+        else:
+            fig.add_trace(go.Scatter(
+                x=self.data[x_col],
+                y=self.data[y_col],
+                mode='lines+markers'
+            ))
+        
         fig.update_layout(
-            hovermode='closest',
-            showlegend=True if color_col else False
-        )
-        return fig
-    
-    def create_correlation_matrix(self, title: str = "") -> go.Figure:
-        """Create correlation matrix"""
-        # Get numeric columns
-        numeric_cols = self.data.select_dtypes(include=[np.number]).columns
-        corr_matrix = self.data[numeric_cols].corr()
-        
-        fig = px.imshow(
-            corr_matrix,
             title=title,
-            template=self.template,
-            height=self.height,
+            xaxis_title=x_col,
+            yaxis_title=y_col,
             width=self.width,
-            color_continuous_scale="RdBu",
-            aspect="auto",
-            labels=dict(color="Correlation")
+            height=self.height,
+            template="plotly_white",
+            hovermode='x unified'
         )
         
-        # Update layout for better readability
-        fig.update_traces(text=corr_matrix.round(2), texttemplate="%{text}")
-        fig.update_layout(
-            xaxis_title="",
-            yaxis_title=""
-        )
-        return fig
+        return fig.to_html(include_plotlyjs=True, full_html=False)
 
 class ChatAnalyzer:
     """Handles chat-based analysis with visualization"""
@@ -190,94 +189,65 @@ class ChatAnalyzer:
             return self.history
     
     def chat(self, message: str, api_key: str) -> Tuple[List[Tuple[str, str]], str]:
-        """Process chat message and generate visualizations"""
-        if self.analyzer.data is None:
-            return [(message, "Please upload a data file first.")], ""
-            
-        if not api_key:
-            return [(message, "Please provide an OpenAI API key.")], ""
-            
+    """Process chat message and generate visualizations"""
+    if self.analyzer.data is None:
+        return [(message, "Please upload a data file first.")], ""
+        
+    if not api_key:
+        return [(message, "Please provide an OpenAI API key.")], ""
+        
+    try:
+        os.environ["OPENAI_API_KEY"] = api_key
+        
+        # Get data context
+        context = self._get_data_context()
+        
+        # Get AI response
+        completion_response = completion(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": self._get_system_prompt()},
+                {"role": "user", "content": f"{context}\n\nUser question: {message}"}
+            ],
+            temperature=0.7
+        )
+        
+        analysis = completion_response.choices[0].message.content
+        
+        # Create visualizations
+        plots_html = ""
         try:
-            os.environ["OPENAI_API_KEY"] = api_key
-            
-            # Get data context
-            context = self._get_data_context()
-            
-            # Get AI response
-            completion_response = completion(
-                model="gpt-4o-mini",
-                messages=[
-                    {"role": "system", "content": self._get_system_prompt()},
-                    {"role": "user", "content": f"{context}\n\nUser question: {message}"}
-                ],
-                temperature=0.7
-            )
-            
-            analysis = completion_response.choices[0].message.content
+            # Extract code blocks
+            import re
+            code_blocks = re.findall(r'```python\n(.*?)```', analysis, re.DOTALL)
             
-            # Create visualizations
-            plot_output = ""
-            try:
-                # Extract code blocks
-                import re
-                code_blocks = re.findall(r'```python\n(.*?)```', analysis, re.DOTALL)
+            for code in code_blocks:
+                # Create namespace for execution
+                namespace = {
+                    'analyzer': self.analyzer,
+                    'df': self.analyzer.data,
+                    'print': lambda x: x
+                }
                 
-                for code in code_blocks:
-                    # Create namespace for execution
-                    namespace = {
-                        'df': self.analyzer.data,
-                        'px': px,
-                        'go': go,
-                        'pd': pd,
-                        'np': np,
-                        'analyzer': self.analyzer
-                    }
-                    
-                    # Execute the code
+                # Execute the code
+                try:
+                    result = eval(code, namespace)
+                    if isinstance(result, str) and ('<div' in result or '<script' in result):
+                        plots_html += f'<div class="plot-container">{result}</div>'
+                except:
                     exec(code, namespace)
-                    
-                    # Look for figure object in namespace
-                    for var in namespace.values():
-                        if isinstance(var, (go.Figure, px.Figure)):
-                            try:
-                                # Try interactive HTML first
-                                html = var.to_html(
-                                    include_plotlyjs=True,
-                                    full_html=False,
-                                    config={
-                                        'displayModeBar': True,
-                                        'responsive': True
-                                    }
-                                )
-                                plot_output += f'''
-                                <div class="plot-container">
-                                    <div style="overflow-x: auto;">{html}</div>
-                                </div>
-                                '''
-                            except Exception as e:
-                                # Fallback to static image
-                                buffer = io.BytesIO()
-                                var.write_image(buffer, format='png')
-                                buffer.seek(0)
-                                image = base64.b64encode(buffer.read()).decode()
-                                plot_output += f'''
-                                <div class="plot-container">
-                                    <img src="data:image/png;base64,{image}" 
-                                         style="max-width: 100%; height: auto;">
-                                </div>
-                                '''
                 
-            except Exception as e:
-                analysis += f"\n\nError creating visualization: {str(e)}"
-            
-            # Update chat history
-            self.history.append((message, analysis))
-            
-            return self.history, plot_output
-            
         except Exception as e:
-            self.history.append((message, f"Error: {str(e)}"))
-            return self.history, ""
+            analysis += f"\n\nError creating visualization: {str(e)}"
+        
+        # Update chat history
+        self.history.append((message, analysis))
+        
+        return self.history, plots_html
+        
+    except Exception as e:
+        self.history.append((message, f"Error: {str(e)}"))
+        return self.history, ""
     
     def _get_data_context(self) -> str:
         """Get current data context for AI"""
@@ -309,51 +279,45 @@ class ChatAnalyzer:
         """
     
     def _get_system_prompt(self) -> str:
-        """Get system prompt for AI"""
-        return """You are a data analysis assistant specialized in creating interactive visualizations.
+    """Get system prompt for AI"""
+    return """You are a data analysis assistant specialized in creating interactive visualizations.
 
 Available visualization functions:
-1. analyzer.create_scatter(x_col, y_col, color_col, title)
-2. analyzer.create_line(x_col, y_cols, title)
-3. analyzer.create_bar(x_col, y_col, color_col, title)
-4. analyzer.create_histogram(column, bins, title)
-5. analyzer.create_box(x_col, y_col, color_col, title)
-6. analyzer.create_correlation_matrix(title)
-
-When analyzing data:
-1. First understand the data type and relationships
-2. Choose appropriate visualizations
-3. Provide insights and analysis
-4. Suggest follow-up analyses
+1. create_histogram(column, bins, title) - For distribution analysis
+2. create_scatter(x_col, y_col, color_col, title) - For relationship analysis
+3. create_box(x_col, y_col, title) - For categorical comparisons
+4. create_line(x_col, y_col, color_col, title) - For trend analysis
 
 Example usage:
 ```python
-# Create scatter plot
-fig = analyzer.create_scatter(
-    x_col='Date',
-    y_col='Value',
-    color_col='Category',
-    title='Value Trends by Category'
+# Create histogram
+result = analyzer.create_histogram(
+    column='Salary',
+    bins=20,
+    title='Salary Distribution'
 )
-print(fig)
+print(result)
 
-# Create multiple visualizations
-fig1 = analyzer.create_histogram(
-    column='Value',
-    bins=30,
-    title='Value Distribution'
+# Create scatter plot
+result = analyzer.create_scatter(
+    x_col='Date',
+    y_col='Salary',
+    color_col='Title',
+    title='Salary Trends by Title'
 )
-print(fig1)
+print(result)
 
-fig2 = analyzer.create_box(
-    x_col='Category',
-    y_col='Value',
-    title='Value Distribution by Category'
+# Create box plot
+result = analyzer.create_box(
+    x_col='Title',
+    y_col='Salary',
+    title='Salary Distribution by Title'
 )
-print(fig2)
+print(result)
 ```
 
-Always wrap code in Python code blocks and print the figures to display them."""
+Always wrap code in Python code blocks and use print() to display the visualizations. 
+Provide analysis and insights about what the visualizations show."""
 
 def create_interface():
     """Create Gradio interface"""