From d6a457ef1d664f7c232da3a70d03f97a49f0f9d6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 21:26:26 -0700 Subject: [PATCH] Improve prepare_single_task trigger checks to linear complexity - Was O(n^2) due to individual channels created for every conditional edge, including the default cond edge created for Command - Now using a single channel per node for all conditional edge / command triggers, reducing to linear complexity - Improves run time on sequential(200) from 1.8s to 0.14s --- libs/langgraph/langgraph/graph/state.py | 11 ++++------- libs/langgraph/langgraph/pregel/io.py | 3 +-- libs/langgraph/tests/test_large_cases.py | 16 ++++++---------- libs/langgraph/tests/test_large_cases_async.py | 16 ++++++---------- 4 files changed, 17 insertions(+), 29 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 11df42b05..a01c3dfcb 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -821,9 +821,11 @@ class CompiledStateGraph(CompiledGraph): input_values = {k: k for k in self.builder.schemas[input_schema]} is_single_input = len(input_values) == 1 and "__root__" in input_values + branch_channel = f"branch:to:{key}" self.channels[key] = EphemeralValue(Any, guard=False) + self.channels[branch_channel] = EphemeralValue(Any, guard=False) self.nodes[key] = PregelNode( - triggers=[], + triggers=[branch_channel], # read state keys and managed values channels=(list(input_values) if is_single_input else input_values), # coerce state dict to schema class (eg. pydantic model) @@ -878,7 +880,7 @@ class CompiledStateGraph(CompiledGraph): if filtered := [p for p in packets if p != END]: writes = [ ( - ChannelWriteEntry(f"branch:{start}:{name}:{p}", start) + ChannelWriteEntry(f"branch:to:{p}", start) if not isinstance(p, Send) else p ) @@ -914,11 +916,6 @@ class CompiledStateGraph(CompiledGraph): if branch.ends else [node for node in self.builder.nodes if node != branch.then] ) - for end in ends: - if end != END: - channel_name = f"branch:{start}:{name}:{end}" - self.channels[channel_name] = EphemeralValue(Any, guard=False) - self.nodes[end].triggers.append(channel_name) # attach then subscriber if branch.then and branch.then != END: diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 30e976d99..55ed413c0 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -14,7 +14,6 @@ from langgraph.constants import ( NULL_TASK_ID, RESUME, RETURN, - SELF, START, TAG_HIDDEN, TASKS, @@ -81,7 +80,7 @@ def map_command( if isinstance(send, Send): yield (NULL_TASK_ID, TASKS, send) elif isinstance(send, str): - yield (NULL_TASK_ID, f"branch:{START}:{SELF}:{send}", START) + yield (NULL_TASK_ID, f"branch:to:{send}", START) else: raise TypeError( f"In Command.goto, expected Send/str, got {type(send).__name__}" diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index f049e4e81..57c9b4071 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -2500,7 +2500,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 2, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ["branch:to:tools"], "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2559,7 +2559,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ["branch:to:tools"], "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2573,7 +2573,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ["branch:to:tools"], "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -6706,7 +6706,7 @@ def test_branch_then( "id": AnyStr(), "name": "tool_two_slow", "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], + "triggers": ["branch:to:tool_two_slow"], }, }, { @@ -10378,9 +10378,7 @@ def test_weather_subgraph( "langgraph_node": "weather_graph", "langgraph_path": [PULL, "weather_graph"], "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], + "langgraph_triggers": ["branch:to:weather_graph"], "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), @@ -10492,9 +10490,7 @@ def test_weather_subgraph( "langgraph_node": "weather_graph", "langgraph_path": [PULL, "weather_graph"], "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], + "langgraph_triggers": ["branch:to:weather_graph"], "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 250cbfe72..95fcb9646 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -2317,7 +2317,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 2, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ["branch:to:tools"], "langgraph_path": ("__pregel_pull", "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2376,7 +2376,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ["branch:to:tools"], "langgraph_path": ("__pregel_pull", "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2390,7 +2390,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_triggers": ["branch:to:tools"], "langgraph_path": ("__pregel_pull", "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -4537,7 +4537,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "tool_two_slow", "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], + "triggers": ["branch:to:tool_two_slow"], }, }, { @@ -7231,9 +7231,7 @@ async def test_weather_subgraph( "langgraph_node": "weather_graph", "langgraph_path": [PULL, "weather_graph"], "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], + "langgraph_triggers": ["branch:to:weather_graph"], "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), @@ -7347,9 +7345,7 @@ async def test_weather_subgraph( "langgraph_node": "weather_graph", "langgraph_path": [PULL, "weather_graph"], "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], + "langgraph_triggers": ["branch:to:weather_graph"], "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(),