Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import pandas as pd
 import matplotlib.pyplot as plt
 import io
+import ast
 from PIL import Image, ImageDraw
 import google.generativeai as genai
 import traceback
@@ -16,33 +17,30 @@ def process_file(file, instructions, api_key):
         file_path = file.name
         df = pd.read_csv(file_path) if file_path.endswith('.csv') else pd.read_excel(file_path)
         
-        # Generate visualization code using Gemini
-        prompt = f"""
-        Analyze the following dataset and instructions:
-        
-        Data columns: {list(df.columns)}
-        Instructions: {instructions}
-        
-        Based on this, create 3 appropriate visualizations. For each visualization, provide:
-        1. A title
-        2. The most suitable plot type (choose from: bar, line, scatter, hist)
-        3. The column to use for the x-axis
-        4. The column to use for the y-axis (use None for histograms)
+        # Generate visualization code
+        response = model.generate_content(f"""
+            Create 3 matplotlib visualization codes based on: {instructions}
+            Data columns: {list(df.columns)}
+            Return Python code as: [('title','plot_type','x','y'), ...]
+            Allowed plot_types: bar, line, scatter, hist
+            Use only DataFrame 'df' and these exact variable names.
+        """)
 
-        Return your response as a Python list of tuples:
-        [
-            ("Title 1", "plot_type1", "x_column1", "y_column1"),
-            ("Title 2", "plot_type2", "x_column2", "y_column2"),
-            ("Title 3", "plot_type3", "x_column3", "y_column3")
-        ]
-        """
+        # Extract code block safely
+        code_block = response.text
+        if '```python' in code_block:
+            code_block = code_block.split('```python')[1].split('```')[0].strip()
+        elif '```' in code_block:
+            code_block = code_block.split('```')[1].strip()
+        
+        print("Generated code block:")
+        print(code_block)
         
-        response = model.generate_content(prompt)
-        plots = eval(response.text)
+        plots = ast.literal_eval(code_block)
         
         # Generate visualizations
         images = []
-        for plot in plots:
+        for plot in plots[:3]:  # Ensure max 3 plots
             fig, ax = plt.subplots(figsize=(10, 6))
             title, plot_type, x, y = plot