Upload agent.py

agents/agent.py

@@ -19,6 +19,16 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
+def strip_answer(answer: str) -> str:
+    """
+    Strip the answer to remove the 'FINAL ANSWER:' prefix and any trailing whitespace.
+    """
+    if "FINAL ANSWER:" in answer:
+        # If the answer contains 'FINAL ANSWER:', we assume it follows the expected format.
+        answer = answer.split("FINAL ANSWER:")[1]
+    return answer.strip()
+
+
 class BasicAgent:
     def __init__(self):
         logger.info("BasicAgent initialized.")
@@ -69,7 +79,8 @@ class SimpleGeminiAgent(BasicAgent):
             f"AdvancedAgent received question (first 50 chars): {question[:50]}..."
         )
         logger.info(f"AdvancedAgent returning answer: {response.text}")
-        return response.text
+        answer = strip_answer(response.text)
+        return answer
 
 
 class DeepResearchGeminiAgent(BasicAgent):
@@ -107,7 +118,8 @@ class DeepResearchGeminiAgent(BasicAgent):
             f"AdvancedAgent received question (first 50 chars): {question[:50]}..."
         )
         logger.info(f"AdvancedAgent returning answer: {final_answer_message.content}")
-        return final_answer_message.content
+        answer = strip_answer(final_answer_message.content)
+        return answer
 
 
 if __name__ == "__main__":