mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-07 18:27:52 +02:00
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
This commit is contained in:
Nuno Campos
@@ -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:
|
||||
|
||||
@@ -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__}"
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(),
|
||||
|
||||
Reference in New Issue
Block a user