add checkpointers

app/agent/graph.py

@@ -59,18 +59,6 @@ async def call_model(state: AgentState) -> Dict[str, List[AIMessage]]:
     return {"messages": [response]}
 
 
-# Define a new graph
-builder = StateGraph(AgentState, input=InputState, config_schema=Configuration)
-
-# Define the two nodes we will cycle between
-builder.add_node(call_model)
-builder.add_node("tools", ToolNode(TOOLS))
-
-# Set the entrypoint as `call_model`
-# This means that this node is the first one called
-builder.add_edge("__start__", "call_model")
-
-
 def route_model_output(state: SQLAgentState) -> Literal["__end__", "tools"]:
     """Determine the next node based on the model's output."""
     last_message = state.messages[-1]
@@ -89,6 +77,20 @@ def route_model_output(state: SQLAgentState) -> Literal["__end__", "tools"]:
     return "tools"
 
 
+# Initialize the checkpointer
+memory = MemorySaver()
+
+# Define a new graph
+builder = StateGraph(AgentState, input=InputState, config_schema=Configuration)
+
+# Define the two nodes we will cycle between
+builder.add_node(call_model)
+builder.add_node("tools", ToolNode(TOOLS))
+
+# Set the entrypoint as `call_model`
+# This means that this node is the first one called
+builder.add_edge("__start__", "call_model")
+
 # Add a conditional edge to determine the next step after `call_model`
 builder.add_conditional_edges(
     "call_model",
@@ -101,9 +103,8 @@ builder.add_conditional_edges(
 # This creates a cycle: after using tools, we always return to the model
 builder.add_edge("tools", "call_model")
 
-# Compile the builder into an executable graph
-memory = MemorySaver()
-graph = builder.compile(name="powersim_agent")
+# Compile the builder into an executable graph WITH checkpointer
+graph = builder.compile(checkpointer=memory, name="powersim_agent")
 
 if __name__ == "__main__":
     import asyncio
@@ -164,4 +165,4 @@ if __name__ == "__main__":
         print(final_response)
 
     # Run the async main function
-    asyncio.run(main())
+    asyncio.run(main())