Update prompts.py

prompts.py

@@ -302,15 +302,30 @@ def validate_ai_response(question, ai_response_dict):
     """
     SAFETY CHECK: Remove plot object if question doesn't ask for visualization
     Call this function AFTER getting AI response and BEFORE using it
+    This is a CRITICAL function - it ensures NO visualizations appear unless explicitly requested
     """
     if not should_allow_visualization(question):
         # Force remove plot if it shouldn't exist
         if "plot" in ai_response_dict and ai_response_dict["plot"] is not None:
-            print(f"⚠️ WARNING: Removing unwanted visualization from AI response")
+            print(f"⚠️ WARNING: Removing unwanted visualization from AI response for question: '{question}'")
             ai_response_dict["plot"] = None
     
     return ai_response_dict
 
+def final_plot_validation(question, plot_object):
+    """
+    LAST LINE OF DEFENSE: Call this right before returning/displaying plot
+    Returns None if visualization should not be shown
+    """
+    if plot_object is None:
+        return None
+    
+    if not should_allow_visualization(question):
+        print(f"🚫 BLOCKED: Preventing visualization for non-visualization query: '{question}'")
+        return None
+    
+    return plot_object
+
 def get_insights_prompt(context_parts, narrative):
     insights_context = "\n".join(context_parts)
     return f"""Based on the analysis results, provide 4-6 informative bullet points that explain key insights, patterns, and findings from the data.
@@ -325,4 +340,66 @@ Guidelines:
 - Each bullet point should be 1-2 complete sentences
 - Focus on meaningful patterns and trends
 - Provide actionable insights
-- Be clear and specific about what the data shows"""
+- Be clear and specific about what the data shows"""
+
+"""
+⚠️⚠️⚠️ CRITICAL: HOW TO USE THESE FUNCTIONS IN YOUR MAIN APP ⚠️⚠️⚠️
+
+In your main application (the file that calls the AI), you MUST use these validation functions:
+
+STEP 1: After getting AI response, validate it immediately:
+```python
+# After getting response from AI
+ai_response = json.loads(response_text)
+
+# CRITICAL: Validate the response to remove unwanted plots
+ai_response = validate_ai_response(user_question, ai_response)
+```
+
+STEP 2: Before displaying/returning plot, do final validation:
+```python
+# Get plot from validated response
+plot_object = ai_response.get("plot")
+
+# FINAL CHECK: Validate plot before using it
+plot_object = final_plot_validation(user_question, plot_object)
+
+# Now use plot_object (it will be None if visualization wasn't requested)
+if plot_object is not None:
+    # Create and display the visualization
+    pass
+else:
+    # No visualization to display
+    pass
+```
+
+EXAMPLE INTEGRATION:
+```python
+def process_question(question, df):
+    # Get AI response
+    response = call_ai_model(question, df)
+    ai_response = json.loads(response)
+    
+    # STEP 1: Validate response
+    ai_response = validate_ai_response(question, ai_response)
+    
+    # Process operations
+    results = process_operations(ai_response["operations"], df)
+    
+    # STEP 2: Validate plot before using
+    plot_spec = final_plot_validation(question, ai_response.get("plot"))
+    
+    # Only create visualization if plot_spec is not None
+    if plot_spec:
+        chart = create_chart(plot_spec, results)
+        return results, chart, ai_response["narrative"]
+    else:
+        return results, None, ai_response["narrative"]
+```
+
+🔴 KEY POINTS:
+1. Call `validate_ai_response()` immediately after parsing AI response
+2. Call `final_plot_validation()` before creating/displaying any chart
+3. These functions will automatically block visualizations when keywords are missing
+4. No changes needed to prompts - validation handles everything
+"""