Update app.py

app.py

@@ -2,12 +2,12 @@ import gradio as gr
 import pandas as pd
 import matplotlib.pyplot as plt
 import io
-import ast
+import json
 from PIL import Image, ImageDraw
 import google.generativeai as genai
 import traceback
 
-def process_file(api_key, file, instructions):
+def process_file(file, instructions, api_key):
     try:
         # Initialize Gemini
         genai.configure(api_key=api_key)
@@ -17,29 +17,36 @@ def process_file(api_key, file, instructions):
         file_path = file.name
         df = pd.read_csv(file_path) if file_path.endswith('.csv') else pd.read_excel(file_path)
         
-        # 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.
-        """)
-
-        # Extract code block safely
-        code_block = response.text.split('```python')[1].split('```')[0].strip()
+        # Generate visualization code using Gemini
+        prompt = f"""
+        Analyze the following dataset and instructions:
+        
+        Data columns: {list(df.columns)}
+        Instructions: {instructions}
         
-        # Print the code block for debugging
-        print("Generated code block:")
-        print(code_block)
+        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)
+        5. A brief explanation of why this visualization is appropriate
+
+        Return your response as a JSON string in this format:
+        [
+            {{"title": "...", "plot_type": "...", "x": "...", "y": "...", "explanation": "..."}},
+            {{"title": "...", "plot_type": "...", "x": "...", "y": "...", "explanation": "..."}},
+            {{"title": "...", "plot_type": "...", "x": "...", "y": "...", "explanation": "..."}}
+        ]
+        """
         
-        plots = ast.literal_eval(code_block)
+        response = model.generate_content(prompt)
+        plots = json.loads(response.text)
         
         # Generate visualizations
         images = []
-        for plot in plots[:3]:  # Ensure max 3 plots
+        for plot in plots:
             fig, ax = plt.subplots(figsize=(10, 6))
-            title, plot_type, x, y = plot
+            title, plot_type, x, y = plot['title'], plot['plot_type'], plot['x'], plot['y']
             
             if plot_type == 'bar':
                 df.plot(kind='bar', x=x, y=y, ax=ax)
@@ -48,21 +55,21 @@ def process_file(api_key, file, instructions):
             elif plot_type == 'scatter':
                 df.plot(kind='scatter', x=x, y=y, ax=ax)
             elif plot_type == 'hist':
-                df[y].hist(ax=ax)
+                df[x].hist(ax=ax)
             
             ax.set_title(title)
             ax.set_xlabel(x)
-            ax.set_ylabel(y)
+            ax.set_ylabel(y if y else 'Frequency')
             plt.tight_layout()
             
             buf = io.BytesIO()
             plt.savefig(buf, format='png')
             buf.seek(0)
             img = Image.open(buf)
-            images.append(img)
+            images.append((img, plot['explanation']))
             plt.close(fig)
 
-        return images if len(images) == 3 else images + [Image.new('RGB', (800, 600), (255,255,255))]*(3-len(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_message = f"Error: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
@@ -70,25 +77,25 @@ def process_file(api_key, file, instructions):
         error_image = Image.new('RGB', (800, 400), (255, 255, 255))
         draw = ImageDraw.Draw(error_image)
         draw.text((10, 10), error_message, fill=(255, 0, 0))
-        return [error_image] * 3
+        return [(error_image, "An error occurred")] * 3
 
-with gr.Blocks(theme=gr.themes.Default(spacing_size="lg")) as demo:
+with gr.Blocks(theme=gr.themes.Default()) as demo:
     gr.Markdown("# Data Analysis Dashboard")
     
     with gr.Row():
-        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", placeholder="Describe the analysis you want...")
     
-    instructions = gr.Textbox(label="Analysis Instructions")
+    api_key = gr.Textbox(label="Gemini API Key", type="password")
     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)]
+    output_images = [gr.Image(label=f"Visualization {i+1}") for i in range(3)]
+    output_texts = [gr.Textbox(label=f"Explanation {i+1}") for i in range(3)]
 
     submit.click(
         process_file,
-        inputs=[api_key, file, instructions],
-        outputs=outputs
+        inputs=[file, instructions, api_key],
+        outputs=output_images + output_texts
     )
 
 if __name__ == "__main__":