Update app.py

app.py

@@ -8,6 +8,7 @@ import warnings
 import google.generativeai as genai
 import os
 from typing import List, Dict, Any, Tuple, Optional
+import re
 
 # --- Configuration & Constants ---
 warnings.filterwarnings('ignore')
@@ -43,8 +44,8 @@ class DataExplorerApp:
             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)
-            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")
+            rows_stat, cols_stat = gr.Textbox("0", interactive=False, show_label=False), gr.Textbox("0", interactive=False, show_label=False)
+            quality_stat, time_cols_stat = gr.Textbox("0%", interactive=False, show_label=False), gr.Textbox("0", interactive=False, show_label=False)
             suggestion_buttons = [gr.Button(visible=False) for _ in range(5)]
             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)
@@ -67,10 +68,8 @@ class DataExplorerApp:
                     with cockpit_page:
                         gr.Markdown("## 📊 Data Cockpit: At-a-Glance Overview")
                         with gr.Row():
-                            with gr.Column(elem_classes="stat-card"): gr.Markdown("<div class='stat-card-title'>Rows</div>"); rows_stat
-                            with gr.Column(elem_classes="stat-card"): gr.Markdown("<div class='stat-card-title'>Columns</div>"); cols_stat
-                            with gr.Column(elem_classes="stat-card"): gr.Markdown("<div class='stat-card-title'>Data Quality</div>"); quality_stat
-                            with gr.Column(elem_classes="stat-card"): gr.Markdown("<div class='stat-card-title'>Date/Time Cols</div>"); time_cols_stat
+                            for title, stat_comp in [("Rows", rows_stat), ("Columns", cols_stat), ("Data Quality", quality_stat), ("Date/Time Cols", time_cols_stat)]:
+                                with gr.Column(elem_classes="stat-card"): gr.Markdown(f"<div class='stat-card-title'>{title}</div>"); stat_comp
                         with gr.Accordion(label="✨ AI Smart Suggestions", open=True): [btn for btn in suggestion_buttons]
                     with deep_dive_page:
                         gr.Markdown("## 🔍 Deep Dive: Manual Dashboard Builder"); gr.Markdown("Construct visualizations to investigate specific relationships.")
@@ -84,15 +83,12 @@ class DataExplorerApp:
             
             # 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)
 
             file_input.upload(self.load_and_process_file, inputs=[file_input], outputs=[
-                state_var, status_output, *pages, rows_stat, cols_stat, quality_stat, time_cols_stat,
-                x_col_dd, y_col_dd, add_plot_btn]).then(lambda: self._switch_page("cockpit", pages), outputs=pages).then(
-                lambda: [gr.update(elem_classes="selected"), gr.update(elem_classes=""), gr.update(elem_classes="")], outputs=nav_buttons)
+                state_var, status_output, *pages, rows_stat, cols_stat, quality_stat, time_cols_stat, x_col_dd, y_col_dd, add_plot_btn])
 
             api_key_input.change(lambda x: gr.update(interactive=bool(x)), inputs=[api_key_input], outputs=[suggestion_btn])
             chat_input.change(lambda x: gr.update(interactive=bool(x.strip())), inputs=[chat_input], outputs=[chat_submit_btn])
@@ -103,9 +99,7 @@ class DataExplorerApp:
             
             suggestion_btn.click(self.get_ai_suggestions, inputs=[state_var, api_key_input], outputs=suggestion_buttons)
             for btn in suggestion_buttons:
-                btn.click(self.handle_suggestion_click, inputs=[btn], outputs=[*pages, chat_input]).then(
-                    lambda: self._switch_page("co-pilot", pages), outputs=pages).then(
-                    lambda: (gr.update(elem_classes=""), gr.update(elem_classes=""), gr.update(elem_classes="selected")), outputs=nav_buttons)
+                btn.click(self.handle_suggestion_click, inputs=[btn], outputs=[*pages, chat_input])
             
             chat_submit_btn.click(self.respond_to_chat, [state_var, api_key_input, chat_input, chatbot], [chatbot, copilot_explanation, copilot_code, copilot_plot, copilot_table]).then(lambda: "", outputs=[chat_input])
             chat_input.submit(self.respond_to_chat, [state_var, api_key_input, chat_input, chatbot], [chatbot, copilot_explanation, copilot_code, copilot_plot, copilot_table]).then(lambda: "", outputs=[chat_input])
@@ -113,12 +107,12 @@ 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 = {"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:
-        return gr.update(visible=plot_type in ['scatter', 'box'])
+    def _update_plot_controls(self, plot_type: str) -> gr.update: return gr.update(visible=plot_type in ['scatter', 'box'])
 
     def load_and_process_file(self, file_obj: Any) -> Tuple[Any, ...]:
         try:
@@ -127,16 +121,14 @@ class DataExplorerApp:
             state = {'df': df, 'metadata': metadata, 'dashboard_plots': []}
             rows, cols, quality = metadata['shape'][0], metadata['shape'][1], metadata['data_quality']
             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'])}",
+            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}"); 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
-        quality = round((df.notna().sum().sum() / (rows * cols)) * 100, 1) if rows * cols > 0 else 0
+        rows, cols, quality = df.shape[0], df.shape[1], round((df.notna().sum().sum() / (df.size)) * 100, 1) if df.size > 0 else 0
         return {'shape': (rows, cols), 'columns': df.columns.tolist(), 'numeric_cols': df.select_dtypes(include=np.number).columns.tolist(),
                 'categorical_cols': df.select_dtypes(include=['object', 'category']).columns.tolist(), 'datetime_cols': df.select_dtypes(include=['datetime64', 'datetime64[ns]']).columns.tolist(),
                 'dtypes_head': df.head(3).to_string(), 'data_quality': quality}
@@ -168,9 +160,7 @@ 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
-        # CORRECTED: Defensive and correct assignment
-        metadata = state.get('metadata', {})
-        columns = metadata.get('columns', [])
+        metadata, columns = state.get('metadata', {}), state.get('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)
@@ -180,24 +170,37 @@ class DataExplorerApp:
     def handle_suggestion_click(self, question: str) -> Tuple[gr.update, ...]:
         return *self._switch_page("co-pilot", [0,1,2,3]), question
 
+    def _sanitize_and_parse_json(self, raw_text: str) -> Dict:
+        """Cleans and parses a JSON string from an LLM response."""
+        # Remove markdown code blocks
+        clean_text = re.sub(r'```json\n?|```', '', raw_text).strip()
+        # Escape single backslashes that are not already escaped
+        clean_text = re.sub(r'(?<!\\)\\(?!["\\/bfnrtu])', r'\\\\', clean_text)
+        return json.loads(clean_text)
+
     def respond_to_chat(self, state: Dict, api_key: str, user_message: str, history: List) -> Any:
         if not user_message.strip(): return history, *[gr.update()]*4
         if not api_key or not state:
-            msg = "I need a Gemini API key and a dataset to work."; history.append((user_message, msg)); return history, *[gr.update(visible=False)]*4
+            history.append((user_message, "I need a Gemini API key and a dataset to work.")); return history, *[gr.update(visible=False)]*4
         
         history.append((user_message, "Thinking... 🤔")); yield history, *[gr.update(visible=False)]*4
         
-        # 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: {dtypes_head}
-        User Question: "{user_message}"
+        metadata, dtypes_head = state.get('metadata', {}), state.get('metadata', {}).get('dtypes_head', 'No metadata available.')
+        prompt = f"""You are 'Chief Data Scientist', an expert AI analyst. Your goal is to answer a user's question about a pandas DataFrame (`df`) by writing and executing Python code.
+        **Instructions:**
+        1. **Analyze:** Understand the user's intent. Infer the best plot type.
+        2. **Plan:** Briefly explain your plan.
+        3. **Code:** Write Python code. Use `fig` for plots (`template='plotly_dark'`) and `result_df` for tables.
+        4. **Insight:** Provide a one-sentence business insight.
+        5. **Respond ONLY with a single JSON object with keys: "plan", "code", "insight".**
+        **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("```", ""))
+            genai.configure(api_key=api_key)
+            # CRITICAL FIX: Use the new sanitizer function
+            response_json = self._sanitize_and_parse_json(genai.GenerativeModel('gemini-1.5-flash').generate_content(prompt).text)
+            
             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})
             
@@ -209,7 +212,7 @@ class DataExplorerApp:
             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 processing the AI response. Please rephrase your question. (Error: {e})")
+            history[-1] = (user_message, f"I encountered an error processing the AI response. Please rephrase your question.\n\n**Details:** `{str(e)}`")
             yield history, *[gr.update(visible=False)]*4
             
     def _safe_exec(self, code_string: str, local_vars: Dict) -> Tuple[Any, ...]:
@@ -224,8 +227,7 @@ if __name__ == "__main__":
         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.")
+        except ImportError: print("Pillow not installed, cannot create dummy bot.png.")
     
     app = DataExplorerApp()
     app.launch()