Update agent.py

agent.py

@@ -148,11 +148,13 @@ def build_graph(provider: str = "groq") -> StateGraph:
             ]
         }
     
-    # 6.2) Node: human → append user question
-    def human_node(state: AgentState, question: str) -> AgentState:
+
+    # 6.2) Node: human → stash the GAIA task ID, then append the question
+    def human_node(state: AgentState, id: str, question: str) -> AgentState:
+        # keep the GAIA task id so we can submit it later
+        state["task_id"] = id
         state["messages"].append(HumanMessage(content=question))
         return state
-
     # 6.3) Node: assistant → call LLM on current messages
     def assistant_node(state: AgentState) -> dict:
         ai_msg = llm.invoke(state["messages"])
@@ -168,6 +170,15 @@ def build_graph(provider: str = "groq") -> StateGraph:
             state["messages"].append(HumanMessage(content=result))
             return {"messages": state["messages"]}
         return tool_node
+    # 6.5) Node: answer → pull out the last assistant reply & format submission dict
+    def answer_node(state: AgentState) -> dict[str,str]:
+        # the GAIA runner will do `.items()` on whatever you return here
+        tid = state["task_id"]
+        # grab the last message (could be a BaseMessage or a raw str)
+        last = state["messages"][-1]
+        text = getattr(last, "content", None) or str(last)
+        return { tid: text }
+
 
     # Instantiate nodes for each tool
     calc_node     = make_tool_node(calculator)
@@ -208,7 +219,12 @@ def build_graph(provider: str = "groq") -> StateGraph:
     g.add_edge("web",      "assistant")
     g.add_edge("wiki",     "assistant")
     g.add_edge("arxiv",    "assistant")
-    # and finally assistant → END when done
-    g.add_edge("assistant", END)
+
+    # register & wire your new answer node
+    g.add_node("answer", answer_node)
+
+    # send assistant → answer → END
+    g.add_edge("assistant", "answer")
+    g.add_edge("answer", END)
 
     return g.compile()