Update app.py

app.py

@@ -4,6 +4,7 @@ import matplotlib.pyplot as plt
 import io
 import base64
 import google.generativeai as genai
+import textwrap
 
 def process_file(api_key, file, instructions):
     # Set up Gemini API
@@ -29,10 +30,9 @@ def process_file(api_key, file, instructions):
     
     User instructions: {instructions if instructions else 'No specific instructions provided.'}
     
-    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.
+    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. Ensure the code is not indented.
 
     Format your response as:
-    ```python
     # Visualization 1
     # Your code here
 
@@ -41,21 +41,24 @@ def process_file(api_key, file, instructions):
 
     # Visualization 3
     # Your code here
-    ```
     """
 
     response = model.generate_content(prompt)
     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()
+    print("Generated code:")
+    print(code)
 
     visualizations = []
     for i, viz_code in enumerate(code.split("# Visualization")[1:4], 1):
         plt.figure(figsize=(10, 6))
         try:
-            exec(viz_code, {'df': df, 'plt': plt})
+            # Remove any leading spaces from each line
+            cleaned_code = '\n'.join(line.strip() for line in viz_code.split('\n') if line.strip())
+            print(f"\nVisualization {i} code:")
+            print(cleaned_code)
+            
+            exec(cleaned_code, {'df': df, 'plt': plt})
             plt.title(f"Visualization {i}")
             
             # Save the plot to a BytesIO object