Update app.py

app.py

@@ -4,7 +4,7 @@ import matplotlib.pyplot as plt
 import io
 import google.generativeai as genai
 from PIL import Image
-import ast  # Added for syntax validation
+import ast
 
 def process_file(api_key, file, instructions):
     genai.configure(api_key=api_key)
@@ -18,27 +18,22 @@ def process_file(api_key, file, instructions):
     except Exception as e:
         return [None]*3
 
-    # Enhanced prompt with strict code requirements
-    prompt = f"""Generate 3 Python code blocks for matplotlib visualizations with these rules:
-    1. Basic Python 3.10 syntax ONLY (NO type annotations)
-    2. Each visualization must:
-       - Start with plt.figure(figsize=(16,9), dpi=120)
-       - Use plt.style.use('ggplot')
-       - Include title, labels, and grid
-       - End with plt.tight_layout()
-    3. Different chart types required (bar, line, scatter)
-    4. Only use these imports: pandas as pd, matplotlib.pyplot as plt
+    # Enhanced prompt with strict variable requirements
+    prompt = f"""Generate 3 matplotlib codes with these rules:
+    1. Use EXACTLY these variables: df (DataFrame), plt (matplotlib)
+    2. NO imports or additional variables
+    3. Start each visualization with:
+       plt.figure(figsize=(16,9), dpi=120)
+       plt.style.use('ggplot')
+    4. End with plt.tight_layout()
     
     Data columns: {list(df.columns)}
-    Sample data: {df.head(3).to_dict()}
+    First 3 rows: {df.head(3).to_dict()}
     User instructions: {instructions or 'None'}
     
     Format EXACTLY as:
     # Visualization 1
-    plt.figure(figsize=(16,9), dpi=120)
-    plt.style.use('ggplot')
-    # Plot code here
-    plt.tight_layout()
+    [code using df and plt]
     """
 
     response = model.generate_content(prompt)
@@ -49,50 +44,30 @@ def process_file(api_key, file, instructions):
         buf = io.BytesIO()
         try:
             # Advanced code cleaning
-            cleaned_code = '\n'.join([
-                line.split('#')[0].strip()  # Remove comments
-                for line in block.split('\n')[1:]  # Skip header line
-                if line.strip() and 
-                not any(s in line for s in ['```', 'Annotation', ': '])
-            ])
+            cleaned_code = '\n'.join(
+                line.replace('data', 'df').split('#')[0].strip()
+                for line in block.split('\n')[1:]
+                if line.strip() and
+                not any(s in line.lower() for s in ['import', 'data =', 'data='])
+            )
             
-            # Syntax validation
+            # Validate syntax
             ast.parse(cleaned_code)
             
-            # Execute code
+            # Execute with controlled environment
+            exec_env = {'df': df, 'plt': plt}
             plt.figure(figsize=(16, 9), dpi=120)
-            exec(cleaned_code, {'df': df, 'plt': plt, 'pd': pd})
+            exec(cleaned_code, exec_env)
             
-            # Save image
             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} Error: {str(e)}")
-            print(f"Problematic Code:\n{cleaned_code}")
+            print(f"Cleaned Code:\n{cleaned_code}")
             visualizations.append(None)
 
     return visualizations + [None]*(3-len(visualizations))
 
-# Gradio interface
-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 CSV/Excel", file_types=[".csv", ".xlsx"])
-    
-    instructions = gr.Textbox(label="Custom Instructions")
-    submit = gr.Button("Generate Visualizations")
-    
-    with gr.Row():
-        outputs = [gr.Image(label=f"Visualization {i+1}") for i in range(3)]
-
-    submit.click(
-        process_file,
-        inputs=[api_key, file, instructions],
-        outputs=outputs
-    )
-
-demo.launch()
+# Rest of the Gradio interface remains the same