how to retrieve the chat history in Langgraph using Checkpointer?

[jayashriv710]

Hi everyone,
I'm having an issue with using TypedDict and Annotated for storing chat history in a class. When I define my AgentState class like this, I can retrieve the chat history without any problems:

from typing import TypedDict, Annotated

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]

However, when I extend the AgentState class to include more fields, I'm unable to retrieve the chat memory

from typing import TypedDict, Annotated

class AgentState(TypedDict):
    input: str
    output: str
    decision: str
    messages: Annotated[list, add_messages]

In this case, I can no longer retrieve the chat memory. Does anyone know why this might be happening and how I can fix it? Any help would be greatly appreciated!

Here is the structure of my StateGraph:

from langchain_ai.langgraph import StateGraph

class AgentState(TypedDict):
    input: str
    output: str
    decision: str
    messages: Annotated[list, add_messages]

workflow = StateGraph(AgentState)

workflow.add_node("analyze", analyze_question)
workflow.add_node("code_agent", answer_code_question)
workflow.add_node("generic_agent", answer_generic_question)
workflow.add_node("rag_agent", answer_rag_question_)

workflow.add_conditional_edges(
    "analyze",
    lambda x: x["decision"],
    {
        "code": "code_agent",
        "rag": "rag_agent",
        "general": "generic_agent"
    }
)

workflow.set_entry_point("analyze")
workflow.add_edge("code_agent", END)
workflow.add_edge("rag_agent", END)
workflow.add_edge("generic_agent", END)

app = workflow.compile(checkpointer=memory)

And this is how I invoke my app:

config = {"configurable": {"thread_id": 5}}
user_input = "can you tell me which food is famous in Japan."

# The config is the second positional argument to stream() or invoke()!
events = app.stream(
    {"messages": [("user", user_input)]}, config, stream_mode="values"
)
for event in events:
    event["messages"][-1].pretty_print()

The Issue
After some conversations, when I ask about which country I'm talking about, the model does not remember the country. Here's an example of such a scenario:

User Input:
 "Can you tell me which food is famous in Japan?"
Response: "Sushi is famous in Japan."
User Input: "Which country am I talking about?"
Response: The model does not remember the country.

My Concern
It seems like the state, specifically the messages key, is not being maintained across different invocations.
Could you please provide insights on how to ensure that the state, especially messages, is preserved and utilized correctly across multiple interactions? Do I need to adjust the way I'm invoking the graph or modify the AgentState definition to maintain the context better?

[Glinte]

MemorySaver does not persist over program runs

[jayashriv710]

Hi Glinte!

Thanks for your response. I've implemented the system with MongoDB, and it works well when I run the LLM alone in the graph—I can get the answers previously I asked for without any issues. However, when I use multiple nodes or agents, I'm unable to retrieve answers for previously asked questions. Is there a way to resolve this issue?

[Glinte]

then no idea, sorry. As one possibility maybe changing ('user', content) to HumanMessage(content) may make it work? I am not even sure whether it is a checkpointer issue or a graph issue, you'll have to check your db and do the debugging yourself to narrow it down.

[hwchase17]

Some things to check:

• You mentioned you're using Mongo Checkpointer - does this also happen if you use the SQLite checkpointer?
• Are you updating the messages key inside the nodes? I noticed you added input/output, so maybe you are updating those instead of the mssages key?
• Without seeing the contents of the nodes, its tough to debug. If you are able to share a fully running reproducible example we can maybe help

[jayashriv710]

Hi Harrison!

Thanks for the message. Here is my work which i implemented with SQLite checkpointer. In both SQLite & MongoDB i am unable to retrieve the previous messages which i asked. Suppose if its worked means what will be limit for these memory? like how many messages we can able to store here?

`from langgraph.graph import MessagesState, StateGraph, START

class MessagesState(TypedDict):
    messages: Annotated[list, add_messages]

def primary_agent(state: MessagesState) -> dict:
    user_query = state["messages"][-1].content

    # Define the prompt template
    prompt_template = PromptTemplate(
        template=(
            "You are a primary agent. You should respond to the user in a friendly manner for the user asking a question.\n"
            "Question: {user_query}\n"
            "If the user is asking about 'healthcare policy related,' you should use RAG.\n"
            "Other than healthcare policy questions, you need to use convo_agent.\n"
            "Your answer (convo_agent / RAG): "
        ),
        input_variables=["user_query"]
    )

    # Define the primary_chain
    primary_chain = (
        {"context": RunnablePassthrough(), "user_query": RunnablePassthrough()}
        | prompt_template
        | llm
        | StrOutputParser()
    )

    # Get the response from the primary_chain
    response = primary_chain.invoke({"user_query": user_query})

    # Determine routing based on response
    if "healthcare" in user_query.lower() or "policy" in user_query.lower():
        response = rag_system(state)
        print('--rag--')
    else:
        response = convo_agent(state)
        print('--convo agent--')

    # Add the AI response to the state
    state["messages"].append(AIMessage(content=response['content']))

    # Return the updated state with new messages
    return {"messages": state["messages"]}
`

`def rag_system(state: MessagesState) -> dict:
    user_input = state["messages"][-1].content
    documents = retriever.invoke(user_input)  # Ensure documents are retrieved or default to an empty list
    context = " ".join(doc.page_content for doc in documents)

    # Define RAG prompt template
    template = """Respond to the user_input using only the provided context. Please ensure the response is professionally formatted and does not exceed 100 words:
    {context}

    Question: {user_input}
    """
    prompt = ChatPromptTemplate.from_template(template)

    # Define RAG chain
    rag_chain = (
        {"context": RunnablePassthrough(), "user_input": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    # Make sure to get the response from the rag_chain
    response = rag_chain.invoke({"context": context, "user_input": user_input})

    return {"content": response}`

`def convo_agent(state: MessagesState) -> dict:
    user_query = state["messages"][-1].content

    # Define a friendly prompt template
    prompt1 = f"""
    You are a helpful and friendly assistant. Engage with the user in a warm and conversational manner. Here is their query:

    {user_query}

    Provide a thoughtful and friendly response to assist them with their question.
    """

    # Create a ChatPromptTemplate from the friendly prompt
    prompt = ChatPromptTemplate.from_template(prompt1)

    # Define the LLM chain
    convo_chain = (
        {"user_query": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    # Get the response from the convo_chain
    response = convo_chain.invoke({"user_query": user_query})

    return {"content": response}

# Define the condition function
def decide_route(state: MessagesState):
    return state.get("route", "")


memory = SqliteSaver.from_conn_string(":memory:")

# Define the workflow
workflow = StateGraph(MessagesState)
workflow.add_node("PrimaryAgent", primary_agent)
workflow.add_node("RAG", rag_system)
workflow.add_node("ConvoAgent", convo_agent)

def decide_to_generate(state: MessagesState):
    user_query = state["messages"][-1].content.lower()
    if "healthcare" in user_query or "policy" in user_query:
        return "RAG"
    return "ConvoAgent"

workflow.add_conditional_edges("PrimaryAgent", decide_to_generate, {
    "RAG": "RAG",
    "ConvoAgent": "ConvoAgent"
})

# Set the entry and finish points
workflow.set_entry_point("PrimaryAgent")
workflow.set_finish_point("RAG")
workflow.set_finish_point("ConvoAgent")  # Adjust as needed

# Compile the application
app = workflow.compile(checkpointer=memory)

# Function to generate a unique thread ID
def get_thread_id_for_user(user_id, thread_id):
    return f"{user_id}_{thread_id}"

# Test the interaction
user_id = "user1"
thread_id = "thread1"
unique_thread_id = get_thread_id_for_user(user_id, thread_id)
config = {"configurable": {"thread_id": unique_thread_id}}`

Here are my output:-

final_state = app.invoke(
    {"messages": [HumanMessage(content="my name is jeni")]},
    config={"configurable": {"thread_id": 42}}
)
final_state["messages"][-1].content

--convo agent--
Hi Jeni! It's great to meet you! How can I assist you today? Whether you have a question or just want to chat, I'm here for you!

final_state = app.invoke(
    {"messages": [HumanMessage(content="do you remember my name?")]},
    config={"configurable": {"thread_id": 42}}
)
final_state["messages"][-1].content

--convo agent--
I’m not sure what your name is, but I’d love to know! What do you like to be called? It helps me feel more connected to you!

[hwchase17]

if you look at

def convo_agent(state: MessagesState) -> dict:
    user_query = state["messages"][-1].content

    # Define a friendly prompt template
    prompt1 = f"""
    You are a helpful and friendly assistant. Engage with the user in a warm and conversational manner. Here is their query:

    {user_query}

    Provide a thoughtful and friendly response to assist them with their question.
    """

    # Create a ChatPromptTemplate from the friendly prompt
    prompt = ChatPromptTemplate.from_template(prompt1)

    # Define the LLM chain
    convo_chain = (
        {"user_query": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    # Get the response from the convo_chain
    response = convo_chain.invoke({"user_query": user_query})

    return {"content": response}

there's a few tings youll want to change:

1/ you are only passing user_query = state["messages"][-1].content into the prompt - so its just passing the content of the last message, it has no information about previous ones! You probably want to pass all of state['messages'] into the prompt in some way

2/ you are responding with return {"content": response} - you probably want to also add in to the list of messages so that you track all messages there (so that when you pass messages into the prompt, as sugggested above, they are present)