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,