Update app.py

app.py

@@ -1,237 +1,29 @@
-import base64
-import io
-import os
-from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
-import json
-
-import gradio as gr
-import numpy as np
-import pandas as pd
-from bokeh.plotting import figure
-from bokeh.layouts import column, row, layout
-from bokeh.models import ColumnDataSource, HoverTool, BoxSelectTool, WheelZoomTool, ResetTool
-from bokeh.embed import components
-from bokeh.resources import CDN
-from litellm import completion
-
-class VisualizationEngine:
-    """Engine for creating interactive Bokeh visualizations"""
-    
-    def __init__(self):
-        self.width = 600
-        self.height = 400
-        self.tools = "pan,box_zoom,wheel_zoom,reset,save,hover"
-    
-    def create_scatter(self, df: pd.DataFrame, x_col: str, y_col: str, 
-                      color_col: Optional[str] = None, title: str = "") -> str:
-        """Create an interactive scatter plot"""
-        source = ColumnDataSource(df)
-        
-        p = figure(width=self.width, height=self.height, title=title, tools=self.tools)
-        
-        if color_col and color_col in df.columns:
-            colors = df[color_col].astype('category').cat.codes
-            scatter = p.scatter(x_col, y_col, source=source, color={'field': color_col, 'transform': 'category10'})
-        else:
-            scatter = p.scatter(x_col, y_col, source=source)
-        
-        p.xaxis.axis_label = x_col
-        p.yaxis.axis_label = y_col
-        
-        hover = p.select(dict(type=HoverTool))
-        hover.tooltips = [(col, f"@{col}") for col in [x_col, y_col] + ([color_col] if color_col else [])]
-        
-        script, div = components(p)
-        return f"{CDN.render()}\n{div}\n{script}"
-    
-    def create_line(self, df: pd.DataFrame, x_col: str, y_cols: List[str], title: str = "") -> str:
-        """Create an interactive line plot"""
-        source = ColumnDataSource(df)
-        
-        p = figure(width=self.width, height=self.height, title=title, tools=self.tools)
-        
-        for y_col in y_cols:
-            p.line(x_col, y_col, line_width=2, source=source, legend_label=y_col)
-        
-        p.xaxis.axis_label = x_col
-        p.yaxis.axis_label = "Values"
-        p.legend.click_policy = "hide"
-        
-        hover = p.select(dict(type=HoverTool))
-        hover.tooltips = [(col, f"@{col}") for col in [x_col] + y_cols]
-        
-        script, div = components(p)
-        return f"{CDN.render()}\n{div}\n{script}"
-    
-    def create_bar(self, df: pd.DataFrame, x_col: str, y_col: str, title: str = "") -> str:
-        """Create an interactive bar plot"""
-        source = ColumnDataSource(df)
-        
-        p = figure(width=self.width, height=self.height, title=title, 
-                  tools=self.tools, x_range=df[x_col].unique().tolist())
-        
-        p.vbar(x=x_col, top=y_col, width=0.9, source=source)
-        
-        p.xaxis.axis_label = x_col
-        p.yaxis.axis_label = y_col
-        p.xgrid.grid_line_color = None
-        
-        hover = p.select(dict(type=HoverTool))
-        hover.tooltips = [(x_col, f"@{x_col}"), (y_col, f"@{y_col}")]
-        
-        script, div = components(p)
-        return f"{CDN.render()}\n{div}\n{script}"
-
-class AnalysisSession:
-    """Maintains state and history for the analysis session"""
-    
-    def __init__(self):
-        self.data: Optional[pd.DataFrame] = None
-        self.chat_history: List[Dict[str, str]] = []
-        self.viz_engine = VisualizationEngine()
-        
-    def add_message(self, role: str, content: str):
-        """Add a message to the chat history"""
-        self.chat_history.append({"role": role, "content": content})
-        
-    def get_context(self) -> str:
-        """Get the current analysis context"""
-        if self.data is None:
-            return "No data loaded yet."
-            
-        context = f"""
-        Current DataFrame Info:
-        - Shape: {self.data.shape}
-        - Columns: {', '.join(self.data.columns)}
-        - Numeric columns: {', '.join(self.data.select_dtypes(include=[np.number]).columns)}
-        - Categorical columns: {', '.join(self.data.select_dtypes(include=['object', 'category']).columns)}
-        """
-        return context
-
-class AnalysisAgent:
-    """Enhanced agent with interactive visualization and chat capabilities"""
-    
-    def __init__(
-        self,
-        model_id: str = "gpt-4o-mini",
-        temperature: float = 0.7,
-    ):
-        self.model_id = model_id
-        self.temperature = temperature
-        self.session = AnalysisSession()
-        
-    def process_query(self, query: str) -> str:
-        """Process a user query and generate response with visualizations"""
-        context = self.session.get_context()
-        
-        messages = [
-            {"role": "system", "content": self._get_system_prompt()},
-            *self.session.chat_history[-5:],  # Include last 5 messages for context
-            {"role": "user", "content": f"{context}\n\nUser query: {query}"}
-        ]
-        
-        try:
-            response = completion(
-                model=self.model_id,
-                messages=messages,
-                temperature=self.temperature,
-            )
-            analysis = response.choices[0].message.content
-            
-            # Extract and execute any code blocks
-            visualizations = []
-            code_blocks = self._extract_code(analysis)
-            
-            for code in code_blocks:
-                try:
-                    # Execute code and capture visualization commands
-                    result = self._execute_visualization(code)
-                    if result:
-                        visualizations.append(result)
-                except Exception as e:
-                    visualizations.append(f"Error creating visualization: {str(e)}")
-            
-            # Add messages to chat history
-            self.session.add_message("user", query)
-            self.session.add_message("assistant", analysis)
-            
-            # Combine analysis and visualizations
-            return analysis + "\n\n" + "\n".join(visualizations)
-            
-        except Exception as e:
-            return f"Error: {str(e)}"
-    
-    def _execute_visualization(self, code: str) -> Optional[str]:
-        """Execute visualization code and return HTML output"""
-        try:
-            # Create a safe namespace with necessary libraries
-            namespace = {
-                'df': self.session.data,
-                'np': np,
-                'pd': pd,
-                'viz': self.session.viz_engine
-            }
-            
-            # Execute the code
-            exec(code, namespace)
-            
-            # Look for visualization result
-            for var in namespace.values():
-                if isinstance(var, str) and ('<script' in var or '<div' in var):
-                    return var
-                    
-            return None
-            
-        except Exception as e:
-            return f"Error executing visualization: {str(e)}"
-    
-    def _get_system_prompt(self) -> str:
-        """Get system prompt with visualization capabilities"""
-        return """You are a data analysis assistant with interactive visualization capabilities.
-        
-Available visualizations:
-1. Scatter plots (viz.create_scatter)
-2. Line plots (viz.create_line)
-3. Bar plots (viz.create_bar)
-
-The following variables are available:
-- df: pandas DataFrame with the current data
-- viz: visualization engine with plotting methods
-- np: numpy library
-- pd: pandas library
-
-When analyzing data:
-1. First understand and explain the data
-2. Create relevant visualizations using the viz engine
-3. Provide insights based on the visualizations
-4. Ask follow-up questions when appropriate
-5. Use markdown for formatting
-
-Example visualization code:
-```python
-# Create scatter plot
-html = viz.create_scatter(df, 'column1', 'column2', title='Analysis')
-print(html)
-
-# Create line plot
-html = viz.create_line(df, 'date_column', ['value1', 'value2'], title='Trends')
-print(html)
-```
-"""
-    
-    @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 create_interface():
-    """Create Gradio interface with chat capabilities"""
+    """Create Gradio interface with proper HTML rendering"""
     
     agent = AnalysisAgent()
     
+    def format_html_output(content: str) -> str:
+        """Format the output to properly render HTML in Gradio"""
+        # Split content into text and HTML parts
+        parts = content.split('<!DOCTYPE html>')
+        
+        if len(parts) == 1:
+            # No HTML content
+            return f'<div style="padding: 20px;">{content}</div>'
+        
+        formatted_parts = []
+        for i, part in enumerate(parts):
+            if i == 0:
+                # Text content
+                if part.strip():
+                    formatted_parts.append(f'<div style="padding: 20px;">{part}</div>')
+            else:
+                # HTML visualization
+                formatted_parts.append(f'<!DOCTYPE html>{part}')
+        
+        return '\n'.join(formatted_parts)
+    
     def process_file(file: gr.File) -> str:
         """Process uploaded file and initialize session"""
         try:
@@ -242,25 +34,29 @@ def create_interface():
             else:
                 return "Error: Unsupported file type"
                 
-            return f"Successfully loaded data: {agent.session.get_context()}"
+            return format_html_output(f"Successfully loaded data: {agent.session.get_context()}")
         except Exception as e:
-            return f"Error loading file: {str(e)}"
+            return format_html_output(f"Error loading file: {str(e)}")
     
     def analyze(file: gr.File, query: str, api_key: str) -> str:
         """Process analysis query"""
         if not api_key:
-            return "Error: Please provide an API key."
+            return format_html_output("Error: Please provide an API key.")
             
         if not file:
-            return "Error: Please upload a file."
+            return format_html_output("Error: Please upload a file.")
             
         try:
             os.environ["OPENAI_API_KEY"] = api_key
-            return agent.process_query(query)
+            result = agent.process_query(query)
+            return format_html_output(result)
         except Exception as e:
-            return f"Error: {str(e)}"
+            return format_html_output(f"Error: {str(e)}")
     
-    with gr.Blocks(title="Interactive Data Analysis Assistant") as interface:
+    with gr.Blocks(css="""
+        .plot-container { margin: 20px 0; }
+        .bokeh-plot { margin: 20px auto; }
+    """) as interface:
         gr.Markdown("""
         # Interactive Data Analysis Assistant
         
@@ -274,7 +70,7 @@ def create_interface():
         """)
         
         with gr.Row():
-            with gr.Column():
+            with gr.Column(scale=1):
                 file = gr.File(
                     label="Upload Data File",
                     file_types=[".csv", ".xlsx", ".xls"]
@@ -290,8 +86,11 @@ def create_interface():
                 )
                 analyze_btn = gr.Button("Analyze")
             
-            with gr.Column():
-                chat_output = gr.HTML(label="Analysis & Visualizations")
+            with gr.Column(scale=2):
+                chat_output = gr.HTML(
+                    label="Analysis & Visualizations",
+                    elem_classes="plot-container"
+                )
         
         # Set up event handlers
         file.change(process_file, inputs=[file], outputs=[chat_output])
@@ -305,15 +104,11 @@ def create_interface():
         gr.Examples(
             examples=[
                 [None, "Show me the distribution of numerical variables"],
-                [None, "Create a correlation analysis with interactive visualizations"],
-                [None, "What are the main trends in the data?"],
-                [None, "Can you identify any interesting patterns?"],
+                [None, "Create an interactive visualization of the relationships between variables"],
+                [None, "Analyze trends in the data over time"],
+                [None, "Compare different categories using interactive charts"],
             ],
             inputs=[file, chat_input]
         )
     
-    return interface
-
-if __name__ == "__main__":
-    interface = create_interface()
-    interface.launch()
+    return interface