Update agent.py

agent.py

@@ -33,19 +33,29 @@ tools = [add,
 
 chat_with_tools = chat.bind_tools(tools)
 
-## Defining our nodes
-def assistant(state: MessagesState):
-    """Assistant node"""
-    return {"messages": [sys_msg] + [chat_with_tools.invoke(state["messages"])]}
-
-# Build graph / nodes
-builder = StateGraph(State)
-builder.add_node("assistant", chat) # Assistant
-builder.add_node("tools", ToolNode(tools)) # Tools
-
-# Logic / edges
-builder.add_edge(START, "assistant")
-builder.add_conditional_edges("assistant", tools_condition)
-builder.add_edge("tools", "assistant")
-
-graph = builder.compile()
+def simple_graph():
+
+    ## Defining our nodes
+    def assistant(state: MessagesState):
+        """Assistant node"""
+        return {"messages": [sys_msg] + [chat_with_tools.invoke(state["messages"])]}
+    
+    # Build graph / nodes
+    builder = StateGraph(State)
+    builder.add_node("assistant", chat) # Assistant
+    builder.add_node("tools", ToolNode(tools)) # Tools
+    
+    # Logic / edges
+    builder.add_edge(START, "assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
+    
+    graph = builder.compile()
+
+    return graph
+
+
+
+if __name__ == "__main__":
+    # Build the graph
+    graph = simple_graph()