Day 4/30: Persistent StateGraph

I still remember the frustration of watching my first LangGraph agent fail miserably after just two turns. I had built a simple support bot designed to answer basic questions about our company's products, but it kept forgetting the context of the conversation. For instance, if a user asked about the pricing of a specific product, the bot would respond correctly on the first turn. However, if the user followed up with a question about the product's features, the bot would act as if it had never heard of the product before, sending the conversation back to square one.

After digging into the code, I realized that the problem lay in how I was handling the conversation state. I was creating a new StateGraph instance on each turn, which meant that the bot was essentially starting from scratch every time it received a new input. This was causing it to lose track of the conversation history and context. To fix this, I needed to find a way to persist the StateGraph across turns, so that the bot could build upon the previous state and maintain a coherent conversation flow.

The key to solving this problem was to understand how StateGraph instances work in LangGraph. A StateGraph represents the current state of the conversation, including the entities, intents, and relationships that have been established so far. By adding nodes and edges to the graph, we can update the state and reflect changes in the conversation. However, when we create a new StateGraph instance, we start with a blank slate, which is why my bot was forgetting the context.

To persist the StateGraph across turns, I used the add_node and add_conditional_edges methods to update the graph based on the user's input. Here's an example of how I did it:

import langgraph as lg

# Create an initial StateGraph instance
state_graph = lg.StateGraph()

# Define a function to update the StateGraph based on user input
def update_state_graph(state_graph, user_input):
    # Add a new node to the graph for the user's input
    input_node = state_graph.add_node(user_input)

    # Add conditional edges to the graph based on the user's intent
    if "price" in user_input:
        state_graph.add_conditional_edges(input_node, "product_pricing")
    elif "features" in user_input:
        state_graph.add_conditional_edges(input_node, "product_features")

    return state_graph

# Simulate a conversation
user_inputs = ["What is the price of product A?", "What are its features?"]

for user_input in user_inputs:
    state_graph = update_state_graph(state_graph, user_input)
    print(state_graph.nodes)  # Print the updated StateGraph nodes

In this example, we create an initial StateGraph instance and define a function update_state_graph to update the graph based on the user's input. We add a new node to the graph for each user input and add conditional edges to reflect the user's intent. By updating the StateGraph instance across turns, we can maintain a coherent conversation flow and ensure that the bot remembers the context.

One practical gotcha to watch out for when working with StateGraph instances is to avoid overwriting the graph with a new instance on each turn. This can cause the bot to lose track of the conversation history and context, as I experienced initially. Instead, make sure to update the existing StateGraph instance by adding new nodes and edges based on the user's input.

As we continue to build more complex agentic AI systems with LangGraph and MCP, we'll explore more advanced techniques for managing conversation state and context. Tomorrow, we'll dive into the world of Model Context Protocol (MCP) and explore how it can help us build more robust and contextualized AI models.