Update app.py

app.py

@@ -321,13 +321,33 @@ class BasicAgent:
         return {
             "messages": [AIMessage(content=final_answer)]
         }
-
-    
+    def format_final_answer(self, question: str, raw_answer: str) -> str:
+        response = self.model.invoke([
+                SystemMessage(content="""
+        You are an exact-match answer formatter.
+        Return only the final answer.
+        No explanation.
+        No markdown.
+        No citations.
+        No extra punctuation.
+        Follow the requested format exactly.
+        If the question asks for a number, return only the number.
+        If the question asks for USD with two decimals, return only the number with two decimals.
+        If the question asks comma-separated values, use ", " between items.
+        If the question asks for algebraic notation, return only the move.
+        If the question says without abbreviations, expand abbreviations.
+        """),
+                HumanMessage(content=f"Question:\n{question}\n\nRaw answer:\n{raw_answer}")
+            ])
+        return response.content.strip()
+        
     def answer_question(self, question: str, task_id: str | None = None) -> str:
         file_info = None
+        q = question.lower()
     
         if task_id:
             info_str = download_task_file.invoke({"task_id": task_id})
+            print(f"[file_info] {info_str}", flush=True)
             try:
                 file_info = json.loads(info_str)
             except Exception:
@@ -337,34 +357,45 @@ class BasicAgent:
             suffix = file_info.get("suffix", "").lower()
             file_path = file_info["file_path"]
     
-            if suffix in [".png", ".jpg", ".jpeg", ".webp"]:
-                return answer_image_question.invoke({
+            is_image_q = any(x in q for x in [
+                "image", "picture", "screenshot", "chess position", "visual", "diagram", "shown in"
+            ])
+            is_audio_q = any(x in q for x in [
+                "audio", "recording", "mp3", "wav", "says", "say in response", "lecture"
+            ])
+            is_python_q = any(x in q for x in [
+                "python code", "attached python", "numeric output from the attached python"
+            ])
+            is_excel_q = any(x in q for x in [
+                "excel", "spreadsheet", "csv", "sales", "table contains"
+            ])
+    
+            if suffix in [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"] or is_image_q:
+                raw = answer_image_question.invoke({
                     "file_path": file_path,
                     "question": question
                 })
+                return self.format_final_answer(question, raw)
     
-            if suffix in [".mp3", ".wav", ".m4a", ".aac", ".flac", ".ogg", ".opus", ".webm"]:
-                return answer_audio_question.invoke({
+            if suffix in [".mp3", ".wav", ".m4a", ".aac", ".flac", ".ogg", ".opus", ".webm"] or is_audio_q:
+                raw = answer_audio_question.invoke({
                     "file_path": file_path,
                     "question": question
                 })
+                return self.format_final_answer(question, raw)
     
-            if suffix == ".py":
-                return answer_python_question.invoke({
+            if suffix == ".py" or is_python_q:
+                raw = answer_python_question.invoke({
                     "file_path": file_path
                 })
+                return self.format_final_answer(question, raw)
     
-            if suffix in [".xlsx", ".xls", ".csv"]:
+            if suffix in [".xlsx", ".xls", ".csv"] or is_excel_q:
                 raw = answer_excel_question.invoke({
                     "file_path": file_path,
                     "question": question
                 })
-                result = self.react_graph.invoke({
-                    "messages": [
-                        HumanMessage(content=f"Question:\n{question}\n\nTool result:\n{raw}\n\nReturn only the final answer.")
-                    ]
-                })
-                return result["messages"][-1].content
+                return self.format_final_answer(question, raw)
     
         user_content = question
         if task_id:
@@ -374,7 +405,6 @@ class BasicAgent:
         return result["messages"][-1].content
 
 
-
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,