AI Agents with LangGraph

Nov 12 2024 · Python 3.12, LangGraph 0.2.x, JupyterLab 4.2.4

Lesson 04: Enhancing Agent Capabilities

Memory Demo

Episode complete

Play next episode

Next
Transcript

Memory Demo

Open the empty memory.ipynb notebook in the Starter project.

To see how the checkpointer works, start by defining your State class and making some node functions that’ll use it:

from typing import TypedDict, Annotated, Sequence
from operator import add

class State(TypedDict):
  messages: Annotated[Sequence[str], add]

def function_1(state):
  return {"messages": ["uno"]}

def function_2(state):
  return {"messages": ["dos"]}

def function_3(state):
  return {"messages": ["tres"]}

Next, create a graph, add your nodes, and connect the nodes with edges:

from langgraph.graph import StateGraph, START, END

graph = StateGraph(State)
graph.add_node("node_1", function_1)
graph.add_node("node_2", function_2)
graph.add_node("node_3", function_3)
graph.add_edge(START, "node_1")
graph.add_edge("node_1", "node_2")
graph.add_edge("node_2", "node_3")
graph.add_edge("node_3", END)

Now you come to the new part where you use MemorySaver. Import MemorySaver. And create a new instance of it. Pass the instance as a parameter to compile the graph. Provide an input message for your state. Set the thread ID to 1. Finally, invoke the app with the input and thread ID.

from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()
app = graph.compile(checkpointer=memory)
user_input = {"messages": ["hola"]}
thread = {"configurable": {"thread_id": "1"}}
app.invoke(user_input, thread)

Run that, and you’ll see:

{'messages': ['hola', 'uno', 'dos', 'tres']}

Hello, one, two, three. So far, so good. Nothing is crashing, but how do you observe if the checkpointer is working? You can loop through the state history with get_state_history:

for state in app.get_state_history(thread):
  print(state)
  print("--")

Run that and see what you get. You have a bunch of StateSnapshots. They start with the newest and go to the oldest. You can see the message content at each step. Also, look at the next property. And the thread ID is always 1. That’s your state history for this thread.

See forum comments
Cinema mode Download course materials from Github
Previous: Memory Next: Structured Output