Update app.py

app.py

@@ -11,265 +11,213 @@ 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
         
-    def create_visualization(self, plot_type: str, **kwargs) -> go.Figure:
-        """Create different types of plotly visualizations"""
+    def create_plot(self, plot_type: str, **kwargs) -> go.Figure:
         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'),
-                color=kwargs.get('color'),
+                self.data, 
+                x=kwargs.get('x'), 
+                y=kwargs.get('y'),
                 title=kwargs.get('title', 'Scatter Plot'),
-                labels=kwargs.get('labels', {}),
-                trendline=kwargs.get('trendline'),
+                color=kwargs.get('color')
             )
-            
         elif plot_type == "line":
             fig = px.line(
-                self.data, x=kwargs.get('x'), y=kwargs.get('y'),
-                color=kwargs.get('color'),
-                title=kwargs.get('title', 'Line Plot'),
-                labels=kwargs.get('labels', {})
+                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'),
-                color=kwargs.get('color'),
-                title=kwargs.get('title', 'Bar Plot'),
-                labels=kwargs.get('labels', {})
+                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'),
-                nbins=kwargs.get('bins', 30),
-                title=kwargs.get('title', 'Histogram'),
-                marginal=kwargs.get('marginal', 'box')
+                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'),
-                color=kwargs.get('color'),
+                self.data, 
+                x=kwargs.get('x'), 
+                y=kwargs.get('y'),
                 title=kwargs.get('title', 'Box Plot')
             )
-            
-        elif plot_type == "violin":
-            fig = px.violin(
-                self.data, x=kwargs.get('x'), y=kwargs.get('y'),
-                color=kwargs.get('color'),
-                box=True,
-                title=kwargs.get('title', 'Violin Plot')
-            )
-            
-        elif plot_type == "correlation":
-            corr = self.data.select_dtypes(include=[np.number]).corr()
-            fig = px.imshow(
-                corr,
-                title=kwargs.get('title', 'Correlation Matrix'),
-                color_continuous_scale="RdBu"
-            )
-            
         else:
             raise ValueError(f"Unknown plot type: {plot_type}")
-        
-        # Update layout for better interactivity
-        fig.update_layout(
-            hovermode='x unified',
-            template='plotly_white',
-            height=500,
-        )
-        
+            
         return fig
 
 class ChatAnalyzer:
-    """Handles chat-based analysis with visualization"""
-    
     def __init__(self):
         self.analyzer = DataAnalyzer()
-        self.chat_history: List[Tuple[str, str]] = []
+        self.history: List[Tuple[str, str]] = []
         
-    def process_file(self, file: gr.File) -> str:
-        """Process uploaded file"""
+    def process_file(self, file: gr.File) -> List[Tuple[str, str]]:
         try:
             if file.name.endswith('.csv'):
                 self.analyzer.data = pd.read_csv(file.name)
             elif file.name.endswith(('.xlsx', '.xls')):
                 self.analyzer.data = pd.read_excel(file.name)
             else:
-                return "Error: Please upload a CSV or Excel file."
+                return [("System", "Error: Please upload a CSV or Excel file.")]
+            
+            info = f"""Data loaded successfully!
+            Shape: {self.analyzer.data.shape}
+            Columns: {', '.join(self.analyzer.data.columns)}"""
             
-            info = f"""
-            Successfully loaded data with shape: {self.analyzer.data.shape}
-            Columns: {', '.join(self.analyzer.data.columns)}
-            """
-            return info
+            self.history = [("System", info)]
+            return self.history
             
         except Exception as e:
-            return f"Error loading file: {str(e)}"
+            self.history = [("System", f"Error loading file: {str(e)}")]
+            return self.history
     
-    def analyze(self, message: str, api_key: str) -> Tuple[str, List[go.Figure]]:
-        """Analyze data based on user message"""
+    def chat(self, message: str, api_key: str) -> Tuple[List[Tuple[str, str]], List[gr.Plot]]:
         if self.analyzer.data is None:
-            return "Please upload a data file first.", []
+            return [(message, "Please upload a data file first.")], []
             
         if not api_key:
-            return "Please provide an OpenAI API key.", []
+            return [(message, "Please provide an OpenAI API key.")], []
             
         try:
             os.environ["OPENAI_API_KEY"] = api_key
             
-            # Prepare context for AI
+            # Get data context
             context = self._get_data_context()
             
             # Get AI response
-            messages = [
-                {"role": "system", "content": self._get_system_prompt()},
-                {"role": "user", "content": f"{context}\n\nUser request: {message}"}
-            ]
-            
-            response = completion(
+            completion_response = completion(
                 model="gpt-4o-mini",
-                messages=messages,
+                messages=[
+                    {"role": "system", "content": self._get_system_prompt()},
+                    {"role": "user", "content": f"{context}\n\nUser question: {message}"}
+                ],
                 temperature=0.7
             )
             
-            analysis = response.choices[0].message.content
+            analysis = completion_response.choices[0].message.content
+            
+            # Create visualizations
+            figures = []
+            try:
+                # Execute any visualization commands in the analysis
+                exec_globals = {
+                    'analyzer': self.analyzer,
+                    'df': self.analyzer.data,
+                    'px': px,
+                    'go': go,
+                    'print': lambda x: figures.append(x) if isinstance(x, go.Figure) else None
+                }
+                
+                # Extract code blocks
+                import re
+                code_blocks = re.findall(r'```python\n(.*?)```', analysis, re.DOTALL)
+                
+                for code in code_blocks:
+                    exec(code, exec_globals)
+                    
+            except Exception as e:
+                analysis += f"\n\nError creating visualization: {str(e)}"
+            
+            # Update chat history
+            self.history.append((message, analysis))
             
-            # Extract visualization commands and create plots
-            figures = self._create_visualizations(analysis)
+            # Convert figures to Gradio plots
+            plots = [gr.Plot(fig) for fig in figures]
             
-            return analysis, figures
+            return self.history, plots
             
         except Exception as e:
-            return f"Error during analysis: {str(e)}", []
+            self.history.append((message, f"Error: {str(e)}"))
+            return self.history, []
     
     def _get_data_context(self) -> str:
-        """Get current data context"""
         df = self.analyzer.data
         numeric_cols = df.select_dtypes(include=[np.number]).columns
         categorical_cols = df.select_dtypes(include=['object', 'category']).columns
         
         return f"""
-        Available Data Information:
+        Data Information:
         - Shape: {df.shape}
         - Numeric columns: {', '.join(numeric_cols)}
         - Categorical columns: {', '.join(categorical_cols)}
-        
-        Basic Statistics:
-        {df.describe().to_string()}
         """
     
     def _get_system_prompt(self) -> str:
-        """Get system prompt"""
-        return """You are a data analysis assistant specialized in creating interactive visualizations using Plotly.
-
-Available plot types:
-1. scatter - for relationships between variables
-2. line - for trends over time
-3. bar - for comparisons between categories
-4. histogram - for distributions
-5. box - for statistical summaries
-6. violin - for distribution comparisons
-7. correlation - for correlation matrix
+        return """You are a data analysis assistant. To create visualizations, use Python code blocks with Plotly.
 
-When creating visualizations:
-1. Specify the plot type and required parameters
-2. Provide insights about the visualization
-3. Suggest follow-up analyses
-4. Use markdown for formatting
-
-Example command format:
+Example commands:
 ```python
 # Create scatter plot
-plot = viz.create_visualization("scatter", x="column1", y="column2", title="Analysis")
-print(plot)
+fig = px.scatter(df, x='column1', y='column2', title='Analysis')
+print(fig)
+
+# Create histogram
+fig = px.histogram(df, x='column', title='Distribution')
+print(fig)
 ```
-"""
-    
-    def _create_visualizations(self, analysis: str) -> List[go.Figure]:
-        """Extract and create visualizations from analysis"""
-        figures = []
-        viz = self.analyzer
-        
-        try:
-            # Execute visualization commands in the analysis
-            exec_globals = {
-                'viz': viz,
-                'print': lambda x: figures.append(x) if isinstance(x, go.Figure) else None
-            }
-            
-            # Extract and execute code blocks
-            import re
-            code_blocks = re.findall(r'```python\n(.*?)```', analysis, re.DOTALL)
-            
-            for code in code_blocks:
-                exec(code, exec_globals)
-                
-        except Exception as e:
-            print(f"Error creating visualizations: {str(e)}")
-            
-        return figures
+
+Provide analysis and insights along with visualizations."""
 
 def create_interface():
-    """Create Gradio interface"""
-    
     analyzer = ChatAnalyzer()
     
-    def chat(message: str, api_key: str) -> Tuple[List[Tuple[str, str]], List[gr.Plot]]:
-        """Handle chat interaction"""
-        response, figures = analyzer.analyze(message, api_key)
-        
-        # Update chat history
-        analyzer.chat_history.append((message, response))
-        
-        # Convert figures to Gradio plots
-        plots = [gr.Plot(fig) for fig in figures]
-        
-        return analyzer.chat_history, plots
-    
     with gr.Blocks() as demo:
         gr.Markdown("""
         # Interactive Data Analysis Chat
-        Upload your data and chat with AI to create interactive visualizations!
+        Upload your data and chat with AI to analyze it!
         """)
         
         with gr.Row():
             with gr.Column(scale=1):
                 file = gr.File(label="Upload Data (CSV or Excel)")
-                api_key = gr.Textbox(label="OpenAI API Key", type="password")
-            
+                api_key = gr.Textbox(
+                    label="OpenAI API Key",
+                    type="password",
+                    placeholder="Enter your API key"
+                )
+                
             with gr.Column(scale=2):
                 chatbot = gr.Chatbot(height=400)
-        
-        with gr.Row():
-            message = gr.Textbox(label="Ask about your data", lines=2)
-            send = gr.Button("Send")
+                message = gr.Textbox(
+                    label="Ask about your data",
+                    placeholder="e.g., Show me a scatter plot of X vs Y"
+                )
+                send = gr.Button("Send")
         
         # Plot output area
-        plot_output = gr.Plot(visible=False)
+        plot_output = gr.Plot(label="Visualization")
+        
+        # Event handlers
+        file.change(
+            analyzer.process_file,
+            inputs=[file],
+            outputs=[chatbot]
+        )
         
-        # Set up event handlers
-        file.change(analyzer.process_file, inputs=[file], outputs=[chatbot])
         send.click(
-            chat,
+            analyzer.chat,
             inputs=[message, api_key],
             outputs=[chatbot, plot_output]
         )
         
+        # Example queries
         gr.Examples(
             examples=[
-                ["Show me a scatter plot of the main numerical variables"],
-                ["Create a correlation matrix of all numerical columns"],
-                ["Analyze the distribution of each variable"],
-                ["Show trends over time if there's temporal data"],
+                ["Show me a scatter plot of the numerical variables"],
+                ["Create a histogram of the distribution"],
+                ["Analyze trends in the data"],
             ],
             inputs=message
         )