Update agent.py

agent.py

@@ -36,6 +36,7 @@ custom_tools = [datetime_tools, transcribe_audio_tool]
 tools = community_tools + custom_tools
 llm_with_tools = llm.bind_tools(tools)
 
+# Prepare tools by name
 tools_by_name = {tool.name: tool for tool in tools}
 
 class MessagesState(TypedDict):  # creates the state (is like the agent's memory at any moment)
@@ -83,11 +84,19 @@ builder.add_edge(START, "llm_call")
 builder.add_conditional_edges(
     "llm_call", 
     should_continue,
-    {
-        # Name returned by should_continue : Name of the next node
-        "Action": "environment",
-        END: END
-})
+    {"Action": "environment", END: END})  # Name returned by should_continue : Name of the next node
 builder.add_edge("environment", "llm_call")
 
-gaia_agent = builder.compile()
+gaia_agent = builder.compile()
+
+# Call the Agent
+def run_gaia(question: str):
+    input_state = {"messages": [HumanMessage(content=question)]}  # prepare the initial user message
+    result = gaia_agent.invoke(input_state, {"recursion_limit": 30})  # invoke with recursion limit to avoid infinite loops
+    final_response = result["messages"][-1].content
+    
+    try:
+        return final_response.split("FINAL ANSWER:")[-1].strip()  #  parse out only what's after "FINAL ANSWER:"
+    except Exception:
+        print("Could not split on 'FINAL ANSWER:'")
+        return final_response