WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-10 11:47:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

chore(prebuilt): restructure tool node and tool injection logic (#5562)

Browse Source 
* 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
This commit is contained in:
Eugene Yurtsev
2025-07-18 09:59:14 -04:00
committed by GitHub
parent dc0f0c5944
commit f63bec8578
4 changed files with 91 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/prebuilt/langgraph/prebuilt/_internal.py
+26
View File
@@ -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."""
libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py
+21 -8
View File
@@ -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:
libs/prebuilt/langgraph/prebuilt/tool_node.py
+32 -16
View File
@@ -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()
}
libs/prebuilt/tests/test_react_agent.py
+12 -10
View File
@@ -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,