Update app.py

app.py

@@ -28,33 +28,6 @@ class CodeExecutionEnvironment:
         except:
             return False
 
-    def extract_and_execute_code(self, text):
-        """Extract code blocks from markdown and execute them"""
-        # Pattern for code blocks
-        code_blocks = re.findall(r'```python(.*?)```', text, re.DOTALL)
-        if not code_blocks:
-            return text, None, []
-
-        all_outputs = []
-        all_figures = []
-        
-        for code in code_blocks:
-            success, output, figures = self.execute_code(code.strip())
-            if success:
-                all_outputs.append(output)
-                all_figures.extend(figures)
-            else:
-                all_outputs.append(f"Error: {output}")
-
-        # Replace code blocks with code + output
-        modified_text = text
-        for i, (code, output) in enumerate(zip(code_blocks, all_outputs)):
-            code_section = f"```python{code}```"
-            output_section = f"\nOutput:\n```\n{output}\n```"
-            modified_text = modified_text.replace(code_section, code_section + output_section)
-
-        return modified_text, "\n".join(all_outputs), all_figures
-
     def execute_code(self, code):
         """Execute code in the managed environment"""
         # Redirect stdout to capture prints
@@ -93,14 +66,26 @@ class CodeExecutionEnvironment:
             exec(code, self.globals_dict)
             output = redirected_output.getvalue()
 
-            # Capture figures
+            # Handle different types of figures
             figures = []
+            
+            # Handle Matplotlib figures
             if plt.get_figs():
                 for i, fig in enumerate(plt.get_figs()):
-                    fig_path = os.path.join(self.figures_dir, f"figure_{len(figures)}.png")
+                    fig_path = os.path.join(self.figures_dir, f"mpl_figure_{len(figures)}.png")
                     fig.savefig(fig_path)
                     figures.append(fig_path)
                 plt.close('all')
+            
+            # Handle Plotly figures
+            if 'fig' in self.globals_dict and 'plotly' in str(type(self.globals_dict['fig'])):
+                fig = self.globals_dict['fig']
+                fig_path = os.path.join(self.figures_dir, f"plotly_figure_{len(figures)}.html")
+                fig.write_html(fig_path)
+                # Also save as image for gallery display
+                img_path = os.path.join(self.figures_dir, f"plotly_figure_{len(figures)}.png")
+                fig.write_image(img_path)
+                figures.append(img_path)
 
             return True, output, figures
 
@@ -109,6 +94,32 @@ class CodeExecutionEnvironment:
         finally:
             sys.stdout = old_stdout
 
+    def extract_and_execute_code(self, text):
+        """Extract code blocks from markdown and execute them"""
+        code_blocks = re.findall(r'```python(.*?)```', text, re.DOTALL)
+        if not code_blocks:
+            return text, None, []
+
+        all_outputs = []
+        all_figures = []
+        
+        for code in code_blocks:
+            success, output, figures = self.execute_code(code.strip())
+            if success:
+                all_outputs.append(output)
+                all_figures.extend(figures)
+            else:
+                all_outputs.append(f"Error: {output}")
+
+        # Replace code blocks with code + output
+        modified_text = text
+        for i, (code, output) in enumerate(zip(code_blocks, all_outputs)):
+            code_section = f"```python{code}```"
+            output_section = f"\nOutput:\n```\n{output}\n```"
+            modified_text = modified_text.replace(code_section, code_section + output_section)
+
+        return modified_text, "\n".join(all_outputs), all_figures
+
 def query_deepseek(prompt: str, api_key: str, system_prompt: str = None):
     """Send a prompt to DeepSeek API"""
     headers = {
@@ -165,6 +176,7 @@ class ChatAndCodeInterface:
             "The user might ask you to analyze data or generate visualizations. "
             "When you write code, wrap it in ```python``` blocks. "
             "You can use any Python library - they will be automatically installed. "
+            "For interactive maps, use plotly.express and ensure you install required dependencies. "
             "\nUser message: " + message
         )
         
@@ -203,8 +215,13 @@ def create_interface():
                 )
                 
             with gr.Column(scale=3):
-                chatbot = gr.Chatbot(height=400)
-                gallery = gr.Gallery(label="Generated Visualizations", columns=2, height=300)
+                chatbot = gr.Chatbot(height=500)
+                gallery = gr.Gallery(
+                    label="Generated Visualizations",
+                    columns=2,
+                    height=400,
+                    object_fit="contain"
+                )
                 
                 with gr.Row():
                     msg = gr.Textbox(
@@ -230,17 +247,18 @@ def create_interface():
         
         Example prompts:
         - "Create a histogram of the numerical columns"
-        - "Analyze the correlation between variables"
-        - "Generate summary statistics and visualize key trends"
+        - "Generate an interactive map of the locations"
+        - "Show the correlation between variables"
+        - "Create a summary dashboard of key metrics"
         
         The assistant will:
         - Generate and execute Python code automatically
-        - Show both code and its output in the chat
-        - Display generated visualizations in the gallery
+        - Handle both static and interactive visualizations
+        - Show code, output, and visualizations in one place
         """)
     
     return demo
 
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch() 
+    demo.launch()