How to manage conversation history in a ReAct Agent¶

Prerequisites

This guide assumes familiarity with the following:

Message history can grow quickly and exceed LLM context window size, whether you're building chatbots with many conversation turns or agentic systems with numerous tool calls. There are several strategies for managing the message history:

message trimming - remove first or last N messages in the history
summarization - summarize earlier messages in the history and replace them with a summary
custom strategies (e.g., message filtering, etc.)

To manage message history in create_react_agent, you need to define a pre_model_hook function or runnable that takes graph state an returns a state update:

Trimming example:

from langchain_core.messages.utils import (
    trim_messages, 
    count_tokens_approximately
)
from langgraph.prebuilt import create_react_agent

# This function will be called every time before the node that calls LLM
def pre_model_hook(state):
    trimmed_messages = trim_messages(
        state["messages"],
        strategy="last",
        token_counter=count_tokens_approximately,
        max_tokens=384,
        start_on="human",
        end_on=("human", "tool"),
    )
    # You can return updated messages either under `llm_input_messages` or 
    # `messages` key (see the note below)
    return {"llm_input_messages": trimmed_messages}

checkpointer = InMemorySaver()
agent = create_react_agent(
    model,
    tools,
    pre_model_hook=pre_model_hook,
    checkpointer=checkpointer,
)

Summarization example:

from langmem.short_term import SummarizationNode
from langchain_core.messages.utils import count_tokens_approximately
from langgraph.prebuilt.chat_agent_executor import AgentState
from langgraph.checkpoint.memory import InMemorySaver
from typing import Any

model = ChatOpenAI(model="gpt-4o")

summarization_node = SummarizationNode(
    token_counter=count_tokens_approximately,
    model=model,
    max_tokens=384,
    max_summary_tokens=128,
    output_messages_key="llm_input_messages",
)

class State(AgentState):
    # NOTE: we're adding this key to keep track of previous summary information
    # to make sure we're not summarizing on every LLM call
    context: dict[str, Any]


checkpointer = InMemorySaver()
graph = create_react_agent(
    model,
    tools,
    pre_model_hook=summarization_node,
    state_schema=State,
    checkpointer=checkpointer,
)

Important

To keep the original message history unmodified in the graph state and pass the updated history only as the input to the LLM, return updated messages under llm_input_messages key
To overwrite the original message history in the graph state with the updated history, return updated messages under messages key

To overwrite the messages key, you need to do the following:

from langchain_core.messages import RemoveMessage
from langgraph.graph.message import REMOVE_ALL_MESSAGES

def pre_model_hook(state):
    updated_messages = ...
    return {
        "messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES), *updated_messages]
        ...
    }

Setup¶

First, let's install the required packages and set our API keys

%%capture --no-stderr
%pip install -U langgraph langchain-openai langmem

import getpass
import os


def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")


_set_env("OPENAI_API_KEY")

Set up LangSmith for LangGraph development

Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here.

Keep the original message history unmodified¶

Let's build a ReAct agent with a step that manages the conversation history: when the length of the history exceeds a specified number of tokens, we will call trim_messages utility that that will reduce the history while satisfying LLM provider constraints.

There are two ways that the updated message history can be applied inside ReAct agent:

Keep the original message history unmodified in the graph state and pass the updated history only as the input to the LLM
Overwrite the original message history in the graph state with the updated history

Let's start by implementing the first one. We'll need to first define model and tools for our agent:

from langchain_openai import ChatOpenAI

model = ChatOpenAI(model="gpt-4o", temperature=0)


def get_weather(location: str) -> str:
    """Use this to get weather information."""
    if any([city in location.lower() for city in ["nyc", "new york city"]]):
        return "It might be cloudy in nyc, with a chance of rain and temperatures up to 80 degrees."
    elif any([city in location.lower() for city in ["sf", "san francisco"]]):
        return "It's always sunny in sf"
    else:
        return f"I am not sure what the weather is in {location}"


tools = [get_weather]

API Reference: ChatOpenAI

Now let's implement pre_model_hook — a function that will be added as a new node and called every time before the node that calls the LLM (the agent node).

Our implementation will wrap the trim_messages call and return the trimmed messages under llm_input_messages. This will keep the original message history unmodified in the graph state and pass the updated history only as the input to the LLM

from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver

from langchain_core.messages.utils import (
    trim_messages,
    count_tokens_approximately,
)


# This function will be added as a new node in ReAct agent graph
# that will run every time before the node that calls the LLM.
# The messages returned by this function will be the input to the LLM.
def pre_model_hook(state):
    trimmed_messages = trim_messages(
        state["messages"],
        strategy="last",
        token_counter=count_tokens_approximately,
        max_tokens=384,
        start_on="human",
        end_on=("human", "tool"),
    )
    return {"llm_input_messages": trimmed_messages}


checkpointer = InMemorySaver()
graph = create_react_agent(
    model,
    tools,
    pre_model_hook=pre_model_hook,
    checkpointer=checkpointer,
)

API Reference: create_react_agent | trim_messages | count_tokens_approximately

from IPython.display import display, Image

display(Image(graph.get_graph().draw_mermaid_png()))

We'll also define a utility to render the agent outputs nicely:

def print_stream(stream, output_messages_key="llm_input_messages"):
    for chunk in stream:
        for node, update in chunk.items():
            print(f"Update from node: {node}")
            messages_key = (
                output_messages_key if node == "pre_model_hook" else "messages"
            )
            for message in update[messages_key]:
                if isinstance(message, tuple):
                    print(message)
                else:
                    message.pretty_print()

        print("\n\n")

Now let's run the agent with a few different queries to reach the specified max tokens limit:

config = {"configurable": {"thread_id": "1"}}

inputs = {"messages": [("user", "What's the weather in NYC?")]}
result = graph.invoke(inputs, config=config)

inputs = {"messages": [("user", "What's it known for?")]}
result = graph.invoke(inputs, config=config)

Let's see how many tokens we have in the message history so far:

messages = result["messages"]
count_tokens_approximately(messages)

You can see that we are close to the max_tokens threshold, so on the next invocation we should see pre_model_hook kick-in and trim the message history. Let's run it again:

inputs = {"messages": [("user", "where can i find the best bagel?")]}
print_stream(graph.stream(inputs, config=config, stream_mode="updates"))

Update from node: pre_model_hook
================================[1m Human Message [0m=================================

What's it known for?
==================================[1m Ai Message [0m==================================

New York City is known for a variety of iconic landmarks, cultural institutions, and vibrant neighborhoods. Some of the most notable features include:

1. **Statue of Liberty**: A symbol of freedom and democracy, located on Liberty Island.
2. **Times Square**: Known for its bright lights, Broadway theaters, and bustling atmosphere.
3. **Central Park**: A large public park offering a natural retreat in the middle of the city.
4. **Empire State Building**: An iconic skyscraper offering panoramic views of the city.
5. **Broadway**: Famous for its world-class theater productions.
6. **Wall Street**: The financial hub of the United States.
7. **Museums**: Including the Metropolitan Museum of Art, Museum of Modern Art (MoMA), and the American Museum of Natural History.
8. **Diverse Cuisine**: A melting pot of cultures offering a wide range of culinary experiences.
9. **Cultural Diversity**: A rich tapestry of cultures and communities from around the world.
10. **Fashion**: A global fashion capital, hosting events like New York Fashion Week.

These are just a few highlights of what makes New York City a unique and vibrant place.
================================[1m Human Message [0m=================================

where can i find the best bagel?



Update from node: agent
==================================[1m Ai Message [0m==================================

New York City is famous for its bagels, and there are several places renowned for serving some of the best. Here are a few top spots where you can find excellent bagels in NYC:

1. **Ess-a-Bagel**: Known for their large, chewy bagels with a variety of spreads and toppings.
2. **Russ & Daughters**: A classic spot offering traditional bagels with high-quality smoked fish and cream cheese.
3. **H&H Bagels**: Famous for their fresh, hand-rolled bagels.
4. **Murray’s Bagels**: Offers a wide selection of bagels and spreads, with a no-toasting policy to preserve freshness.
5. **Absolute Bagels**: Known for their authentic, fluffy bagels and a variety of cream cheese options.
6. **Tompkins Square Bagels**: Offers creative bagel sandwiches and a wide range of spreads.
7. **Bagel Hole**: Known for their smaller, denser bagels with a crispy crust.

Each of these places has its own unique style and flavor, so it might be worth trying a few to find your personal favorite!

You can see that the pre_model_hook node now only returned the last 3 messages, as expected. However, the existing message history is untouched:

updated_messages = graph.get_state(config).values["messages"]
assert [(m.type, m.content) for m in updated_messages[: len(messages)]] == [
    (m.type, m.content) for m in messages
]

Overwrite the original message history¶

Let's now change the pre_model_hook to overwrite the message history in the graph state. To do this, we’ll return the updated messages under messages key. We’ll also include a special RemoveMessage(REMOVE_ALL_MESSAGES) object, which tells create_react_agent to remove previous messages from the graph state:

from langchain_core.messages import RemoveMessage
from langgraph.graph.message import REMOVE_ALL_MESSAGES


def pre_model_hook(state):
    trimmed_messages = trim_messages(
        state["messages"],
        strategy="last",
        token_counter=count_tokens_approximately,
        max_tokens=384,
        start_on="human",
        end_on=("human", "tool"),
    )
    # NOTE that we're now returning the messages under the `messages` key
    # We also remove the existing messages in the history to ensure we're overwriting the history
    return {"messages": [RemoveMessage(REMOVE_ALL_MESSAGES)] + trimmed_messages}


checkpointer = InMemorySaver()
graph = create_react_agent(
    model,
    tools,
    pre_model_hook=pre_model_hook,
    checkpointer=checkpointer,
)

API Reference: RemoveMessage

Now let's run the agent with the same queries as before:

config = {"configurable": {"thread_id": "1"}}

inputs = {"messages": [("user", "What's the weather in NYC?")]}
result = graph.invoke(inputs, config=config)

inputs = {"messages": [("user", "What's it known for?")]}
result = graph.invoke(inputs, config=config)
messages = result["messages"]

inputs = {"messages": [("user", "where can i find the best bagel?")]}
print_stream(
    graph.stream(inputs, config=config, stream_mode="updates"),
    output_messages_key="messages",
)

Update from node: pre_model_hook
================================[1m Remove Message [0m================================


================================[1m Human Message [0m=================================

What's it known for?
==================================[1m Ai Message [0m==================================

New York City is known for a variety of iconic landmarks, cultural institutions, and vibrant neighborhoods. Some of the most notable features include:

1. **Statue of Liberty**: A symbol of freedom and democracy, located on Liberty Island.
2. **Times Square**: Known for its bright lights, Broadway theaters, and bustling atmosphere.
3. **Central Park**: A large public park offering a natural oasis amidst the urban environment.
4. **Empire State Building**: An iconic skyscraper offering panoramic views of the city.
5. **Broadway**: Famous for its world-class theater productions and musicals.
6. **Wall Street**: The financial hub of the United States, located in the Financial District.
7. **Museums**: Including the Metropolitan Museum of Art, Museum of Modern Art (MoMA), and the American Museum of Natural History.
8. **Diverse Cuisine**: A melting pot of cultures, offering a wide range of international foods.
9. **Cultural Diversity**: Known for its diverse population and vibrant cultural scene.
10. **Brooklyn Bridge**: An iconic suspension bridge connecting Manhattan and Brooklyn.

These are just a few highlights, as NYC is a city with endless attractions and activities.
================================[1m Human Message [0m=================================

where can i find the best bagel?



Update from node: agent
==================================[1m Ai Message [0m==================================

New York City is famous for its bagels, and there are several places renowned for serving some of the best. Here are a few top spots where you can find delicious bagels in NYC:

1. **Ess-a-Bagel**: Known for its large, chewy bagels and a wide variety of spreads and toppings. Locations in Midtown and the East Village.

2. **Russ & Daughters**: A historic appetizing store on the Lower East Side, famous for its bagels with lox and cream cheese.

3. **Absolute Bagels**: Located on the Upper West Side, this spot is popular for its fresh, fluffy bagels.

4. **Murray’s Bagels**: Known for its traditional, hand-rolled bagels. Located in Greenwich Village.

5. **Tompkins Square Bagels**: Offers a wide selection of bagels and creative cream cheese flavors. Located in the East Village.

6. **Bagel Hole**: A small shop in Park Slope, Brooklyn, known for its classic, no-frills bagels.

7. **Leo’s Bagels**: Located in the Financial District, known for its authentic New York-style bagels.

Each of these places has its own unique style and flavor, so it might be worth trying a few to find your personal favorite!

You can see that the pre_model_hook node returned the last 3 messages again. However, this time, the message history is modified in the graph state as well:

updated_messages = graph.get_state(config).values["messages"]
assert (
    # First 2 messages in the new history are the same as last 2 messages in the old
    [(m.type, m.content) for m in updated_messages[:2]]
    == [(m.type, m.content) for m in messages[-2:]]
)

Summarizing message history¶

Finally, let's apply a different strategy for managing message history — summarization. Just as with trimming, you can choose to keep original message history unmodified or overwrite it. The example below will only show the former.

We will use the SummarizationNode from the prebuilt langmem library. Once the message history reaches the token limit, the summarization node will summarize earlier messages to make sure they fit into max_tokens.

from langmem.short_term import SummarizationNode
from langgraph.prebuilt.chat_agent_executor import AgentState
from typing import Any

model = ChatOpenAI(model="gpt-4o")
summarization_model = model.bind(max_tokens=128)

summarization_node = SummarizationNode(
    token_counter=count_tokens_approximately,
    model=summarization_model,
    max_tokens=384,
    max_summary_tokens=128,
    output_messages_key="llm_input_messages",
)


class State(AgentState):
    # NOTE: we're adding this key to keep track of previous summary information
    # to make sure we're not summarizing on every LLM call
    context: dict[str, Any]


checkpointer = InMemorySaver()
graph = create_react_agent(
    model,
    tools,
    pre_model_hook=summarization_node,
    state_schema=State,
    checkpointer=checkpointer,
)

config = {"configurable": {"thread_id": "1"}}
inputs = {"messages": [("user", "What's the weather in NYC?")]}

result = graph.invoke(inputs, config=config)

inputs = {"messages": [("user", "What's it known for?")]}
result = graph.invoke(inputs, config=config)

inputs = {"messages": [("user", "where can i find the best bagel?")]}
print_stream(graph.stream(inputs, config=config, stream_mode="updates"))

Update from node: pre_model_hook
================================[1m System Message [0m================================

Summary of the conversation so far: The user asked about the current weather in New York City. In response, the assistant provided information that it might be cloudy, with a chance of rain, and temperatures reaching up to 80 degrees.
================================[1m Human Message [0m=================================

What's it known for?
==================================[1m Ai Message [0m==================================

New York City, often referred to as NYC, is known for its:

1. **Landmarks and Iconic Sites**:
   - **Statue of Liberty**: A symbol of freedom and democracy.
   - **Central Park**: A vast green oasis in the middle of the city.
   - **Empire State Building**: Once the tallest building in the world, offering stunning views of the city.
   - **Times Square**: Known for its bright lights and bustling atmosphere.

2. **Cultural Institutions**:
   - **Broadway**: Renowned for theatrical performances and musicals.
   - **Metropolitan Museum of Art** and **Museum of Modern Art (MoMA)**: World-class art collections.
   - **American Museum of Natural History**: Known for its extensive exhibits ranging from dinosaurs to space exploration.

3. **Diverse Neighborhoods and Cuisine**:
   - NYC is famous for having a melting pot of cultures, reflected in neighborhoods like Chinatown, Little Italy, and Harlem.
   - The city offers a wide range of international cuisines, from street food to high-end dining.

4. **Financial District**:
   - Home to Wall Street, the New York Stock Exchange (NYSE), and other major financial institutions.

5. **Media and Entertainment**:
   - Major hub for television, film, and media, with numerous studios and networks based there.

6. **Fashion**:
   - Often referred to as one of the "Big Four" fashion capitals, hosting events like New York Fashion Week.

7. **Sports**:
   - Known for its passionate sports culture with teams like the Yankees (MLB), Mets (MLB), Knicks (NBA), and Rangers (NHL).

These elements, among others, contribute to NYC's reputation as a vibrant and dynamic city.
================================[1m Human Message [0m=================================

where can i find the best bagel?



Update from node: agent
==================================[1m Ai Message [0m==================================

Finding the best bagel in New York City can be subjective, as there are many beloved spots across the city. However, here are some renowned bagel shops you might want to try:

1. **Ess-a-Bagel**: Known for its chewy and flavorful bagels, located in Midtown and Stuyvesant Town.

2. **Bagel Hole**: A favorite for traditionalists, offering classic and dense bagels, located in Park Slope, Brooklyn.

3. **Russ & Daughters**: A legendary appetizing store on the Lower East Side, famous for their bagels with lox.

4. **Murray’s Bagels**: Located in Greenwich Village, known for their fresh and authentic New York bagels.

5. **Absolute Bagels**: Located on the Upper West Side, they’re known for their fresh, fluffy bagels with a variety of spreads.

6. **Tompkins Square Bagels**: In the East Village, famous for their creative cream cheese options and fresh bagels.

7. **Zabar’s**: A landmark on the Upper West Side known for their classic bagels and smoked fish.

Each of these spots offers a unique take on the classic New York bagel experience, and trying several might be the best way to discover your personal favorite!

You can see that the earlier messages have now been replaced with the summary of the earlier conversation!

How to manage conversation history in a ReAct Agent¶

Setup¶

Keep the original message history unmodified¶

Overwrite the original message history¶

Summarizing message history¶

Comments