How to update graph state with tool execution result

[jiwonmik]

class HumanInputRun(BaseTool):
    """Tool that asks user for input."""

    name: str = "human"
    description: str = (
        "You can ask a human for getting answer for the clarifying question. "
        "The input should be a question and the answer options for the human."
    )
    prompt_func: Callable[[str], None] = Field(default_factory=lambda: _print_func)
    input_func: Callable = Field(default_factory=lambda: input)

    def _run(
        self,
        question: str,
        answer_options: List[str],
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        """Use the Human input tool."""
        self.prompt_func(question, answer_options)
        answer = self.input_func()
        return {"clarifying_question_answer": answer, "clarifying_question": question, "answer_options": answer_options}
    
tool = HumanInputRun(input_func=get_input)
tools = [tool]
tool_node = ToolNode(tools)

I have this tool and I want to update the state with what this tool returns.

However, it only returns Toolmessage with the return dict in content
(e.g., messages=[..., Toolmessage(content='{"clarifying_question_answer": answer, "clarifying_question": question, "answer_options": answer_options}'])

What I want to do is to update the graph state with this {"clarifying_question_answer": answer, "clarifying_question": question, "answer_options": answer_options}.

How can I do that?

[hwchase17]

if you want to update the state directly, you should probably add this as a node, not as a tool. currently I dont think we have great way for tools to update state - they generate a tool message

[alessnet]

@hwchase17 Just to clarify, does it mean it's ok to pass graph state but it's not possible to update state? I checked this doc Does this doc say it allows reading state but not updating ?

I might be wrong, but just by looking at the source code it seems safe to modify the injected State from the tool call, because the state is simply passed from the ToolNode to the BaseTool without copy.

class GetTableDetailsTool(BaseTool):
    """...."""
    async def _arun(
        self,
        table_logical_name: str,
        
        state: Annotated[State, InjectedState]
    ) -> str:
        
       metadata = await get_entity_metadata(table_logical_name)
       state["visited_tables"][table_logical_name] = metadata # here we modify the state
       return "some string"

However it would be nice to have a confirmation this is indeed a feature and not a private implementation detail.
Can someone confirm?

Thank you :)

[markmcd]

However it would be nice to have a confirmation this is indeed a feature and not a private implementation detail.

I believe it's covered by Hyrum's Law

[vbarda]

State should only be modified via the state updates returned from the individual nodes (including the nodes with tool execution results, like ToolNode). We strongly recommend against manually mutating the state. Will post an update here as soon as we have a solution for this

[alessnet]

ok thanks for the clarification :)

It would have been easy (but not clean) to just return a special object (instead of a string) to inform Langgraph that the tool is returning a State modification as well. e.g.

return ToolResult("result of the tool", state={"new_state": "xxx"})

However, the BaseTool class from langchain_core always coerse the return of the tool to a string, even before Langgraph's ToolNode could do anything about it.

So maybe we could inject a kind of "State Modifier" in the tool call using the Annotated[] typehint ?

def _run(
        self,
        question: str,
        answer_options: List[str],

        state_modifier: Annotated[StateModifier, InjectedStateModifier], # <-----
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        """Use the Human input tool."""
        self.prompt_func(question, answer_options)
        answer = self.input_func()

        # state change would be accepted only if the tool didn't crash
        state_modifier.push({"clarifying_question_answer": answer, "clarifying_question": question, "answer_options": answer_options})

        return answser 
    

[rajasblack]

State should only be modified via the state updates returned from the individual nodes (including the nodes with tool execution results, like ToolNode). We strongly recommend against manually mutating the state. Will post an update here as soon as we have a solution for this

It will be interesting to see the solution. But it would be great to be able to update the state within a tool.

[chensylz]

I have the same problem.

[kaushik143]

@hwchase17 Is there any plan to add capability to update state from tools too in future? Not sure if ToolsMessage will have all the information returned from tools.

[NathanAP]

This should be addressed because ToolMessage's content is just a string. You probably want to return something else more complex/complete/normalized from your tools.

[humbroll]

Just fyi, I ended up implementing follow up node to address this issue.

[NathanAP]

I actually did something close to that. Since the flow is Node A > Tool > Node A > ... > Node B, we can try to find the last ToolMessage in messages with that tool name and do whatever we want. This is a terrible workout IMO but it worked...

[humbroll]

Alternatively, you can connect a follow-up node immediately after the tool nodes to handle post-processing based on the tool name. I found this to be the most straightforward approach. I also discussed this issue with Harrison in person, and he recommended overriding the ToolNode class, which might be the best solution.

[alessnet]

That's too much hassle... I have finished the implementation I have suggested in my previous comment and it work like a charm. I just had to override ToolNode to make him inject an additional (optional) parameter which allow any tool to modify the state. No need for follow-up nodes or messages list manipulation.

[sgaseretto]

@alessnet do you have any gist on how you implemented it?

[MrAlekhin]

I had to implement something unconventional, but it works.

I created a tool using the @tool(response_format="content_and_artifact") directive. The tool only receives parameters and returns them as a tuple, like this: return "some_tool_message", params. After that, I created a separate node that always executes after the tool call to handle the actual functionality. Essentially, the tool serves as an entry trigger, while the node handles the real logic.

Honestly, this approach feels odd, but it works. I haven’t found a way to modify the state directly within the tool call itself. All my attempts so far have been unsuccessful. Here’s the code example:

# My tool
@tool(response_format="content_and_artifact")
def tool_a(state: Annotated[InputState, InjectedState], config: RunnableConfig, **other_params) -> Tuple[str, dict]:
    tool_params = {
        'param1': 1,
        'param2': 2
    }
    return "some_message", tool_params

# My node
def tool_a_node(state: State):
    # Modify state here
    # Send AI messages from here
    pass  # Add your functionality

# Mapping data from the tool trigger to the actual node for processing
def map_data_from_tool_to_node(state: State):
    tool_message = state["messages"][-1]
    tool_data = tool_message.artifact
    Send("tool_a_node", tool_data)

# Adding the tool and connecting it to the node
builder.add_node("tool_a", ToolNode([tool_a]))
builder.add_conditional_edges("tool_a", map_data_from_tool_to_node, ["tool_a_node"])

Important

If anyone knows a better way to modify the state directly from the tool call, I’d appreciate the insights. So far, this workaround is the best I could come up with.

[MrAlekhin]

@vbarda mentioned that state mutation should come from ToolNode returns, while @alessnet suggested a direct state modification approach. I’m not sure which method is better—does anyone have thoughts or recommendations? For me personally, the tool call needs to serve as a trigger for the workflow rather than directly returning the tool message to the state as the first step. Instead, I need the tool message to be posted as the last step of the workflow.

[artificial-feelings]

Had the same question and I am not sure if this was solved already, but what is wrong with InjectedState and the approach explained here with examples:
https://langchain-ai.github.io/langgraph/reference/prebuilt/#langgraph.prebuilt.tool_node.InjectedState

I believe we can pass state directly to the tool. The rest is the same with default ToolNode.

from langgraph.prebuilt import InjectedState, ToolNode


class AgentState(TypedDict):
    messages: List[BaseMessage]
    foo: str

@tool
def state_tool(x: int, state: Annotated[dict, InjectedState]) -> str:
    '''Do something with state.'''
    if len(state["messages"]) > 2:
        return state["foo"] + str(x)
    else:
        return "not enough messages"

[falmanna]

Setter Getter approach is seems to be working well with me.
This is langgraph.ts, but should be the same with python.

The only thing missing is I would like to trigger the reducer instead of doing Object.assign, but not sure how yet.

async function toolNodeWithState(state: typeof GlobalState.State) {
    const tools = await generateTools(
	    context,
	    () => state,
	    (newState) => Object.assign(state, newState),
    );
    const toolNodeWithConfig = new ToolNode(tools);
    state.messages = (await toolNodeWithConfig.invoke(state)).messages;
    return state;
}

[alessnet]

@MrAlekhin I'm not doing direct state modification. I'm respecting the langgraph philosophy: you just "return" the state modifications you want your tool to apply (more below)

@sgaseretto this is how I made the modification:

First you need to define "InjectedStateModifier". That's the annotation used to tell langchain that the tool parameter is injected by us and not the AI. The actual data being injected would be something similar to my StateModifier.

class InjectedStateModifier(InjectedToolArg):
    pass


class StateModifier:
    def __init__(self,
                state_schema: Type,
                parsed_state_schema: Dict[str, Callable]
                ):
        
        self._state_schema = state_schema
        self._parsed_state_schema = parsed_state_schema
        self._state: Dict[str, Any] = {}
        
        
    def push_statechange(self, s: Dict[str, Any]):
        _inject_tool_results(self._parsed_state_schema, self._state, s) # a method that merge tool states respecting the operators (if any)

    @property
    def pending_state(self):
        return {**self._state}

Look at how the out of the box InjectedState or InjectedStore works for a example.

Then modify the ToolNode class to create a StateModifier and inject that to all tool calls. Simply mimic the existing code _inject_store() to know how to inject that class to the parameter:

  async def _afunc(
      self,
      input: Union[
          list[AnyMessage],
          dict[str, Any],
          BaseModel,
      ],
      config: RunnableConfig,
      *,
      store: BaseStore,
  ) -> Any:
      statemodifier: StateModifier = StateModifier(self._state_schema, self._parsed_state_schema) # <---
      tool_calls, output_type = self._parse_input(input, store, statemodifier)
      if output_type == "list":
          raise Exception('output_type == "list" is not supported')
      
      outputs = await asyncio.gather(
          *(self._arun_one(call, config) for call in tool_calls)
      )

      statemodifier.push_statechange({self.messages_key: outputs}) # <---
      return statemodifier.pending_state

The last line is what the ToolNode would actually return to Langgraph : a Message + any state returned by the tool.

In the end, it is super simple to use once that ToolNode is modified:

class StateForTool(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    author: str
    elem: Annotated[list[int], operator.add]

class GetTableDetailsToolM(BaseModel):
    table_logical_name: str = Field(..., description="")
    state: Annotated[StateForTool, InjectedState] = Field()
    store: Annotated[Any, InjectedStore, SkipValidation] = Field()
    statemodifier: Annotated[Any, InjectedStateModifier, SkipValidation] = Field()


class GetTableDetailsTool(BaseTool):
    name: str = "get_table_details"
    description: str = """..."""
    args_schema: Type[BaseModel] = GetTableDetailsToolM


    async def _arun(
        self,
        table_logical_name: str,

        state: Annotated[StateForTool, InjectedState],
        store: Annotated[BaseStore, InjectedStore],
        statemodifier: Annotated[StateModifier, InjectedStateModifier],
        config: RunnableConfig,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:

        # this is the state this tool want to push
        statemodifier.push_statechange({
            "elem": 69,
            "author": "Alessandro"
        })

        return "This is what my tool returns"

In the above example, the tool can decide to push some state change on its own. Please note that if there are multiple tool calls at the same time modifying "elem", they will be correctly aggregated because my implementation also check for Annotated fields operators.

Once ToolNode is patched, you can easily modify state without having to deal with followup nodes or messages list inspection.

[MrAlekhin]

@alessnet Thank you man! That should do the trick 👍 Way more elegant 🪄

[PepijnBoers]

What about using Command?

https://langchain-ai.github.io/langgraph/how-tos/update-state-from-tools/

from langgraph.prebuilt.chat_agent_executor import AgentState
from langgraph.types import Command
from langchain_core.tools import tool
from langchain_core.tools.base import InjectedToolCallId
from langchain_core.messages import ToolMessage
from langchain_core.runnables import RunnableConfig

from typing_extensions import Any, Annotated


class State(AgentState):
    # user provided
    last_name: str
    # updated by the tool
    user_info: dict[str, Any]


@tool
def lookup_user_info(
    tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig
):
    """Use this to look up user information to better assist them with their questions."""
    user_id = config.get("configurable", {}).get("user_id")
    if user_id is None:
        raise ValueError("Please provide user ID")

    if user_id not in USER_ID_TO_USER_INFO:
        raise ValueError(f"User '{user_id}' not found")

    user_info = USER_ID_TO_USER_INFO[user_id]
    return Command(
        update={
            # update the state keys
            "user_info": user_info,
            # update the message history
            "messages": [
                ToolMessage(
                    "Successfully looked up user information", tool_call_id=tool_call_id
                )
            ],
        }
    )

[vbarda]

Yup, this is the correct way of doing this -- it was just recently added!

[vbarda]

Hi folks! We have now added support for updating the state from the tools using the new Command type. Please see this how-to guide for reference https://langchain-ai.github.io/langgraph/how-tos/update-state-from-tools/. Let me know if you have any questions / feedback

[MrAlekhin]

@vbarda Christmas came early 😂🙏🎅

[humbroll]

Indeed, what a gift! @vbarda

[lucamoroz]

Hi @vbarda , thank you for the update!
I noticed that to update the state from a @tool using Command, the state must be a subclass of AgentState.
Updating a custom state that inherits from TypedDict or MessagesState leads to KeyError. Is this intentional?

[vbarda]

@lucamoroz yes, the state was subclassed from AgentState in the example only because the prebuilt create_react_agent expects it.

if you are building your own agent implementation, you don't need it and can just use TypedDict / MessagesState. the key logic for handling state updates from tools actually lives in the prebuilt ToolNode which is used by create_react_agent.

it is also not necessary to use ToolNode - you can implement your own tool calling node that returns a list of tool outputs (including Command objects)

hope this helps