Implement direct matplotlib figure display with st.pyplot()

- Added streamlit import to code execution environment
- Updated plotting requirements to use st.pyplot(fig) for direct display
- Should provide better resolution and faster rendering than file-based approach
- Fixed immediate user message display without rerun delays

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

app.py

@@ -494,7 +494,10 @@ def show_custom_response(response):
             with st.expander("📋 View Generated Code", expanded=False):
                 st.code(response["gen_code"], language="python")
         
-        # Try to display image if content is a file path
+        # Check if this is a plot response (plots are now displayed directly via st.pyplot)
+        is_plot_response = isinstance(content, str) and "Plot displayed successfully" in content
+        
+        # Try to display image if content is a file path (for backward compatibility)
         try:
             if isinstance(content, str) and (content.endswith('.png') or content.endswith('.jpg')):
                 if os.path.exists(content):

new_system_prompt.txt

@@ -33,10 +33,9 @@ DATA SAFETY:
 - Use .dropna() to remove missing values before analysis
 
 PLOTTING REQUIREMENTS:
-- Create plots for visualization requests: plt.figure(figsize=(12, 7))
-- Save plots with ultra-high resolution: filename = f"plot_{uuid.uuid4().hex[:8]}.png"; plt.savefig(filename, dpi=1200, bbox_inches='tight', facecolor='white', edgecolor='none')
-- Close plots: plt.close()
-- Store filename: answer = filename
+- Create plots for visualization requests: fig, ax = plt.subplots(figsize=(12, 7))
+- Display plots directly: st.pyplot(fig); plt.close()
+- Store success message: answer = "Plot displayed successfully"
 - For non-plots: answer = "text result"
 
 BASIC ERROR PREVENTION:

src.py

@@ -166,7 +166,8 @@ def ask_question(model_name, question):
             'uuid': __import__('uuid'),
             'calendar': __import__('calendar'),
             'np': __import__('numpy'),
-            'df': df  # <-- pass your DataFrame here
+            'df': df,  # <-- pass your DataFrame here
+            'st': __import__('streamlit')  # Add streamlit for st.pyplot()
         }
 
         # allow user to inject more globals (optional)