From 7bd8616b1e7c4074c30d6e87e4cec32bfe926dbb Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Mon, 19 May 2025 14:07:42 -0400 Subject: [PATCH] feature: Implement `post_model_hook` and `HumanInterruptNode` (#4583) --- .../langgraph/prebuilt/chat_agent_executor.py | 88 ++++++-- libs/prebuilt/langgraph/prebuilt/interrupt.py | 167 ++++++++++++++- libs/prebuilt/langgraph/prebuilt/tool_node.py | 17 +- .../tests/test_interrupt_tool_node.py | 191 ++++++++++++++++++ libs/prebuilt/tests/test_react_agent.py | 141 +++++++++++++ 5 files changed, 579 insertions(+), 25 deletions(-) create mode 100644 libs/prebuilt/tests/test_interrupt_tool_node.py diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index a8597d9f0..855d35e7b 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -247,6 +247,7 @@ def create_react_agent( Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]] ] = None, pre_model_hook: Optional[RunnableLike] = None, + post_model_hook: Optional[RunnableLike] = None, state_schema: Optional[StateSchemaType] = None, config_schema: Optional[Type[Any]] = None, checkpointer: Optional[Checkpointer] = None, @@ -321,6 +322,12 @@ def create_react_agent( ... } ``` + post_model_hook: An optional node to add after the `agent` node (i.e., the node that calls the LLM). + Useful for implementing human-in-the-loop, guardrails, validation, or other post-processing. + Post-model hook must be a callable or a runnable that takes in current graph state and returns a state update. + + !!! Note + Only available with `version="v2"`. state_schema: An optional state schema that defines graph state. Must have `messages` and `remaining_steps` keys. Defaults to `AgentState` that defines those two keys. @@ -591,6 +598,10 @@ def create_react_agent( workflow.set_entry_point(entrypoint) + if post_model_hook is not None: + workflow.add_node("post_model_hook", post_model_hook) + workflow.add_edge("agent", "post_model_hook") + if response_format is not None: workflow.add_node( "generate_structured_response", @@ -598,7 +609,10 @@ def create_react_agent( generate_structured_response, agenerate_structured_response ), ) - workflow.add_edge("agent", "generate_structured_response") + if post_model_hook is not None: + workflow.add_edge("post_model_hook", "generate_structured_response") + else: + workflow.add_edge("agent", "generate_structured_response") return workflow.compile( checkpointer=checkpointer, @@ -610,17 +624,24 @@ def create_react_agent( ) # Define the function that determines whether to continue or not - def should_continue(state: StateSchema) -> Union[str, list]: + def should_continue(state: StateSchema) -> Union[str, list[Send]]: messages = _get_state_value(state, "messages") last_message = messages[-1] # If there is no function call, then we finish if not isinstance(last_message, AIMessage) or not last_message.tool_calls: - return END if response_format is None else "generate_structured_response" + if post_model_hook is not None: + return "post_model_hook" + elif response_format is not None: + return "generate_structured_response" + else: + return END # Otherwise if there is, we continue else: if version == "v1": return "tools" elif version == "v2": + if post_model_hook is not None: + return "post_model_hook" tool_calls = [ tool_node.inject_tool_args(call, state, store) # type: ignore[arg-type] for call in last_message.tool_calls @@ -649,6 +670,14 @@ def create_react_agent( # This means that this node is the first one called workflow.set_entry_point(entrypoint) + agent_paths = ["tools", END] + post_model_hook_paths = [entrypoint, "tools", END] + + # Add a post model hook node if post_model_hook is provided + if post_model_hook is not None: + workflow.add_node("post_model_hook", post_model_hook) + agent_paths.append("post_model_hook") + # Add a structured output node if response_format is provided if response_format is not None: workflow.add_node( @@ -657,19 +686,52 @@ def create_react_agent( generate_structured_response, agenerate_structured_response ), ) - workflow.add_edge("generate_structured_response", END) - should_continue_destinations = ["tools", "generate_structured_response"] - else: - should_continue_destinations = ["tools", END] + if post_model_hook is not None: + post_model_hook_paths.append("generate_structured_response") + else: + agent_paths.append("generate_structured_response") + + if post_model_hook is not None: + + def post_model_hook_router(state: StateSchema) -> Union[str, list[Send]]: + """Route to the next node after post_model_hook. + + Routes to one of: + * "tools": if there are pending tool calls without a corresponding message. + * "generate_structured_response": if no pending tool calls exist and response_format is specified. + * END: if no pending tool calls exist and no response_format is specified. + """ + + messages = _get_state_value(state, "messages") + tool_messages = [ + m.tool_call_id for m in messages if isinstance(m, ToolMessage) + ] + last_ai_message = next( + m for m in reversed(messages) if isinstance(m, AIMessage) + ) + pending_tool_calls = [ + c for c in last_ai_message.tool_calls if c["id"] not in tool_messages + ] + + if pending_tool_calls: + return [Send("tools", [tool_call]) for tool_call in pending_tool_calls] + elif isinstance(messages[-1], ToolMessage): + return entrypoint + elif response_format is not None: + return "generate_structured_response" + else: + return END + + workflow.add_conditional_edges( + "post_model_hook", + post_model_hook_router, # type: ignore[arg-type] + path_map=post_model_hook_paths, + ) - # We now add a conditional edge workflow.add_conditional_edges( - # First, we define the start node. We use `agent`. - # This means these are the edges taken after the `agent` node is called. "agent", - # Next, we pass in the function that will determine which node is called next. - should_continue, - path_map=should_continue_destinations, + should_continue, # type: ignore[arg-type] + path_map=agent_paths, ) def route_tool_responses(state: StateSchema) -> str: diff --git a/libs/prebuilt/langgraph/prebuilt/interrupt.py b/libs/prebuilt/langgraph/prebuilt/interrupt.py index 2ab3ee0d1..0b5778d45 100644 --- a/libs/prebuilt/langgraph/prebuilt/interrupt.py +++ b/libs/prebuilt/langgraph/prebuilt/interrupt.py @@ -1,11 +1,12 @@ -from typing import ( - Literal, - Optional, - Union, -) +from copy import deepcopy +from typing import Any, Literal, Optional, Union, cast +from langchain_core.messages import ToolCall, ToolMessage from typing_extensions import TypedDict +from langgraph.types import Command, interrupt +from langgraph.utils.runnable import RunnableCallable + class HumanInterruptConfig(TypedDict): """Configuration that defines what actions are allowed for a human interrupt. @@ -92,3 +93,159 @@ class HumanResponse(TypedDict): type: Literal["accept", "ignore", "response", "edit"] args: Union[None, str, ActionRequest] + + +class InterruptToolNode(RunnableCallable): + """Prebuilt post model hook node used to enable common patterns for tool interrupts. + + For any tools with specified policies, an interrupt will be raised when the LLM returns + a tool call for said tool. The interrupt policy will be used to determine what sort of resume logic is allowed. + Any of the following resume patterns are supported: + + * accept: the tool call is executed as planned + * edit: the args for the tool call are edited and then the tool call is executed + * response: text response/feedback is fed back into the LLM + * ignore: the current tool call is ignored / skipped + + Args: + **interrupt_policy: a mapping of tool names to [`HumanInterruptConfig`][prebuilt.interrupt.HumanInterruptConfig] dictionaries + specifying which interrupt patterns to enable for said tool. + + Example: + ```python + from langgraph.prebuilt import create_react_agent + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.prebuilt.interrupt import HumanInterruptConfig, InterruptToolNode + from langgraph.types import Command + + + def book_hotel(hotel_name: str) -> str: + '''Book a room at the provided hotel.''' + # Some hotel API calls, a sensitive / expensive operation + return f"Booked a hotel at {hotel_name}." + + + agent = create_react_agent( + "openai:gpt-4.1", + tools=[book_hotel], + prompt="You are a hotel booking assistant.", + post_model_hook=InterruptToolNode( + book_hotel=HumanInterruptConfig( + allow_accept=True, + allow_edit=True, + allow_ignore=True, + allow_respond=True, + ) + ), + checkpointer=InMemorySaver(), + ) + + config = {"configurable": {"thread_id": 1}} + + response = agent.invoke( + {"messages": [{"role": "user", "content": "please book a hotel at the hilton inn in boston."}]}, + config=config, + ) + + response = agent.invoke(Command(resume={"type": "accept"}), config=config) + ``` + """ + + def __init__(self, **interrupt_policy: HumanInterruptConfig): + super().__init__(self._func, self._afunc) + self.interrupt_policy = interrupt_policy + + def _interrupt( + self, + tool_call: ToolCall, + interrupt_config: HumanInterruptConfig, + ) -> Union[ToolCall, ToolMessage]: + """Interrupt before a tool call and ask for human input.""" + call_id = tool_call["id"] + tool_name = tool_call["name"] + + request = HumanInterrupt( + action_request=ActionRequest( + action=tool_name, + args=tool_call["args"], + ), + config=interrupt_config, + description=f"Please review tool call for `{tool_name}` before execution.", + ) + response = interrupt([request]) + + # resume provided by agent inbox as a list + response = response[0] if isinstance(response, list) else response + + try: + response_type = response.get("type") + except AttributeError: + raise TypeError( + f"Unexpected resume value: {response}." + f"Expected a dict with `'type'` key." + ) + + if response_type == "accept" and interrupt_config["allow_accept"]: + return tool_call + elif response_type == "edit" and interrupt_config["allow_edit"]: + return ToolCall( + args=cast(ActionRequest, response)["args"]["args"], + name=tool_name, + id=call_id, + type="tool_call", + ) + elif response_type == "response" and interrupt_config["allow_respond"]: + return ToolMessage( + content=cast(str, response["args"]), + name=tool_name, + tool_call_id=call_id, + status="error", + ) + elif response_type == "ignore" and interrupt_config["allow_ignore"]: + return ToolMessage( + content=f"User ignored the tool call for `{tool_name}` with id {call_id}", + name=tool_name, + tool_call_id=call_id, + status="success", + ) + + allowed_types = [ + type_name + for type_name, is_allowed in { + "accept": interrupt_config["allow_accept"], + "edit": interrupt_config["allow_edit"], + "response": interrupt_config["allow_respond"], + "ignore": interrupt_config["allow_ignore"], + }.items() + if is_allowed + ] + + raise ValueError( + f"Unexpected human response: {response}. " + f"Expected one with `'type'` in {allowed_types} based on {tool_name}'s interrupt configuration." + ) + + def _func(self, input: dict[str, Any]) -> Command: + ai_msg = input["messages"][-1] + tool_calls: list[ToolCall] = deepcopy(ai_msg.tool_calls) or [] + tool_messages: list[ToolMessage] = [] + + for idx, tool_call in enumerate(tool_calls): + if interrupt_config := self.interrupt_policy.get(tool_call["name"]): + interrupt_result = self._interrupt( + tool_call=tool_call, interrupt_config=interrupt_config + ) + + if isinstance(interrupt_result, ToolMessage): + tool_messages.append(interrupt_result) + else: + tool_calls[idx] = interrupt_result + + updated_ai_msg = ai_msg.copy(update={"tool_calls": tool_calls}) + + # conditional routing logic for post_model_hook will direct to the tools node + # or agent node depending on if there are pending tool calls + return {"messages": [updated_ai_msg, *tool_messages]} + + async def _afunc(self, input: dict[str, Any]) -> Command: + return self._func(input) diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index b4ec0a569..03bf1c3e9 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -431,22 +431,25 @@ class ToolNode(RunnableCallable): return tool_calls, input_type else: input_type = "list" - message: AnyMessage = input[-1] + messages = input elif isinstance(input, dict) and (messages := input.get(self.messages_key, [])): input_type = "dict" - message = messages[-1] - elif messages := getattr(input, self.messages_key, None): + elif messages := getattr(input, self.messages_key, []): # Assume dataclass-like state that can coerce from dict input_type = "dict" - message = messages[-1] else: raise ValueError("No message found in input") - if not isinstance(message, AIMessage): - raise ValueError("Last message is not an AIMessage") + try: + latest_ai_message = next( + m for m in reversed(messages) if isinstance(m, AIMessage) + ) + except StopIteration: + raise ValueError("No AIMessage found in input") tool_calls = [ - self.inject_tool_args(call, input, store) for call in message.tool_calls + self.inject_tool_args(call, input, store) + for call in latest_ai_message.tool_calls ] return tool_calls, input_type diff --git a/libs/prebuilt/tests/test_interrupt_tool_node.py b/libs/prebuilt/tests/test_interrupt_tool_node.py new file mode 100644 index 000000000..e305d1201 --- /dev/null +++ b/libs/prebuilt/tests/test_interrupt_tool_node.py @@ -0,0 +1,191 @@ +import pytest +from langchain_core.messages import ToolMessage +from langchain_core.runnables import RunnableConfig + +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.prebuilt import create_react_agent +from langgraph.prebuilt.interrupt import HumanInterruptConfig, InterruptToolNode +from langgraph.types import Command +from tests.model import FakeToolCallingModel + + +def hello_tool(name: str) -> str: + """Return a greeting for the provided person.""" + return f"Hello, {name}!" + + +post_model_hook = InterruptToolNode( + hello_tool=HumanInterruptConfig( + allow_accept=True, + allow_edit=True, + allow_ignore=True, + allow_respond=True, + ) +) + +default_model = FakeToolCallingModel( + tool_calls=[ + [ + { + "name": "hello_tool", + "args": {"name": "lady gaga"}, + "id": "some-random-id", + } + ] + ] +) + + +def test_interrupt_surfaced( + request: pytest.FixtureRequest, + sync_checkpointer: BaseCheckpointSaver, +) -> None: + agent = create_react_agent( + default_model, + [hello_tool], + checkpointer=sync_checkpointer, + post_model_hook=post_model_hook, + ) + config: RunnableConfig = {"configurable": {"thread_id": "1"}} + result = agent.invoke({"messages": [("user", "Say hi to lady gaga!")]}, config) + + interrupt_data = result["__interrupt__"] + assert interrupt_data[0].value == [ + { + "action_request": {"action": "hello_tool", "args": {"name": "lady gaga"}}, + "config": { + "allow_accept": True, + "allow_edit": True, + "allow_ignore": True, + "allow_respond": True, + }, + "description": "Please review tool call for `hello_tool` before execution.", + } + ] + + response = agent.invoke(Command(resume={"type": "accept"}), config=config) + tool_message: ToolMessage = response["messages"][-2] + assert tool_message.content == "Hello, lady gaga!" + assert tool_message.name == "hello_tool" + + +@pytest.mark.parametrize( + "resume, expected_content", + [ + ({"type": "accept"}, "Hello, lady gaga!"), + ( + {"type": "ignore"}, + "User ignored the tool call for `hello_tool` with id some-random-id", + ), + ( + { + "type": "edit", + "args": {"action": "hello_tool", "args": {"name": "bruno mars"}}, + }, + "Hello, bruno mars!", + ), + ], +) +def test_interrupt_resume_variants( + request: pytest.FixtureRequest, + sync_checkpointer: BaseCheckpointSaver, + resume: dict, + expected_content: str, +) -> None: + agent = create_react_agent( + default_model, + [hello_tool], + checkpointer=sync_checkpointer, + post_model_hook=post_model_hook, + ) + + config: RunnableConfig = {"configurable": {"thread_id": "1"}} + agent.invoke({"messages": [("user", "Say hi to lady gaga!")]}, config) + + response = agent.invoke(Command(resume=resume), config=config) + tool_message: ToolMessage = response["messages"][-2] + assert tool_message.name == "hello_tool" + assert tool_message.content == expected_content + + if resume["type"] == "edit": + ai_msg = response["messages"][-1] + assert ai_msg.tool_calls == [ + { + "name": "hello_tool", + "args": {"name": "lady gaga"}, + "id": "some-random-id", + "type": "tool_call", + } + ] + + +def test_resume_with_response( + request: pytest.FixtureRequest, + sync_checkpointer: BaseCheckpointSaver, +) -> None: + model = FakeToolCallingModel( + tool_calls=[ + [ + { + "name": "hello_tool", + "args": {"name": "lady gaga"}, + "id": "some-random-id", + } + ], + [ + { + "name": "hello_tool", + "args": {"name": "bruno mars"}, + "id": "some-random-id-2", + } + ], + ] + ) + + agent = create_react_agent( + model, + [hello_tool], + checkpointer=sync_checkpointer, + post_model_hook=post_model_hook, + ) + + config: RunnableConfig = {"configurable": {"thread_id": "1"}} + agent.invoke({"messages": [("user", "Say hi to lady gaga!")]}, config) + + # Provide user response + agent.invoke( + Command( + resume={ + "type": "response", + "args": "actually, please say hello to bruno mars", + } + ), + config=config, + ) + + # Accept the updated call + response = agent.invoke(Command(resume={"type": "accept"}), config=config) + + assert len(response["messages"]) == 6 + tool_message: ToolMessage = response["messages"][-2] + assert tool_message.name == "hello_tool" + assert tool_message.content == "Hello, bruno mars!" + + +def test_resume_with_type_not_allowed(sync_checkpointer: BaseCheckpointSaver) -> None: + agent = create_react_agent( + default_model, + [hello_tool], + checkpointer=sync_checkpointer, + post_model_hook=post_model_hook, + ) + config: RunnableConfig = {"configurable": {"thread_id": "1"}} + agent.invoke({"messages": [("user", "Say hi to lady gaga!")]}, config) + + with pytest.raises(ValueError) as exc_info: + agent.invoke(Command(resume={"type": "not-allowed"}), config=config) + + assert ( + str(exc_info.value) + == "Unexpected human response: {'type': 'not-allowed'}. Expected one with `'type'` in ['accept', 'edit', 'response', 'ignore'] based on hello_tool's interrupt configuration." + ) diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index c115fe675..649b9dd6e 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -1399,3 +1399,144 @@ def test_pre_model_hook() -> None: AIMessage(content="Hello!", id="1"), ] } + + +def test_post_model_hook() -> None: + class FlagState(AgentState): + flag: bool + + model = FakeToolCallingModel(tool_calls=[]) + + def post_model_hook(state: FlagState) -> dict[str, bool]: + return {"flag": True} + + pmh_agent = create_react_agent( + model, [], post_model_hook=post_model_hook, state_schema=FlagState + ) + + assert "post_model_hook" in pmh_agent.nodes + + result = pmh_agent.invoke({"messages": [HumanMessage("hi?")], "flag": False}) + assert result["flag"] is True + + events = list(pmh_agent.stream({"messages": [HumanMessage("hi?")], "flag": False})) + assert events == [ + { + "agent": { + "messages": [ + AIMessage( + content="hi?", + additional_kwargs={}, + response_metadata={}, + id="1", + ) + ] + } + }, + {"post_model_hook": {"flag": True}}, + ] + + +def test_post_model_hook_with_structured_output() -> None: + class WeatherResponse(BaseModel): + temperature: float = Field(description="The temperature in fahrenheit") + + tool_calls = [[{"args": {}, "id": "1", "name": "get_weather"}]] + + def get_weather(): + """Get the weather""" + return "The weather is sunny and 75°F." + + expected_structured_response = WeatherResponse(temperature=75) + model = FakeToolCallingModel( + tool_calls=tool_calls, structured_response=expected_structured_response + ) + + class State(AgentState): + flag: bool + structured_response: WeatherResponse + + def post_model_hook(state: State) -> Union[dict[str, bool], Command]: + return {"flag": True} + + agent = create_react_agent( + model, + [get_weather], + response_format=WeatherResponse, + post_model_hook=post_model_hook, + state_schema=State, + ) + + assert "post_model_hook" in agent.nodes + assert "generate_structured_response" in agent.nodes + + response = agent.invoke( + {"messages": [HumanMessage("What's the weather?")], "flag": False} + ) + assert response["flag"] is True + assert response["structured_response"] == expected_structured_response + + events = list( + agent.stream({"messages": [HumanMessage("What's the weather?")], "flag": False}) + ) + assert "generate_structured_response" in events[-1] + assert events == [ + { + "agent": { + "messages": [ + AIMessage( + content="What's the weather?", + additional_kwargs={}, + response_metadata={}, + id="2", + tool_calls=[ + { + "name": "get_weather", + "args": {}, + "id": "1", + "type": "tool_call", + } + ], + ) + ] + } + }, + {"post_model_hook": {"flag": True}}, + { + "tools": { + "messages": [ + _AnyIdToolMessage( + content="The weather is sunny and 75°F.", + name="get_weather", + tool_call_id="1", + ), + ] + } + }, + { + "agent": { + "messages": [ + AIMessage( + content="What's the weather?-What's the weather?-The weather is sunny and 75°F.", + additional_kwargs={}, + response_metadata={}, + id="3", + tool_calls=[ + { + "name": "get_weather", + "args": {}, + "id": "1", + "type": "tool_call", + } + ], + ) + ] + } + }, + {"post_model_hook": {"flag": True}}, + { + "generate_structured_response": { + "structured_response": WeatherResponse(temperature=75.0) + } + }, + ]