Building Conversational Chains with Message History

The Need for Memory

As someone who has built chatbots before, you're already familiar with the core challenge of managing conversational state. In a typical implementation, you would manually maintain a list of messages, append the user's input and the AI's response after each turn, and then inject this entire history back into the prompt for the next API call. This ensures the model has the necessary context to understand follow-up questions.

LangChain abstracts this entire process with a convenient wrapper: RunnableWithMessageHistory. This component automatically handles loading the correct conversation history, passing it to your chain, and saving the new messages after the chain has run.

The diagram below illustrates the data flow. The RunnableWithMessageHistory component wraps your core logic (your "Runnable"). When invoked, it uses a session_id to load the appropriate message history, runs your chain, saves the new turn, and returns the output.

This diagram shows how RunnableWithMessageHistory orchestrates the flow of data. It takes inputs and a session_id, uses the ID to load message history, executes your chain, saves the response back to the history, and finally produces the output.

To understand this mechanism, let's start with a simple example before integrating it into our more complex RAG chain.

Implementing In-Memory History

RunnableWithMessageHistory requires two main things:

1. The chain (Runnable) to wrap.
2. A function that returns a chat history object based on a unique identifier (a session_id).

We will begin by using a simple in-memory store, which is perfect for development and testing. The following tutorial provides a clear, modern example of this pattern.

Let's apply this pattern. We'll add the necessary components to our src/lcel_quickstart/chain.py file.

First, add the required imports at the top of your file. We need MessagesPlaceholder to add a spot for the history in our prompt, and the history-related classes.

# src/lcel_quickstart/chain.py (updated imports)

# ... (keep existing imports)
from langchain_core.prompts import MessagesPlaceholder
from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

Next, let's create the history factory. This is the mechanism for providing distinct conversation histories for different users or sessions. We'll use a simple dictionary as a global store.

# src/lcel_quickstart/chain.py (add history factory)

store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    """
    Retrieves a chat history object for a given session ID.
    If the session ID doesn't exist, a new in-memory history is created.
    """
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]

Now, let's create a new chain-building function, get_conversational_rag_chain, that incorporates this history. We'll modify the prompt template to accept the chat history and then wrap our RAG chain with RunnableWithMessageHistory.

The prompt needs to be updated to handle a list of messages. We will use ChatPromptTemplate.from_messages and include a MessagesPlaceholder. This tells LangChain where to inject the conversation history.

# src/lcel_quickstart/chain.py (new conversational RAG chain function)

def get_conversational_rag_chain():
    """
    Builds and returns a conversational RAG chain with memory.
    """
    retriever = get_retriever()
    model = ChatOpenAI(model="gpt-3.5-turbo")

    # New prompt with a placeholder for chat history
    prompt = ChatPromptTemplate.from_messages([
        ("system", "Answer the user's questions based on the context below. If the answer is not in the context, say you don't know."),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{question}"),
        ("system", "Context:\n{context}")
    ])

    # The sub-chain for generating an answer, now accepts chat_history
    answer_generator = (
        {
            "context": itemgetter("documents") | RunnableLambda(format_docs),
            "question": itemgetter("question"),
            "chat_history": itemgetter("chat_history"),
        }
        | prompt
        | model
        | StrOutputParser()
    )

    # The base RAG chain that retrieves documents and generates an answer
    # Note: It now expects a dictionary with 'question' and 'chat_history'
    base_rag_chain = (
        RunnableParallel(
            {
                "documents": itemgetter("question") | retriever,
                "question": itemgetter("question"),
                "chat_history": itemgetter("chat_history"),
            }
        ).assign(answer=answer_generator)
    )

    # Wrap the base chain with history management
    conversational_rag_chain = RunnableWithMessageHistory(
        base_rag_chain,
        get_session_history,
        input_messages_key="question",
        history_messages_key="chat_history",
        output_messages_key="answer", # Specifies which output to save as AIMessage
    )

    return conversational_rag_chain

Let's carefully examine the changes:

1. Prompt: The prompt now includes MessagesPlaceholder(variable_name="chat_history"). This is where the list of HumanMessage and AIMessage objects will be injected.
2. answer_generator: The dictionary passed to this sub-chain now uses itemgetter to pull chat_history from its input and pass it along to the prompt.
3. base_rag_chain: This is the core logic before memory is added. Crucially, its input is now a dictionary { "question": ..., "chat_history": ... }. The RunnableParallel step is updated to pass the question to the retriever while also passing both question and chat_history through to the next step.
4. RunnableWithMessageHistory Wrapper: This is the final step.
   • It wraps our base_rag_chain.
   • It's configured to use our get_session_history factory.
   • input_messages_key="question" tells it that the question field in the input dictionary is the new human message.
   • history_messages_key="chat_history" tells it to inject the loaded history into the chat_history field of the input for the wrapped chain.
   • output_messages_key="answer" informs it that the answer field in the output dictionary should be saved as the AI's response in the history.

Testing the Conversational Chain

With the new chain defined, let's update the main execution block to test it. The key difference is that we now invoke the chain with a dictionary input and provide a config object containing the session_id.

# src/lcel_quickstart/chain.py (updated main block)

if __name__ == '__main__':
    conversational_chain = get_conversational_rag_chain()
    
    session_id = "test_session_123"
    config = {"configurable": {"session_id": session_id}}

    print("--- Conversational RAG Chain Demonstration ---")
    
    # First question
    query1 = "What is LCEL?"
    print(f"\nUser: {query1}")
    result1 = conversational_chain.invoke({"question": query1}, config=config)
    print(f"AI: {result1['answer']}")

    # Follow-up question
    query2 = "What is its main benefit?"
    print(f"\nUser: {query2}")
    result2 = conversational_chain.invoke({"question": query2}, config=config)
    print(f"AI: {result2['answer']}")
    
    # You can inspect the history directly
    print("\n--- Inspecting History ---")
    history = store[session_id]
    print(history.messages)
    print("--------------------------")

Run the script from your terminal:

uv run python src/lcel_quickstart/chain.py

You should see the AI answer the first question, and then correctly interpret the follow-up question "What is its main benefit?" by using the context from the first turn. When you inspect the history, you'll see the full conversation stored as a sequence of HumanMessage and AIMessage objects. Try running it with a different session_id to confirm that conversation histories are kept separate.

Conclusion

In this lesson, you have successfully added memory to our RAG application. By wrapping your chain with RunnableWithMessageHistory and providing a simple history factory, you've enabled it to handle multi-turn conversations without needing to manage the state manually.

The key takeaways are:

• RunnableWithMessageHistory is the standard LCEL component for making any chain conversational.
• It automates the loading, passing, and saving of chat history based on a session_id provided in the config.
• The wrapped chain must be designed to accept a history object, typically by including a MessagesPlaceholder in its prompt.
• An in-memory dictionary is a straightforward way to manage session histories for development, but for production, you would swap this with a persistent store like Redis or a database.

While our chain is now conversational, it has a subtle but critical flaw. When you ask a follow-up question like "what is its main benefit?", the retriever receives this exact string. This query is too ambiguous on its own and may not retrieve the correct documents. The only reason it works here is because the LLM uses the chat history at the final generation step to make sense of it. For a truly robust conversational RAG system, we need to make the retrieval step itself context-aware.

In the next lesson, we will address this by creating a preliminary chain that rephrases the user's follow-up question into a standalone query before it hits the retriever.