Update app.py

app.py

@@ -4,6 +4,7 @@ import matplotlib.pyplot as plt
 import io
 import base64
 import google.generativeai as genai
+import traceback
 
 def process_file(api_key, file, instructions):
     # Set up Gemini API
@@ -31,13 +32,16 @@ def process_file(api_key, file, instructions):
     
     Suggest 3 ways to visualize this data. For each visualization:
     1. Describe the visualization type and what it will show.
-    2. Provide Python code using matplotlib to create the visualization.
+    2. Provide Python code using matplotlib to create the visualization. Ensure the code is complete and executable.
     3. Explain why this visualization is useful for understanding the data.
     
     Format your response as:
     Visualization 1:
     Description: ...
-    Code: ...
+    Code:
+    ```python
+    # Your code here
+    ```
     Explanation: ...
     
     Visualization 2:
@@ -54,21 +58,35 @@ def process_file(api_key, file, instructions):
     for i, suggestion in enumerate(suggestions[1:4], 1):  # Process only the first 3 visualizations
         parts = suggestion.split("Code:")
         description = parts[0].strip()
-        code = parts[1].split("Explanation:")[0].strip()
+        code_parts = parts[1].split("Explanation:")
+        code = code_parts[0].strip()
+        explanation = code_parts[1].strip() if len(code_parts) > 1 else "No explanation provided."
+        
+        # Extract code from markdown code block if present
+        if "```python" in code and "```" in code:
+            code = code.split("```python")[1].split("```")[0].strip()
         
         # Execute the code
         plt.figure(figsize=(10, 6))
-        exec(code)
-        plt.title(f"Visualization {i}")
-        
-        # Save the plot to a BytesIO object
-        buf = io.BytesIO()
-        plt.savefig(buf, format='png')
-        buf.seek(0)
-        img_str = base64.b64encode(buf.getvalue()).decode()
-        plt.close()
-        
-        visualizations.append((f"data:image/png;base64,{img_str}", description, code))
+        try:
+            exec(code)
+            plt.title(f"Visualization {i}")
+            
+            # Save the plot to a BytesIO object
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png')
+            buf.seek(0)
+            img_str = base64.b64encode(buf.getvalue()).decode()
+            plt.close()
+            
+            visualizations.append((f"data:image/png;base64,{img_str}", description, code, explanation))
+        except Exception as e:
+            error_message = f"Error in visualization {i}: {str(e)}\n{traceback.format_exc()}"
+            visualizations.append((None, description, code, error_message))
+
+    # Ensure we always return 3 visualizations, even if there were errors
+    while len(visualizations) < 3:
+        visualizations.append((None, "Error", "No code generated", "An error occurred"))
 
     return visualizations
 
@@ -95,13 +113,18 @@ with gr.Blocks() as demo:
         code2 = gr.Code(language="python", label="Code 2")
         code3 = gr.Code(language="python", label="Code 3")
 
+    with gr.Row():
+        expl1 = gr.Textbox(label="Explanation/Error 1")
+        expl2 = gr.Textbox(label="Explanation/Error 2")
+        expl3 = gr.Textbox(label="Explanation/Error 3")
+
     submit.click(
         fn=process_file,
         inputs=[api_key, file, instructions],
         outputs=[
-            output1, desc1, code1,
-            output2, desc2, code2,
-            output3, desc3, code3
+            output1, desc1, code1, expl1,
+            output2, desc2, code2, expl2,
+            output3, desc3, code3, expl3
         ],
         show_progress=True,
     )