Update app.py

app.py

@@ -6,7 +6,6 @@ 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')
 
@@ -16,34 +15,23 @@ def process_file(api_key, file, instructions):
     else:
         df = pd.read_excel(file.name)
 
-    # Enhanced prompt with strict code requirements
-    prompt = f"""Generate 3 matplotlib visualization codes for this data:
-    Columns: {list(df.columns)}
-    First 3 rows: {df.head(3).to_dict()}
+    # Updated prompt with valid matplotlib styles
+    prompt = f"""Generate 3 matplotlib codes with these requirements:
+    1. Start with: plt.figure(figsize=(16,9), dpi=120)
+    2. Use one of these styles: ggplot, bmh, dark_background, fast
+    3. Include: title, labels, grid, legend if needed
+    4. Different chart types (bar, line, scatter, etc)
     
-    Requirements:
-    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}
+    Data columns: {list(df.columns)}
+    Sample data: {df.head(3).to_dict()}
+    User instructions: {instructions or 'None'}
     
     Format exactly as:
     # Visualization 1
-    [complete code]
-    
-    # Visualization 2
-    [complete code]
-    
-    # Visualization 3
-    [complete code]
+    plt.figure(figsize=(16,9), dpi=120)
+    plt.style.use('ggplot')  # Example valid style
+    [code]
+    plt.tight_layout()
     """
 
     response = model.generate_content(prompt)
@@ -52,21 +40,22 @@ 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')
             
-            # Execute generated code
-            exec(cleaned_code, {'df': df, 'plt': plt})
+            # Execute code with safe environment
+            exec_env = {
+                'df': df,
+                'plt': plt,
+                'pd': pd
+            }
+            exec(cleaned_code, exec_env)
             
-            # Save HD image
             plt.tight_layout()
             plt.savefig(buf, format='png', bbox_inches='tight')
             plt.close()
@@ -77,7 +66,6 @@ def process_file(api_key, file, instructions):
             print(f"Visualization {i} Error: {str(e)}")
             visualizations.append(None)
 
-    # Return exactly 3 images, filling with None if needed
     return visualizations + [None]*(3-len(visualizations))
 
 # Gradio interface
@@ -88,7 +76,7 @@ with gr.Blocks() as demo:
         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="Instructions (optional)")
+    instructions = gr.Textbox(label="Custom Instructions")
     submit = gr.Button("Generate Visualizations")
     
     with gr.Row():