Fixing Memory Limitations: A Migration from LangChain to LangGraph

Solving LangChain's short-term memory issue with LangGraph's memory persistence

Introduction

I was facing lots of issues with memory using ConversationalBufferMemory with LangChain, so I found a working alternative using LangGraph.

Previously, I created a chatbot with a reAct Agent for Apple products and ran into problems with the agent forgetting the users inputs once it has finished a run. It seems to be a common issue that the memory is not persisted across agent runs using ConversationBufferMemory and LangChain in a streamlit application. According to Google and this article on memory on Langchain- “The conversational buffer memory in LangChain agents doesn't persist across separate agent runs because it's designed to be a short-term memory for each individual interaction. Each time the agent is instantiated or invoked, a new memory buffer is created, discarding any previous interactions.”

Case in point below:

Despite implementing the fix from stackoverflow link above, I was still not able to to get the chatbot to remember across runs.

Thats when I got the idea to use LangGraph persistence to share the context across multiple interactions and enable multi-turn chats

Recap

Langchain is a framework for developing Large Language Model Applications. It allows for a structured way to build, deploy and manage LLM-based applications by offering tools and components that connect LLMs to external data sources, APIs, and other services.

LangGraph is a framework built on top of LangChain, that enables the creation of complex, stateful, and cyclical workflows for AI agents using graph-based architectures.

The main components include:

• Graph-based Architecture:

  LangGraph represents workflows as directed graphs, where nodes represent functions, tools, or Large Language Models (LLMs), and edges define the connections and flow between them.

• Stateful Execution:

  LangGraph manages state across multiple turns of interaction, making it suitable for applications that require maintaining context and memory of past actions.

• Dynamic Control Flow:

  It allows for conditional routing between nodes based on the results of previous computations, enabling complex decision-making processes within the workflow.

• Extensibility:

  LangGraph extends the core Runnerble API from LangChain, making it easy to integrate with existing LangChain components and tools.

• State Management:

  LangGraph provides mechanisms for persisting and managing the state of the application, which is crucial for long-running or multi-session workflows.

Taking advantage of the Stateful Execution, I managed to resolve the memory problem for LangChain on my chatbot.

Migration

This was not a simple task as I had to create a new environment and reinstall the packages for LangGraph instead of LangChain as well as to refactor the chatbot code to LangGraph to add the individual nodes instead of simply instantiating a reAct Agent as in the case of LangChain.

Before:

llm = ChatOpenAI(model_name="gpt-4o-mini")

self.agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, memory=agent_memory)

After:

class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]

class Agent:
    def __init__(self, model, checkpointer):
        self.system = system_prompt
        graph = StateGraph(AgentState)
        graph.add_node("llm", self.call_openai)
        graph.add_node("action", self.take_action)
        graph.add_conditional_edges("llm", self.exists_action, {True: "action", False: END})
        graph.add_edge("action", "llm")
        graph.set_entry_point("llm")

Luckily, as LangGraph is based on LangChain, the tools which were created before could be used for the LangGraph agent as well.

I chose this approach instead of using the create_react_agent function in LangGraph as I would be able to add new nodes such as a decision-making node in the future.

Memory:

LangGraph’s checkpointer allowed for persistent checkpointing. If you supply a checkpointer during graph compilation and specify a thread_id when invoking the graph, LangGraph will automatically persist the state after each step. When you call the graph again with the same thread_id, it retrieves the saved state, enabling the chatbot to resume from where it previously stopped.

class Agent:
    def __init__(self, model, checkpointer):
        self.system = system_prompt
        graph = StateGraph(AgentState)
        graph.add_node("llm", self.call_openai)
        graph.add_node("action", self.take_action)
        graph.add_conditional_edges("llm", self.exists_action, {True: "action", False: END})
        graph.add_edge("action", "llm")
        graph.set_entry_point("llm")
        self.graph = graph.compile(checkpointer=checkpointer)
        self.model = model

class Chatbot:
    def __init__(self, agent: Agent, thread_id: str = "user-123"):
        self.agent = agent
        self.thread = {"configurable": {"thread_id": thread_id}}
        self.history: List = []

    def send(self, user_input: str) -> List[str]:
        self.history.append(HumanMessage(content=user_input))
        responses = []

        for event in self.agent.graph.stream({"messages": self.history}, self.thread):
            for output in event.values():
                if isinstance(output, dict) and "messages" in output:
                    new_msgs = output["messages"]
                    self.history.extend(new_msgs)
                    for msg in new_msgs:
                        if hasattr(msg, "content") and msg.content:
                            responses.append(msg.content)
        return responses

This allows for multi-turn conversations

And it allowed me to save the memory in a database:

conn = sqlite3.connect("chat_memory.db", check_same_thread=False)  
memory = SqliteSaver(conn)

TL;DR

I solved my issue with LangChain’s ConversationalBufferMemory not persisting across agent runs by migrating to LangGraph, which allowed for multi-turn conversations using its memory persistence feature.

On another topic, LangGraph seems to be better for more complex workflows such as incorporating human in the loop decision-making, or having an additional node to direct the query to different tools based on logic.