Update app.py

app.py

@@ -42,8 +42,8 @@ SYSTEM_INSTRUCTION = (
     "1. **Reasoning:** Explain your plan, the steps taken, and the insights derived from the data. Format this in Markdown. "
     "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 output: `st.dataframe(...)` for results, `st.bar_chart()`, `st.line_chart()`, or `st.pyplot()` for plots. "
-    "   - If you need custom plotting, you can assume `import matplotlib.pyplot as plt` is available and use `plt.figure()`, `plt.bar()`, etc., followed by `st.pyplot(plt)` to display the plot. "
+    "   - Use Streamlit functions for simple outputs: `st.dataframe(...)`, `st.bar_chart()`, `st.line_chart()`. "
+    "   - For custom, complex plots, you MUST use Matplotlib: start with `plt.figure()`, then use `plt.` commands (like `plt.hist()`), and finish with `st.pyplot(plt)` to display the plot. "
     "   - Ensure the code is self-contained and ready to execute."
 )
 
@@ -53,7 +53,7 @@ SYSTEM_INSTRUCTION = (
 def chat_with_gemini(prompt, context):
     """Sends a prompt and data context to the Gemini model for structured analysis (reasoning + code)."""
     
-    # FIX: Correctly prepend 'models/' to the model name in the URL path
+    # Correctly prepend 'models/' to the model name in the URL path
     url = f"{GEMINI_BASE}/models/{CHAT_MODEL}:generateContent?key={GEMINI_API_KEY}"
     
     # Construct the full prompt including the data context
@@ -125,7 +125,6 @@ if uploaded:
             df = st.session_state.df # Local variable for code execution context
             
             # Summarize dataset for context sent to the LLM
-            # FIX: Using to_string() instead of to_markdown() to avoid the 'tabulate' dependency error.
             context = f"Dataset Columns: {', '.join(df.columns.astype(str))}\n\nFirst 5 rows of data:\n{df.head(5).to_string(index=False)}"
 
             st.markdown("---")
@@ -166,6 +165,9 @@ if uploaded:
                         # Executing the code within the local scope
                         exec(code, globals(), local_scope)
                         
+                        # FIX: Explicitly close all Matplotlib figures to prevent cross-run contamination
+                        plt.close('all') 
+                        
                         st.success("Code execution complete. Results are displayed above.")
 
                     except Exception as e: