Update agent.py

agent.py

@@ -41,13 +41,12 @@ tools = [
 # Load system prompt
 system_prompt = """
 You are a general AI assistant. I will ask you a question.
-Report your thoughts, and finish your answer with the following template: 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.
+Report your thoughts, and finish your answer with only the answer, no extra text, no prefix, and no explanation.
+Your 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 a comma to write your number, nor 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 on whether the element to be put in the list is a number or a string.
+Format your output as: Answers (answers): [{"task_id": ..., "submitted_answer": ...}]
 """
 
 
@@ -117,9 +116,16 @@ def build_graph():
     # --- Nodes ---
     def assistant(state: MessagesState):
         """Assistant node"""
-        # Prepend the system message to the state
         messages_with_system_prompt = [sys_msg] + state["messages"]
-        return {"messages": [llm_with_tools.invoke(messages_with_system_prompt)]}
+        llm_response = llm_with_tools.invoke(messages_with_system_prompt)
+        # Extract the answer text (strip any "FINAL ANSWER:" if present)
+        answer_text = llm_response.content
+        if answer_text.strip().lower().startswith("final answer:"):
+            answer_text = answer_text.split(":", 1)[1].strip()
+        # Get task_id from state or set a placeholder
+        task_id = state.get("task_id", "1")  # Replace with actual logic if needed
+        formatted = f'Answers (answers): [{{"task_id": "{task_id}", "submitted_answer": "{answer_text}"}}]'
+        return {"messages": [formatted]}
 
     def retriever_node(state: MessagesState):
         """