Update agent.py

agent.py

@@ -60,6 +60,7 @@ def tool_node(state: MessagesState):
         result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"]))
     return {"messages": result}
 
+# Conditional edge function to route to the tool node or end based upon whether the LLM made a tool call
 def should_continue(state: MessagesState) -> Literal["Action", END]:
     """Decide if we should continue the loop or stop based upon whether the LLM made a tool call"""
 
@@ -69,3 +70,24 @@ def should_continue(state: MessagesState) -> Literal["Action", END]:
         return "Action"
     # Otherwise, we stop (reply to the user)
     return END
+
+# Build workflow
+builder = StateGraph(MessagesState)
+
+# Add nodes
+builder.add_node("llm_call", llm_call)
+builder.add_node("environment", tool_node)
+
+# Add edges to connect nodes
+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
+})
+builder.add_edge("environment", "llm_call")
+
+gaia_agent = builder.compile()