Update app.py

app.py

@@ -64,7 +64,7 @@ agent = initialize_agent(
     agent_kwargs={
         "system_message": SystemMessage(content="You are a general AI assistant. I will ask you a question. "
             "Report your thoughts, and finish your answer with the following template: "
-            "[YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. "
+            "FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. "
             "If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. "
             "If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. "
             "If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.")
@@ -77,8 +77,19 @@ class BasicAgent:
         self.agent = agent
 
     def __call__(self, question: str) -> str:
-        print(f"\nQuestion: {question}")
-        return self.agent.run(question)
+
+        def extract_model_answer(full_response: str) -> str:
+            marker = "FINAL ANSWER:"
+            idx = full_response.find(marker)
+            if idx == -1:
+                return full_response.strip()  # fallback: return whole response
+            return full_response[idx + len(marker):].strip().rstrip(".")
+            
+        response = self.agent.run(question)
+        model_answer = response["messages"][-1].content
+        fixed_answer = extract_model_answer(model_answer)
+        reasoning_trace = "\n".join([f"{m.type}: {m.content}" for m in messages + response["messages"]])
+        return fixed_answer
 
     
 def run_and_submit_all( profile: gr.OAuthProfile | None):