Update app.py

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
+import json
 
 from typing import List, TypedDict, Annotated, Optional
 from langchain_openai import ChatOpenAI
@@ -26,8 +27,11 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 
-class AgentState(TypedDict):
+class AgentState(TypedDict, total=False):
     messages: Annotated[list[AnyMessage], add_messages]
+    verification_status: str
+    verification_notes: str
+    verify_retries: int
     
 class BasicAgent:
     def __init__(self):
@@ -69,20 +73,36 @@ class BasicAgent:
         # Define nodes: these do the work
         builder.add_node("assistant", self.assistant)
         builder.add_node("tools", ToolNode(self.tools))
+        builder.add_node("verify_answer", self.verify_answer)
+        builder.add_node("retry_with_feedback", self.retry_with_feedback)
         builder.add_node("final_process", self.clean_answer)
         
         # Define edges: these determine how the control flow moves
         builder.add_edge(START, "assistant")
+        
         builder.add_conditional_edges(
             "assistant",
             self.route_after_assistant,
             {
                 "tools": "tools",
-                "final_process": "final_process",
+                "verify_answer": "verify_answer",
             },
         )
+        
         builder.add_edge("tools", "assistant")
+        
+        builder.add_conditional_edges(
+            "verify_answer",
+            self.route_after_verify,
+            {
+                "retry": "retry_with_feedback",
+                "final_process": "final_process",
+            },
+        )
+        
+        builder.add_edge("retry_with_feedback", "assistant")
         builder.add_edge("final_process", END)
+        
         self.react_graph = builder.compile()
 
     def __call__(self, question: str, task_id: str | None = None) -> str:
@@ -124,6 +144,141 @@ class BasicAgent:
         return {
             "messages": [self.chat_with_tools.invoke([sys_msg] + state["messages"])],
         }
+    def _safe_parse_json(self, text: str) -> dict:
+    """
+    Parse JSON from model output safely.
+    The model may sometimes wrap JSON with extra text.
+    """
+    text = text.strip()
+
+    try:
+        return json.loads(text)
+    except Exception:
+        pass
+
+    start = text.find("{")
+    end = text.rfind("}")
+
+    if start != -1 and end != -1 and end > start:
+        try:
+            return json.loads(text[start:end + 1])
+        except Exception:
+            pass
+
+    return {
+        "verdict": "PASS",
+        "revised_answer": text,
+        "issue": "JSON parse failed, using raw verifier output.",
+    }
+
+
+    def verify_answer(self, state: AgentState):
+        messages = state["messages"]
+    
+        question = messages[0].content
+        candidate_answer = messages[-1].content
+        retry_count = state.get("verify_retries", 0)
+    
+        # 只拿最近一部分上下文，避免 prompt 过长
+        recent_context = []
+        for m in messages[-8:]:
+            role = m.__class__.__name__
+            content = getattr(m, "content", "")
+            recent_context.append(f"{role}:\n{content}")
+    
+        context_text = "\n\n---\n\n".join(recent_context)
+    
+        response = self.model.invoke([
+                SystemMessage(content="""
+        You are a strict answer verifier for an exact-match QA benchmark.
+        
+        Your job:
+        1. Check whether the candidate answer satisfies the question.
+        2. Check whether the answer format is exactly what the question asks.
+        3. Use the available conversation/tool evidence only.
+        4. Do not introduce unsupported new facts.
+        5. If the answer is probably correct but badly formatted, revise the format.
+        6. If the answer is unsupported, clearly wrong, says file/audio/image is unavailable, or ignores a required source/date constraint, request a retry.
+        
+        Return only valid JSON with this schema:
+        {
+          "verdict": "PASS" or "RETRY",
+          "revised_answer": "final answer if PASS, otherwise empty string",
+          "issue": "short reason"
+        }
+        
+        Important exact-match formatting rules:
+        - Remove explanations, markdown, citations, and extra punctuation.
+        - If the question asks for a number, return only the number.
+        - If the question asks for a name part, return only that part.
+        - If the question asks comma-separated values, use comma + space.
+        - If the question says without abbreviations, expand abbreviations.
+        - If the candidate says no file/image/audio is available but the question has a task_id or attachment, use RETRY.
+        - If the candidate ignores a date/source constraint, use RETRY.
+        """),
+                HumanMessage(content=f"""
+        Original question:
+        {question}
+        
+        Candidate answer:
+        {candidate_answer}
+        
+        Recent conversation and tool context:
+        {context_text}
+        """)
+            ])
+    
+        data = self._safe_parse_json(response.content)
+    
+        verdict = str(data.get("verdict", "PASS")).upper().strip()
+        revised_answer = str(data.get("revised_answer", candidate_answer)).strip()
+        issue = str(data.get("issue", "")).strip()
+    
+        max_verify_retries = 1
+    
+        if verdict == "RETRY" and retry_count < max_verify_retries:
+            return {
+                "verification_status": "retry",
+                "verification_notes": issue,
+                "verify_retries": retry_count + 1,
+            }
+    
+        # 如果 verifier 通过，或者已经重试过一次还不行，就进入最终清洗
+        final_candidate = revised_answer if revised_answer else candidate_answer
+    
+        return {
+            "messages": [AIMessage(content=final_candidate)],
+            "verification_status": "pass",
+            "verification_notes": issue,
+            "verify_retries": retry_count,
+        }
+    
+    
+    def route_after_verify(self, state: AgentState):
+        if state.get("verification_status") == "retry":
+            return "retry"
+    
+        return "final_process"
+    
+    
+    def retry_with_feedback(self, state: AgentState):
+        notes = state.get("verification_notes", "")
+    
+        return {
+            "messages": [
+                HumanMessage(content=f"""
+    Your previous answer failed verification.
+    
+    Verifier notes:
+    {notes}
+    
+    Please answer the original question again.
+    Use tools if needed.
+    Pay attention to source/date constraints and exact output format.
+    Return only the final answer.
+    """.strip())
+            ]
+        }
     
     def clean_answer(self, state: AgentState):
         messages = state["messages"]
@@ -132,42 +287,77 @@ class BasicAgent:
         raw_answer = messages[-1].content
     
         response = self.model.invoke([
-            SystemMessage(content=(
-                "You are an answer formatter. "
-                "Extract only the final answer required by the question. "
-                "No explanation. No markdown. No extra words."
-            )),
-            HumanMessage(content=f"Question:\n{question}\n\nRaw answer:\n{raw_answer}")
-        ])
+                SystemMessage(content="""
+        You are an exact-match answer formatter.
+        
+        Rules:
+        - Return only the final answer.
+        - No explanation.
+        - No markdown.
+        - No citations.
+        - No extra punctuation.
+        - 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 unless $ is explicitly requested.
+        - If the question asks comma-separated values, use ", " between items.
+        - If the question asks for a first name, surname, city, country code, or algebraic notation, return only that.
+        - If the question says "without abbreviations", expand abbreviations.
+        - Preserve required capitalization when obvious.
+        """),
+                HumanMessage(content=f"Question:\n{question}\n\nRaw answer:\n{raw_answer}")
+            ])
     
         final_answer = response.content.strip()
     
         return {
             "messages": [AIMessage(content=final_answer)]
         }
+
     
-    def route_after_assistant(self, state: AgentState):
-        last_message = state["messages"][-1]
+    def answer_question(self, question: str, task_id: str | None = None) -> str:
+        file_info = None
     
-        if getattr(last_message, "tool_calls", None):
-            return "tools"
+        if task_id:
+            info_str = download_task_file.invoke({"task_id": task_id})
+            try:
+                file_info = json.loads(info_str)
+            except Exception:
+                file_info = None
     
-        return "final_process"
-
-    def answer_question(self, question: str, task_id: str | None = None) -> str:
-        user_content = question
+        if file_info and "file_path" in file_info:
+            suffix = file_info.get("suffix", "").lower()
+            file_path = file_info["file_path"]
     
+            if suffix in [".png", ".jpg", ".jpeg", ".webp"]:
+                return answer_image_question.invoke({
+                    "file_path": file_path,
+                    "question": question
+                })
+    
+            if suffix in [".mp3", ".wav", ".m4a"]:
+                return answer_audio_question.invoke({
+                    "file_path": file_path,
+                    "question": question
+                })
+    
+            if suffix == ".py":
+                return answer_python_question.invoke({
+                    "file_path": file_path
+                })
+    
+            if suffix in [".xlsx", ".xls", ".csv"]:
+                raw = answer_excel_question.invoke({
+                    "file_path": file_path,
+                    "question": question
+                })
+                return self.format_final_answer(question, raw)
+    
+        user_content = question
         if task_id:
-            user_content += (
-                f"\n\nTask ID: {task_id}. "
-                "If the question mentions an attached file, spreadsheet, code file, audio, or image, "
-                "use the available tools with this task_id."
-            )
-    
-        messages = [HumanMessage(content=user_content)]
-        result = self.react_graph.invoke({"messages": messages})
-        return result["messages"][-1].content
+            user_content += f"\n\nTask ID: {task_id}."
     
+        result = self.react_graph.invoke({"messages": [HumanMessage(content=user_content)]})
+        return result["messages"][-1].content
+
 
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):