Update app.py

app.py

@@ -1,4 +1,7 @@
+import base64
+import io
 import os
+from dataclasses import dataclass
 from typing import List, Optional, Tuple, Dict, Any
 import json
 
@@ -11,61 +14,149 @@ from plotly.subplots import make_subplots
 from litellm import completion
 
 class DataAnalyzer:
+    """Handles data analysis and visualization"""
+    
     def __init__(self):
         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,
+            title=title,
+            template=self.template,
+            height=self.height,
+            width=self.width
+        )
+        fig.update_layout(
+            hovermode='closest',
+            showlegend=True if color_col else False
+        )
+        return fig
+    
+    def create_line(self, x_col: str, y_cols: List[str], title: str = "") -> go.Figure:
+        """Create line plot"""
+        fig = go.Figure()
         
-    def create_plot(self, plot_type: str, **kwargs) -> str:
-        """Create plot and return HTML string"""
-        if self.data is None:
-            raise ValueError("No data loaded")
-            
-        if plot_type == "scatter":
-            fig = px.scatter(
-                self.data, 
-                x=kwargs.get('x'), 
-                y=kwargs.get('y'),
-                title=kwargs.get('title', 'Scatter Plot'),
-                color=kwargs.get('color')
-            )
-        elif plot_type == "line":
-            fig = px.line(
-                self.data, 
-                x=kwargs.get('x'), 
-                y=kwargs.get('y'),
-                title=kwargs.get('title', 'Line Plot')
-            )
-        elif plot_type == "bar":
-            fig = px.bar(
-                self.data, 
-                x=kwargs.get('x'), 
-                y=kwargs.get('y'),
-                title=kwargs.get('title', 'Bar Plot')
-            )
-        elif plot_type == "histogram":
-            fig = px.histogram(
-                self.data, 
-                x=kwargs.get('x'),
-                title=kwargs.get('title', 'Distribution')
-            )
-        elif plot_type == "box":
-            fig = px.box(
-                self.data, 
-                x=kwargs.get('x'), 
-                y=kwargs.get('y'),
-                title=kwargs.get('title', 'Box Plot')
+        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'
+                )
             )
-        else:
-            raise ValueError(f"Unknown plot type: {plot_type}")
         
-        # Convert plot to HTML string
-        return fig.to_html(include_plotlyjs=True, full_html=False)
+        fig.update_layout(
+            title=title,
+            template=self.template,
+            height=self.height,
+            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
+        )
+        return fig
+    
+    def create_histogram(self, column: str, bins: int = 30, 
+                        title: str = "") -> go.Figure:
+        """Create histogram"""
+        fig = px.histogram(
+            self.data,
+            x=column,
+            nbins=bins,
+            title=title,
+            template=self.template,
+            height=self.height,
+            width=self.width,
+            marginal="box"  # Add box plot on the margin
+        )
+        fig.update_layout(
+            hovermode='closest',
+            showlegend=False
+        )
+        return fig
+    
+    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
+        )
+        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,
+            width=self.width,
+            color_continuous_scale="RdBu",
+            aspect="auto",
+            labels=dict(color="Correlation")
+        )
+        
+        # Update layout for better readability
+        fig.update_traces(text=corr_matrix.round(2), texttemplate="%{text}")
+        fig.update_layout(
+            xaxis_title="",
+            yaxis_title=""
+        )
+        return fig
 
 class ChatAnalyzer:
+    """Handles chat-based analysis with visualization"""
+    
     def __init__(self):
         self.analyzer = DataAnalyzer()
         self.history: List[Tuple[str, str]] = []
-        
+    
     def process_file(self, file: gr.File) -> List[Tuple[str, str]]:
+        """Process uploaded file and initialize analyzer"""
         try:
             if file.name.endswith('.csv'):
                 self.analyzer.data = pd.read_csv(file.name)
@@ -74,9 +165,22 @@ class ChatAnalyzer:
             else:
                 return [("System", "Error: Please upload a CSV or Excel file.")]
             
+            # Convert date columns to datetime
+            date_cols = self.analyzer.data.select_dtypes(include=['object']).columns
+            for col in date_cols:
+                try:
+                    self.analyzer.data[col] = pd.to_datetime(self.analyzer.data[col])
+                except:
+                    continue
+            
             info = f"""Data loaded successfully!
             Shape: {self.analyzer.data.shape}
-            Columns: {', '.join(self.analyzer.data.columns)}"""
+            Columns: {', '.join(self.analyzer.data.columns)}
+            
+            Numeric columns: {', '.join(self.analyzer.data.select_dtypes(include=[np.number]).columns)}
+            Date columns: {', '.join(self.analyzer.data.select_dtypes(include=['datetime64']).columns)}
+            Categorical columns: {', '.join(self.analyzer.data.select_dtypes(include=['object']).columns)}
+            """
             
             self.history = [("System", info)]
             return self.history
@@ -86,6 +190,7 @@ 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.")], ""
             
@@ -111,120 +216,318 @@ class ChatAnalyzer:
             analysis = completion_response.choices[0].message.content
             
             # Create visualizations
-            plot_html = ""
+            plot_output = ""
             try:
-                # Execute any visualization commands in the analysis
-                exec_globals = {
-                    'analyzer': self.analyzer,
-                    'df': self.analyzer.data,
-                    'px': px,
-                    'go': go,
-                    'print': lambda x: x  # Will be used to capture plot output
-                }
-                
                 # Extract code blocks
                 import re
                 code_blocks = re.findall(r'```python\n(.*?)```', analysis, re.DOTALL)
                 
                 for code in code_blocks:
-                    result = exec(code, exec_globals)
-                    if isinstance(result, (go.Figure, px.Figure)):
-                        plot_html += result.to_html(include_plotlyjs=True, full_html=False)
+                    # Create namespace for execution
+                    namespace = {
+                        'df': self.analyzer.data,
+                        'px': px,
+                        'go': go,
+                        'pd': pd,
+                        'np': np,
+                        'analyzer': self.analyzer
+                    }
                     
+                    # Execute the code
+                    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_html
+            return self.history, plot_output
             
         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"""
         df = self.analyzer.data
         numeric_cols = df.select_dtypes(include=[np.number]).columns
-        categorical_cols = df.select_dtypes(include=['object', 'category']).columns
+        date_cols = df.select_dtypes(include=['datetime64']).columns
+        categorical_cols = df.select_dtypes(include=['object']).columns
+        
+        # Get basic statistics
+        stats = df[numeric_cols].describe().to_string() if len(numeric_cols) > 0 else "No numeric columns"
         
         return f"""
         Data Information:
         - Shape: {df.shape}
         - Numeric columns: {', '.join(numeric_cols)}
+        - Date columns: {', '.join(date_cols)}
         - Categorical columns: {', '.join(categorical_cols)}
+        
+        Basic Statistics:
+        {stats}
+        
+        Available visualization functions:
+        - analyzer.create_scatter(x_col, y_col, color_col, title)
+        - analyzer.create_line(x_col, y_cols, title)
+        - analyzer.create_bar(x_col, y_col, color_col, title)
+        - analyzer.create_histogram(column, bins, title)
+        - analyzer.create_box(x_col, y_col, color_col, title)
+        - analyzer.create_correlation_matrix(title)
         """
     
     def _get_system_prompt(self) -> str:
-        return """You are a data analysis assistant. To create visualizations, use Python code blocks with Plotly.
+        """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)
 
-Example commands:
+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
+
+Example usage:
 ```python
 # Create scatter plot
-fig = px.scatter(df, x='column1', y='column2', title='Analysis')
-fig  # Return the figure object
+fig = analyzer.create_scatter(
+    x_col='Date',
+    y_col='Value',
+    color_col='Category',
+    title='Value Trends by Category'
+)
+print(fig)
+
+# Create multiple visualizations
+fig1 = analyzer.create_histogram(
+    column='Value',
+    bins=30,
+    title='Value Distribution'
+)
+print(fig1)
 
-# Create histogram
-fig = px.histogram(df, x='column', title='Distribution')
-fig  # Return the figure object
+fig2 = analyzer.create_box(
+    x_col='Category',
+    y_col='Value',
+    title='Value Distribution by Category'
+)
+print(fig2)
 ```
 
-Provide analysis and insights along with visualizations."""
+Always wrap code in Python code blocks and print the figures to display them."""
 
 def create_interface():
+    """Create Gradio interface"""
+    
     analyzer = ChatAnalyzer()
     
-    with gr.Blocks(css="footer {visibility: hidden}") as demo:
+    # Custom CSS for better visualization display
+    css = """
+    .container { max-width: 1200px; margin: auto; }
+    .plot-container {
+        margin: 20px 0;
+        padding: 15px;
+        border: 1px solid #e0e0e0;
+        border-radius: 8px;
+        background: white;
+        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+    }
+    .chat-message {
+        margin-bottom: 15px;
+        padding: 10px;
+        border-radius: 8px;
+        background: #f8f9fa;
+    }
+    .title {
+        text-align: center;
+        margin-bottom: 20px;
+    }
+    .footer {
+        text-align: center;
+        margin-top: 20px;
+        font-size: 0.9em;
+        color: #666;
+    }
+    """
+    
+    with gr.Blocks(css=css, title="Interactive Data Analysis Chat") as demo:
         gr.Markdown("""
         # Interactive Data Analysis Chat
-        Upload your data and chat with AI to analyze it!
+        
+        Upload your data and chat with AI to analyze it! Features:
+        - Interactive visualizations with zoom, pan, and hover capabilities
+        - Natural language analysis and insights
+        - Statistical analysis and summaries
+        - Trend detection and pattern analysis
+        
+        Start by uploading a CSV or Excel file.
         """)
         
         with gr.Row():
             with gr.Column(scale=1):
-                file = gr.File(label="Upload Data (CSV or Excel)")
+                file = gr.File(
+                    label="Upload Data (CSV or Excel)",
+                    file_types=[".csv", ".xlsx", ".xls"],
+                    elem_classes="file-upload"
+                )
                 api_key = gr.Textbox(
                     label="OpenAI API Key",
                     type="password",
-                    placeholder="Enter your API key"
+                    placeholder="Enter your API key",
+                    elem_classes="api-input"
                 )
                 
             with gr.Column(scale=2):
-                chatbot = gr.Chatbot(height=400)
-                message = gr.Textbox(
-                    label="Ask about your data",
-                    placeholder="e.g., Show me a scatter plot of X vs Y"
+                chatbot = gr.Chatbot(
+                    height=400,
+                    elem_classes="chat-message",
+                    show_label=False
                 )
-                send = gr.Button("Send")
+                with gr.Row():
+                    message = gr.Textbox(
+                        label="Ask about your data",
+                        placeholder="e.g., Show me trends in the data",
+                        lines=2,
+                        elem_classes="message-input",
+                        scale=4
+                    )
+                    send = gr.Button(
+                        "Send",
+                        scale=1,
+                        elem_classes="send-button"
+                    )
         
-        # Plot output area (using HTML for interactive plots)
-        plot_output = gr.HTML(label="Visualization")
+        # Plot output area
+        plot_output = gr.HTML(
+            label="Visualizations",
+            elem_classes="plot-container",
+            visible=True  # Always show container even when empty
+        )
         
         # Event handlers
         file.change(
-            analyzer.process_file,
+            fn=analyzer.process_file,
             inputs=[file],
-            outputs=[chatbot]
+            outputs=[chatbot],
+            api_name="upload"
+        )
+        
+        # Handle both click and enter key
+        msg_handler = send.click(
+            fn=analyzer.chat,
+            inputs=[message, api_key],
+            outputs=[chatbot, plot_output],
+            api_name="chat"
         )
         
-        send.click(
-            analyzer.chat,
+        message.submit(
+            fn=analyzer.chat,
             inputs=[message, api_key],
             outputs=[chatbot, plot_output]
         )
         
+        # Clear message after sending
+        msg_handler.then(
+            fn=lambda: "",
+            inputs=[],
+            outputs=[message]
+        )
+        
         # Example queries
         gr.Examples(
             examples=[
-                ["Show me a scatter plot of the numerical variables"],
-                ["Create a histogram of the distribution"],
-                ["Analyze trends in the data"],
+                ["Show me a scatter plot of the main numerical variables and explain any patterns you see"],
+                ["Create a correlation analysis with heatmap and highlight the strongest relationships"],
+                ["Show the distribution of values using histograms and describe the shapes"],
+                ["Create box plots to analyze categories and identify any outliers"],
+                ["Show trends over time using line plots and explain the patterns"],
+                ["Generate a comprehensive analysis with multiple visualizations"],
+                ["Compare the distribution across different categories using appropriate plots"],
+                ["Identify and visualize any seasonal patterns or cycles in the data"],
             ],
-            inputs=message
+            inputs=[message],
+            label="Example Analysis Queries"
+        )
+        
+        # Tips section
+        gr.Markdown("""
+        ### Tips for better analysis:
+        1. **Data Preparation**: Upload clean CSV or Excel files with clear column names
+        2. **Specific Questions**: Ask clear, specific questions about your data
+        3. **Interactive Features**: Use zoom, pan, and hover on visualizations
+        4. **Follow-up Questions**: Ask for deeper analysis of interesting patterns
+        5. **Multiple Views**: Request different visualization types for better insights
+        
+        ### Available Visualization Types:
+        - Scatter plots for relationships
+        - Line plots for trends
+        - Bar charts for comparisons
+        - Histograms for distributions
+        - Box plots for statistical summaries
+        - Correlation matrices for relationship analysis
+        """)
+        
+        # Footer
+        gr.Markdown("""
+        <div class="footer">
+        Built with Gradio • Powered by OpenAI • Interactive Visualizations
+        </div>
+        """)
+        
+        # Theme customization
+        demo.theme = gr.themes.Soft(
+            primary_hue="blue",
+            secondary_hue="gray",
+            neutral_hue="gray",
+            text_size=gr.themes.sizes.text_md
         )
     
     return demo
 
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch()
+    demo.launch(
+        share=False,  # Set to True to create a public link
+        debug=True,   # Set to False in production
+        show_error=True,  # Show detailed error messages
+        server_port=7860  # Specify port number
+    )