Update app.py

app.py

@@ -4,7 +4,6 @@ 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
@@ -30,46 +29,33 @@ def process_file(api_key, file, instructions):
     
     User instructions: {instructions if instructions else 'No specific instructions provided.'}
     
-    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. Ensure the code is complete and executable.
-    3. Explain why this visualization is useful for understanding the data.
-    
+    Generate Python code for 3 different visualizations using matplotlib. Each visualization should be unique and provide insights into the data. Do not include any explanations or descriptions, only the Python code for each visualization.
+
     Format your response as:
-    Visualization 1:
-    Description: ...
-    Code:
     ```python
+    # Visualization 1
+    # Your code here
+
+    # Visualization 2
+    # Your code here
+
+    # Visualization 3
     # Your code here
     ```
-    Explanation: ...
-    
-    Visualization 2:
-    ...
-    
-    Visualization 3:
-    ...
     """
 
     response = model.generate_content(prompt)
-    suggestions = response.text.split("Visualization")
+    code = response.text.strip()
+
+    # Extract code from markdown code block if present
+    if "```python" in code and "```" in code:
+        code = code.split("```python")[1].split("```")[0].strip()
 
     visualizations = []
-    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 = 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
+    for i, viz_code in enumerate(code.split("# Visualization")[1:4], 1):
         plt.figure(figsize=(10, 6))
         try:
-            exec(code)
+            exec(viz_code, {'df': df, 'plt': plt})
             plt.title(f"Visualization {i}")
             
             # Save the plot to a BytesIO object
@@ -79,14 +65,14 @@ def process_file(api_key, file, instructions):
             img_str = base64.b64encode(buf.getvalue()).decode()
             plt.close()
             
-            visualizations.append((f"data:image/png;base64,{img_str}", description, code, explanation))
+            visualizations.append(f"data:image/png;base64,{img_str}")
         except Exception as e:
-            error_message = f"Error in visualization {i}: {str(e)}\n{traceback.format_exc()}"
-            visualizations.append((None, description, code, error_message))
+            print(f"Error in visualization {i}: {str(e)}")
+            visualizations.append(None)
 
-    # Ensure we always return 3 visualizations, even if there were errors
+    # Ensure we always return 3 visualizations
     while len(visualizations) < 3:
-        visualizations.append((None, "Error", "No code generated", "An error occurred"))
+        visualizations.append(None)
 
     return visualizations
 
@@ -102,30 +88,11 @@ with gr.Blocks() as demo:
         output1 = gr.Image(label="Visualization 1")
         output2 = gr.Image(label="Visualization 2")
         output3 = gr.Image(label="Visualization 3")
-    
-    with gr.Row():
-        desc1 = gr.Textbox(label="Description 1")
-        desc2 = gr.Textbox(label="Description 2")
-        desc3 = gr.Textbox(label="Description 3")
-    
-    with gr.Row():
-        code1 = gr.Code(language="python", label="Code 1")
-        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, expl1,
-            output2, desc2, code2, expl2,
-            output3, desc3, code3, expl3
-        ],
+        outputs=[output1, output2, output3],
         show_progress=True,
     )