Update app.py

app.py

@@ -14,6 +14,13 @@ import google.generativeai as genai
 genai.configure(api_key="AIzaSyBO3-HG-WcITn58PdpK7mMyvFQitoH00qA")  # Replace with your actual Gemini API key
 
 # Load Gemini model
+def output_guard(answer):
+    # Check if answer is empty or too short
+    if not answer or len(answer.strip()) < 20:
+        print("Output guard triggered: answer too short or empty.")
+        return False
+    # You can add more checks here if needed
+    return True
 
 
 import re
@@ -137,14 +144,22 @@ class WebFallbackQA(dspy.Program):
         return {"answer": answer, "retrieved_context": "Tavily"}
 
 
+
 class MathRouter(dspy.Program):
     def forward(self, question):
         print("Routing question:", question)
         if not is_valid_math_question(question):
-            return dspy.Output(answer="❌ Only math questions are accepted. Please rephrase.", retrieved_context="")
+            return {"answer": "❌ Only math questions are accepted. Please rephrase.", "retrieved_context": ""}
+        
         result = MathRetrievalQA().forward(question)
-        #return result if result.answer else WebFallbackQA().forward(question)
+        
+        # Apply output guard here
+        if not output_guard(result["answer"]):
+            print("Output guard failed, falling back to web search.")
+            return WebFallbackQA().forward(question)
+        
         return result if result["answer"] else WebFallbackQA().forward(question)
+
 router = MathRouter()
 
 # === Feedback Storage ===