diff --git a/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py b/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py index 17b9652bb..a4556e232 100644 --- a/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py @@ -31,7 +31,7 @@ from langgraph.managed import IsLastStep, RemainingSteps from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.prebuilt.tool_node import ToolNode from langgraph.store.base import BaseStore -from langgraph.types import Checkpointer +from langgraph.types import Checkpointer, Send from langgraph.utils.runnable import RunnableCallable StructuredResponse = Union[dict, BaseModel] @@ -248,6 +248,9 @@ def create_react_agent( interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, debug: bool = False, + tool_call_parallelism: Literal[ + "single_tool_node", "parallel_tool_nodes" + ] = "single_tool_node", ) -> CompiledGraph: """Creates a graph that works with a chat model that utilizes tool calling. @@ -297,6 +300,15 @@ def create_react_agent( Should be one of the following: "agent", "tools". This is useful if you want to return directly or run additional processing on an output. debug: A flag indicating whether to enable debug mode. + tool_call_parallelism: Determines what state is sent to the tool node. + Can be one of: + + - `"single_tool_node"`: The tool node processes a single message. All tool + calls in the message are executed in parallel within the tool node. + - `"parallel_tool_nodes"`: The tool node processes a tool call. + Tool calls are distributed across multiple instances of the tool + node using the [Send](https://langchain-ai.github.io/langgraph/concepts/low_level/#send) + API. Returns: A compiled LangChain runnable that can be used for chat interactions. @@ -572,6 +584,11 @@ def create_react_agent( TimeoutError: Timed out at step 2 ``` """ + if tool_call_parallelism not in ("single_tool_node", "parallel_tool_nodes"): + raise ValueError( + f"Invalid version {tool_call_parallelism}. Supported versions are " + "'single_tool_node' and 'parallel_tool_nodes'." + ) if state_schema is not None: required_keys = {"messages", "remaining_steps"} @@ -747,7 +764,7 @@ def create_react_agent( ) # Define the function that determines whether to continue or not - def should_continue(state: AgentState) -> str: + def should_continue(state: AgentState) -> Union[str, list]: messages = state["messages"] last_message = messages[-1] # If there is no function call, then we finish @@ -755,7 +772,14 @@ def create_react_agent( return END if response_format is None else "generate_structured_response" # Otherwise if there is, we continue else: - return "tools" + if tool_call_parallelism == "single_tool_node": + return "tools" + elif tool_call_parallelism == "parallel_tool_nodes": + tool_calls = [ + tool_node.inject_tool_args(call, state, store) # type: ignore[arg-type] + for 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) diff --git a/libs/langgraph/langgraph/prebuilt/tool_node.py b/libs/langgraph/langgraph/prebuilt/tool_node.py index 581a89997..8175c9cda 100644 --- a/libs/langgraph/langgraph/prebuilt/tool_node.py +++ b/libs/langgraph/langgraph/prebuilt/tool_node.py @@ -135,6 +135,8 @@ class ToolNode(RunnableCallable): If multiple tool calls are requested, they will be run in parallel. The output will be a list of ToolMessages, one for each tool call. + Tool calls can also be passed directly as a list of `ToolCall` dicts. + Args: tools: A sequence of tools that can be invoked by the ToolNode. name: The name of the ToolNode in the graph. Defaults to "tools". @@ -168,10 +170,28 @@ class ToolNode(RunnableCallable): return {"messages": result} ``` + Tool calls can also be passed directly to a ToolNode. This can be useful when using + the Send API, e.g., in a conditional edge: + + ```python + def example_conditional_edge(state: dict) -> List[Send]: + tool_calls = state["messages"][-1].tool_calls + # If tools rely on state or store variables (whose values are not generated + # directly by a model), you can inject them into the tool calls. + tool_calls = [ + tool_node.inject_tool_args(call, state, store) + for call in last_message.tool_calls + ] + return [Send("tools", [tool_call]) for tool_call in tool_calls] + ``` + Important: - - The state MUST contain a list of messages. - - The last message MUST be an `AIMessage`. - - The `AIMessage` MUST have `tool_calls` populated. + - The input state can be one of the following: + - A dict with a messages key containing a list of messages. + - A list of messages. + - A list of tool calls. + - If operating on a message list, the last message must be an `AIMessage` with + `tool_calls` populated. """ name: str = "ToolNode" @@ -276,7 +296,7 @@ class ToolNode(RunnableCallable): def _run_one( self, call: ToolCall, - input_type: Literal["list", "dict"], + input_type: Literal["list", "dict", "tool_calls"], config: RunnableConfig, ) -> ToolMessage: if invalid_tool_message := self._validate_tool_call(call): @@ -331,7 +351,7 @@ class ToolNode(RunnableCallable): async def _arun_one( self, call: ToolCall, - input_type: Literal["list", "dict"], + input_type: Literal["list", "dict", "tool_calls"], config: RunnableConfig, ) -> ToolMessage: if invalid_tool_message := self._validate_tool_call(call): @@ -392,10 +412,15 @@ class ToolNode(RunnableCallable): BaseModel, ], store: Optional[BaseStore], - ) -> Tuple[list[ToolCall], Literal["list", "dict"]]: + ) -> Tuple[list[ToolCall], Literal["list", "dict", "tool_calls"]]: if isinstance(input, list): - input_type = "list" - message: AnyMessage = input[-1] + if isinstance(input[-1], dict) and input[-1].get("type") == "tool_call": + input_type = "tool_calls" + tool_calls = input + return tool_calls, input_type + else: + input_type = "list" + message: AnyMessage = input[-1] elif isinstance(input, dict) and (messages := input.get(self.messages_key, [])): input_type = "dict" message = messages[-1] @@ -410,7 +435,7 @@ class ToolNode(RunnableCallable): raise ValueError("Last message is not an AIMessage") 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 message.tool_calls ] return tool_calls, input_type @@ -490,7 +515,7 @@ class ToolNode(RunnableCallable): } return tool_call - def _inject_tool_args( + def inject_tool_args( self, tool_call: ToolCall, input: Union[ @@ -500,6 +525,21 @@ class ToolNode(RunnableCallable): ], store: Optional[BaseStore], ) -> ToolCall: + """Injects the state and store into the tool call. + + Tool arguments with types annotated as `InjectedState` and `InjectedStore` are + ignored in tool schemas for generation purposes. This method injects them into + tool calls for tool invocation. + + Args: + tool_call (ToolCall): The tool call to inject state and store into. + input (Union[list[AnyMessage], dict[str, Any], BaseModel]): The input state + to inject. + store (Optional[BaseStore]): The store to inject. + + Returns: + ToolCall: The tool call with injected state and store. + """ if tool_call["name"] not in self.tools_by_name: return tool_call @@ -509,11 +549,14 @@ class ToolNode(RunnableCallable): return tool_call_with_store def _validate_tool_command( - self, command: Command, call: ToolCall, input_type: Literal["list", "dict"] + self, + command: Command, + call: ToolCall, + input_type: Literal["list", "dict", "tool_calls"], ) -> Command: if isinstance(command.update, dict): # input type is dict when ToolNode is invoked with a dict input (e.g. {"messages": [AIMessage(..., tool_calls=[...])]}) - if input_type != "dict": + if input_type not in ("dict", "tool_calls"): raise ValueError( f"Tools can provide a dict in Command.update only when using dict with '{self.messages_key}' key as ToolNode input, " f"got: {command.update} for tool '{call['name']}'" diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 69751dc1a..3c4a518b2 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -73,6 +73,8 @@ from tests.messages import _AnyIdHumanMessage, _AnyIdToolMessage pytestmark = pytest.mark.anyio +REACT_TOOL_CALL_PARALLELISM = ["single_tool_node", "parallel_tool_nodes"] + class FakeToolCallingModel(BaseChatModel): tool_calls: Optional[list[list[ToolCall]]] = None @@ -148,13 +150,21 @@ class FakeToolCallingModel(BaseChatModel): @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_no_prompt(request: pytest.FixtureRequest, checkpointer_name: str) -> None: +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +def test_no_prompt( + request: pytest.FixtureRequest, checkpointer_name: str, tool_call_parallelism: str +) -> None: checkpointer: BaseCheckpointSaver = request.getfixturevalue( "checkpointer_" + checkpointer_name ) model = FakeToolCallingModel() - agent = create_react_agent(model, [], checkpointer=checkpointer) + agent = create_react_agent( + model, + [], + checkpointer=checkpointer, + tool_call_parallelism=tool_call_parallelism, + ) inputs = [HumanMessage("hi?")] thread = {"configurable": {"thread_id": "123"}} response = agent.invoke({"messages": inputs}, thread, debug=True) @@ -316,7 +326,8 @@ def test_runnable_prompt(): assert response == expected_response -def test_prompt_with_store(): +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +def test_prompt_with_store(tool_call_parallelism: str): def add(a: int, b: int): """Adds a and b""" return a + b @@ -336,7 +347,13 @@ def test_prompt_with_store(): model = FakeToolCallingModel() # test state modifier that uses store works - agent = create_react_agent(model, [add], prompt=prompt, store=in_memory_store) + agent = create_react_agent( + model, + [add], + prompt=prompt, + store=in_memory_store, + tool_call_parallelism=tool_call_parallelism, + ) response = agent.invoke( {"messages": [("user", "hi")]}, {"configurable": {"user_id": "1"}} ) @@ -344,7 +361,11 @@ def test_prompt_with_store(): # test state modifier that doesn't use store works agent = create_react_agent( - model, [add], prompt=prompt_no_store, store=in_memory_store + model, + [add], + prompt=prompt_no_store, + store=in_memory_store, + tool_call_parallelism=tool_call_parallelism, ) response = agent.invoke( {"messages": [("user", "hi")]}, {"configurable": {"user_id": "2"}} @@ -395,7 +416,8 @@ async def test_prompt_with_store_async(): @pytest.mark.parametrize("tool_style", ["openai", "anthropic"]) -def test_model_with_tools(tool_style: str): +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +def test_model_with_tools(tool_style: str, tool_call_parallelism: str): model = FakeToolCallingModel(tool_style=tool_style) @dec_tool @@ -409,7 +431,11 @@ def test_model_with_tools(tool_style: str): return f"Tool 2: {some_val}" # check valid agent constructor - agent = create_react_agent(model.bind_tools([tool1, tool2]), [tool1, tool2]) + agent = create_react_agent( + model.bind_tools([tool1, tool2]), + [tool1, tool2], + tool_call_parallelism=tool_call_parallelism, + ) result = agent.nodes["tools"].invoke( { "messages": [ @@ -576,7 +602,8 @@ def test__infer_handled_types() -> None: not IS_LANGCHAIN_CORE_030_OR_GREATER, reason="Pydantic v1 is required for this test to pass in langchain-core < 0.3", ) -def test_react_agent_with_structured_response() -> None: +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +def test_react_agent_with_structured_response(tool_call_parallelism: str) -> None: class WeatherResponse(BaseModel): temperature: float = Field(description="The temperature in fahrenheit") @@ -592,7 +619,10 @@ def test_react_agent_with_structured_response() -> None: ) for response_format in (WeatherResponse, ("Meow", WeatherResponse)): agent = create_react_agent( - model, [get_weather], response_format=response_format + model, + [get_weather], + response_format=response_format, + tool_call_parallelism=tool_call_parallelism, ) response = agent.invoke({"messages": [HumanMessage("What's the weather?")]}) assert response["structured_response"] == expected_structured_response @@ -600,7 +630,7 @@ def test_react_agent_with_structured_response() -> None: assert response["messages"][-2].content == "The weather is sunny and 75°F." -# tools for testing Too +# tools for testing ToolNode def tool1(some_val: int, some_other_val: str) -> str: """Tool 1 docstring.""" if some_val == 0: @@ -731,6 +761,47 @@ async def test_tool_node(): assert tool_message.tool_call_id == "some 3" +async def test_tool_node_tool_call_input(): + # Single tool call + tool_call_1 = { + "name": "tool1", + "args": {"some_val": 1, "some_other_val": "foo"}, + "id": "some 0", + "type": "tool_call", + } + result = ToolNode([tool1]).invoke([tool_call_1]) + assert result["messages"] == [ + ToolMessage(content="1 - foo", tool_call_id="some 0", name="tool1"), + ] + + # Multiple tool calls + tool_call_2 = { + "name": "tool1", + "args": {"some_val": 2, "some_other_val": "bar"}, + "id": "some 1", + "type": "tool_call", + } + result = ToolNode([tool1]).invoke([tool_call_1, tool_call_2]) + assert result["messages"] == [ + ToolMessage(content="1 - foo", tool_call_id="some 0", name="tool1"), + ToolMessage(content="2 - bar", tool_call_id="some 1", name="tool1"), + ] + + # Test with unknown tool + tool_call_3 = tool_call_1.copy() + tool_call_3["name"] = "tool2" + result = ToolNode([tool1]).invoke([tool_call_1, tool_call_3]) + assert result["messages"] == [ + ToolMessage(content="1 - foo", tool_call_id="some 0", name="tool1"), + ToolMessage( + content="Error: tool2 is not a valid tool, try one of [tool1].", + name="tool2", + tool_call_id="some 0", + status="error", + ), + ] + + async def test_tool_node_error_handling(): def handle_all(e: Union[ValueError, ToolException, ValidationError]): return TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e)) @@ -999,7 +1070,8 @@ def test_tool_node_incorrect_tool_name(): assert tool_message.tool_call_id == "some 0" -def test_tool_node_node_interrupt(): +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +def test_tool_node_node_interrupt(tool_call_parallelism: str): def tool_normal(some_val: int) -> str: """Tool docstring.""" return "normal" @@ -1045,13 +1117,14 @@ def test_tool_node_node_interrupt(): checkpointer = MemorySaver() config = {"configurable": {"thread_id": "1"}} agent = create_react_agent( - model, [tool_interrupt, tool_normal], checkpointer=checkpointer + model, + [tool_interrupt, tool_normal], + checkpointer=checkpointer, + tool_call_parallelism=tool_call_parallelism, ) result = agent.invoke({"messages": [HumanMessage("hi?")]}, config) - assert result["messages"] == [ - _AnyIdHumanMessage( - content="hi?", - ), + expected_messages = [ + _AnyIdHumanMessage(content="hi?"), AIMessage( content="hi?", id="0", @@ -1070,7 +1143,14 @@ def test_tool_node_node_interrupt(): }, ], ), + _AnyIdToolMessage(content="normal", name="tool_normal", tool_call_id="2"), ] + if tool_call_parallelism == "single_tool_node": + # Interrupt blocks second tool result + assert result["messages"] == expected_messages[:-1] + elif tool_call_parallelism == "parallel_tool_nodes": + assert result["messages"] == expected_messages + state = agent.get_state(config) assert state.next == ("tools",) task = state.tasks[0] @@ -1082,7 +1162,8 @@ def test_tool_node_node_interrupt(): not IS_LANGCHAIN_CORE_030_OR_GREATER, reason="Langchain core 0.3.0 or greater is required", ) -async def test_tool_node_command(): +@pytest.mark.parametrize("input_type", ["dict", "tool_calls"]) +async def test_tool_node_command(input_type: str): from langchain_core.tools.base import InjectedToolCallId @dec_tool @@ -1159,19 +1240,15 @@ async def test_tool_node_command(): """Add two numbers""" return a + b - result = ToolNode([add, transfer_to_bob]).invoke( - { - "messages": [ - AIMessage( - "", - tool_calls=[ - {"args": {"a": 1, "b": 2}, "id": "1", "name": "add"}, - {"args": {}, "id": "2", "name": "transfer_to_bob"}, - ], - ) - ] - } - ) + tool_calls = [ + {"args": {"a": 1, "b": 2}, "id": "1", "name": "add", "type": "tool_call"}, + {"args": {}, "id": "2", "name": "transfer_to_bob", "type": "tool_call"}, + ] + if input_type == "dict": + input_ = {"messages": [AIMessage("", tool_calls=tool_calls)]} + elif input_type == "tool_calls": + input_ = tool_calls + result = ToolNode([add, transfer_to_bob]).invoke(input_) assert result == [ { @@ -1646,7 +1723,8 @@ async def test_tool_node_command_list_input(): not IS_LANGCHAIN_CORE_030_OR_GREATER, reason="Langchain core 0.3.0 or greater is required", ) -def test_react_agent_update_state(): +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +def test_react_agent_update_state(tool_call_parallelism: str): from langchain_core.tools.base import InjectedToolCallId class State(AgentState): @@ -1684,6 +1762,7 @@ def test_react_agent_update_state(): state_schema=State, prompt=prompt, checkpointer=checkpointer, + tool_call_parallelism=tool_call_parallelism, ) config = {"configurable": {"thread_id": "1"}} # run until interrpupted @@ -1699,6 +1778,87 @@ def test_react_agent_update_state(): assert tool_message.name == "get_user_name" +@pytest.mark.skipif( + not IS_LANGCHAIN_CORE_030_OR_GREATER, + reason="Langchain core 0.3.0 or greater is required", +) +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +def test_react_agent_parallel_tool_calls(tool_call_parallelism: str): + human_assistance_execution_count = 0 + + @dec_tool + def human_assistance(query: str) -> str: + """Request assistance from a human.""" + nonlocal human_assistance_execution_count + human_response = interrupt({"query": query}) + human_assistance_execution_count += 1 + return human_response["data"] + + get_weather_execution_count = 0 + + @dec_tool + def get_weather(location: str) -> str: + """Use this tool to get the weather.""" + nonlocal get_weather_execution_count + get_weather_execution_count += 1 + return "It's sunny!" + + checkpointer = MemorySaver() + tool_calls = [ + [ + {"args": {"location": "sf"}, "id": "1", "name": "get_weather"}, + {"args": {"query": "request help"}, "id": "2", "name": "human_assistance"}, + ], + [], + ] + model = FakeToolCallingModel(tool_calls=tool_calls) + agent = create_react_agent( + model, + [human_assistance, get_weather], + checkpointer=checkpointer, + tool_call_parallelism=tool_call_parallelism, + ) + config = {"configurable": {"thread_id": "1"}} + query = "Get user assistance and also check the weather" + message_types = [] + for event in agent.stream( + {"messages": [("user", query)]}, config, stream_mode="values" + ): + message_types.append([message.type for message in event["messages"]]) + + if tool_call_parallelism == "single_tool_node": + assert message_types == [ + ["human"], + ["human", "ai"], + ] + elif tool_call_parallelism == "parallel_tool_nodes": + assert message_types == [ + ["human"], + ["human", "ai"], + ["human", "ai", "tool"], + ] + + # Resume + message_types = [] + for event in agent.stream( + Command(resume={"data": "Hello"}), config, stream_mode="values" + ): + message_types.append([message.type for message in event["messages"]]) + + assert message_types == [ + ["human", "ai"], + ["human", "ai", "tool", "tool"], + ["human", "ai", "tool", "tool", "ai"], + ] + + if tool_call_parallelism == "single_tool_node": + assert human_assistance_execution_count == 1 + assert get_weather_execution_count == 2 + elif tool_call_parallelism == "parallel_tool_nodes": + assert human_assistance_execution_count == 1 + assert get_weather_execution_count == 1 + + def my_function(some_val: int, some_other_val: str) -> str: return f"{some_val} - {some_other_val}" @@ -1887,6 +2047,49 @@ def test_tool_node_inject_state(schema_: Type[T]) -> None: assert tool_message.content == "hi?" +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +def test_create_react_agent_inject_vars(tool_call_parallelism: str) -> None: + class AgentStateExtraKey(AgentState): + foo: int + + store = InMemoryStore() + namespace = ("test",) + store.put(namespace, "test_key", {"bar": 3}) + + def tool1( + some_val: int, + state: Annotated[dict, InjectedState], + store: Annotated[BaseStore, InjectedStore()], + ) -> str: + """Tool 1 docstring.""" + store_val = store.get(namespace, "test_key").value["bar"] + return some_val + state["foo"] + store_val + + tool_call = { + "name": "tool1", + "args": {"some_val": 1}, + "id": "some 0", + "type": "tool_call", + } + model = FakeToolCallingModel(tool_calls=[[tool_call], []]) + agent = create_react_agent( + model, + [tool1], + state_schema=AgentStateExtraKey, + store=store, + tool_call_parallelism=tool_call_parallelism, + ) + input_message = HumanMessage("hi") + result = agent.invoke({"messages": [input_message], "foo": 2}) + assert result["messages"] == [ + input_message, + 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"), + ] + assert result["foo"] == 2 + + @pytest.mark.skipif( not IS_LANGCHAIN_CORE_030_OR_GREATER, reason="Langchain core 0.3.0 or greater is required", @@ -2020,7 +2223,8 @@ def test_tool_node_messages_key() -> None: ] -async def test_return_direct() -> None: +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +async def test_return_direct(tool_call_parallelism: str) -> None: @dec_tool(return_direct=True) def tool_return_direct(input: str) -> str: """A tool that returns directly.""" @@ -2044,7 +2248,11 @@ async def test_return_direct() -> None: tool_calls=first_tool_call, ) model = FakeToolCallingModel(tool_calls=[first_tool_call, []]) - agent = create_react_agent(model, [tool_return_direct, tool_normal]) + agent = create_react_agent( + model, + [tool_return_direct, tool_normal], + tool_call_parallelism=tool_call_parallelism, + ) # Test direct return for tool_return_direct result = agent.invoke( @@ -2138,7 +2346,8 @@ def test_inspect_react() -> None: inspect.getclosurevars(agent.nodes["agent"].bound.func) -def test_react_with_subgraph_tools() -> None: +@pytest.mark.parametrize("tool_call_parallelism", REACT_TOOL_CALL_PARALLELISM) +def test_react_with_subgraph_tools(tool_call_parallelism: str) -> None: class State(TypedDict): a: int b: int @@ -2187,7 +2396,12 @@ def test_react_with_subgraph_tools() -> None: ) checkpointer = MemorySaver() tool_node = ToolNode([addition, multiplication], handle_tool_errors=False) - agent = create_react_agent(model, tool_node, checkpointer=checkpointer) + agent = create_react_agent( + model, + tool_node, + checkpointer=checkpointer, + tool_call_parallelism=tool_call_parallelism, + ) result = agent.invoke( {"messages": [HumanMessage(content="What's 2 + 3 and 2 * 3?")]}, config={"configurable": {"thread_id": "1"}},