Update app.py

app.py

@@ -13,79 +13,101 @@ def process_file(api_key, file, instructions):
 
     try:
         df = pd.read_csv(file.name) if file.name.endswith('.csv') else pd.read_excel(file.name)
+        if df.empty:
+            raise ValueError("Uploaded file contains no data")
     except Exception as e:
-        print(f"File Error: {str(e)}")
-        return [None]*3
+        print(f"Data Error: {str(e)}")
+        return [generate_error_image(str(e))]*3
 
-    # Enhanced prompt with strict syntax rules
-    prompt = f"""Generate 3 Python code blocks for matplotlib visualizations with:
-    1. Perfect Python syntax
-    2. No markdown or incomplete lines
-    3. Each block must start with:
+    # Enhanced prompt with strict plotting requirements
+    prompt = f"""Generate 3 matplotlib codes with these rules:
+    1. Use ONLY these variables: df (DataFrame), plt
+    2. Each visualization MUST:
+       - Plot actual data from df
+       - Include title, axis labels, and data labels if needed
+       - Use clear color schemes
+       - Avoid empty plots
+    3. Code structure:
        plt.figure(figsize=(16,9), dpi=120)
        plt.style.use('ggplot')
-    4. Use ONLY df and plt variables
-    5. End with plt.tight_layout()
+       # Plotting code using df columns: {list(df.columns)}
+       plt.tight_layout()
     
-    Columns: {list(df.columns)}
-    Sample: {df.head(3).to_dict()}
-    Instructions: {instructions or 'None'}
+    Sample data: {df.head(3).to_dict()}
+    User instructions: {instructions or 'General insights'}
     
     Format EXACTLY as:
     # Visualization 1
-    [code]
-    
-    # Visualization 2
-    [code]
-    
-    # Visualization 3
-    [code]
+    [complete code]
     """
 
     try:
         response = model.generate_content(prompt)
         code_blocks = re.split(r'# Visualization \d+', response.text)[1:4]
     except Exception as e:
-        print(f"Gemini Error: {str(e)}")
-        return [None]*3
+        return [generate_error_image("API Error")]*3
 
     visualizations = []
     for i, block in enumerate(code_blocks, 1):
-        buf = io.BytesIO()
         try:
-            # Advanced code cleaning
-            cleaned_code = '\n'.join(
-                line.strip().replace("'", "").replace('"', '')  # Remove stray quotes
-                for line in block.split('\n') 
-                if line.strip() and 
-                not any(c in line for c in ['`', '```', 'Annotation']) and
-                re.match(r'^[a-zA-Z0-9_().=, <>:]+$', line)  # Basic syntax validation
-            )
+            # Advanced code sanitization
+            cleaned_code = sanitize_code(block, df.columns)
             
-            # Add missing parentheses check
-            cleaned_code = re.sub(r'plt.style.use\([\'"]ggplot$', 
-                                'plt.style.use("ggplot")', cleaned_code)
-            
-            # Syntax validation
+            # Validate and execute
             ast.parse(cleaned_code)
-            
-            # Execute code
-            exec_env = {'df': df, 'plt': plt}
-            plt.figure(figsize=(16, 9), dpi=120)
-            exec(cleaned_code, exec_env)
-            
-            plt.savefig(buf, format='png', bbox_inches='tight')
-            plt.close()
-            visualizations.append(Image.open(buf))
+            img = execute_plot_code(cleaned_code, df)
+            visualizations.append(img)
         except Exception as e:
             print(f"Visualization {i} Error: {str(e)}")
-            print(f"Cleaned Code:\n{cleaned_code}")
-            visualizations.append(None)
+            visualizations.append(generate_error_image(f"Plot {i} Error"))
+
+    return visualizations + [generate_error_image("Not Generated")]*(3-len(visualizations))
+
+def sanitize_code(code_block, columns):
+    """Clean and validate generated code"""
+    replacements = {
+        r"'y_axis'": f"'{columns[1]}'" if len(columns) > 1 else "'Value'",
+        r"'x_axis'": f"'{columns[0]}'",
+        r"data": "df",
+        r"plt.legend\(\)": ""  # Remove empty legend calls
+    }
+    
+    cleaned = []
+    for line in code_block.split('\n'):
+        line = line.strip()
+        if not line or line.startswith('`'):
+            continue
+            
+        # Apply replacements
+        for pattern, replacement in replacements.items():
+            line = re.sub(pattern, replacement, line)
+            
+        cleaned.append(line)
+    
+    return '\n'.join(cleaned)
+
+def execute_plot_code(code, df):
+    """Safely execute plotting code"""
+    buf = io.BytesIO()
+    plt.figure(figsize=(16, 9), dpi=120)
+    plt.style.use('ggplot')
+    
+    try:
+        exec(code, {'df': df, 'plt': plt})
+        plt.tight_layout()
+        plt.savefig(buf, format='png', bbox_inches='tight')
+        buf.seek(0)
+        return Image.open(buf)
+    finally:
+        plt.close()
 
-    return visualizations + [None]*(3-len(visualizations))
+def generate_error_image(message):
+    """Create error indication image"""
+    img = Image.new('RGB', (1920, 1080), color=(73, 109, 137))
+    return img
 
 # Gradio interface
-with gr.Blocks() as demo:
+with gr.Blocks(theme=gr.themes.Default(spacing_size="lg")) as demo:
     gr.Markdown("# Professional Data Visualizer")
     
     with gr.Row():
@@ -93,7 +115,7 @@ with gr.Blocks() as demo:
         file = gr.File(label="Upload Data File", file_types=[".csv", ".xlsx"])
     
     instructions = gr.Textbox(label="Visualization Instructions")
-    submit = gr.Button("Generate", variant="primary")
+    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)]