Update app.py

app.py

@@ -6,39 +6,45 @@ import google.generativeai as genai
 from PIL import Image
 
 def process_file(api_key, file, instructions):
+    # Configure Gemini with precise model version
     genai.configure(api_key=api_key)
     model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
 
-    try:
-        if file.name.endswith('.csv'):
-            df = pd.read_csv(file.name)
-        else:
-            df = pd.read_excel(file.name)
-    except Exception as e:
-        return [f"File Error: {str(e)}"] * 3
+    # Read file
+    if file.name.endswith('.csv'):
+        df = pd.read_csv(file.name)
+    else:
+        df = pd.read_excel(file.name)
 
-    # Properly terminated prompt string
-    prompt = f"""Generate exactly 3 distinct matplotlib visualizations for:
+    # Enhanced prompt with strict code requirements
+    prompt = f"""Generate 3 matplotlib visualization codes for this data:
     Columns: {list(df.columns)}
-    Data types: {dict(df.dtypes)}
-    Sample data: {df.head(3).to_dict()}
+    First 3 rows: {df.head(3).to_dict()}
     
     Requirements:
-    1. 1920x1080 resolution (figsize=(16,9), dpi=120)
-    2. Professional styling (seaborn, grid, proper labels)
-    3. Diverse chart types (include at least 1 advanced visualization)
+    1. Each visualization must start with:
+       plt.figure(figsize=(16,9), dpi=120)
+       plt.style.use('seaborn')
+    2. Include complete plotting code with:
+       - Title
+       - Axis labels
+       - Legend if needed
+       - plt.tight_layout()
+    3. Different chart types (bar, line, scatter, etc)
+    4. No explanations - only valid Python code
     
-    User instructions: {instructions or 'None provided'}
+    User instructions: {instructions}
     
-    Format response strictly as:
+    Format exactly as:
     # Visualization 1
-    plt.figure(figsize=(16,9), dpi=120)
-    [code]
-    plt.tight_layout()
+    [complete code]
     
     # Visualization 2
-    ...
-    """  # Closing triple quotes added here
+    [complete code]
+    
+    # Visualization 3
+    [complete code]
+    """
 
     response = model.generate_content(prompt)
     code_blocks = response.text.split("# Visualization ")[1:4]
@@ -46,42 +52,47 @@ def process_file(api_key, file, instructions):
     visualizations = []
     for i, block in enumerate(code_blocks, 1):
         try:
+            # Clean and validate code
+            cleaned_code = '\n'.join([
+                line.strip() for line in block.split('\n') 
+                if line.strip() and not line.startswith('```')
+            ])
+            
+            # Create HD figure
+            buf = io.BytesIO()
             plt.figure(figsize=(16, 9), dpi=120)
             plt.style.use('seaborn')
             
-            cleaned_code = '\n'.join([line.strip() for line in block.split('\n')[1:] if line.strip()])
+            # Execute generated code
             exec(cleaned_code, {'df': df, 'plt': plt})
-            plt.title(f"Visualization {i}", fontsize=14)
+            
+            # Save HD image
             plt.tight_layout()
-
-            buf = io.BytesIO()
-            plt.savefig(buf, format='png', dpi=120, bbox_inches='tight')
+            plt.savefig(buf, format='png', bbox_inches='tight')
             plt.close()
+            
             buf.seek(0)
             visualizations.append(Image.open(buf))
         except Exception as e:
-            print(f"Visualization {i} failed: {str(e)}")
-            visualizations.append(Image.new('RGB', (1920, 1080), color=(73, 109, 137)))
-
-    while len(visualizations) < 3:
-        visualizations.append(Image.new('RGB', (1920, 1080), color=(73, 109, 137)))
+            print(f"Visualization {i} Error: {str(e)}")
+            visualizations.append(None)
 
-    return visualizations[:3]
+    # Return exactly 3 images, filling with None if needed
+    return visualizations + [None]*(3-len(visualizations))
 
 # Gradio interface
-with gr.Blocks(theme=gr.themes.Default(spacing_size="lg")) as demo:
-    gr.Markdown("# **HD Data Visualizer**  📊✨")
+with gr.Blocks() as demo:
+    gr.Markdown("# Data Visualization Tool")
     
     with gr.Row():
-        api_key = gr.Textbox(label="🔑 Gemini API Key", type="password")
-        file = gr.File(label="📁 Upload Dataset", file_types=[".csv", ".xlsx"])
+        api_key = gr.Textbox(label="Gemini API Key", type="password")
+        file = gr.File(label="Upload CSV/Excel", file_types=[".csv", ".xlsx"])
     
-    instructions = gr.Textbox(label="💡 Custom Instructions (optional)", 
-                            placeholder="E.g.: Focus on time series patterns...")
-    submit = gr.Button("🚀 Generate Visualizations", variant="primary")
+    instructions = gr.Textbox(label="Instructions (optional)")
+    submit = gr.Button("Generate Visualizations")
     
     with gr.Row():
-        outputs = [gr.Image(label=f"Visualization {i+1}", width=600) for i in range(3)]
+        outputs = [gr.Image(label=f"Visualization {i+1}") for i in range(3)]
 
     submit.click(
         process_file,