Langgraph State Updates with Tools

Using tools to update state was tricky to figure out. This is a quick walkthrough of some key pieces using langgraph.

Overview

At a high level, let’s assume we have a ReAct agent that calls various tools in a langgraph graph and a tool needs to be able to update state. Typically the tool just returns a result given the input parameters. However, the ability to update state leads to more possible options such as intermediate caching, sharing and updating knowledge across agents, etc.

Agent Overview

To start, an agent might be something like the following:

from typing import Literal
from langgraph.types import Command
from langgraph.prebuilt import create_react_agent

def general_agent(state) -> Command[Literal['__end__']]:
    # Define your system prompt
    SYSTEM_PROMPT = "You are a helpful assistant that can use tools to update state."
    
    model = get_model()  # e.g., ChatDatabricks(endpoint=...)

    react_agent = create_react_agent(
        model=model,
        prompt=SYSTEM_PROMPT,
        tools=[
            tool_a,
            tool_b,
            tool_z
        ],
        state_schema=CustomState # subclass langgraph.graph.MessagesState and add in params as needed
    )
    output = react_agent.invoke(state)
    new_message = output['messages'][-1]

    return Command(
        goto='__end__',
        update={
            'messages': new_message,
            'other_data': output['other_data']
        }
    )

This simple agent is given a list of tools and some state and makes tool calls as needed. The important piece comes with how those tools are constructed.

Tools

To set up a state-modifying tool, you need to give it access to the state and tool call id.

import copy
from typing import Annotated, Literal, Any
from langchain_core.tools import tool, InjectedToolCallId
from langchain_core.messages import ToolMessage
from langgraph.prebuilt import InjectedState
from langgraph.types import Command
from src.your_module import CustomState  # Adjust import path as needed
@tool
def tool_example(
    regular_params: Any,
    state: Annotated[CustomState, InjectedState],
    tool_call_id: Annotated[str, InjectedToolCallId]
) -> Command:

    # Create a copy of state to read
    state_copy = copy.deepcopy(state['key_to_copy'])

    # Mutate state as needed
    state_copy['some_key'] = 'new_value'  # Example mutation

    # Return final state
    confirmation_message = "State updated successfully"
    return Command(update={
        'key_in_state_to_update': state_copy,
        'messages': [ToolMessage(confirmation_message, tool_call_id=tool_call_id)]
    })

The key here is that the tool is returning a command back to the agent, with an update to both a key in state and the messages list itself with the tool message.

Reducer Strategies

When tools update state in parallel, you need reducers to handle conflicting updates. Reducers take two values (left and right) and return the resolved value. Here are common strategies:

Replace Reducer - Takes the most recent update:

def replace_reducer(left, right):
    return right

Then in your state you’d do something like:

from langgraph.graph import MessagesState
from typing import Annotated

class CustomState(MessagesState):
    data_that_could_conflict: Annotated[dict, replace_reducer]

Summary

This was a very quick and dirty overview of some of the pieces to keep in mind where tools update state. Hopefully it helps you avoid some key struggles I had wiring this all together.