From 4b7ec256e9929ba856a01db47f4961fc265fa504 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 28 Apr 2025 11:07:09 -0400 Subject: [PATCH 1/2] prebuilt: switch to executing parallel tool calls via Send by default --- libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py | 9 ++++++++- libs/prebuilt/tests/test_react_agent.py | 8 ++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index 667e9961a..6b3d72480 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -277,7 +277,7 @@ def create_react_agent( interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, debug: bool = False, - version: Literal["v1", "v2"] = "v1", + version: Literal["v1", "v2"] = "v2", name: Optional[str] = None, ) -> CompiledGraph: """Creates a graph that works with a chat model that utilizes tool calling. @@ -930,6 +930,13 @@ def create_react_agent( break if m.name in should_return_direct: return END + + # handle a case of parallel tool calls where + # the tool w/ `return_direct` was executed in a different `Send` + if isinstance(m, AIMessage) and m.tool_calls: + if any(call["name"] in should_return_direct for call in m.tool_calls): + return END + return entrypoint if should_return_direct: diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index eea4444e9..d2d6960d3 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -1078,7 +1078,9 @@ async def test_return_direct(version: str) -> None: ), ] model = FakeToolCallingModel(tool_calls=[second_tool_call, []]) - agent = create_react_agent(model, [tool_return_direct, tool_normal]) + agent = create_react_agent( + model, [tool_return_direct, tool_normal], version=version + ) result = agent.invoke( {"messages": [HumanMessage(content="Test normal", id="hum1")]} ) @@ -1107,7 +1109,9 @@ async def test_return_direct(version: str) -> None: ), ] model = FakeToolCallingModel(tool_calls=[both_tool_calls, []]) - agent = create_react_agent(model, [tool_return_direct, tool_normal]) + agent = create_react_agent( + model, [tool_return_direct, tool_normal], version=version + ) result = agent.invoke({"messages": [HumanMessage(content="Test both", id="hum2")]}) assert result["messages"] == [ HumanMessage(content="Test both", id="hum2"), From 78f809bc52cb903a14177f55304ed269e63cfd01 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 28 Apr 2025 11:55:00 -0400 Subject: [PATCH 2/2] update other tests --- libs/langgraph/tests/test_large_cases.py | 126 ++++++------------ .../langgraph/tests/test_large_cases_async.py | 49 ++++--- 2 files changed, 73 insertions(+), 102 deletions(-) diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index b94585548..0206b4802 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -39,6 +39,7 @@ from langgraph.types import ( interrupt, ) from tests.agents import AgentAction, AgentFinish +from tests.any_int import AnyInt from tests.any_str import AnyDict, AnyStr, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_SYNC, @@ -2447,13 +2448,15 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ] } - assert [ + events = [ c for c in app.stream( {"messages": [HumanMessage(content="what is weather in sf")]}, stream_mode="messages", ) - ] == [ + ] + + assert events[:3] == [ ( _AnyIdAIMessageChunk( content="", @@ -2495,8 +2498,8 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 2, "langgraph_node": "tools", - "langgraph_triggers": ("branch:to:tools",), - "langgraph_path": (PULL, "tools"), + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), "langgraph_checkpoint_ns": AnyStr("tools:"), }, ), @@ -2545,6 +2548,9 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: "ls_model_type": "chat", }, ), + ] + + assert events[3:5] == UnsortedSequence( ( _AnyIdToolMessage( content="result for another", @@ -2554,8 +2560,8 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ("branch:to:tools",), - "langgraph_path": (PULL, "tools"), + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), "langgraph_checkpoint_ns": AnyStr("tools:"), }, ), @@ -2568,11 +2574,13 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ("branch:to:tools",), - "langgraph_path": (PULL, "tools"), + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), "langgraph_checkpoint_ns": AnyStr("tools:"), }, ), + ) + assert events[5:] == [ ( _AnyIdAIMessageChunk( content="answer", @@ -2603,12 +2611,17 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: model.i = 0 # reset the model - assert ( - app.invoke( - {"messages": [HumanMessage(content="what is weather in sf")]}, - stream_mode="updates", - )[0]["agent"]["messages"] - == [ + invoke_updates_events = app.invoke( + {"messages": [HumanMessage(content="what is weather in sf")]}, + stream_mode="updates", + ) + + stream_updates_events = [ + *app.stream({"messages": [HumanMessage(content="what is weather in sf")]}) + ] + + for output in (invoke_updates_events, stream_updates_events): + assert output[:3] == [ { "agent": { "messages": [ @@ -2657,6 +2670,8 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ] } }, + ] + assert output[3:5] == UnsortedSequence( { "tools": { "messages": [ @@ -2665,6 +2680,12 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: name="search_api", tool_call_id="tool_call234", ), + ] + } + }, + { + "tools": { + "messages": [ _AnyIdToolMessage( content="result for a third one", name="search_api", @@ -2673,79 +2694,10 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ] } }, - {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, - ][0]["agent"]["messages"] - ) - - assert [ - *app.stream({"messages": [HumanMessage(content="what is weather in sf")]}) - ] == [ - { - "agent": { - "messages": [ - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _AnyIdToolMessage( - content="result for query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - } - }, - { - "agent": { - "messages": [ - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another"}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - _AnyIdToolMessage( - content="result for a third one", - name="search_api", - tool_call_id="tool_call567", - ), - ] - } - }, - {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, - ] + ) + assert output[5:] == [ + {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}} + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 3062682f7..0acdb5ca3 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -35,7 +35,8 @@ from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, Pregel from langgraph.store.memory import InMemoryStore from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter -from tests.any_str import AnyDict, AnyStr +from tests.any_int import AnyInt +from tests.any_str import AnyDict, AnyStr, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, REGULAR_CHECKPOINTERS_ASYNC, @@ -2269,13 +2270,15 @@ async def test_prebuilt_tool_chat() -> None: ] } - assert [ + events = [ c async for c in app.astream( {"messages": [HumanMessage(content="what is weather in sf")]}, stream_mode="messages", ) - ] == [ + ] + + assert events[:3] == [ ( _AnyIdAIMessageChunk( content="", @@ -2301,7 +2304,7 @@ async def test_prebuilt_tool_chat() -> None: "langgraph_step": 1, "langgraph_node": "agent", "langgraph_triggers": ("branch:to:agent",), - "langgraph_path": ("__pregel_pull", "agent"), + "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", @@ -2317,8 +2320,8 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 2, "langgraph_node": "tools", - "langgraph_triggers": ("branch:to:tools",), - "langgraph_path": ("__pregel_pull", "tools"), + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), "langgraph_checkpoint_ns": AnyStr("tools:"), }, ), @@ -2360,13 +2363,16 @@ async def test_prebuilt_tool_chat() -> None: "langgraph_step": 3, "langgraph_node": "agent", "langgraph_triggers": ("branch:to:agent",), - "langgraph_path": ("__pregel_pull", "agent"), + "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", "ls_model_type": "chat", }, ), + ] + + assert events[3:5] == UnsortedSequence( ( _AnyIdToolMessage( content="result for another", @@ -2376,8 +2382,8 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ("branch:to:tools",), - "langgraph_path": ("__pregel_pull", "tools"), + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), "langgraph_checkpoint_ns": AnyStr("tools:"), }, ), @@ -2390,11 +2396,13 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ("branch:to:tools",), - "langgraph_path": ("__pregel_pull", "tools"), + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), "langgraph_checkpoint_ns": AnyStr("tools:"), }, ), + ) + assert events[5:] == [ ( _AnyIdAIMessageChunk( content="answer", @@ -2403,7 +2411,7 @@ async def test_prebuilt_tool_chat() -> None: "langgraph_step": 5, "langgraph_node": "agent", "langgraph_triggers": ("branch:to:agent",), - "langgraph_path": ("__pregel_pull", "agent"), + "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", @@ -2412,12 +2420,13 @@ async def test_prebuilt_tool_chat() -> None: ), ] - assert [ + stream_updates_events = [ c async for c in app.astream( {"messages": [HumanMessage(content="what is weather in sf")]} ) - ] == [ + ] + assert stream_updates_events[:3] == [ { "agent": { "messages": [ @@ -2466,6 +2475,8 @@ async def test_prebuilt_tool_chat() -> None: ] } }, + ] + assert stream_updates_events[3:5] == UnsortedSequence( { "tools": { "messages": [ @@ -2474,6 +2485,12 @@ async def test_prebuilt_tool_chat() -> None: name="search_api", tool_call_id="tool_call234", ), + ] + } + }, + { + "tools": { + "messages": [ _AnyIdToolMessage( content="result for a third one", name="search_api", @@ -2482,7 +2499,9 @@ async def test_prebuilt_tool_chat() -> None: ] } }, - {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, + ) + assert stream_updates_events[5:] == [ + {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}} ]