Commit Graph
6 Commits
Author SHA1 Message Date
8cea8ae1de fix(prebuilt): update ToolNode to allow Command update to remove all messages (#5678)
## Description

Previously, when a tool returned `Command` to update the graph's state,
the `_validate_tool_command` method in `ToolNode` would raise a
`ValueError` if the `messages_update` list contained only a
`RemoveMessage(id=REMOVE_ALL_MESSAGES)` object. This was because the
validation logic expected a matching `ToolMessage` for the tool call and
did not account for this specific state-clearing scenario.

This commit modifies the validation logic to check if the
`messages_update` list contains a single
`RemoveMessage(id=REMOVE_ALL_MESSAGES)` element. If this condition is
met, the `ToolMessage` validation is bypassed, allowing a tool to clear
the entire message history without causing a validation error.

A new test case, `test_tool_node_command_remove_all_messages`, has been
added to `tests/test_tool_node.py` to verify this change and prevent
future regressions.

## Example

Here is a self-contained example that illustrates the problem and the
fix. Without this change, the code block for `Example 2` would raise a
`ValueError`.

```python
from typing import Annotated, List

from langchain_core.messages import (
    AIMessage,
    AnyMessage,
    HumanMessage,
    RemoveMessage,
    ToolMessage,
)
from langchain_core.tools import InjectedToolCallId, tool
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph, add_messages
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langgraph.prebuilt import InjectedState, ToolNode
from langgraph.types import Command
from pydantic import BaseModel, Field


# Agent state tracks current and all messages
class AgentState(BaseModel):
    messages: Annotated[List[AnyMessage], add_messages] = Field(
        default_factory=list, description="Current conversation messages."
    )
    all_messages: Annotated[List[AnyMessage], add_messages] = Field(
        default_factory=list, description="All messages, including removed ones."
    )


# Tool to clear history if long enough, otherwise returns a warning
@tool
def clear_history_tool(
    state: Annotated[AgentState, InjectedState],
    tool_call_id: Annotated[str, InjectedToolCallId],
):
    """Clears message history if it's long enough."""
    if len(state.messages) < 3:
        return Command(
            update={
                "messages": [
                    ToolMessage(
                        "History is not long enough to be cleared. Please try again.",
                        tool_call_id=tool_call_id,
                    )
                ]
            }
        )
    else:
        return Command(
            update={
                "messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)],
                "all_messages": state.messages
                + [
                    ToolMessage(
                        "History has been successfully cleared.",
                        tool_call_id=tool_call_id,
                    )
                ],
            }
        )


# Bind the tool to the model
model = ChatOpenAI(model="gpt-4o-mini").bind_tools([clear_history_tool])


def model_node(state: AgentState):
    return {"messages": [model.invoke(state.messages)]}


# Build the agent graph
graph_builder = StateGraph(AgentState)
graph_builder.add_node("model", model_node)
graph_builder.add_node("tools", ToolNode([clear_history_tool]))
graph_builder.set_entry_point("model")
graph_builder.add_edge("model", "tools")
graph_builder.add_edge("tools", END)
graph = graph_builder.compile()


def print_messages(header, messages):
    print(f"\n{header}")
    for message in messages:
        message.pretty_print()


### Example 1: Not enough history to clear
state_1 = AgentState(
    messages=[HumanMessage(content="Please clear my message history.")]
)
output_1 = graph.invoke(state_1)
print_messages("First call: State 'messages'", output_1["messages"])
print_messages("First call: State 'all_messages'", output_1["all_messages"])

### Example 2: History is cleared
state_2 = AgentState(
    messages=[
        HumanMessage(content="Will this PR get merged?"),
        AIMessage(content="Maybe, if it's good enough."),
        HumanMessage(content="Please clear my message history."),
    ]
)
# Without the changes in this PR, the following line will raise a ValueError
output_2 = graph.invoke(state_2)
print_messages("Second call: State 'messages'", output_2["messages"])
print_messages("Second call: State 'all_messages'", output_2["all_messages"])
```

### Outputs

*Without the changes in this PR:*

```
First call: State 'messages'
================================ Human Message =================================

Please clear my message history.
================================== Ai Message ==================================
Tool Calls:
  clear_history_tool (ba421ac3-1e1a-4208-a8f6-c5500ee0abcc)
 Call ID: ba421ac3-1e1a-4208-a8f6-c5500ee0abcc
  Args:
================================= Tool Message =================================
Name: clear_history_tool

History is not long enough to be cleared. Please try again.

First call: State 'all_messages'


Traceback (most recent call last):
  File "main.py", line 114, in <module>
    output_2 = graph.invoke(state_2)
               ^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/langgraph/pregel/__init__.py", line 2844, in invoke
    for chunk in self.stream(
  File ".venv/lib/python3.11/site-packages/langgraph/pregel/__init__.py", line 2534, in stream
    for _ in runner.tick(
  File ".venv/lib/python3.11/site-packages/langgraph/prebuilt/tool_node.py", line 241, in _func
    outputs = [
              ^
  File ".venv/lib/python3.11/concurrent/futures/_base.py", line 619, in result_iterator
    yield _result_or_cancel(fs.pop())
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/concurrent/futures/_base.py", line 317, in _result_or_cancel
    return fut.result(timeout)
           ^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/concurrent/futures/_base.py", line 449, in result
    return self.__get_result()
           ^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/concurrent/futures/_base.py", line 401, in __get_result
    raise self._exception
  File ".venv/lib/python3.11/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/langchain_core/runnables/config.py", line 555, in _wrapped_fn
    return contexts.pop().run(fn, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/langgraph/prebuilt/tool_node.py", line 353, in _run_one
    return self._validate_tool_command(response, call, input_type)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/langgraph/prebuilt/tool_node.py", line 616, in _validate_tool_command
    raise ValueError(
ValueError: Expected to have a matching ToolMessage in Command.update for tool 'clear_history_tool', got: [RemoveMessage(content='', additional_kwargs={}, response_metadata={}, id='__remove_all__')]. Every tool call (LLM requesting to call a tool) in the message history MUST have a corresponding ToolMessage. You can fix it by modifying the tool to return `Command(update={"messages": [ToolMessage("Success", tool_call_id=tool_call_id), ...]}, ...)`.
```

*With the changes in this PR:*

```
First call: State 'messages'
================================ Human Message =================================

Please clear my message history.
================================== Ai Message ==================================
Tool Calls:
  clear_history_tool (ba421ac3-1e1a-4208-a8f6-c5500ee0abcc)
 Call ID: ba421ac3-1e1a-4208-a8f6-c5500ee0abcc
  Args:
================================= Tool Message =================================
Name: clear_history_tool

History is not long enough to be cleared. Please try again.

First call: State 'all_messages'


Second call: State 'messages'

Second call: State 'all_messages'
================================ Human Message =================================

Will this PR get merged?
================================== Ai Message ==================================

Maybe, if it's good enough.
================================ Human Message =================================

Please clear my message history.
================================== Ai Message ==================================
Tool Calls:
  clear_history_tool (499b1be3-6df1-493f-85e5-8d7e429dead8)
 Call ID: 499b1be3-6df1-493f-85e5-8d7e429dead8
  Args:
================================= Tool Message =================================
Name: clear_history_tool

History has been successfully cleared.
```

## Twitter handle

[@samuelpullely](https://x.com/samuelpullely)

---------

Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
2025-07-29 20:30:42 +00:00
Sydney RunkleandGitHub d1710e2eac change[langgraph]: clean up Interrupt interface for v1 (#5405) 2025-07-09 14:03:08 -04:00
Nuno Campos bc17c3522b Remove old checkpoint test fixtures
- Now all tests fully migrated to more recent sync_checkpointer and async_checkpointer fixtures for parametrising on checkpointer
- Use sync/async_store fixtures where tests used only in memory store
- Remove unused "should snapshot" check for older versions of langchain core no longer tested against
2025-05-24 14:27:13 -07:00
Vadym BardaandGitHub 4095f0a927 prebuilt: only combine Command.PARENT for Send gotos in ToolNode (#4019) 2025-03-25 15:47:46 -04:00
Vadym BardaandGitHub b97cda4290 prebuilt: add support for multiple Command(graph=Command.PARENT) returned by tools (#4003) 2025-03-25 13:51:05 -04:00
Vadym BardaandGitHub 24c13c211e langgraph: separate prebuilt into a standalone package (#3589) 2025-02-26 18:33:07 -05:00