Update app.py

app.py

@@ -3,7 +3,7 @@ import pandas as pd
 import matplotlib.pyplot as plt
 import io
 import ast
-from PIL import Image
+from PIL import Image, ImageDraw
 import google.generativeai as genai
 
 def process_file(api_key, file, instructions):
@@ -13,30 +13,25 @@ def process_file(api_key, file, instructions):
         model = genai.GenerativeModel('gemini-pro')
         
         # Read uploaded file
-        file_path = file.name  # Get full file path
-        if file_path.endswith('.csv'):
-            df = pd.read_csv(file_path)
-        else:
-            df = pd.read_excel(file_path)
+        file_path = file.name
+        df = pd.read_csv(file_path) if file_path.endswith('.csv') else pd.read_excel(file_path)
         
-        # Generate visualization code based on instructions
-        columns = list(df.columns)
+        # Generate visualization code
         response = model.generate_content(f"""
             Create 3 matplotlib visualization codes based on: {instructions}
-            Data columns: {columns}
-            Return only Python code as: [('title','plot_type','x','y'), ...]
+            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.
         """)
+
+        # Extract code block safely
+        code_block = response.text.split('```python')[1].split('```')[0].strip()
+        plots = ast.literal_eval(code_block)
         
-        # Parse and validate generated code
-        plots = ast.literal_eval(response.text.split('```')[-2].strip('python\n '))
-        if len(plots) != 3:
-            raise ValueError("Exactly 3 visualizations required")
-            
-        # Generate plots
+        # Generate visualizations
         images = []
-        for plot in plots:
+        for plot in plots[:3]:  # Ensure max 3 plots
             fig = plt.figure()
             title, plot_type, x, y = plot
             
@@ -56,7 +51,7 @@ def process_file(api_key, file, instructions):
             images.append(Image.open(buf))
             plt.close()
 
-        return images
+        return images if len(images) == 3 else images + [Image.new('RGB', (800, 600), (255,255,255))]*(3-len(images))
 
     except Exception as e:
         error_image = Image.new('RGB', (800, 100), (255, 255, 255))
@@ -71,4 +66,17 @@ with gr.Blocks(theme=gr.themes.Default(spacing_size="lg")) as demo:
         api_key = gr.Textbox(label="Gemini API Key", type="password")
         file = gr.File(label="Upload Dataset", file_types=[".csv", ".xlsx"])
     
-    instructions = gr.Textbox(label="Analysis Instructions
+    instructions = gr.Textbox(label="Analysis Instructions")
+    submit = gr.Button("Generate Insights", variant="primary")
+    
+    with gr.Row():
+        outputs = [gr.Image(label=f"Visualization {i+1}", width=600) for i in range(3)]
+
+    submit.click(
+        process_file,
+        inputs=[api_key, file, instructions],
+        outputs=outputs
+    )
+
+if __name__ == "__main__":
+    demo.launch()