update agent and prompt

agent.py

@@ -10,7 +10,11 @@ from smolagents import (
     WikipediaSearchTool,
 )
 
-from prompt import calc_agent_prompt, web_search_agent_prompt
+from prompt import (
+    calc_agent_prompt,
+    gen_GAIA_answer_formatter_prompt,
+    web_search_agent_prompt,
+)
 
 if pathlib.Path(".env").exists():
     dotenv.load_dotenv(".env")
@@ -50,6 +54,41 @@ calc_agent = CodeAgent(
     tools=[]
 )
 
+class ExtractFailedException(Exception):
+    """Custom exception for failed extraction of formatted answer."""
+    pass
+
+class GAIAAnswerFormatter:
+
+    model = OpenAIServerModel(
+        model_id="gpt-4.1-mini",
+        api_key=OPENAI_API_KEY,
+    )
+
+    def extract_formatted_answer(self, llm_response: str) -> str:
+        import re 
+
+        match = re.search(r'<formated_answer>(.*?)</formated_answer>', llm_response, re.DOTALL)
+        if match:
+            return match.group(1).strip()
+        raise ExtractFailedException(
+            "Failed to extract formatted answer from the LLM response."
+        )
+
+    def __call__(self, question: str, answer: str) -> str:
+        message = [{
+            "role": "user",
+            "content": gen_GAIA_answer_formatter_prompt(question, answer),
+        }]
+        response = self.model.generate(messages=message)
+        try:
+            return self.extract_formatted_answer(response.content)
+        except ExtractFailedException:
+            return answer
+
+
+answer_formatter = GAIAAnswerFormatter()
+
 class Agent:
     def __init__(self):
         self.agent = CodeAgent(
@@ -61,7 +100,8 @@ class Agent:
 
 
     def __call__(self, question: str) -> str:
-        return self.agent.run(question)
+        answer = self.agent.run(question)
+        return answer_formatter(question, answer)
 
 
 if __name__ == "__main__":

prompt.py

@@ -288,4 +288,35 @@ Now Begin!""",
         pre_messages="",
         post_messages="",
     ),
-)
+)
+
+def gen_GAIA_answer_formatter_prompt(question: str, answer: str) -> str:
+    return f"""You are a GAIA answer format validator and formatter. Your task is to check if an agent's answer meets GAIA benchmark requirements and reformat it if necessary.
+
+GAIA FORMAT REQUIREMENTS:
+- YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings
+- If asked for a number: don't use commas to write numbers, no units like $ or % unless specified
+- If asked for a string: no articles (a, an, the), no abbreviations (e.g. for cities), write digits in plain text unless specified
+- If asked for a comma separated list: apply above rules for each element
+
+ORIGINAL QUESTION:
+{question}
+
+AGENT'S ANSWER:
+{answer}
+
+INSTRUCTIONS:
+1. First, analyze what type of answer the question is asking for (number, string, or list)
+2. Check if the agent's answer meets GAIA format requirements
+3. If the answer is already correctly formatted, return it as is
+4. If the answer needs reformatting, extract the core information and reformat according to GAIA rules
+5. Provide your final answer wrapped in XML tags: <formated_answer>YOUR FORMATTED ANSWER</formated_answer>
+
+Remember:
+- Keep only essential information
+- Remove unnecessary words, articles, and explanations
+- Follow the specific formatting rules for numbers, strings, or lists
+- If the agent's answer contains multiple pieces of information, extract only what the question specifically asks for
+- The content inside <formated_answer> tags should be the clean, formatted answer without any additional text
+
+Now analyze and format the answer:"""