Update app.py

app.py

@@ -5,8 +5,8 @@ import numpy as np
 import requests
 import json
 import time 
-import matplotlib.pyplot as plt # Explicitly imported for use in exec() scope
-import seaborn as sns
+import matplotlib.pyplot as plt 
+import seaborn as sns 
 
 # --- CONFIG ---
 # Note: GEMINI_API_KEY is retrieved from environment variables/secrets.
@@ -44,13 +44,13 @@ SYSTEM_INSTRUCTION = (
     "2. **Code:** If the question requires calculation, aggregation, or visualization, you MUST generate Python code to execute against the 'df' DataFrame. "
     "   - The DataFrame is already loaded as a variable named 'df'. Do NOT redefine it. "
     "   - Use Streamlit functions for simple outputs: `st.dataframe(...)`, `st.bar_chart()`, `st.line_chart()`. "
-    "   - For **ALL** custom, complex plots, you MUST follow this strict Matplotlib sequence: **Start with `plt.figure()`, use `plt.` commands for plotting, and explicitly end with `st.pyplot(plt)`** to display the output. "
+    "   - For **ALL** custom, complex plots, you MUST follow this strict Matplotlib sequence: **Start with `plt.figure()`, use `plt.` or `sns.` commands for plotting, and explicitly end with `st.pyplot(plt)`** to display the output. "
+    "   - **CRITICAL GUARDRAIL:** When generating code that uses logical conditions (e.g., in `if` statements or for complex filters) on Pandas Series or NumPy arrays, you **MUST** resolve ambiguity by using `.any()` or `.all()`. Do NOT compare a series directly to a single boolean value."
     "   - Ensure the code is self-contained and ready to execute."
 )
 
 # --- Helper Functions ---
 
-# Function to chat with the Gemini API and enforce structured JSON output
 def chat_with_gemini(prompt, context):
     """Sends a prompt and data context to the Gemini model for structured analysis (reasoning + code)."""
     
@@ -76,10 +76,9 @@ def chat_with_gemini(prompt, context):
     for attempt in range(max_retries):
         try:
             r = requests.post(url, headers={'Content-Type': 'application/json'}, data=json.dumps(payload))
-            r.raise_for_status() # This will raise an exception for 4xx or 5xx errors
+            r.raise_for_status() 
             data = r.json()
             
-            # The JSON output is a string inside the 'text' part
             json_str = data["candidates"][0]["content"]["parts"][0]["text"]
             return json.loads(json_str)
 
@@ -161,7 +160,8 @@ if uploaded:
                             'st': st, 
                             'pd': pd,
                             'np': np,
-                            'plt': plt, # Explicitly include pyplot here
+                            'plt': plt,
+                            'sns': sns, 
                         }
                         
                         # Append a neutral statement to the code to prevent implicit Streamlit display of the last value