Update app.py

app.py

@@ -5,116 +5,98 @@ import io
 import google.generativeai as genai
 from PIL import Image
 import ast
+import re
 
 def process_file(api_key, file, instructions):
-    # Configure Gemini API with correct model
     genai.configure(api_key=api_key)
     model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
 
     try:
-        # Read uploaded file
-        if file.name.endswith('.csv'):
-            df = pd.read_csv(file.name)
-        else:
-            df = pd.read_excel(file.name)
+        df = pd.read_csv(file.name) if file.name.endswith('.csv') else pd.read_excel(file.name)
     except Exception as e:
         print(f"File Error: {str(e)}")
-        return [None, None, None]
+        return [None]*3
 
-    # Enhanced prompt with strict coding rules
-    prompt = f"""Generate 3 distinct matplotlib visualization codes with these rules:
-    1. Use ONLY these variables: df (existing DataFrame), plt
-    2. No imports or additional data loading
-    3. Each visualization must:
-       - Start with: plt.figure(figsize=(16,9), dpi=120)
-       - Use plt.style.use('ggplot')
-       - Include title, axis labels, and grid
-       - End with plt.tight_layout()
-    4. Different chart types (bar, line, scatter, etc)
+    # 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:
+       plt.figure(figsize=(16,9), dpi=120)
+       plt.style.use('ggplot')
+    4. Use ONLY df and plt variables
+    5. End with plt.tight_layout()
     
-    Dataset columns: {list(df.columns)}
-    Sample data: {df.head(3).to_dict()}
-    User instructions: {instructions or 'None'}
+    Columns: {list(df.columns)}
+    Sample: {df.head(3).to_dict()}
+    Instructions: {instructions or 'None'}
     
     Format EXACTLY as:
     # Visualization 1
-    plt.figure(figsize=(16,9), dpi=120)
-    plt.style.use('ggplot')
-    # Visualization code using df
-    plt.tight_layout()
+    [code]
+    
+    # Visualization 2
+    [code]
+    
+    # Visualization 3
+    [code]
     """
 
     try:
         response = model.generate_content(prompt)
-        code_blocks = response.text.split("# Visualization ")[1:4]
+        code_blocks = re.split(r'# Visualization \d+', response.text)[1:4]
     except Exception as e:
         print(f"Gemini Error: {str(e)}")
-        return [None, None, None]
+        return [None]*3
 
     visualizations = []
     for i, block in enumerate(code_blocks, 1):
         buf = io.BytesIO()
         try:
-            # Clean and validate generated code
-            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 ='])
-            ])
+            # 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
+            )
+            
+            # Add missing parentheses check
+            cleaned_code = re.sub(r'plt.style.use\([\'"]ggplot$', 
+                                'plt.style.use("ggplot")', cleaned_code)
             
-            # Syntax check
+            # Syntax validation
             ast.parse(cleaned_code)
             
-            # Execute code in controlled environment
+            # Execute code
             exec_env = {'df': df, 'plt': plt}
             plt.figure(figsize=(16, 9), dpi=120)
             exec(cleaned_code, exec_env)
             
-            # Save HD 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)
 
-    # Ensure exactly 3 outputs
     return visualizations + [None]*(3-len(visualizations))
 
 # Gradio interface
-with gr.Blocks(theme=gr.themes.Default(spacing_size="lg")) as demo:
-    gr.Markdown("# 🔍 Data Visualization Generator")
+with gr.Blocks() as demo:
+    gr.Markdown("# Professional Data Visualizer")
     
     with gr.Row():
-        api_key = gr.Textbox(
-            label="🔑 Gemini API Key",
-            type="password",
-            placeholder="Enter your API key here"
-        )
-        file = gr.File(
-            label="📁 Upload Dataset",
-            file_types=[".csv", ".xlsx"],
-            type="filepath"
-        )
-    
-    instructions = gr.Textbox(
-        label="💡 Custom Instructions",
-        placeholder="E.g.: Compare sales trends across regions..."
-    )
+        api_key = gr.Textbox(label="Gemini API Key", type="password")
+        file = gr.File(label="Upload Data File", file_types=[".csv", ".xlsx"])
     
-    submit = gr.Button("🚀 Generate Visualizations", variant="primary")
+    instructions = gr.Textbox(label="Visualization Instructions")
+    submit = gr.Button("Generate", variant="primary")
     
     with gr.Row():
-        outputs = [
-            gr.Image(
-                label=f"Visualization {i+1}",
-                width=600,
-                height=400
-            ) for i in range(3)
-        ]
+        outputs = [gr.Image(label=f"Visualization {i+1}", width=600) for i in range(3)]
 
     submit.click(
         process_file,
@@ -122,5 +104,4 @@ with gr.Blocks(theme=gr.themes.Default(spacing_size="lg")) as demo:
         outputs=outputs
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()