Update app.py

app.py

@@ -35,38 +35,28 @@ class DataExplorerApp:
         with gr.Blocks(theme=gr.themes.Glass(primary_hue="indigo", secondary_hue="blue"), css=CSS, title="AI Data Explorer Pro") as demo:
             state_var = gr.State({})
             
-            # --- Component Definition (Define-Then-Place Pattern) ---
-            # Sidebar
-            ### CRITICAL FIX HERE ###
-            # Buttons are defined individually with their value (text) to fix the "Run" bug.
+            # Component Definition
             cockpit_btn = gr.Button("📊 Data Cockpit", elem_classes="selected", elem_id="cockpit")
             deep_dive_btn = gr.Button("🔍 Deep Dive Builder", elem_id="deep_dive")
             copilot_btn = gr.Button("🤖 Chief Data Scientist", elem_id="co-pilot")
-            
             file_input = gr.File(label="📁 Upload CSV File", file_types=[".csv"])
             status_output = gr.Markdown("Status: Awaiting data...")
             api_key_input = gr.Textbox(label="🔑 Gemini API Key", type="password", placeholder="Enter key to enable AI...")
             suggestion_btn = gr.Button("Get Smart Suggestions", variant="secondary", interactive=False)
-            
-            # Cockpit
             rows_stat, cols_stat = gr.Textbox("0", interactive=False, show_label=False, elem_classes="stat-card-value"), gr.Textbox("0", interactive=False, show_label=False, elem_classes="stat-card-value")
             quality_stat, time_cols_stat = gr.Textbox("0%", interactive=False, show_label=False, elem_classes="stat-card-value"), gr.Textbox("0", interactive=False, show_label=False, elem_classes="stat-card-value")
             suggestion_buttons = [gr.Button(visible=False) for _ in range(5)]
-            
-            # Deep Dive
             plot_type_dd = gr.Dropdown(['histogram', 'bar', 'scatter', 'box'], label="Plot Type", value='histogram')
             x_col_dd = gr.Dropdown([], label="X-Axis / Column", interactive=False)
             y_col_dd = gr.Dropdown([], label="Y-Axis (for Scatter/Box)", visible=False, interactive=False)
             add_plot_btn, clear_plots_btn = gr.Button("Add to Dashboard", variant="primary", interactive=False), gr.Button("Clear Dashboard", interactive=False)
             dashboard_plots = [gr.Plot(visible=False) for _ in range(MAX_DASHBOARD_PLOTS)]
-
-            # Co-pilot
             chatbot = gr.Chatbot(height=500, label="Conversation", show_copy_button=True, avatar_images=(None, "bot.png"))
             copilot_explanation, copilot_code = gr.Markdown(visible=False, elem_classes="explanation-block"), gr.Code(language="python", visible=False, label="Executed Code")
             copilot_plot, copilot_table = gr.Plot(visible=False, label="Generated Visualization"), gr.Dataframe(visible=False, label="Generated Table", wrap=True)
             chat_input, chat_submit_btn = gr.Textbox(label="Your Question", placeholder="e.g., 'What is the relationship between age and salary?'", scale=4), gr.Button("Ask AI", variant="primary", interactive=False)
             
-            # --- Layout Arrangement ---
+            # Layout Arrangement
             with gr.Row():
                 with gr.Column(scale=1, elem_classes="sidebar"):
                     gr.Markdown("## 🚀 AI Explorer Pro", elem_id="app-title"); cockpit_btn; deep_dive_btn; copilot_btn; gr.Markdown("---")
@@ -92,8 +82,9 @@ class DataExplorerApp:
                         with gr.Accordion("AI's Detailed Response", open=True): copilot_explanation; copilot_code; copilot_plot; copilot_table
                         with gr.Row(): chat_input; chat_submit_btn
             
-            # --- Event Handlers Registration ---
+            # Event Handlers Registration
             pages, nav_buttons = [welcome_page, cockpit_page, deep_dive_page, copilot_page], [cockpit_btn, deep_dive_btn, copilot_btn]
+            
             for i, btn in enumerate(nav_buttons):
                 btn.click(lambda id=btn.elem_id: self._switch_page(id, pages), outputs=pages).then(
                     lambda i=i: [gr.update(elem_classes="selected" if j==i else "") for j in range(len(nav_buttons))], outputs=nav_buttons)
@@ -122,9 +113,8 @@ class DataExplorerApp:
 
     def launch(self): self.demo.launch(debug=True)
 
-    # --- Backend Logic Methods ---
     def _switch_page(self, page_id: str, all_pages: List) -> List[gr.update]:
-        visibility = {"cockpit": 1, "deep_dive": 2, "co-pilot": 3}
+        visibility = {"welcome":0, "cockpit":1, "deep_dive":2, "co-pilot":3}
         return [gr.update(visible=i == visibility.get(page_id, 0)) for i in range(len(all_pages))]
     
     def _update_plot_controls(self, plot_type: str) -> gr.update:
@@ -136,12 +126,13 @@ class DataExplorerApp:
             metadata = self._extract_dataset_metadata(df)
             state = {'df': df, 'metadata': metadata, 'dashboard_plots': []}
             rows, cols, quality = metadata['shape'][0], metadata['shape'][1], metadata['data_quality']
-            updates = (state, f"✅ **{os.path.basename(file_obj.name)}** loaded.", *self._switch_page("cockpit", [0,1,2,3]),
-                       f"{rows:,}", f"{cols}", f"{quality}%", f"{len(metadata['datetime_cols'])}",
-                       gr.update(choices=metadata['columns'], interactive=True), gr.update(choices=metadata['columns'], interactive=True), gr.update(interactive=True))
-            return updates
+            page_updates = self._switch_page("cockpit", [0,1,2,3])
+            return (state, f"✅ **{os.path.basename(file_obj.name)}** loaded.", *page_updates,
+                    f"{rows:,}", f"{cols}", f"{quality}%", f"{len(metadata['datetime_cols'])}",
+                    gr.update(choices=metadata['columns'], interactive=True), gr.update(choices=metadata['columns'], interactive=True), gr.update(interactive=True))
         except Exception as e:
-            gr.Error(f"File Load Error: {e}"); return {}, f"❌ Error: {e}", *self._switch_page("welcome", [0,1,2,3]), "0", "0", "0%", "0", gr.update(choices=[], interactive=False), gr.update(choices=[], interactive=False), gr.update(interactive=False)
+            gr.Error(f"File Load Error: {e}"); page_updates = self._switch_page("welcome", [0,1,2,3]);
+            return {}, f"❌ Error: {e}", *page_updates, "0", "0", "0%", "0", gr.update(choices=[], interactive=False), gr.update(choices=[], interactive=False), gr.update(interactive=False)
 
     def _extract_dataset_metadata(self, df: pd.DataFrame) -> Dict[str, Any]:
         rows, cols = df.shape
@@ -151,22 +142,21 @@ class DataExplorerApp:
                 'dtypes_head': df.head(3).to_string(), 'data_quality': quality}
 
     def add_plot_to_dashboard(self, state: Dict, x_col: str, y_col: Optional[str], plot_type: str) -> List[Any]:
-        if len(state.get('dashboard_plots', [])) >= MAX_DASHBOARD_PLOTS:
-            gr.Warning(f"Dashboard is full. Max {MAX_DASHBOARD_PLOTS} plots."); return [state, gr.update(), *self._get_plot_updates(state)]
-        if not x_col: gr.Warning("Please select an X-axis column."); return [state, gr.update(), *self._get_plot_updates(state)]
-        df = state['df']
-        title = f"{plot_type.capitalize()}: {y_col} by {x_col}" if y_col and plot_type in ['box', 'scatter'] else f"Distribution of {x_col}"
+        dashboard_plots = state.get('dashboard_plots', [])
+        if len(dashboard_plots) >= MAX_DASHBOARD_PLOTS:
+            gr.Warning(f"Dashboard is full. Max {MAX_DASHBOARD_PLOTS} plots."); return [state, gr.update(interactive=True), *self._get_plot_updates(state)]
+        if not x_col: gr.Warning("Please select an X-axis column."); return [state, gr.update(interactive=True), *self._get_plot_updates(state)]
+        df, title = state.get('df'), f"{plot_type.capitalize()}: {y_col} by {x_col}" if y_col and plot_type in ['box', 'scatter'] else f"Distribution of {x_col}"
         try:
+            fig=None;
             if plot_type == 'histogram': fig = px.histogram(df, x=x_col, title=title)
             elif plot_type == 'box': fig = px.box(df, x=x_col, y=y_col, title=title)
             elif plot_type == 'scatter': fig = px.scatter(df, x=x_col, y=y_col, title=title, trendline="ols")
-            elif plot_type == 'bar':
-                counts = df[x_col].value_counts().nlargest(20)
-                fig = px.bar(counts, x=counts.index, y=counts.values, title=f"Top 20 Categories for {x_col}", labels={'index': x_col, 'y': 'Count'})
+            elif plot_type == 'bar': fig = px.bar(df[x_col].value_counts().nlargest(20), title=f"Top 20 for {x_col}")
             if fig:
-                fig.update_layout(template="plotly_dark"); state['dashboard_plots'].append(fig); gr.Info(f"Added '{title}' to dashboard.")
+                fig.update_layout(template="plotly_dark"); dashboard_plots.append(fig); gr.Info(f"Added '{title}' to dashboard.")
             return [state, gr.update(interactive=True), *self._get_plot_updates(state)]
-        except Exception as e: gr.Error(f"Plotting Error: {e}"); return [state, gr.update(), *self._get_plot_updates(state)]
+        except Exception as e: gr.Error(f"Plotting Error: {e}"); return [state, gr.update(interactive=True), *self._get_plot_updates(state)]
     
     def _get_plot_updates(self, state: Dict) -> List[gr.update]:
         plots = state.get('dashboard_plots', [])
@@ -178,7 +168,10 @@ class DataExplorerApp:
     def get_ai_suggestions(self, state: Dict, api_key: str) -> List[gr.update]:
         if not api_key: gr.Warning("API Key is required."); return [gr.update(visible=False)]*5
         if not state: gr.Warning("Please load data first."); return [gr.update(visible=False)]*5
-        metadata = state['metadata']; prompt = f"From columns {metadata['columns']}, generate 4 impactful analytical questions. Return ONLY a JSON list of strings."
+        # CORRECTED: Defensive and correct assignment
+        metadata = state.get('metadata', {})
+        columns = metadata.get('columns', [])
+        prompt = f"From columns {columns}, generate 4 impactful analytical questions. Return ONLY a JSON list of strings."
         try:
             genai.configure(api_key=api_key); suggestions = json.loads(genai.GenerativeModel('gemini-1.5-flash').generate_content(prompt).text)
             return [gr.Button(s, visible=True) for s in suggestions] + [gr.Button(visible=False)] * (5 - len(suggestions))
@@ -194,23 +187,29 @@ class DataExplorerApp:
         
         history.append((user_message, "Thinking... 🤔")); yield history, *[gr.update(visible=False)]*4
         
-        metadata, prompt = state['metadata'], f"""You are 'Chief Data Scientist', an expert AI analyst...
+        # CRITICAL FIX: Correctly and safely get metadata before using it.
+        metadata = state.get('metadata', {})
+        dtypes_head = metadata.get('dtypes_head', 'No metadata available.')
+        
+        prompt = f"""You are 'Chief Data Scientist', an expert AI analyst...
         Respond ONLY with a single JSON object with keys: "plan", "code", "insight".
-        Metadata: {metadata['dtypes_head']}
+        Metadata: {dtypes_head}
         User Question: "{user_message}"
         """
         try:
             genai.configure(api_key=api_key); response_json = json.loads(genai.GenerativeModel('gemini-1.5-flash').generate_content(prompt).text.strip().replace("```json", "").replace("```", ""))
             plan, code, insight = response_json.get("plan"), response_json.get("code"), response_json.get("insight")
             stdout, fig, df_result, error = self._safe_exec(code, {'df': state['df'], 'px': px, 'pd': pd})
+            
             history[-1] = (user_message, f"**Plan:** {plan}")
             explanation = f"**Insight:** {insight}"
             if stdout: explanation += f"\n\n**Console Output:**\n```\n{stdout}\n```"
             if error: gr.Error(f"AI Code Execution Failed: {error}")
+            
             yield (history, gr.update(visible=bool(explanation), value=explanation), gr.update(visible=bool(code), value=code),
                    gr.update(visible=bool(fig), value=fig), gr.update(visible=bool(df_result is not None), value=df_result))
         except Exception as e:
-            history[-1] = (user_message, f"I encountered an error. Please rephrase your question. (Error: {e})")
+            history[-1] = (user_message, f"I encountered an error processing the AI response. Please rephrase your question. (Error: {e})")
             yield history, *[gr.update(visible=False)]*4
             
     def _safe_exec(self, code_string: str, local_vars: Dict) -> Tuple[Any, ...]:
@@ -221,17 +220,12 @@ class DataExplorerApp:
         except Exception as e: return None, None, None, str(e)
 
 if __name__ == "__main__":
-    # Create dummy files for avatar images if they don't exist
     if not os.path.exists("bot.png"):
-        from PIL import Image
-        img = Image.new('RGB', (100, 100), color = 'darkblue')
-        img.save('bot.png')
-    if not os.path.exists("user.png"):
-        from PIL import Image
-        img = Image.new('RGB', (100, 100), color = 'gray')
-        img.save('user.png')
-    if not os.path.exists("workflow.png"):
-        gr.Warning("workflow.png not found. The welcome screen image will be missing.")
-
+        try:
+            from PIL import Image
+            Image.new('RGB', (1, 1)).save('bot.png')
+        except ImportError:
+            print("Pillow not installed, cannot create dummy bot.png. Please create it manually.")
+    
     app = DataExplorerApp()
     app.launch()