From f63bec8578b561bc949686d3b7a228880e33e683 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Fri, 18 Jul 2025 09:59:14 -0400 Subject: [PATCH] chore(prebuilt): restructure tool node and tool injection logic (#5562) * Cleaning up the underlying tool injection logic which is happening in multiple locations. * State was being injected into the ToolCall via Send in two places in create react agent and the logic doesn't belong there, the actual injection should be happening inside the ToolNode where there's awareness of what run time parameters the tool accepts. Change is required to unblock: https://github.com/langchain-ai/langgraph/pull/5537 --- libs/prebuilt/langgraph/prebuilt/_internal.py | 26 ++++++++++ .../langgraph/prebuilt/chat_agent_executor.py | 29 +++++++---- libs/prebuilt/langgraph/prebuilt/tool_node.py | 48 ++++++++++++------- libs/prebuilt/tests/test_react_agent.py | 22 +++++---- 4 files changed, 91 insertions(+), 34 deletions(-) create mode 100644 libs/prebuilt/langgraph/prebuilt/_internal.py diff --git a/libs/prebuilt/langgraph/prebuilt/_internal.py b/libs/prebuilt/langgraph/prebuilt/_internal.py new file mode 100644 index 000000000..a5d66211c --- /dev/null +++ b/libs/prebuilt/langgraph/prebuilt/_internal.py @@ -0,0 +1,26 @@ +from typing import Any, Literal, TypedDict + +from langchain_core.messages import ToolCall + + +class ToolCallWithContext(TypedDict): + """ToolCall with additional context for graph state. + + This is an internal data-structure meant to help the ToolNode accept + tools calls with additional context (e.g. state) when dispatched using the + `Send` API. + + The Send API is used in create_react_agent to be able to distribute the tool + calls in parallel and support human-in-the-loop workflows where graph execution + may be paused for an indefinite time. + """ + + tool_call: ToolCall + __type: Literal["tool_call_with_context"] + """Type to parameterize the payload. + + Using "__" as a prefix to be defensive against potential name collisions with + regular user state. + """ + state: Any + """The state is provided as additional context.""" diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index 6425c7c3c..0b474f259 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -39,6 +39,7 @@ from langgraph.graph import END, StateGraph from langgraph.graph.message import add_messages from langgraph.graph.state import CompiledStateGraph from langgraph.managed import IsLastStep, RemainingSteps +from langgraph.prebuilt._internal import ToolCallWithContext from langgraph.prebuilt.tool_node import ToolNode from langgraph.store.base import BaseStore from langgraph.types import Checkpointer, Send @@ -650,11 +651,17 @@ def create_react_agent( 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 + return [ + Send( + "tools", + ToolCallWithContext( + __type="tool_call_with_context", + tool_call=tool_call, + state=state, + ), + ) + for tool_call in last_message.tool_calls ] - return [Send("tools", [tool_call]) for tool_call in tool_calls] # Define a new graph workflow = StateGraph(state_schema or AgentState, config_schema=config_schema) @@ -733,11 +740,17 @@ def create_react_agent( ] if pending_tool_calls: - pending_tool_calls = [ - tool_node.inject_tool_args(call, state, store) # type: ignore[arg-type] - for call in pending_tool_calls + return [ + Send( + "tools", + ToolCallWithContext( + __type="tool_call_with_context", + tool_call=tool_call, + state=state, + ), + ) + for tool_call in 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: diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 5980d0492..da84cc520 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -71,6 +71,7 @@ from pydantic import BaseModel from typing_extensions import Annotated, get_args, get_origin from langgraph.errors import GraphBubbleUp +from langgraph.prebuilt._internal import ToolCallWithContext from langgraph.store.base import BaseStore from langgraph.types import Command, Send from langgraph.utils.runnable import RunnableCallable @@ -358,7 +359,8 @@ class ToolNode(RunnableCallable): *, store: Optional[BaseStore], ) -> Any: - tool_calls, input_type = self._parse_input(input, store) + tool_calls, input_type = self._parse_input(input) + tool_calls = [self.inject_tool_args(call, input, store) for call in tool_calls] config_list = get_config_list(config, len(tool_calls)) input_types = [input_type] * len(tool_calls) with get_executor_for_config(config) as executor: @@ -379,7 +381,8 @@ class ToolNode(RunnableCallable): *, store: Optional[BaseStore], ) -> Any: - tool_calls, input_type = self._parse_input(input, store) + tool_calls, input_type = self._parse_input(input) + tool_calls = [self.inject_tool_args(call, input, store) for call in tool_calls] outputs = await asyncio.gather( *(self._arun_one(call, input_type, config) for call in tool_calls) ) @@ -436,18 +439,19 @@ class ToolNode(RunnableCallable): input_type: Literal["list", "dict", "tool_calls"], config: RunnableConfig, ) -> ToolMessage: + """Run a single tool call synchronously.""" if invalid_tool_message := self._validate_tool_call(call): return invalid_tool_message - try: - input = {**call, **{"type": "tool_call"}} - response = self.tools_by_name[call["name"]].invoke(input, config) + call_args = {**call, **{"type": "tool_call"}} + response = self.tools_by_name[call["name"]].invoke(call_args, config) # GraphInterrupt is a special exception that will always be raised. # It can be triggered in the following scenarios: # (1) a NodeInterrupt is raised inside a tool # (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool - # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool + # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph + # called as a tool # (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture) except GraphBubbleUp as e: raise e @@ -491,6 +495,7 @@ class ToolNode(RunnableCallable): input_type: Literal["list", "dict", "tool_calls"], config: RunnableConfig, ) -> ToolMessage: + """Run a single tool call asynchronously.""" if invalid_tool_message := self._validate_tool_call(call): return invalid_tool_message @@ -548,16 +553,25 @@ class ToolNode(RunnableCallable): dict[str, Any], BaseModel, ], - store: Optional[BaseStore], ) -> Tuple[list[ToolCall], Literal["list", "dict", "tool_calls"]]: + input_type: Literal["list", "dict", "tool_calls"] if isinstance(input, list): if isinstance(input[-1], dict) and input[-1].get("type") == "tool_call": input_type = "tool_calls" - tool_calls = input + tool_calls = cast(list[ToolCall], input) return tool_calls, input_type else: input_type = "list" messages = input + elif ( + isinstance(input, dict) and input.get("__type") == "tool_call_with_context" + ): + # mypy will not be able to type narrow correctly since the signature + # for input contains dict[str, Any]. We'd need to type dict[str, Any] + # before we can apply correct typing. + input = cast(ToolCallWithContext, input) # type: ignore[assignment] + input_type = "tool_calls" + return [input["tool_call"]], input_type elif isinstance(input, dict) and (messages := input.get(self.messages_key, [])): input_type = "dict" elif messages := getattr(input, self.messages_key, []): @@ -573,10 +587,7 @@ class ToolNode(RunnableCallable): except StopIteration: raise ValueError("No AIMessage found in input") - tool_calls = [ - self.inject_tool_args(call, input, store) - for call in latest_ai_message.tool_calls - ] + tool_calls = [call for call in latest_ai_message.tool_calls] return tool_calls, input_type def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]: @@ -618,15 +629,20 @@ class ToolNode(RunnableCallable): required_fields_str = ", ".join(f for f in required_fields if f) err_msg += f" State should contain fields {required_fields_str}." raise ValueError(err_msg) - if isinstance(input, dict): + + if isinstance(input, dict) and input.get("__type") == "tool_call_with_context": + state = input["state"] + else: + state = input + + if isinstance(state, dict): tool_state_args = { - tool_arg: input[state_field] if state_field else input + tool_arg: state[state_field] if state_field else state for tool_arg, state_field in state_args.items() } - else: tool_state_args = { - tool_arg: getattr(input, state_field) if state_field else input + tool_arg: getattr(state, state_field) if state_field else state for tool_arg, state_field in state_args.items() } diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index 8c110126b..d31c67256 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -5,6 +5,7 @@ from functools import partial from typing import ( Annotated, List, + Literal, Optional, Type, TypeVar, @@ -509,7 +510,7 @@ class CustomStatePydantic(AgentStatePydantic): @pytest.mark.parametrize("state_schema", [CustomState, CustomStatePydantic]) def test_react_agent_update_state( sync_checkpointer: BaseCheckpointSaver, - version: str, + version: Literal["v1", "v2"], state_schema: StateSchemaType, ) -> None: @dec_tool @@ -557,7 +558,7 @@ def test_react_agent_update_state( version=version, ) config = {"configurable": {"thread_id": "1"}} - # run until interrpupted + # Run until interrupted agent.invoke({"messages": [("user", "what's my name")]}, config) # supply the value for the interrupt response = agent.invoke(Command(resume="Archibald"), config) @@ -781,8 +782,9 @@ class AgentStateExtraKeyPydantic(AgentStatePydantic): "state_schema", [AgentStateExtraKey, AgentStateExtraKeyPydantic] ) def test_create_react_agent_inject_vars( - version: str, state_schema: StateSchemaType + version: Literal["v1", "v2"], state_schema: StateSchemaType ) -> None: + """Test that the agent can inject state and store into tool functions.""" store = InMemoryStore() namespace = ("test",) store.put(namespace, "test_key", {"bar": 3}) @@ -817,15 +819,14 @@ def test_create_react_agent_inject_vars( model = FakeToolCallingModel(tool_calls=[[tool_call], []]) agent = create_react_agent( model, - [tool1], + ToolNode([tool1], handle_tool_errors=False), state_schema=state_schema, store=store, version=version, ) - input_message = HumanMessage("hi") - result = agent.invoke({"messages": [input_message], "foo": 2}) + result = agent.invoke({"messages": [{"role": "user", "content": "hi"}], "foo": 2}) assert result["messages"] == [ - input_message, + _AnyIdHumanMessage(content="hi"), AIMessage(content="hi", tool_calls=[tool_call], id="0"), _AnyIdToolMessage(content="6", name="tool1", tool_call_id="some 0"), AIMessage("hi-hi-6", id="1"), @@ -1580,13 +1581,14 @@ def test_create_react_agent_inject_vars_with_post_model_hook( "type": "tool_call", } - def post_model_hook(state: dict) -> None: - return + def post_model_hook(state: dict) -> dict: + """Post model hook is injecting a new foo key.""" + return {"foo": 2} model = FakeToolCallingModel(tool_calls=[[tool_call], []]) agent = create_react_agent( model, - [tool1], + ToolNode([tool1], handle_tool_errors=False), state_schema=state_schema, store=store, post_model_hook=post_model_hook,